Optimizing Personalized Federated Learning through Adaptive Layer-Wise Learning

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 $D_{1},D_{2},\ldots,D_{N}$, which has different data classes and sizes. These datasets exhibit heterogeneity, characterized by non-IID (non-independently and identically distributed) (Zhao et al. 2018). The goal of pFL can be defined as:

where $\theta_{k}$ is the model parameters of client $k$, $m_{k}$ is the data size of $D_{k}$, $M$ is the whole data size of all clients, $L_{k}(\theta_{k})$ is the loss function of client $k$.

Figure 1 depicts the workflow of FLAYER during the training of the $k$-th client in the $t$-th iteration. Initially, the global model is downloaded to each client for local aggregation, FLAYER implements an adaptive aggregation strategy, dynamically adjusting the aggregation weights between the global and local models based on the inference accuracy of the local model. This allows models, whether local or global, with higher performance to have a greater influence on the initialization of the local head layer. During the local training phase, FLAYER leverages a layer-wise adaptive learning rate scheme. This strategy dynamically adjusts the learning rates according to each layer’s position within the network and the gradient of that layer. By optimizing the learning step for each layer, FLAYER enhances local model personalization performance and accelerates convergence. In the final phase of updating the model to the central server, FLAYER employs a layer-wise masking strategy. This selective approach uploads only the essential parameters from each local model, preventing the global averaging process from diluting the crucial information captured by the local models, thus enhancing the global representation. Algorithm 1 presents the overall FL process in FLAYER.

Building on previous findings (Yosinski et al. 2014; Zhu, Hong, and Zhou 2021) that base layers of 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 $D_{k}$ from the previous iteration. The process is defined as follows:

Here, $\theta_{k}$ is the local model parameter matrix of the $k$-th client, $\tilde{\theta}_{k}$ denotes the local model parameter matrix of the $k$-th client after initialization, $L$ is the total number of layers, $s$ is the number of layers in the head for personalization, and $\theta_{g}^{(1:L-s,t-1)}$ 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. $\theta_{k}^{(L-s+1:L,t-1)}$ 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 $A_{k,g}$ and $A_{k,l}$ control the influence of global and local parameters, respectively. At iteration ${t-1}$, the local model inference accuracy on dataset $D_{k}$ sets the local weight $A_{k,l}^{t-1}$, with $(1-A_{k,l})$ adjusting for global influence. A lower local accuracy increases reliance on the global model through $A_{k,g}$, providing stability in early training phases. As the accuracy of the local model increases, the initialization of the head layer becomes more dependent on the local model, thereby incorporating more localized and personalized information.

After the model initialization, the local model trains on its local dataset $D_{k}$:

$\eta$ represents the learning rate, $\nabla_{\tilde{\theta}_{k}}\mathcal{L}_{k}(\tilde{\theta}_{k}^{t},D_{k})$ is the gradient of the loss function $\mathcal{L}_{k}$ with respect to the parameters $\tilde{\theta}_{k}$ evaluated using local dataset $D_{k}$.

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:

• •

  $\eta$ 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.

In non-IID settings, averaging updated client parameters on the server side can dilute important information during aggregation. To address this, we propose a layer-wise binary masking scheme for server aggregation, aimed at preserving critical information and ensuring a high-quality global representation. The core idea is to selectively upload parameters from each layer based on their significance and layer position, distinguishing between general features in early layers and more complex, client-specific features in deeper layers (Luo et al. 2021). In detail, our strategy prioritizes uploading parameters with high significance, typically those with greater changes in early layers, as these are likely to capture essential patterns and features common to all client datasets, thereby enhancing the model’s generalization ability. For deeper layers, we employ a more inclusive approach, uploading a greater proportion of parameters to capture a wide range of client-specific details and complex features essential for the model’s performance on localized tasks. The proportion of parameters uploaded from each layer, denoted as $UP^{i}$ for layer $i$, is calculated based on the layer’s position within the network architecture using the following formula:

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:

To manage which parameters are uploaded from each layer $i$ of a local model on client $k$, we use a binary mask matrix, $M_{k}^{t}$, which has the same dimensions as the parameter matrix $\hat{\theta}_{k}^{t}$. Initially, all elements of this matrix are initialized to one. Each item $m_{j,k}^{(i,t)}$ in $M_{k}^{(i,t)}$ is determined by whether the corresponding parameter $\theta_{j,k}^{(i,t)}$ in $\Delta{\theta}_{k}^{(i,t)}$ belongs to the subset of parameters with the highest weight changes, $S_{k}^{(i,t)}$. This setup uses the following rule:

Finally, we obtain the essential parameter $\theta_{k}^{t}$ required for uploading by multiplying the $\hat{\theta}_{k}^{t}$ with the binary mask $M_{k}^{t}$.

