pMixFed: Efficient Personalized Federated Learning through Adaptive Layer-Wise Mixup

The limitations of full model personalization methods with global and fully independent local models are discussed in Section 2. Partial PFL methods improve personalization by providing more flexibility through allowing clients to choose which parts of their models to be personalized based on their specific needs and constraints for improved performance. Let $L_{k}^{t}$ be a partial local model $k$ in round $t$ which is partitioned into two parts $\langle L_{l,k}^{t};L_{g,k}^{t}\rangle$, where $l,g\subseteq\{1,\dots M\}$ are the personalized and global layers, and $M$ is the number of layers. We can integrate both personalized and generalized layers in a local model $L_{k}$ as:

Among different partitioning strategies for partial PFL [34], the most popular technique is to assign local personalized layers $L_{l,k}^{t}$ to final layers and allow the base layers $L_{g,k}^{t}$ to share the knowledge similar to FedPer [4]. This choice aligns with insights from MAML ³³3Model-Agnostic Meta-Learning algorithm, suggesting that initial layers keeps general and broad information while personalized characteristics manifest prominently in the higher layers. Accordingly, we would have:

For simplicity $L_{g}=L_{g,k}$ and $L_{l}=L_{g,k}$ where $\{1\leq g\leq s\leq l\leq M\}$ and $s$ is the split(cut) layer. The objective in solving Equation 3.1 is to find the optimal $s$ (cut layer) which minimizes the personalization objective: $\sum_{k=1}^{K}\frac{1}{|\mathcal{D}_{k}|}\mathcal{F}_{k}(\theta_{k})=L_{l}^{(t)}(L_{g}^{(t)}(x_{k}))$.

In partial models, after several rounds of local training, both personalized and global layers of local model are updated. This update could be synchronous like FedSim or asynchronous as in FedAlt. The Personalized layers will be frozen until the next communication round, $L_{l}^{(t+1)}=L^{\prime}_{l}$ and the global layers will be sent to the server for global model aggregation : $G^{(t+1)}\leftarrow\sum_{k=1}^{K}\frac{|\mathcal{D}_{k}|}{|\mathcal{D}|}L^{\prime}_{g}(x_{k})$. In the next broadcasting phase, the shared layers of the local model will be updated as $L_{g}^{(t+1)}\leftarrow G^{(t+1)}$. Our goal is to benefit from Mixup to improve personalization.

Mixup is a data augmentation technique for enhancing model generalization [61] based on learning to generalize on linear combinations of training examples. Variations of Mixup have consistently excelled in vision tasks, contributing to improved robustness, generalization, and adversarial privacy. Mixup creates augmented samples as:

where $x_{i}$ and $x_{j}$ are two input samples, $y_{i}$ and $y_{j}$ are the corresponding labels, $\lambda$, $\lambda\sim\text{Beta}(\alpha,\alpha),\lambda\in[0,1]$, is the degree of interpolation between the two samples. Mixup relates data points belonging to different classes which has been shown to be successful in mitigating overfitting and improving model generalization [49, 13, 62]. There are many variants of Mixup that has been developed to address specific challenges and enhance its effectiveness. For instance, AlignMixup [48] improves local spatial alignment by introducing transformations that better preserve semantic consistency between input pairs. Manifold Mixup [49] extends the concept to hidden layer representations, acting as a powerful regularization technique by training deep neural networks (DNNs) on linear combinations of intermediate features. CutMix [60] replaces patches between images, blending visual information while retaining spatial structure. Remix [7] addresses class imbalance by assigning higher weights to minority classes during the mixing process, enhancing the robustness of the trained model on imbalanced datasets. AdaMix [13] dynamically optimizes the mixing distributions, reducing overlaps and improving training efficiency. This linear interpolation also serves as a regularization technique that shapes smoother decision boundaries, thereby enhancing the ability of a trained model to generalize to unseen data. Mixup can also increase robustness against adversarial attacks [62, 5]; and improves performance against noise, corrupted labels, and uncertainty as it relaxes the dependency on specific information [13].

