Adaptive Self-Distillation for Minimizing Client Drift in Heterogeneous Federated Learning
M Yashwanth, Gaurav Kumar Nayak, Arya Singh, Yogesh Simmhan, Anirban Chakraborty
TL;DR
Federated Learning suffers from client drift when data are non-iid, especially with label distribution imbalance. The authors propose Adaptive Self-Distillation (ASD), a per-sample KL-divergence regularizer whose weights depend on the global model's prediction entropy and the client's label distribution, requiring no auxiliary data and easily adding to existing FL methods. They provide theoretical analysis showing reduced gradient dissimilarity and evidence of improved generalization via Hessian flatness, complemented by extensive experiments on CIFAR-10/100 and Tiny-Imagenet where ASD yields consistent gains across baselines. The work demonstrates a practical, low-overhead, plug-and-play approach to robust Federated Learning under label heterogeneity, improving convergence and performance without additional communication.
Abstract
Federated Learning (FL) is a machine learning paradigm that enables clients to jointly train a global model by aggregating the locally trained models without sharing any local training data. In practice, there can often be substantial heterogeneity (e.g., class imbalance) across the local data distributions observed by each of these clients. Under such non-iid label distributions across clients, FL suffers from the 'client-drift' problem where every client drifts to its own local optimum. This results in slower convergence and poor performance of the aggregated model. To address this limitation, we propose a novel regularization technique based on adaptive self-distillation (ASD) for training models on the client side. Our regularization scheme adaptively adjusts to each client's training data based on the global model's prediction entropy and the client-data label distribution. We show in this paper that our proposed regularization (ASD) can be easily integrated atop existing, state-of-the-art FL algorithms, leading to a further boost in the performance of these off-the-shelf methods. We theoretically explain how incorporation of ASD regularizer leads to reduction in client-drift and empirically justify the generalization ability of the trained model. We demonstrate the efficacy of our approach through extensive experiments on multiple real-world benchmarks and show substantial gains in performance when the proposed regularizer is combined with popular FL methods.