Optimizing Personalized Federated Learning Through Adaptive Layer-Wise Learning

1 Introduction

Federated learning (FL) enables collaborative model training across diverse, decentralized data sources while preserving the confidentiality and integrity of each dataset. It is widely used in mobile applications like private face recognition [Niu and Deng, 2022], predictive text, speech recognition, and image annotation [Song et al., 2022]. However, data from mobile devices frequently exhibits non-independent and identically distributed (non-IID) characteristics due to variations in user behavior, device types, or regional differences [Zhu et al., 2021a]. This data heterogeneity poses significant challenges for typical FL algorithms, as the trained global model may struggle to adapt to the specific needs of individual clients, resulting in poor inference performance and slow convergence.

Personalized Federated Learning (pFL) tackles this issue by tailoring learning to each client [Tan et al., 2022]. Instead of producing a single global model, pFL generates client-specific models to better capture local data nuances. Customization can be implemented via diverse strategies, including model-wise [Luo and Wu, 2022], layer-wise [Ma et al., 2022], or element-wise [Zhang et al., 2023b] aggregation, each offering varying degrees of granularity and control.

Model-wise aggregation methods like APPLE [Luo and Wu, 2022], FedAMP [Huang et al., 2021] and Ditto [Li et al., 2021b] aggregate entire client models using learned weights, enabling broad global knowledge integration but lacking granularity for data-specific variations. This approach also incurs high computational overhead for model selection, hindering scalability. Layer-wise methods address some of these challenges by allowing local aggregation in layer units. Here, a "layer unit" can be a single neural network layer or a block comprising multiple layers. Representative approaches include FedPer [Arivazhagan et al., 2019], FedRep [Collins et al., 2021] and pFedLA [Ma et al., 2022]. These methods allow for more targeted adaptation of different parts of the network. For example, FedRep uses the global model to construct the lower layers (i.e., layers toward the input layer, also termed as base layers) of each local model and the higher layers (i.e., layers toward the model's output layer, also termed as head layers) are built solely on local data for personalization. This design strikes a balance between global knowledge sharing and local personalization. Element-wise aggregation, such as FedALA [Zhang et al., 2023b], assigns per-parameter weights for finer control, but introduces prohibitive computational/memory costs in large models. Its pre-defined static weights limit adaptability to dynamic data environments. Moreover, all above pFL methods use a constant learning rate across all clients during training, without considering the differing learning needs of various layers for each client. This may lead to overpersonalization for each client, resulting in the loss of previously aggregated global information.

To address these gaps, we introduce FLAYER, a novel layer-wise optimization framework for pFL that effectively balances global representation learning with local customization. Figure 1 outlines FLAYER 's local learning process, comprising three key stages: local model initialization (incorporating both local and global insights into the head layers), adaptive layer-wise training (mitigating gradient vanishing and facilitating faster convergence), and layer-specific parameter uploading via masking (ensuring that only essential updates are shared to preserve global features). Through this layered, adaptive approach, FLAYER effectively balances personalized learning with shared global knowledge while reducing computational overhead. We evaluate FLAYER by applying it to both image [Krizhevsky et al., 2009; Chrabaszcz et al., 2017] and text [Zhang et al., 2015] classification tasks using four widely adopted benchmarks. The results show that FLAYER outperforms eight other pFL methods in inference accuracy and computational cost. This paper makes the following contributions:

• A new performance-guided layer-wise aggregation method allows clients to dynamically incorporate both local and global information in a cost-effective manner;
• A new layer-specific adaptive learning rate scheme for pFL to steer the personalization and speed up convergence;
• A new layer-wise masking technique for selectively uploading essential parameters to the central server to improve global representation.

2 Background and Overview

2.1 ProblemDefinition

pFL learns personalized models cooperatively among clients. Consider a scenario where we have n clients, and each client processes their distinct private training data denoted as D1,D2, . . . ,DN, which has different data classes and sizes. These datasets exhibit heterogeneity, characterized by nonIID [Zhao et al., 2018]. The goal of pFL can be defined as:

where θk is the model parameters of client k, mk is the data size of Dk, M is the whole data size of all clients, Lk(θk) is the loss function of client k.

2.2 Overview of FLAYER

Figure 1 illustrates FLAYER workflow for the k-th client in iteration t. Each client begins by downloading the global model and applying an adaptive aggregation strategy that dynamically weights global and local models based on local inference accuracy, allowing stronger performers to guide head-layer initialization. During local training, a layer-wise adaptive learning rate optimizes updates for each layer according to its position and gradient, improving personalization and convergence. Finally, a layerwise masking strategy selectively uploads only essential parameters to the server, preserving critical local insights and enhancing the global representation. Algorithm 1 details the entire FL process.

3 Methodology

3.1 Performance-guided Layer-wise Local Aggregation

Building on previous findings [Yosinski et al., 2014; Zhu et al., 2021b] that base layers of deep neural network (DNN) capture generalized features while head layers encode task-specific features, we introduce a differentiated update strategy for these layers. In our approach, during the local model initialization stage, the base layers of each client's model are directly updated using parameters from the global model. This ensures the consistent refinement of generalized features across all clients. For the head layers, we use a dynamic aggregation method to integrate both global and local parameters. This integration is tailored based on the performance of the k-th client's model on its local dataset Dk from the previous iteration. The process is defined as follows:

Here, θk is the local model parameter matrix of the k-th client, ˜ θk denotes the local model parameter matrix of the kth client after initialization, L is the total number of layers, s is the number of layers in the head for personalization, and θ(1:L-s,t-1) g represents the lower L - s layers in the base part of the global model at iteration t - 1, which are used to update the base layers in the local model, and all the clients share the same base layers. θ(L-s+1:L,t-1) k denotes the head layers of the k-th local model. We aggregate the head layers from the local and global models to initialize the head layers for the local model. The aggregation weights Ak,g and Ak,l control the influence of global and local parameters, respectively. At iteration t -1, the local model inference accuracy on dataset Dk sets the local weight At-1 k,l , with (1 - Ak,l) adjusting for global influence. The dynamic weighting for head layers helps each client gradually "specialize" its head layers, without getting stuck in unproductive local optima. When local performance is low, leaning on the global head acts like a safety net, pulling the client's parameters closer to the collective knowledge. As local performance improves, the client can afford to rely more on its own specialized parameters. FLAYER ensures convergence by aligning foundational features (base layers) and employing performance-driven head adjustments (head layers), the method fosters stable updates and mitigates divergence, enabling steady convergence.

3.2 Adaptive Layer-specific Learning Rate

After the model initialization, the local model trains on its local dataset Dk:ˆ

η represents the learning rate, ∇˜ θk Lk( ˜ θtk ,Dk) is the gradient of the loss function Lk with respect to the parameters ˜ θk evaluated using local dataset Dk.

The existing pFL methods [Luo and Wu, 2022; Li et al., 2021b; Huang et al., 2021; Arivazhagan et al., 2019; Collins et al., 2021; Zhang et al., 2023b] typically employ a fixed learning rate. However, in the context of FL with non-IID data, we observe that the learning rate is a critical hyperparameter that significantly impacts both the performance and convergence speed of local models (see Section Ablation). Previous work [Singh et al., 2015] pioneered layer-wise learning rate adjustments, primarily to mitigate the vanishing gradient issue in the lower layers of DNNs within a non-distributed training context. However, this approach is not well-suited for the pFL context. Inspired by [Luo et al., 2021] and our observation (see Section Layer Similarity), we note that the first layer among all local models, showing the highest similarity, captures universal features and thus requires a smaller learning rate with more gradual adjustments. In contrast, deeper layers, exhibiting greater divergence, deal with more complex, client-specific features and benefit from larger learning steps, which aids local model personalization. Building on these insights, we implement an adaptive learning rate scheme for pFL that integrates layer positional information with the corresponding gradient:

where:

• η is the base learning rate.
• g(i,t) represents the gradient vector of the i-th layer at iteration t.
• ∥g(i,t) ∥2 is the L2 norm (Euclidean norm) of the gradient.
• L is the total number of layers.

This approach applies lower learning rates to the initial layers, which capture universal features, and larger rates to deeper layers, which handle client-specific features. Empirical results (see Section Ablation) demonstrate that the adaptive learning rate scheme significantly improves convergence speed and model accuracy compared to fixed learning rate strategies.

3.3 Layer-wise Sparse BinaryMasking

In non-IID settings, naively averaging updated client parameters on the server can overshadow essential local information. To mitigate this issue, we propose a layer-wise binary masking scheme that selectively uploads parameters to preserve critical knowledge and maintain a robust global representation. Our approach leverages the observation that early layers capture broadly shared features across clients, whereas deeper layers learn more specialized, client-specific representations [Luo et al., 2021]. Specifically, we prioritize parameters in the initial layers that undergo larger changes, as they likely embody fundamental patterns vital for model generalization. In contrast, we upload a more extensive set of parameters for deeper layers to retain a diverse range of complex, client-specific features necessary for personalized performance. Formally, we denote the upload proportion for the i-th layer as UP i , computing it based on the layer's position within the network:

where UP i is constrained to be at least 0.1 to ensure that every layer contributes to the global model aggregation, but not more than 1, reflecting a full update contribution.

To identify and select significant weights for sharing, we calculate the absolute weight fluctuation value of the local model within each layer after local training:

Then, to focus on parameters that have undergone notable changes, we identify the top UP i parameters from each layer i in k-th client model based on their fluctuation values after the t-th training round:

Algorithm 1 FLAYER

Input: N clients, ρ: client joining ratio, L: loss function, Θ0g : initial global model, η: base local learning rate, s: the hyperparameter of FLAYER.

Output: Well-performing local models ˜Θ1, . . . , ˜ΘN

• 1: Server sends Θ0g to all clients to initialize local models.
• 2: for iteration t = 1, . . . , T do
• 3: Server samples a subset C t of clients according to ρ.
• 4: Server sends Θt-1 g to |C t | clients.
• 5: for Client k ∈ C t in parallel do
• 6: Client k initializes local model ˜Θtk by Equation (2).
• 7: Client k obtains ˆΘtk by Equation (3) - (4).
• 8: ▷ Local model training
• 9: Client k obtains masked Θtk by Equation (5) - (9).
• 10: Client k sends Θtk to the server. ▷ Uploading
• 11: end for
• 12: Server-side Aggregation:
• 13: Server obtains Θtg by Θtg ← ∑k∈Ct nk ∑j∈Ct nj Θtk by Equation (10).
• 14: end for
• 15: return ˜Θ1, . . . , ˜ΘN

To manage which parameters are uploaded from each layer i of a local model on client k, we use a binary mask matrix, M tk , which has the same dimensions as the parameter matrix ˆ θtk . Initially, all elements of this matrix are initialized to one. Each item m(i,t) j,k in M (i,t) k is determined by whether the corresponding parameter θ(i,t) j,k in ∆θ(i,t) k belongs to the subset of parameters with the highest weight changes, S(i,t) k . This setup uses the following rule:

Finally, we obtain the essential parameter θtk required for uploading by multiplying the ˆ θtk with the binary mask M tk .

During the server aggregation in round t, each client k ∈ C t provides updated parameters θtk and has a local dataset of size nk. Let Θt-1 g denote the global parameters from the previous round and I be the indicator function. The condition ∑k∈Ct nkI(θtk [i] ̸= 0) > 0 checks if at least one client provides a non-zero update for parameter i, and ensures that we only average over those clients that have a non-zero parameter, preventing division by the total number of clients when many of them do not contribute to that parameter. The aggregation process is defined as:

FLAYER employs selective weight sharing to balance global knowledge with client-specific adaptation, distinguishing it from FedMask [Li et al., 2021a] and FedSelect [Tamirisa et al., 2024]. While FedMask, which also achieves personalization using a heterogeneous binary mask with a small overhead. However, FedMask does not consider the unique characteristics of different layers, failing to capture layer-specific information. Moreover, FedMask's binary parameter aggregation is insufficient for complex tasks, such as CIFAR-100. FLAYER applies a layerwise masking strategy informed by parameter significance and layer position. This preserves universal features in early layers and ensures rich adaptation in deeper layers. Moreover, unlike FedSelect, which uses the Lottery Ticket Hypothesis to incrementally identify and expand "winning" subnetworks, which adds complexity due to evolving subnetworks, FLAYER retains a stable global structure, thus reducing computational overhead while maintaining finegrained personalization. This approach effectively integrates fundamental representations and diverse local insights, resulting in higher accuracy and more stable convergence in non-IID settings.

4 Evaluation Setup

4.1 Platforms andWorkloads

To evaluate the performance of FLAYER, we use a fourlayer CNN [McMahan et al., 2017] and ResNet-18 [He et al., 2016] for CV tasks, training them on three benchmark datasets: CIFAR-10, CIFAR-100 [Krizhevsky et al., 2009], and Tiny-ImageNet [Chrabaszcz et al., 2017]. For the NLP task, we train fastText [Joulin et al., 2017] on the AG News dataset [Zhang et al., 2015]. We use the Dirichlet distribution Dir(β) with β = 0.1 [Lin et al., 2020; Wang et al., 2020] to model a high level of heterogeneity across client data. Following FedAvg, we use a batch size of 10 and a single epoch of local model training per iteration. We execute the training process five times for each task and calculate the geometric mean of training latency and inference accuracy until convergence. Our experiments consider 20 clients. The number of layers in the head for CNN, ResNet-18, and fastText is 1, 2 and 1, respectively. Following FedALA, we set a base learning rate of 0.1 for ResNet-18 and fastText and 0.005 for CNN during local training. All experiments were conducted on a multi-core server with a 24-core 5.7GHz Intel i9-12900K CPU and an NVIDIA RTX A5000 GPU with 24GB of GPU memory.

4.2 Competitive Baselines

We compare FLAYER with eight other pFL methods alongside FedAvg, including model-wise aggregation methods APPLE, Ditto and FedAMP, layer-wise aggregation methods FedPer, FedRep, GPFL [Zhang et al., 2023a] and FedCP [Zhang et al., 2023c], and element-wise FedALA

on four popular benchmark datasets in inference accuracy. In addition, we also evaluate the performance of FLAYER in terms of the computation cost, hyperparameter, layer similarity, data heterogeneity, scalability, and applicability.

5 Experimental Results

5.1 Overall Performance

Inference accuracy. Table 1 compares the inference accuracy of FLAYER with eight other pFL methods in CV and NLP domains with Dir(0.1). APPLE gives the highest accuracy in the model-wise category, but with a high computation cost. FedPer uses a simple local aggregation strategy, utilizing global base layers and local head layers to initialize the local model, improving accuracy by an average of 17.66% over FedAvg. Building on this, GPFL and FedCP seek to integrate global and personalized features:

GPFL adjusts base layers, while FedCP focuses on head layers. Although both outperform FedRep and FedPer with a 4-layer CNN, they fail on deeper models like ResNet18. In contrast, FedRep extends FedPer by alternating head and base layer training, boosting accuracy by 19.28% over FedAvg. FedALA enhances FedPer by incorporating global information into local head initialization, achieving a 21.85% improvement. Previous pFL methods acknowledge the distinct roles of base and head layers in non-IID settings and integrate global and local information only at initialization. However, they fail to adapt these layers effectively during local training, limiting the capture of on-demand local information and slowing convergence. In contrast, FLAYER integrates global and local knowledge in a layer-wise manner throughout initialization, local training, and model updating, achieving the highest test accuracy among all pFL methods and improving over FedAvg by 23.31%.

Computation cost. Table 2 compares the computation cost of FLAYER with other pFL methods and FedAvg, measured by the training time required for convergence. Except for CIFAR-10 with CNN, where FedAMP and GPFL bring the lower training cost (but with poor inference accuracy) than FLAYER, our approach gives the lowest computation cost across all other tasks, reducing total training cost by an average of 58.9% (up to 80.1%) compared to FedAvg. Specifically, model-wise methods like APPLE and Ditto involve complex calculations leading to high overhead. FedRep trains the base and head layers separately, which incurs significant training costs. FLAYER effectively incorporates both local and global information across all layers, resulting in fewer rounds needed for convergence compared to FedALA, with an average reduction of 52.7% in total training time.

5.2 Evaluation on Personalization Layers

Table 3 shows inference accuracy for a 4-layer CNN and ResNet-18 with varying sizes (termed as s) of the head layers. For ResNet-18, the highest inference accuracy is achieved with s set to 2, focusing personalization on the final two layers. For the 4-layer CNN, s is set to 1, with the remaining layers updated using the global model.

5.3 Layer Similarity

We measure Centered Kernel Alignment (CKA) [Kornblith et al., 2019] similarity across clients' model layers to understand how pFL methods balance personalization and generalization on non-IID data. Figure 2 compares FedAvg, APPLE, FedRep, FedALA, GPFL, FedCP, and FLAYER on CIFAR-10. After training, base-layer (layer 1, 2, 3) similarity increases for all methods, indicating effective global feature extraction. In contrast, the head layers remain more distinct, focusing on client-specific patterns. FLAYER achieves moderate similarity in both base and head layers, suggesting a balanced integration of global and local features. Additionally, under FedAvg, the last layer is a little higher than that of the second-to-last layer due to FedAvg aggregating local models by averaging their parameters. The last layer, having fewer parameters (especially in classification tasks with a moderate number of classes), experiences less variability and is more directly influenced by the averaging process.

5.4 Evaluation on Data Heterogeneity

We examine the impact of statistical heterogeneity on FLAYER and other eight pFL methods, using 20 clients under three settings: β = 0.5, β = 0.1, and β = 0.01. Here, smaller β values indicate higher data heterogeneity, with β = 0.1 serving as our baseline. As summarized in Table 4, FLAYER and FedALA deliver comparable performance when heterogeneity is low (β = 0.5). However, under more challenging scenarios (β = 0.01, and β = 0.1), FLAYER consistently outperforms all competing methods, achieving the highest accuracy. Additionally, we observe that some pFL methods, such as Ditto and FedAMP, exhibit lower accuracy than FedAvg under β = 0.5. This outcome may occur because they rely too heavily on local information, overshadowing the benefits of aggregated global

(a) ResNet-18 at round 1

(b) ResNet-18 until convergence

Figure 3: Learning curves of FLAYER and five baselines on CIFAR-100 using ResNet-18 model.

knowledge and thus leading to suboptimal generalization. Figure 3 presents the learning curves of FedAvg, Ditto, GPFL, FedPer, FedALA, and our approach. FLAYER converges more rapidly and achieves higher accuracy than all baselines. Moreover, as the degree of non-IID data distribution increases, all pFL methods tend to improve in accuracy, although some, such as FedRep and GFPL, require more communication rounds to converge. In contrast, FLAYER 's dynamic learning capability enables it to reach convergence even under these challenging conditions rapidly.

5.5 Scalability

To evaluate the scalability, we vary the number of clients from 20 to 100 on CIFAR-100 with ResNet-18 under the Dir(0.1) heterogeneity setting. Table 4 compares the average inference accuracy between FLAYER and other pFL methods. We can see that FLAYER consistently outperforms others across various scales of client quantity. However, we notice that the accuracy is higher with 50 clients compared to 20 clients for both FedALA and FLAYER. With 50 clients, the learning algorithm has access to more diverse training data than with 20 clients, improving the generalization ability and overall accuracy of the learned model. While a decrease in accuracy across all methods is observed as the client count increases from 50 to 100, FLAYER experiences a decline of less than 2%. In contrast, the APPLE method shows a significant drop in performance, with a 9.6% decrease in inference accuracy in the same scenario. This underlines the efficiency of FLAYER in managing larger numbers of clients, particularly in scenarios characterized by increased scalability demands.

5.6 Applicability

Our evaluation so far applied FLAYER to FedAvg. We now apply FLAYER to other FL methods to evaluate the generalization ability of our approach. Note that FLAYER does not replace the foundational architectures of an FL method. Table 4 reports the inference accuracy and improvements achieved after applying our approach to an underlying FL method. FLAYER improves the accuracy of all pFL methods, boosting the accuracy by 4.90% to 11.97%.

5.7 Ablation Study

Figure 4 illustrates how each component impacts both convergence speed and final accuracy. Figure 4(a) shows that Agg. Only contributes the least to convergence speed yet still improves accuracy over standard methods. For example, FedPer requires 184 iterations to converge and achieves an accuracy of 54.26%, Agg. Only reaches 55.11% accuracy in just 63 iterations, underscoring the benefit of even partial global aggregation over purely local strategies. Meanwhile, LR. Only delivers the highest overall accuracy, whereas Masking Only excels in accelerating convergence. Figure 4(b) indicates that combining the LR and Masking schemes yields only about a 1% drop in accuracy compared to using all three components. Notably, FLAYER achieves better accuracy at 53 iterations than its counterpart without Aggregation, suggesting that Agg.Only can be introduced later in training to boost performance further. These findings also highlight that LR. Only and Masking Only schemes are highly effective in non-IID settings.

6 RelatedWork

Previous pFL methods for managing non-IID data issues typically fall into two categories: personalizing the global model and customizing individual models for each client. Our work focuses on learning personalized models, where pFL tailors individual models to each client's data, primarily via weighted aggregation for local model adaptation.

Model-wise aggregation. Train personalized models for each client by combining clients' models using weighted aggregation. For example, FedFomo [Zhang et al., 2020] employs a distance metric for weighted aggregation, while APPLE [Luo and Wu, 2022] introduces an adaptive mechanism to balance global and local objectives. FedAMP [Huang et al., 2021] uses attention functions for clientspecific models, and Ditto [Li et al., 2021b] incorporates a proximal term for personalized models. However, existing model aggregation methods may overlook complex variations and unique characteristics in client data, leading to suboptimal personalization.

Layer-wise aggregation. Customizing different layers has been explored by FedPer [Arivazhagan et al., 2019] and FedRep [Collins et al., 2021], and extended by approaches like GPFL [Zhang et al., 2023a] and FedCP [Zhang et al., 2023c], which aim to incorporate global and personalized information within the base or head layers. pFedLA [Ma et al., 2022] further uses hypernetworks to learn layer-wise aggregation weights at a high computational cost. However, all these methods overlook the influence of diverse local data on both base and head layers, constraining their potential to improve accuracy.

Element-wise aggregation. This is the most fine-grained local aggregation approach, aggregating at the parameter level. FedALA [Zhang et al., 2023b] introduces an elementlevel aggregation weight matrix in the head layers, enhancing accuracy across various tasks. However, extra computation is required for weight calculation and does not account for the distinct roles and learning abilities of different layers during training.

7 Conclusion

We have presented FLAYER, a new layer-wise pFL approach to optimize FL in the face of non-IID data. FLAYER adaptively adjusts the aggregation weights and learning rate and selects layer-wise masking to effectively incorporate local and global information throughout all network layers. Experimental results show that FLAYER achieves the best inference accuracy and significantly reduces computational overheads compared to existing pFL methods.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China (No.62372281); Shaanxi Province 2024 Innovation Capacity Support Program (2024ZC-KJXX-014). For any correspondence, please contact Jie Ren (E-mail: renjie@snnu.edu.cn).