[Paper Review] No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data

Mi Luo, Fei Chen+4 more|arXiv (Cornell University)|Jun 9, 2021

Privacy-Preserving Technologies in DataComputer Science130 citations

TL;DR

The paper shows that classifier bias is the main bottleneck in non-IID federated learning and introduces CCVR, a post-training calibration method using virtual representations drawn from an estimated Gaussian Mixture Model to significantly improve accuracy on standard benchmarks.

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.

Motivation & Objective

• Investigate how data heterogeneity affects representations across neural network layers in federated learning.
• Identify the layer most affected by non-IID data and diagnose why the classifier underperforms.
• Propose a simple, privacy-preserving classifier calibration method to mitigate bias.
• Demonstrate that post-training calibration with virtual representations yields state-of-the-art results on common benchmarks.

Proposed method

• Empirical analysis of layer-wise representations using Centered Kernel Alignment (CKA) across local models trained with FedAvg on non-IID data.
• Observation that the classifier layer exhibits the lowest cross-client similarity and biased weight norms.
• Evaluation of regularization strategies (clsnorm, clsprox) and post-calibration using IID samples as baseline comparisons.
• Proposal of Classifier Calibration with Virtual Representations (CCVR) which estimates per-class feature distributions and samples virtual features from a Gaussian mixture model to retrain only the classifier.
• Privacy-preserving design: only local Gaussian statistics are uploaded; CCVR is post-hoc and can be combined with existing FL methods.

Experimental results

Research questions

• RQ1 Does non-IID data cause layer-wise representation divergence, particularly in the classifier, more than other layers?
• RQ2 Can debiasing the classifier during training or after training improve global accuracy in federated learning under non-IID conditions?
• RQ3 Is a post-hoc, privacy-preserving calibration using synthetic (virtual) representations sufficient to achieve state-of-the-art performance?
• RQ4 Under what data regimes and hyperparameters does CCVR provide the largest gains?

Key findings

• Classifier representations diverge more for deeper layers under non-IID data, with the classifier showing the strongest bias across clients.
• Post-training calibration of the classifier yields significant accuracy gains across CIFAR-10, CIFAR-100, and CINIC-10, often surpassing other regularization methods.
• Regularizing only the classifier during training helps in light heterogeneity but degrades with higher heterogeneity; post-hoc calibration remains robust.
• CCVR, which generates virtual representations from a class-conditioned Gaussian mixture model, provides substantial improvements across baseline FL methods (e.g., FedAvg, FedProx, MOON).
• The effectiveness of CCVR correlates with the separability of the learned feature distributions, as measured by GMM Wasserstein separability, and benefits more when representations are already strong.