Our goal is to leverage the well-established benefits of Mixup in the context of personalized FL. While Mixup has previously been employed in FL frameworks, such as XORMixup [40], FEDMIX [57], and FedMix [51], prior studies have primarily focused on using Mixup for data augmentation or data averaging. We propose pMixFed by integrating Mixup on the model parameter space, rather using it on the feature space. We apply Mixup between the parameters of the global and the local models in a layer-wise manner for more customized and adaptive PFL. Our approach eliminates the need for static and rigid partitioning strategies. Specifically, during both the broadcasting and aggregation stages of our partial PFL framework, we generate mixed model weights using a an interpolation strategy which is illustrated in Figure 2. Mixup offers the flexibility in combining models by introducing a mix degree for each layer $\lambda_{i}$, which changes gradually according to $\mu$, i.e., the mix factor. Parameter $\mu$ is also updated adaptively in each communication round and for each client according to the test accuracy of the global and the local model during the evaluation phase of FL. $\mu$ is computed as follows:

$\delta=(t/T)$, where $t$, and $T$ are the current communication round, and the overall number of communication rounds respectively. Also, $Acc$ calculated as: $Acc=(Acc_{k}^{t}/Acc_{overall}^{(t-1)})$ in the broadcasting phase and average test accuracy of the previous global model $\left(Acc=G^{t}(x,\theta^{t}),y\right)$ on all local test sets $x=\{x_{1},x_{2},\ldots,x_{K}\}$. , in the aggregation stage. $b$ is the offset parameter for sigmoid function which is set to $2$ after several experiments. More detailed discussion on parameter $\mu$’s rule of update is discussed in section 5 and the reason behind this formulation is discussed in the Appendix 8.

As shown in Figure 2, Mixup is applied in two distinct stages of FL. Firstly, when transferring shared knowledge to local models, the local model $L_{k}$ is mixed up with the current global model $G$ according to the dynamic mixing factor $\mu$, which determines the change ratio of $\lambda_{i}$ (layer-wise Mixup degree in Eq. (4)) across different layers. $\lambda_{i}$ gradually is changed from $1\rightarrow 0$ as we move from the head to the base layer. $\lambda_{i}=1$ means sharing the 100% of the global model and $\lambda_{i}=0$ means that the corresponding layer in local model is frozen and will not be mixed up with the global model $G$. Calculation of the Mixup degree of layer $i$ $\lambda_{i}$ at both broadcasting and aggregation stages is performed as follows:

where $n$ is total number of layers, $i$ is the current layer number starting from the base layer to the head as $i=0\rightarrow i=n$. $\mu$ is the mix factor which will be adaptively updated in each communication round $t$ and for each local model $k$ according to Eq. 5.

Datasets: We used three datasets widely used federated learning: MNIST [23], CIFAR-10, and CIFAR-100 [3], We followed the setup in [32] and [33] to simulate heterogeneous, non-IID data distributions across clients for both train and test datasets. The maximum number of classes per user is set to $S=5$ for CIFAR-10 and MNIST, and $S=50$ for CIFAR-100. Experiments were conducted across varying heterogeneous settings, including small-scale ($N=10$) and large-scale ($N=100$) client populations, with different client participation rates $C=[100\%,10\%]$ to measure the effects of stragglers.
Baselines and backbone: We compare two version of our method (i) pMixFed: an adaptive and dynamic mixup-based PFL approach, and (ii) pMixFed-Dynamic: a dynamic-only mixup variant where the parameter $\mu$ is fixed across communication rounds, against several baselines. These baselines include: FedAvg [32], FedAlt [34], FedSim [34], FedBABU [33]. Additionally, we compare against full model personalization methods, Ditto [28], Per-FedAvg [12] , and LG-FedAvg [29]. For all experiments, the number of local training epochs was set to $r=5$, number of communication rounds was fixed at $T=100$, and the batch size was 32. The Adam optimizer has been used while the learning rate, for both global and local updates, was $lr=0.001$ across all communication rounds. The average personalized test accuracy across individual client’s data in the final communication round is reported in Table 2 and 1. Figure 3 presents the training accuracy versus communication rounds for CIFAR10 and CIFAR100 datasets.

