Source-linked AI summary
FedKD: Communication Efficient Federated Learning via Knowledge Distillation
Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, Xing Xie
TL;DR
Federated learning incurs heavy communication costs when large model updates are exchanged over many rounds. FedKD shares only a collaboratively updated student model, combines it with an adaptive mutual teacher–student distillation framework, and compresses exchanged gradients using dynamic SVD approximation. Experiments across benchmark tasks report substantially reduced communication costs with competitive results.
Problem
Repeated exchange of large model updates makes federated learning communication-intensive, especially when models contain many parameters.
Method
FedKD locally updates a large teacher, collaboratively shares a small student, adapts mutual distillation to prediction correctness, and compresses gradients using dynamic SVD approximation.
Results
FedKD largely reduces federated-learning communication costs while achieving competitive performance across benchmark datasets and tasks.
Takeaways & Limitations
Sharing only the student model and dynamically approximating gradients provides a communication-efficient federated-learning approach with competitive results.
Takeaways & Limitations
The approach assumes private raw data remains on each client and factorizes different parameter matrices independently in practice.
Abstract
from arXiv · showhide
Federated learning is widely used to learn intelligent models from decentralized data. In federated learning, clients need to communicate their local model updates in each iteration of model learning. However, model updates are large in size if the model contains numerous parameters, and there usually needs many rounds of communication until model converges. Thus, the communication cost in federated learning can be quite heavy. In this paper, we propose a communication efficient federated learning method based on knowledge distillation. Instead of directly communicating the large models between clients and server, we propose an adaptive mutual distillation framework to reciprocally learn a student and a teacher model on each client, where only the student model is shared by different clients and updated collaboratively to reduce the communication cost. Both the teacher and student on each client are learned on its local data and the knowledge distilled from each other, where their distillation intensities are controlled by their prediction quality. To further reduce the communication cost, we propose a dynamic gradient approximation method based on singular value decomposition to approximate the exchanged gradients with dynamic precision. Extensive experiments on benchmark datasets in different tasks show that our approach can effectively reduce the communication cost and achieve competitive results.
1 Introduction
Federated learning faces heavy communication costs because large model updates must be exchanged repeatedly. FedKD reduces this burden by collaboratively sharing only a small student model, using adaptive mutual distillation and dynamic gradient approximation.
- Motivation: Large model updates exchanged over many federated-learning rounds create enormous communication costs and hinder deploying models such as BERT.The server aggregates client updates and redistributes global updates until convergence.
- FedKD framework: FedKD keeps a large teacher model local while sharing and collaboratively updating only a small student model across clients.Teacher and student models are distilled from each other using local data and prediction correctness.
- Evaluation: Extensive benchmark experiments across different tasks report effective communication-cost reduction and competitive results.The paper identifies these experiments as evidence of the approach’s effectiveness and efficiency.
- FedKD framework: Adaptive mutual knowledge distillation makes teacher and student learn from each other while accounting for their prediction correctness.This mechanism combines local-data supervision with knowledge transferred between the two models.
- Communication compression: Dynamic gradient approximation based on SVD compresses communicated gradients with dynamic precision to further reduce communication cost.The method factorizes gradients before exchange and reconstructs them at the server and clients.
2 Related Work
Federated learning keeps raw data decentralized but still requires repeated model-update communication, creating a large-model communication problem. Knowledge distillation offers model compression, though conventional distillation generally assumes centralized data, motivating federated alternatives such as FedKD.
- 2.1 Federated Learning: Federated learning coordinates decentralized client models through repeated upload, aggregation, and redistribution of model updates without moving raw user data.This privacy-aware setup supports collaborative learning across user devices and servers.
- 2.1 Federated Learning: Large models make repeated federated model-update exchange expensive, motivating methods that reduce communicated parameters while retaining competitive performance.The paper presents knowledge distillation as one remedy for this communication burden.
- 2.2 Knowledge Distillation: Knowledge distillation transfers knowledge from a large teacher to a small student for model compression, but conventional methods usually rely on centralized data storage.Prior approaches distill from outputs, hidden states, or attention matrices, while federated methods also use gradient compression and codistillation.
3 FedKD
FedKD reduces federated-learning communication by collaboratively sharing only a small student model while each client retains a large teacher model. It combines adaptive mutual distillation with dynamically compressed student gradients to reduce communication while preserving learning quality.
- 3.1 Problem Definition: FedKD assigns each client a large local teacher and a smaller shared student model, collaboratively updating only the student across clients.Both models are learned from labeled local data and knowledge distilled from each other.
- 3.2 Federated Knowledge Distillation: Teacher models update locally, while student gradients are uploaded, aggregated centrally, and returned for collaborative student-model learning.The local gradients are encrypted before upload because raw gradients may contain private information.
- 3.2 Federated Knowledge Distillation: Adaptive mutual distillation weights knowledge transfer according to teacher and student prediction quality, weakening distillation when predictions are unreliable.The framework uses output soft labels, hidden states and attention heatmaps, alongside task-specific supervision.
- 3.3 Dynamic Gradients Approximation: SVD-based dynamic gradient approximation factorizes each student gradient into smaller matrices and retains singular values according to an energy threshold.Compression reduces exchanged-gradient size when PK + K^2 + KQ < PQ, where K is the number of retained singular values.
- 3.3 Dynamic Gradients Approximation: The energy threshold increases dynamically during training, using rough gradient approximations early and more accurate approximations near convergence.Tstart and Tend control the threshold's start and end values.
- 3.4 Complexity Analysis: FedKD communication costs O(R|Θs|/ρ), versus O(R|Θt|) for directly learning the large model, because |Θs| ≪ |Θt| and ρ > 1.Its total computational cost is O(RD|Θt| + RD|Θs| + RPQ^2), and the paper reports efficiency in both communication and computation.
4 Experiments
Experiments evaluate FedKD on personalized news recommendation and adverse drug reaction tweet detection using decentralized user data. Results examine performance, adaptive distillation, client count, threshold choices, and SVD-based gradient compression.
- Datasets and Experimental Settings: Experiments use MIND for personalized news recommendation and SMM4H for adverse drug reaction tweet detection.MIND contains impression logs from 1 million users; SMM4H experiments use 16,694 crawled tweets.
- Datasets and Experimental Settings: Training data are randomly divided into four folds, with each fold stored locally on a different client.This setup simulates decentralized private data across clients.
- Performance Evaluation: FedKD achieves performance comparable with UniLM (Fed) and UniLM (Cen) while using much less communication than the original UniLM model.The authors also report better performance and lower communication cost than FedSGD and FedDropout.
- Effectiveness of Adaptive Mutual Distillation: Mutual distillation improves both teacher and student models, especially the teacher, which is used for inference during testing.Knowledge transfer helps the student imitate the teacher and helps the teacher overcome limited local labeled data.
- Effectiveness of Adaptive Mutual Distillation: Adaptive mutual and hidden losses improve performance, while removing adaptive loss weighting produces suboptimal results.Adaptive weighting is intended to account for prediction correctness, distill higher-quality knowledge, and mitigate overfitting.
- Influence of Client Number: FedKD performance remains similar or slightly improves with more clients, whereas UniLM4/2 performance slightly declines.FedKD can learn from multiple teacher models, while vanilla FedAvg may sacrifice performance when learning across limited local data.
- Influence of Tstart and Tend: The selected thresholds Tstart = 0.95 and Tend = 0.98 balance communication cost and model performance.Lower Tstart reduces communication but performance drops quickly below 0.95; Tend = 0.98 similarly provides a favorable trade-off.
- Analysis of Dynamic Gradient Approximation: UniLM gradient matrices are low-rank, and more singular values must be retained later in training to maintain the same energy threshold.The method therefore increases retained precision as training progresses to preserve model accuracy.
5 Conclusion
FedKD combines adaptive mutual distillation with dynamic gradient approximation to reduce federated-learning communication while maintaining promising model performance.
- FedKD reciprocally distills a large teacher model and a small student model on each client, with distillation intensity controlled by prediction correctness.
- Only the small student model is shared across clients and learned collaboratively, while the large teacher model is updated locally.
- Dynamic gradient approximation further reduces communication cost by compressing communicated gradients.
- Experiments on two benchmark datasets for different tasks show that FedKD largely reduces communication cost while keeping promising model performance.
A.1 Comparison with Additional Baselines
The additional-baseline comparison evaluates FedKD on MIND against multiple established news-recommendation models, including recurrent, attention-based, and personalized approaches.
- The MIND comparison includes EBNR, DKN, NPA, NAML, LSTUR, NRMS, and FIM as additional baseline methods.
- EBNR is described as embedding-based news recommendation with GRU, while DKN uses a deep knowledge network for news recommendation.
- The listed baselines also cover personalized attention, attentive multi-view learning, long short-term user interest, and multi-head self-attention.
A.2 Additional Results on SMM4H
Additional SMM4H results are presented in Figures 8–11, where the reported phenomena resemble those observed on MIND; the baseline models use centralized data storage.
- Additional SMM4H results are shown in Figures 8–11.
- The SMM4H experiments observe phenomena similar to the results on MIND.
- The baselines are trained on a centralized data storage.
- Figure 8 examines the influence of mutual distillation on the student and teacher models.
- Figure 9 examines the effect of adaptive mutual distillation.
A.3 Algorithm Workflow
The FedKD workflow is summarized in Algorithm 1.
- Algorithm 1 summarizes the workflow of FedKD.
- The algorithm provides a workflow-level summary of the proposed FedKD method.
- The workflow summary is presented as Algorithm 1 in the paper.
A.4 Experimental Environment
The experiments use a Linux server with four Tesla V100 GPUs and parallel training via Horovod. Additional evaluations examine client number and the effects of Tstart and Tend on performance and communication cost.
- The experimental environment runs Ubuntu 16.04 and Python 3.6 on a server with four Tesla V100 GPUs.The server provides 32GB memory per GPU.
- Figure 10 evaluates the influence of client number.
- Figure 11 evaluates how Tstart and Tend affect model performance and communication cost.
- Horovod enables parallel model training across the four GPUs, with each GPU representing a platform.
A.5 Model Initialization
The section describes student-model initialization from UniLM layers and reports training settings, optimization steps, gradient approximation, and training durations for FedKD variants.
- Model Initialization: FedKD4 and FedKD2 initialize the student model with UniLM token embeddings and its first four or two layers, respectively.The hidden dimension is unchanged because UniLMv2 models with other hidden dimensions are unavailable.
- Training Cost: FedKD4 and FedKD2 require about 66 and 57 hours on MIND, and 12 and 10.5 minutes on SMM4H, respectively.The complete hyperparameter settings are listed in Table 6.
- Training Procedure: The training procedure sets teacher and student learning rates, client number, and the Tstart and Tend hyperparameters before parallel client execution.
- Gradient Approximation: Clients upload encrypted SVD components, while the server decrypts and reconstructs client gradients before forming global gradients.
- Gradient Approximation: The procedure continues until local models converge, after which clients reconstruct the global gradient.
- Training Procedure: Clients compute task and distillation losses, then derive local teacher and student gradients for return.