Adaptive Self-Distillation for Minimizing Client Drift in Heterogeneous Federated Learning

TL;DR

This paper introduces an adaptive self-distillation regularization technique for federated learning that reduces client drift caused by data heterogeneity, improving convergence and model performance.

Contribution

The paper proposes a novel ASD regularization method that adaptively mitigates client drift in federated learning, compatible with existing algorithms.

Findings

ASD reduces client drift in heterogeneous FL settings.

Combining ASD with existing methods improves accuracy.

Extensive experiments show substantial performance gains.

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.