We evaluate two versions of our proposed method: pMixFed, which utilizes both adaptive and dynamic layer-wise mixup degree updates, and pMixFed-Dynamic, where the mixup coefficient $\mu_{i}$ is fixed across all training rounds and gradually decreases from $1\rightarrow 0$ from the head to the base layer ($i=M\rightarrow 1$) following a Sigmoid function designed for each client. Our results show that the final accuracy of pMixFed outperforms baseline methods in most cases. As shown in Figure 3, the accuracy curve of pMixFed exhibits smoother and faster convergence, which may suggest the potential for early stopping in FL settings. Previous studies on mixup also suggest that linear interpolation between features provides the most benefit in the early training phase [65]. Moreover, Figure 3 highlights that partial models with a hard split are highly sensitive to hyperparameter selection and different distribution settings, which can lead to training instability. This issue appears to be addressed more effectively in pMixFed, as discussed further in Section 4. According to the results in Table 2 and Table 1, Ditto demonstrates relatively robust performance across different heterogeneity settings. However, its effectiveness diminishes as the model depth increases, such as with MobileNet, and under larger client populations ($N=100$). On the other hand, while FedAlt and FedSim report above-baseline results, they consistently fail during training. Adjusting hyperparameters did not resolve this issue.

Adaptive Robustness to Performance Degradation: During our experiments, we observed the algorithm’s ability to adapt and recover from performance degradation, especially in challenging scenarios such as gradient vanishing, adding new users, or introducing unseen incoming data. For instance, in complex settings with larger models, such as MobileNet on the CIFAR-100 dataset, partial PFL models like FedSim and FedAlt experience sudden accuracy drops due to zero gradients or the incorporation of new participants into the cohort. We believe the reason behind this phenomenon is due to the: 1-local and global model update discrepancy in partial models with strict cut in the middle depicted in Figure 1. This degradation is mitigated by the adaptive mixup coefficient, which dynamically adjusts the degree of personalization based on the local model’s performance during both the broadcasting and aggregation stages. Specifically, if the global model $G^{(t)}$ lacks sufficient strength, the mixup coefficient $\mu^{(t)}$ is reduced, decreasing the influence of global model. 2- Catastrophic forgetting which is addressed in pMixFed by keeping the historical models $H\big{|}_{1}^{T}G^{t}$ in the aggregation process as discussed in section 4.2.2. Figure 4 illustrates that even applying mixup only to the shared layers of the same partial PFL models (FedAlt and Fedsim) enhances resilience against sudden accuracy drops, maintaining model performance over time. In both experiments, the mixup degree for personalized layers is set to $\lambda_{i}=0$ for all clients similar to FedSim and FedAlt algorithms.

Random vs gradual mix factor from $\beta$ distribution: In this paper, we have explored different designs for calculating g mix factor $\mu_{k}$. The value of $\lambda$ in Eq. 4, naturally sampled from a $\beta\neq(\alpha,\alpha)$ distribution [61] which is on the interval [0,1]. We have also experimented the random $\lambda_{i}$ using $\beta$ distribution with different $\alpha$. If $\alpha=1$, the $\beta$ distribution is uniform meaning that the $\lambda$ would be sampled uniformly from [0,1]. Moreover $\alpha>1$, The $\lambda$ would be more in between, creating a more mixed output between $L_{k}$ and $G$. On the contrary, if $\alpha<1$ the mixed model tend to choose just one of the global and local models where $\lambda=1or\lambda=0$. The effects of different $\alpha$ on mixup degree $\lambda$ is discussed in Appendix.

