Source-linked AI summary
Layer-wised Model Aggregation for Personalized Federated Learning
Xiaosong Ma, Jie Zhang, Song Guo, Wenchao Xu
TL;DR
Non-IID clients differ across layers, but prior personalized federated learning methods rely on whole-model similarity and can suffer slower convergence or inadequate personalization. pFedLA learns layer-wise aggregation weights with client-specific server hypernetworks, and experiments report superior performance across four classification tasks while retaining a communication-efficient variant with negligible performance loss.
Problem
Existing personalized federated learning methods use whole-model distances or losses, overlooking layer-specific utility and limiting personalization for heterogeneous non-IID clients.
Method
pFedLA uses a dedicated server-side hypernetwork for each client to generate and update aggregation weights for each model layer.
Results
Superior performance over state-of-the-art approaches was demonstrated across four image-classification tasks and different data distributions.
Takeaways & Limitations
Layer-wise aggregation can reinforce collaboration among similar clients and provide adequate personalization over non-IID datasets.
Abstract
from arXiv · showhide
Personalized Federated Learning (pFL) not only can capture the common priors from broad range of distributed data, but also support customized models for heterogeneous clients. Researches over the past few years have applied the weighted aggregation manner to produce personalized models, where the weights are determined by calibrating the distance of the entire model parameters or loss values, and have yet to consider the layer-level impacts to the aggregation process, leading to lagged model convergence and inadequate personalization over non-IID datasets. In this paper, we propose a novel pFL training framework dubbed Layer-wised Personalized Federated learning (pFedLA) that can discern the importance of each layer from different clients, and thus is able to optimize the personalized model aggregation for clients with heterogeneous data. Specifically, we employ a dedicated hypernetwork per client on the server side, which is trained to identify the mutual contribution factors at layer granularity. Meanwhile, a parameterized mechanism is introduced to update the layer-wised aggregation weights to progressively exploit the inter-user similarity and realize accurate model personalization. Extensive experiments are conducted over different models and learning tasks, and we show that the proposed methods achieve significantly higher performance than state-of-the-art pFL methods.
1. Introduction
Non-IID client data makes one shared model inadequate, while existing personalized aggregation overlooks layer-specific utility. pFedLA addresses this gap with layer-wise aggregation and dedicated hypernetworks, achieving superior performance across four image-classification tasks.
- Non-IID data can make a shared global model slow convergence or harm inference when it deviates from clients’ local distributions.
- Existing personalized methods commonly aggregate client models using whole-model distances, overlooking differences in layer utility.Shallow layers emphasize local feature extraction, whereas deeper layers extract global features.
- A dedicated server-side hypernetwork learns cross-client layer weights during personalized federated training.
- The paper explicitly compares layer-wise and model-wise personalization to reveal benefits of layer-wise aggregation for heterogeneous clients.
- pFedLA performs personalized federated aggregation at layer granularity to exploit inter-client similarities and produce accurate models for non-IID clients.
- Extensive experiments on four image-classification tasks demonstrated pFedLA’s superior performance over state-of-the-art approaches.
2. Related Work
Prior personalized federated learning methods modify data, optimize single models, or combine multiple models, but generally treat each client model as a whole. The paper instead motivates fine-grained layer-wise aggregation and uses hypernetworks to evaluate layer importance.
- Personalized federated learning methods include data-based approaches and model-based approaches that customize model structures or parameters.
- Data-based methods may require server knowledge of clients’ statistical distributions, potentially creating privacy concerns.
- Model-based methods include single-model techniques such as fine-tuning, regularization, mixture, meta-learning, and parameter decomposition.
- Existing approaches treat each client model as a whole entity rather than considering layer-wise utility for personalized aggregation.
- Hypernetworks map target-task embeddings to parameters of another neural network and are used here to evaluate model-layer importance.
3. Method
pFedLA uses client-specific hypernetworks to generate layer-wise aggregation weights, producing personalized models while preserving the federated learning objective and private local training. Its heuristic variant retains selected layers locally to reduce communication with negligible performance reduction.
- pFedLA Algorithm: pFedLA applies dedicated server-side hypernetworks to evaluate layer importance and generate client-specific aggregation weights for personalized models.Each hypernetwork takes a client embedding vector as input and outputs an aggregation weight matrix.
- pFedLA Algorithm: The aggregation weight matrix assigns distinct weights to different client models for each layer rather than applying one identical weight across all layers.For layer n, client i aggregates parameters using the layer-specific weights α_i^{ln,j}.
- pFedLA Algorithm: Client models are formed by weighted aggregation of intermediate parameters after local training, followed by private-data local updates.The algorithm initializes model, hypernetwork, and embedding parameters, aggregates models at the server, and returns local updates from each client.
- pFedLA Algorithm: pFedLA transforms optimization over personalized client parameters into optimization of client embeddings and hypernetwork parameters.The objective jointly optimizes V, the client embedding vectors, and Ψ, the hypernetwork parameters.
- HeurpFedLA: HeurpFedLA retains selected top-k layers without aggregation and reduces server-to-client transmission while maintaining pFL model performance with negligible reduction.The remaining layers follow the same operations as pFedLA, and retained-layer parameters need not be transmitted from the server.
4. Evaluation
Experiments evaluate pFedLA and HeurpFedLA across four non-IID image-classification datasets, comparing personalization, weight evolution, layer importance, and communication efficiency. pFedLA generally outperforms baselines, while HeurpFedLA trades retained layers for lower communication costs with limited accuracy impact.
- Experimental Setup: Experiments cover EMNIST, FashionMNIST, CIFAR10, and CIFAR100 under two non-IID data-distribution scenarios.The study uses both class-assignment heterogeneity and imbalanced class distributions, with 70% training and 30% test data.
- Experimental Setup: pFedLA is compared with FedAvg, Local Training, Per-Fedavg, pFedMe, pFedHN, FedBN, FedRep, and FedFomo.The baselines represent local, meta-learning, regularization, hypernetwork, batch-normalization, representation, and distance-based aggregation approaches.
- Performance Evaluation: pFedLA provides superior performance over baselines on the four datasets with different data distributions in most cases, while HeurpFedLA incurs negligible performance reduction.The reported averages are obtained after 600 rounds for 10 clients and 2500 rounds for 100 clients.
- Analysis of Weight Evolution: Aggregation weights become larger between clients with more overlapping classes, while each client’s self-weight is highest in the layer-wise heatmaps.This pattern is observed for adjacent client IDs designed to represent progressively greater data similarity.
- Analysis of Communication Efficiency: Layer importance differs by dataset: the first fully connected layer is most weighted on CIFAR10, whereas the classifier layer is most weighted on EMNIST.Retaining highly weighted layers locally can compromise CIFAR10 accuracy but may increase EMNIST accuracy, consistent with the cited FedRep observation.
- Effect of k: Retaining different numbers of top-k layers does not significantly affect accuracy, enabling a trade-off between training efficiency and communication cost.HeurpFedLA selects retained local layers according to available communication bandwidth.
5. Conclusion
pFedLA performs personalized model aggregation layer by layer for non-IID federated learning, reinforcing collaboration among similar clients and outperforming conventional model-wise approaches. An improved version reduces communication overhead with negligible performance loss, supporting large-scale FL scenarios.
- pFedLA performs personalized model aggregation in a layer-wise manner over non-IID datasets.
- Layer-wise aggregation progressively reinforces collaboration among similar clients and generates adequate personalization.
- pFedLA outperforms conventional model-wise approaches in personalized federated learning.
- An improved pFedLA version reduces communication overhead with negligible performance loss.
- Evaluations on four classification tasks demonstrate pFedLA's feasibility and superior performance.