Adopting selective weight sharing, FLAYER enhances the global model representation. Our approach differs from FedMask (Li et al. 2021a), 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. In our approach, early layers, which capture universal features, are updated only with the most critical changes during server aggregation, preserving a robust foundation for all clients and preventing the dilution of essential base features. Conversely, deeper layers, which capture complex and client-specific features, receive updates from a larger proportion of parameters. This ensures the global model incorporates a diverse set of features, enhancing its generalization ability.

To evaluate the performance of FLAYER, we use a four-layer 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, Hinton et al. 2009), and Tiny-ImageNet (Chrabaszcz, Loshchilov, and Hutter 2017). For the NLP task, we train fastText (Joulin et al. 2017) on the AG News dataset (Zhang, Zhao, and LeCun 2015). We use the Dirichlet distribution $Dir(\beta)$ with $\beta=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 24 GB of GPU memory.

We compare FLAYER with six other pFL methods alongside FedAvg, including model-wise aggregation methods APPLE, Ditto and FedAMP, layer-wise aggregation methods FedPer and FedRep, 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.

Table 4 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.

To analyze how pFL methods perform across layers on non-IID datasets, we measure the Centered Kernel Alignment (CKA) (Kornblith et al. 2019) similarity of features from the same layer of different clients’ models using identical test samples. This analysis helps to evaluate the balance between personalization and generalization of different pFL methods. Figure 2 presents the CKA similarity across 20 clients for methods like FedAvg, APPLE, FedRep, FedALA, and FLAYER on CIFAR-10, highlighting changes from the initial round to the training convergence. We observe that after training, the similarity of the base layers in both the CNN and ResNet-18 improves across all FL methods, indicating that the global model effectively captures common features shared by different clients. The deeper layers show lower similarity, with the head layers exhibiting the least, reflecting their focus on localized, client-specific data. Additionally, the simpler structure of the 4-layer CNN results in higher similarity across all layers compared to ResNet-18, suggesting it is less capable of capturing specialized features. FLAYER achieves moderate similarity levels in the base and head layers, suggesting that it balances well between integrating global patterns and adapting to local specifics, thereby enhancing overall model performance.

We also evaluate the impact of statistical heterogeneity on FLAYER and other SOTA pFL methods using 20 clients. Specifically, we set two degrees of heterogeneity on CIFAR-100. The first scenario is $\beta=0.01$, where the smaller the value of $\beta$, the greater the heterogeneity of the setting. We use $\beta=0.1$ as the baseline performance. Table 5 reports the performance impact under these varying degrees of heterogeneity. FLAYER consistently outperforms all other SOTA pFL methods across all heterogeneous settings, delivering the highest accuracy.

To evaluate the scalability of our approach, we vary the number of clients from 20 to 100 when applying ResNet-18 to CIFAR-100, setting the heterogeneity parameter $Dir(0.1)$. Table 5 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. 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.

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 5 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%.

Figure 3 presents the accuracy of three strategies in FLAYER when used alone: Adaptive Aggregation (Agg.) Only, Adaptive Learning Rate (LR) Only, and Masking Only. The results show that Adaptive LR Only achieves the highest accuracy on both ResNet-18 and CNN when trained on the CIFAR-100. For CNN, the full FLAYER exhibits a convergence speed similar to Adaptive LR Only, suggesting that the learning rate is the most influential factor for CNN performance. While Masking Only shows a convergence trend comparable to FLAYER on ResNet-18, it converges more slowly on CNN compared to the other two strategies. The Masking Only benefits deeper network structures by prioritizing critical parameters involved in residual connections, thereby preserving the integrity of these connections and enhancing performance and convergence in deeper networks like ResNet-18. Adaptive Agg. Only brings the least benefit and has the slowest convergence speed when used alone. However, its role is essential in incorporating local and global information in the head layers before training, which lays a solid foundation for the effectiveness of adaptive learning rate and masking strategies. During ResNet-18 training, accuracy improves significantly around the 50th iteration, aligning with the trend observed in Adaptive Agg. Only. While Adaptive LR Only is crucial for performance enhancements, particularly in CNNs, the combined approach of Adaptive Agg., Adaptive LR, and Masking within FLAYER offers a balanced and synergistic strategy that leverages the strengths of each scheme.

Global model personalization aims to adjust the global model to suit diverse client data distributions, creating a model that universally benefits all clients. This typically involves training the global model on varied data and local adaptations for each client’s specific data. Previous studies have explored strategies to mitigate data heterogeneity and improve the global model’s generalization (Pillutla et al. 2022; Zhang et al. 2023a).

Personalized model learning tailors individual models to each client’s data, emphasizing local adaptation. This approach often employs weighted aggregation methods for local model personalization.