Among the various partitioning strategies for partial PFL discussed in the main paper and in [34], one of the most widely adopted techniques is assigning higher layers of local model $L_{l,k}^{t}$ to personalization while allowing the base layers $L_{g,k}^{t}$ to be shared across clients as the global model [33, 45, 4]. This design aligns with insights from the Model-Agnostic Meta-Learning (MAML) algorithm [10], which demonstrates that lower layers generally retain task-agnostic, generalized features, while the higher layers capture task-specific, personalized characteristics. Accordingly, in this work, we designate the head of the model as containing personalized information, while the base layers represent generalized information shared across clients ⁴⁴4It should be noted that our method is fully adaptable to different partial PFL designs as well discussed in [34]..
Broadcasting: To achieve a nuanced and gradual transition in the mixing process between the global model and the local model, we define the mix degree $\lambda_{i}$ as follows. The local model’s head $L_{n}$ remains frozen, and the head of the global model is excluded from being shared with the local model. The mix degree for each layer $\lambda_{i}$ increases incrementally based on the mix factor $\mu$, such that as we move toward the base layers, the personalization impact decreases. This dynamic behavior is represented as:

where $\lambda_{i}$ controls the degree of mixing at layer $i$. This process is visually illustrated in Figure 8.

Aggregation: The aggregation stage focuses on preserving the generalized information from the history of previous global models to mitigate the catastrophic forgetting problem. In contrast to the broadcasting process, the primary goal here is to retain the generalized information of the previous global model, which encapsulates the history of all prior models $H\big{|}_{1}^{t-1}{G}$. Therefore, the base layers should predominantly be shared from the global model, particularly during the first round of training, where the updated local models are still underdeveloped. Nonetheless, as we transition to the head, it’s less prominent to transfer knowledge from the previous global models. Figure 8 illustrates how the mixup degree transitions from the head to the base layers.

To enable online, adaptive updates of the mix degree $\lambda_{i}$, we implemented several algorithms, as detailed in Section 5.4 Ablation Study. We observed a clear relationship between the mixup degree and the similar behavior to lr, which is further discussed in Section 9.4.1. Inspired by learning rate schedulers, we introduced the mix-factor $\mu$ to adaptively update the layer-wise mixup degree based on the current communication round $\delta=t/T$ and the relative performance of the current local model $L_{k}^{t}$ compared to the global model $G^{t}$. The Sigmoid function in Figure 8 illustrates how $\mu$ evolves with respect to $\delta$ and accuracy ($Acc$). The best results were achieved when setting $b=1$ and using the square of $Acc$ as an exponent. The rationale behind this approach is that a more experienced, better-performing model should share more information. Specifically, if the local model accuracy $Acc_{l}$ significantly exceeds that of the global model ($Acc_{G}$) in the current round ($t$) such that $Acc\gg 1$, less information is shared from the global model. Conversely, when $0<Acc<1$, the global model dominates the parameter updates in both the global and local models. Figure 7 shows the distribution of the calculated $\mu$ across different communication rounds for client 0 in the broadcasting stage. In the broadcasting stage, higher $Acc_{l}$ values result in freezing more layers for personalization, leading to a decrease in $\mu$. Conversely, during the aggregation phase, if the global model accuracy $Acc_{G}$ outperforms the updated local model accuracy $Acc_{l^{\prime}}^{(t+1)}$, more base layers are shared. This is especially important during the first round of training, where local models are less stable. As demonstrated in Section 5.3, our proposed method adapts dynamically to performance drops and high-distribution complexity, adjusting the mixup degree as needed. This adaptability is effective even when applied partially to partial PFL methods such as FedAlt, and FedSim.

