Source-linked AI summary
Parameterized Knowledge Transfer for Personalized Federated Learning
Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wencao Xu, Feijie Wu
TL;DR
Statistical heterogeneity and parameter aggregation limit personalized federated learning when clients use different architectures. KT-pFL transfers personalized soft predictions through a learned knowledge coefficient matrix, and experiments report improved efficiency and reduced communication overhead across heterogeneous settings.
Problem
Personalized federated learning methods that aggregate model parameters require identical model structures and sizes, limiting heterogeneous model settings.
Method
KT-pFL maintains personalized server-side soft predictions formed from clients’ local soft predictions and jointly trains the knowledge coefficient matrix with model parameters.
Results
Experiments across heterogeneous models and data distributions report significant performance gains, improved training efficiency, and reduced communication overhead.
Takeaways & Limitations
Parameterized group knowledge transfer supports personalized federated training without requiring homogeneous client model structures.
Takeaways & Limitations
The experiments configure the public dataset size |D_r| to 3000 because very large public datasets slow convergence and increase client computation time.
Abstract
from arXiv · showhide
In recent years, personalized federated learning (pFL) has attracted increasing attention for its potential in dealing with statistical heterogeneity among clients. However, the state-of-the-art pFL methods rely on model parameters aggregation at the server side, which require all models to have the same structure and size, and thus limits the application for more heterogeneous scenarios. To deal with such model constraints, we exploit the potentials of heterogeneous model settings and propose a novel training framework to employ personalized models for different clients. Specifically, we formulate the aggregation procedure in original pFL into a personalized group knowledge transfer training algorithm, namely, KT-pFL, which enables each client to maintain a personalized soft prediction at the server side to guide the others' local training. KT-pFL updates the personalized soft prediction of each client by a linear combination of all local soft predictions using a knowledge coefficient matrix, which can adaptively reinforce the collaboration among clients who own similar data distribution. Furthermore, to quantify the contributions of each client to others' personalized training, the knowledge coefficient matrix is parameterized so that it can be trained simultaneously with the models. The knowledge coefficient matrix and the model parameters are alternatively updated in each round following the gradient descent way. Extensive experiments on various datasets (EMNIST, Fashion\_MNIST, CIFAR-10) are conducted under different settings (heterogeneous models and data distributions). It is demonstrated that the proposed framework is the first federated learning paradigm that realizes personalized model training via parameterized group knowledge transfer while achieving significant performance gain comparing with state-of-the-art algorithms.
1 Introduction
Statistical heterogeneity limits conventional federated learning and motivates personalized models, but parameter aggregation restricts clients to identical architectures. KT-pFL addresses this constraint through parameterized personalized knowledge transfer and reports efficiency benefits across heterogeneous settings.
- Non-IID client datasets can prevent a conventional global model from generalizing well to each client.
- Existing personalized federated learning methods still require identical model structure and size for parameter aggregation.
- KT-pFL maintains a personalized server-side soft prediction for each client and updates it by combining all clients’ local soft predictions.
- The knowledge coefficient matrix reinforces collaboration among clients with similar data distributions and is trained alongside model parameters using alternating updates.
- KT-pFL improves training efficiency and reduces communication overhead compared with traditional parameter-transfer approaches.
- KT-pFL is presented as the first framework for personalized knowledge transfer in federated learning and is evaluated with theoretical guarantees and extensive experiments.
2 Related Work
Prior personalized federated learning methods generally assume homogeneous models and organize personalization through single, multiple, or absent global models. Knowledge-distillation approaches support heterogeneous architectures, but commonly combine soft predictions heuristically rather than learning their weights jointly.
- Homogeneous-model personalized federated learning methods use single, multiple, or no global models at the server.
- Single global model: Single-global-model methods include local fine-tuning, regularization, hybrid local and global models, and meta-learning.
- Multiple or no global models: Multiple-global-model methods cluster clients by similarity, while some methods omit a global model to address heterogeneity.
- Existing personalized methods that aggregate model parameters require identical model structure and size, preventing personalized architectures for heterogeneous clients.
- Heterogeneous FL and Knowledge Distillation: Knowledge-distillation methods aggregate local soft predictions instead of parameters, but most use averaged or heuristic teacher combinations.
- Heterogeneous FL and Knowledge Distillation: KT-pFL updates soft-prediction weights together with model parameters during every federated learning iteration.
3 Problem Formulation
The formulation targets personalized models with different structures under private, non-IID client data. KT-pFL replaces parameter aggregation with personalized soft-prediction transfer, using knowledge coefficients and regularization to coordinate local training.
- The framework considers N clients with private datasets and aims to collaboratively train personalized models that may have different structures.
- Conventional federated learning minimizes total empirical loss with a unified model structure, which cannot represent clients using unique models.
- Each client’s personalized loss adds a KL-divergence term that transfers collaborative knowledge from a personalized teacher.
- The knowledge coefficient c_mn estimates client m’s contribution to client n and weights the personalized aggregated knowledge.
- In each communication round, clients train locally, send soft predictions on public data, receive personalized predictions, and perform local distillation.
- The soft prediction is computed from logits with a temperature-scaled softmax, where T is the temperature hyperparameter.
- The regularization term discourages clients with different data distributions from assigning excessively large knowledge coefficients and eliminating collaboration.
4 KT-pFL Algorithm
KT-pFL alternates local model updates with server-side updates of a knowledge coefficient matrix, using personalized soft predictions to transfer knowledge across clients. Its analysis establishes conditions under which personalized ensemble models match or outperform local-only and average ensemble models.
- Alternating optimization: KT-pFL alternates optimization of model parameters and the knowledge coefficient matrix across communication rounds.The server fixes one variable while optimizing the other through gradient descent.
- Client update: Each round combines private-data local training with public-data distillation from personalized soft predictions.Clients train on private mini-batches, then distill knowledge using public-data mini-batches before returning local parameters.
- Server aggregation: The server forms each client’s personalized soft prediction by aggregating clients’ collaborative knowledge according to the knowledge coefficient matrix.The resulting personalized predictions guide subsequent local training, while coefficient vectors are maintained per client.
- Performance guarantee: Under suitable coefficient matrices, the personalized ensemble model has expected loss no larger than the local-only model, and can be strictly better.The theorem gives LDn(personalized ensemble) ≤ LDn(local-only) and identifies cases with strict inequality.
- Data setting: The method assumes public data are accessible to all clients, and it can operate with either labeled or unlabeled public datasets.The public dataset supports the distillation stage and need not be labeled.
- Performance guarantee: Under suitable coefficient matrices, the personalized ensemble model also has expected loss no larger than the average ensemble model, with strict improvement possible.The comparison extends the guarantee beyond local-only training to average ensemble predictions.
5 Evaluations
Experiments evaluate KT-pFL across image datasets, Non-IID distributions, heterogeneous and homogeneous models, client scales, and hyperparameter settings. KT-pFL achieves strong accuracy, supports larger federated systems, and reduces communication overhead through soft-prediction transfer.
- Experimental setup: Experiments use EMNIST, Fashion_MNIST, and CIFAR-10 under two Non-IID settings, with 75% training and 25% testing data.The settings assign clients either two classes or all classes with client-specific class proportions.
- Experimental setup: Four lightweight architectures—LeNet, AlexNet, ResNet-18, and ShuffleNetV2—are assigned across 20 clients to evaluate heterogeneous models.Five clients use each model structure.
- Homogeneous FL: KT-pFL is also compared with FedAvg and state-of-the-art personalized federated learning methods under homogeneous CNN models.The baselines include Per-FedAvg, Fedavg-Local FT, pFedMe, FedAMP, FedFomo, and FedHN.
- Heterogeneous FL: KT-pFL obtains comparable or better accuracy than competing distillation methods in both Non-IID settings, with the best performance attributed to personalized aggregation of local soft predictions.The comparison includes FedMD, FedDF, pFedDF, Sim-pFL, and TopK-pFL.
- Scalability: KT-pFL works in both small-scale and large-scale federated systems, including experiments with 100 clients and partial knowledge-coefficient updates.The large-scale setting uses the same client-selection mechanism as the baselines.
- Hyperparameter effects: Larger local-epoch values benefit convergence, whereas larger distillation-step values do not always improve performance and require moderation.The experiments examine the trade-off between local computation and global communication.
- Hyperparameter effects: Increasing public-data size improves average test accuracy, but very large public datasets slow convergence and increase client computation time.The experiments configure the public-data size as 3000.
- Communication overhead: Soft-prediction-based KT-pFL has far lower communication overhead than conventional parameter-based personalized federated learning.The evaluation records data, model parameters, and soft predictions without compression.
Broader Impact
KT-pFL addresses heterogeneous client needs by removing homogeneous model constraints, reducing communication overhead, and improving training efficiency without extra client-side computation.
- KT-pFL supports customized neural architectures across devices with different computation, communication, and storage capacities.
- The framework can significantly reduce communication overhead during federated training.
- KT-pFL improves training efficiency through a parameterized update mechanism without additional computation overhead at the client side.