No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data (2106.05001v2)

Abstract: A central challenge in training classification models in the real-world federated system is learning with non-IID data. To cope with this, most of the existing works involve enforcing regularization in local optimization or improving the model aggregation scheme at the server. Other works also share public datasets or synthesized samples to supplement the training of under-represented classes or introduce a certain level of personalization. Though effective, they lack a deep understanding of how the data heterogeneity affects each layer of a deep classification model. In this paper, we bridge this gap by performing an experimental analysis of the representations learned by different layers. Our observations are surprising: (1) there exists a greater bias in the classifier than other layers, and (2) the classification performance can be significantly improved by post-calibrating the classifier after federated training. Motivated by the above findings, we propose a novel and simple algorithm called Classifier Calibration with Virtual Representations (CCVR), which adjusts the classifier using virtual representations sampled from an approximated gaussian mixture model. Experimental results demonstrate that CCVR achieves state-of-the-art performance on popular federated learning benchmarks including CIFAR-10, CIFAR-100, and CINIC-10. We hope that our simple yet effective method can shed some light on the future research of federated learning with non-IID data.

• The paper identifies significant classifier bias in federated learning with non-IID data and proposes the CCVR method for effective post-training calibration using privacy-preserving virtual representations.
• Analysis shows classifier layers are particularly susceptible to bias in non-IID settings, and calibrating the classifier with aggregated statistics substantially improves model performance.
• Experimental results demonstrate that the CCVR method achieves state-of-the-art performance on benchmarks and can enhance existing federated learning algorithms like FedAvg and FedProx.

Classifier Calibration in Federated Learning with Non-IID Data

The paper "No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data" presents an investigation into the training of classification models in federated learning systems with non-IID data distributions. The authors focus on the representation learning dynamics across different layers of a neural network and address classifier bias—a significant challenge in federated learning environments. Their novel method, Classifier Calibration with Virtual Representations (CCVR), proposes a pragmatic solution by calibrating classifiers post-training to improve model performance.

Key Insights from the Study

The authors embark on an experimental analysis to explore how non-IID data distribution affects neural networks, emphasizing representation similarity across layers. Two surprising observations are highlighted: first, the classifier layer exhibits greater bias compared to other layers; second, post-training classifier calibration yields notable performance improvements. They apply Centered Kernel Alignment (CKA) to measure feature similarity, discovering that classifier features present the lowest similarity across client models. This indicates significant discrepancies emerging from non-IID data, suggesting that debiasing the classifier could be crucial for performance enhancement.

Proposed Method: CCVR

The paper introduces the CCVR algorithm, which post-processes the global model by adjusting the classifier using virtual representations. These representations are drawn from an approximated Gaussian Mixture Model (GMM) in the feature space, ensuring privacy by not relying on actual training data. The server computes aggregate feature statistics from client-side updates, using them to generate synthetic features. This approach preserves privacy and achieves significant improvements without altering the training process itself.

Experimental Validation

Extensive evaluations on benchmarks such as CIFAR-10, CIFAR-100, and CINIC-10 demonstrate the efficacy of CCVR, establishing it as a state-of-the-art method. The paper presents clear evidence that classifiers trained with non-IID data can be substantially biased, but calibration with as few as a limited number of IID samples can mitigate this issue. The comparison further shows that methods like FedAvg and FedProx achieve enhanced performance with CCVR, highlighting its broad applicability.

Practical and Theoretical Implications

Practically, CCVR provides a method that can be easily integrated into existing federated learning workflows, offering model performance improvements without demanding extensive change in the federated setup. Theoretically, the work sheds light on the inherent limitations of federated learning frameworks dealing with non-IID data, emphasizing the importance of understanding layer-wise representations. It challenges the focus solely on federated aggregation methodologies, redirecting attention towards classifier-specific adjustments.

Future Directions

The research opens several avenues for future exploration. The calibrated approach might inform better initialization strategies for federated settings or inspire alternative architectures adapted to highly heterogeneous data distributions. The paper also prompts a broader consideration of privacy-preserving training processes beyond classification tasks, potentially influencing the development of federated approaches for other neural network architectures like LSTMs and Transformers.

In conclusion, the paper adeptly addresses a critical gap in federated learning by focusing on classifier calibration and offers a scalable, privacy-aware solution that could significantly influence subsequent AI research and applications in this domain.