pMixFed is capable of handling variable model sizes across different clients. The global model, $M_{g}$, retains the maximum number of layers from all clients, i.e., $M_{G}=\max(M_{1},M_{2},\dots,M_{N})$. During the matching process between the global model $G_{i}$ and the local model $L_{i}$, if a layer block from the local model does not match a corresponding global layer, we set $\lambda_{i}=0$, meaning that the layer block will neither participate in the broadcasting nor aggregation processes. So the layers are participating according to their existing participation rate. For instance, if only 40% of clients have more than 4 layers, the generalization degree (in both broadcasting and aggregation stages), will be less than $0.4$ in each training round due to different participation rates.

We first analyze the local training updates that occur between communication rounds.

This lemma provides a bound on how the local training process improves the loss function. The bound depends on the learning rate $\eta_{l}$, the smoothness constant $L_{1}$, and the variance of the gradient estimates $\sigma^{2}$.

Next, we consider the effect of aggregating local models at the server after each communication round.

This lemma shows that, during aggregation, the global loss increases by a term proportional to the global learning rate $\eta_{g}$ and the parameter variation $\delta^{2}$.

With the above lemmas, we can now derive the convergence properties of pMixFed.

This theorem demonstrates that the global loss decreases with each communication round, with the convergence depending on the local and global learning rates, the gradient norms, and the parameter variation.

This final theorem indicates that pMixFed achieves non-convex convergence with a rate of $\mathcal{O}(1/T)$, demonstrating that the algorithm improves over time as the number of communication rounds increases.

In FedSGD, the gradients are aggregated and the server will be update the global model according to the aggregated gradients. the FedSGD is sometimes preferred over FedAvg due to its potentially faster convergence. However, it lacks robustness in heterogeneous environments. pMixFed leverages the faster convergence characteristics of FedSGD by incorporating early stopping mechanisms, facilitated by the use of mixup. As demonstrated in Section 5, the mixup factor $\lambda$ functions analogously to an SGD update at the server, even though pMixFed aggregates model weights rather than gradients, similar to FedAvg. Section 9.4 provides a more detailed explanation of this mechanism.
the pMixFed algorithm combines the advantages of FedSGD and FedAvg by aggregating model weights rather than gradients, while still ensuring convergence even in heterogeneous data settings. The incorporation of the mixup mechanism enhances stability, providing faster convergence rates compared to FedSGD, particularly in non-convex settings. Considering local SGD steps $r$ and $\eta_{l}$ as the local learning rate we can further expand the above formulation:

Since in FedSGD, the global model $G^{(t)}$ is fully shared with each local model $L_{k}^{(t)}$ in each communication round t, and $\sum_{i=1}^{k}\Omega_{k}=1$, we can update the problem as a constrained optimization problem.

The right term of the equation is the history of global model gradients which could be written as $H\big{|}_{2}^{r-1}\nabla{G^{t}}=a.G^{t}-b.G^{t+1}$. The problem of the FedSGD algorithm is that the gradients are too small to keep the state of the previous global model, leading to catastrophic forgetting in partial models that use gradients in their aggregation stage. We hypothesize that the mixup coefficient $\lambda$ acts similarly to the learning rate $\eta$, suggesting that $\lambda$ plays a role analogous to $\eta$ in FedSGD.
FedSGD Update Rule : In FedSGD, the global model $G^{(t)}$ is updated by aggregating the gradients from all clients:

where $\Omega_{k}=\frac{|\mathcal{D}_{k}|}{|\mathcal{D}|}$ is the weight associated with client $k$. Since $L_{k}^{(t+1)}=L_{k}^{(t)}-\eta_{l}\nabla L_{k}^{(t+1)}$, we can rewrite the update as:

pMixFed Update Rule: In pMixFed, the global model update incorporates mixup coefficients:

Assuming $\lambda_{k,i}=\lambda_{k}$ for all layers $i$, and considering that $G^{(t)}$ is sent to all clients at round $t$, we simplify the update to:

Dataset: We used three widely used federated learning datasets: MNIST [23], CIFAR-10, CIFAR-100 [3], and EMNIST [8]. CIFAR-10 consists of 50,000 images of size $32\times 32$ for training and 10,000 images for testing. CIFAR-100, on the other hand, consists of 100 classes, with 500 $32\times 32$ images per class for training and 100 images per class for testing. MNIST contains 10 labels and includes 60,000 samples of $28\times 28$ grayscale images for training and 10,000 for testing. For creating heterogenity, we followed the
Training Details: For evaluation, We have reported the average test accuracy of the global model [59] for different approaches. The final global model at the last communication round is saved and used during the evaluation. The global model is then personalized according to each baseline’s personalization or fine-tuning algorithm for $r=4$ local epochs and $T=50$. For FedAlt, the local model is reconstructed from the global model and fine-tuned on the test data. For FedSim, both the global and local models are fine-tuned partially but simultaneously. In the case of FedBABU, the head (fully connected layers) remains frozen during local training, while the body is updated. Since we could not directly apply pFedHN in our platform setting, we adapted their method using the same hyper parameters discussed above and employed hidden layers with 100 units for the hypernetwork and 16 kernels. The local update process for LG-FedAvg, FedAvg, and Per-FedAvg simply involves updating all layers jointly during the fine-tuning process. The global learning rate for the methods that need sgd update in the global server e.g., FedAvg, has been set from $lr_{global}=[1e-3,1e-4,and1e-5]$.It should be noted that due to the performance drop for some methods (FedAlt , FedSim) in round 10 or 40 in some settings, we’ve reported the highest accuracy achieved. Also this is the reason the accuracy curves are illustrated for 39 rounds instead of 50. ⁵⁵5As discussed in the main paper, the change of hyper-parameters such as lr, batch size, momentum and even changing the optimizer to adam didn’t help with the performance drop in most cases.

In the evaluation of the effect of learning rate and mixup, the average test accuracy⁶⁶6Classification accuracy using softmax is measured on cold-start clients $|D^{ts}_{k\cap\text{unseen}}|,k\subset\{1,\dots,M\}$, where $D^{ts}\neq D^{tr}$. These clients have not participated in the federation at any point during training. FedAlt and FedSim perform poorly on cold-start users or unseen clients, highlighting their limited generalization capability. The test accuracy of pFedMix, while affected under a 10% participation rate, benefits significantly from seeing more clients. Increased client participation directly improves accuracy, as observed in previous studies [34].

The Beta distribution $\beta(\alpha,\alpha)$ is defined on the interval [0,1], where the parameter $\alpha$ controls the shape of the distribution. The value of $\lambda$, used in Eq. 4, is naturally sampled from this distribution. By varying the parameter $\alpha$, we can adjust how much mixing occurs between the global model $G$ and the local model $L_{k}$.

• •

  Uniform Distribution: When $\alpha=1$, the Beta distribution becomes uniform over [0,1]. In this case, $\lambda$ is sampled uniformly across the entire interval, meaning that each model, $G$ and $L_{k}$, has an equal probability of being weighted more or less in the mixup process. This leads to a broad exploration of different combinations of global and local models, allowing for a wide range of mixed models.

• •

  Concentrated Mixup ( $\alpha>1$): When $\alpha>1$, the Beta distribution is concentrated around the center of the interval [0,1]. As a result, the mixup factor $\lambda$ is more likely to be closer to 0.5, leading to more balanced combinations of the global and local models. This results in outputs that are more ”mixed,” with neither model dominating the mixup process. Such a setup can enhance the robustness of the combined model, as it prevents extreme weighting of either model, creating smoother interpolations between them.

• •

  Extremal Mixup ( $\alpha<1$): In contrast, when $\alpha<1$, the Beta distribution becomes U-shaped, with more probability mass near 0 and 1. This means that $\lambda$ tends to be either very close to 0 or very close to 1, favoring one model over the other in the mixup process. When $\lambda\approx 0$, the local model $L_{k}$ is chosen almost exclusively, and when $\lambda\approx 1$, the global model $G$ is predominantly selected. This form of mixup creates a more deterministic selection between global and local models, with less mixing occurring.