No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data

Mi Luo; Fei Chen; Dapeng Hu; Yifan Zhang; Jian Liang; Jiashi Feng

Source-linked AI summary

No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data

Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, Jiashi Feng

arXiv:2106.05001v2Jun 9, 2021cs.LGcs.CVcs.DCstat.ML

TL;DR

Federated learning with non-IID data lacks a deep understanding of how heterogeneity affects different network layers. The paper analyzes layer representations, identifies classifier bias, and proposes CCVR, which calibrates the classifier using virtual GMM-sampled representations; experiments show significant gains across image benchmarks, including up to 10.41% on CINIC-10.

Problem
Existing federated learning methods address non-IID data through local regularization, server aggregation, data sharing, or personalization, but provide limited understanding of heterogeneity’s effects across deep-network layers.
Method
The paper analyzes cross-client layer representations and introduces CCVR, which post-calibrates the classifier with virtual features sampled from an estimated Gaussian Mixture Model without accessing real images.
Results
CCVR significantly improves baseline accuracies across CIFAR-10, CIFAR-100, and CINIC-10, with gains up to 10.41% on CINIC-10.
Takeaways & Limitations
Classifier bias is a central source of performance degradation under non-IID federated training, and post-training calibration can improve models without modifying the original training process or transmitting raw representations.
Takeaways & Limitations
Classifier calibration alone cannot match IID accuracy in highly heterogeneous settings because its performance depends strongly on representation quality.

Abstract

from arXiv · show

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.

1 Introduction

Federated learning with non-IID data can degrade deep classification models, and layer-wise analysis identifies the classifier as especially heterogeneous and biased. This motivates CCVR, a privacy-preserving post-training calibration method using virtual representations, which improves benchmark performance.

Non-IID client data can cause unstable, slow convergence and suboptimal or detrimental model performance.
Layer-wise CKA analysis finds the classifier has the lowest feature similarity across local models, indicating greater classifier heterogeneity.
Post-calibrating the classifier significantly improves accuracy with only a small fraction of IID data, but direct deployment can violate federated-learning privacy rules.
CCVR rectifies the classifier after federated training by generating virtual representations from an approximated Gaussian Mixture Model.
CCVR achieves significant accuracy improvements over several federated algorithms on CIFAR-10, CIFAR-100, and CINIC-10 benchmarks.

2 Related Work

Related work addresses federated non-IID data through client-drift mitigation, aggregation changes, data sharing, and personalization. The paper instead decouples representation learning from classifier calibration, making CCVR compatible with existing approaches.

Federated-learning research on non-IID data commonly pursues client-drift mitigation, aggregation schemes, data sharing, or personalized models.
Client Drift Mitigation: Client-drift methods modify local objectives or weighting to reduce inconsistency caused by heterogeneous local optima.
Aggregation Scheme: Aggregation methods modify server-side model fusion to address oscillation and harmful effects from element-wise weight averaging.
Data Sharing: Data-sharing methods provide public, synthesized, or condensed samples to supplement absent classes, but may violate federated-learning privacy rules.
CCVR decouples representation learning from classifier calibration and can be combined with most existing federated-learning approaches.

3 Heterogeneity in Federated Learning: The Devil Is in Classifier

The study examines how non-IID federated training affects different layers and finds the classifier is especially heterogeneous and biased. Classifier regularization and post-training calibration are evaluated, with calibration producing strong improvements even from limited IID data.

3.1 Problem Setup: The experiment measures layer-wise CKA similarity across local models trained on CIFAR-10 with Dirichlet-partitioned non-IID data.Ten clients train a seven-layer convolutional network with FedAvg for 100 communication rounds and 10 local epochs per round.
3.2 A Closer Look at Classification Model: Classifier Bias: Deeper layers exhibit lower cross-client representation similarity, indicating greater heterogeneity toward the classifier.The classifier is the last compared layer and shows the lowest feature similarity under non-IID training.
3.2 A Closer Look at Classification Model: Classifier Bias: Classifier weight norms become biased toward classes with more local training samples, with stronger bias after non-IID training than IID training.The bias is observed both at the initial stage and more heavily at the end of training under non-IID data.
3.3 Classifier Regularization and Calibration: Post-calibration with IID samples fixes the feature extractor and optimizes the classifier using cross-entropy loss, but requires IID raw features from heterogeneous clients.The authors state that this IID-sample calibration cannot be applied directly in a real federated learning system.
3.3 Classifier Regularization and Calibration: Classifier weight normalization helps under light heterogeneity but can provide less benefit or cause damage as heterogeneity increases, whereas proximal regularization gives consistently minor improvements.Both approaches are compared with classifier post-calibration on CIFAR-10 across three heterogeneity levels.
3.3 Classifier Regularization and Calibration: Post-calibrating the FedAvg classifier with all training samples significantly improves performance across all tested heterogeneity levels.Calibration remains consistently effective even when using only 1/50 of the training data, and is also tested on FedProx models.

4 Classifier Calibration with Virtual Representations

CCVR calibrates the classifier after federated training using virtual feature representations generated from class-wise Gaussian statistics, without accessing real images. Clients upload local feature means and covariances, which the server aggregates to generate virtual features for classifier retraining.

Classifier Calibration with Virtual Representations: CCVR runs on the server after federated training and recalibrates the global model’s classifier while retaining its learned feature extractor.The method uses virtual features rather than real images.
Feature Distribution Estimation: CCVR approximates each class’s feature distribution as a Gaussian and aggregates client statistics into global means and covariances.The Gaussian assumption is applied in the feature space of the learned extractor.
Feature Distribution Estimation: Each client extracts class-specific features and computes local means and covariances before uploading Gaussian statistics to the server.The uploaded statistics replace raw representations in the server-side calibration process.
Virtual Representations Generation: The server samples Mc virtual features for each class from N(µc, Σc), assigning every generated feature its class label.The number of samples can reflect each class’s fraction of the overall data.
Classifier Re-Training: The classifier is retrained with cross-entropy on virtual representations, producing a calibrated classifier composed with the pretrained feature extractor.This yields the final model g_eϕ ◦ f_bθ.
Privacy Protection: CCVR provides basic privacy protection because clients upload Gaussian statistics instead of raw representations and can be combined with additional privacy techniques.The paper describes CCVR as a post-hoc method.

5 Experiment

Experiments evaluate CCVR on CIFAR-10, CIFAR-100, and CINIC-10 across federated learning baselines and heterogeneous partitions. CCVR consistently improves baseline accuracy, with especially large gains when classifier bias is severe and representations are more separable.

Experiment Setup: CCVR is evaluated on CIFAR-10, CIFAR-100, and CINIC-10 against FedAvg, FedProx, FedAvgM, and MOON.CINIC-10 introduces distribution shifts because it combines samples from ImageNet and CIFAR-10.
Can classifier calibration improve performance of federated learning?: CCVR increases accuracy for all baseline methods, with gains reaching 10.41% on CINIC-10 without modifying federated training.Calibration can be applied as post-processing to the trained global model.
In what situation does CCVR work best?: CIFAR-100 shows subtler improvement because its difficult 100-class task can produce representations with low separability, while CCVR approaches the classifier-calibration upper bound.The result indicates effective classifier correction despite a poor feature extractor.
In what situation does CCVR work best?: On CINIC-10, FedAvg features separate some classes well, but classifier bias causes systematic misclassification that CCVR alleviates.After calibration, errors are concentrated nearer decision boundaries rather than dominated by particular classes.
How to forecast the performance of classifier calibration?: Mean Sliced Wasserstein distance between GMM mixtures is positively correlated with the accuracy upper bound of classifier calibration.The paper proposes this distance as a way to assess simulated-representation quality and forecast calibration performance.
How many virtual features to generate?: Even a few virtual features can significantly increase classification accuracy across CIFAR-10 heterogeneity settings.The number of virtual features per class is varied over {0, 50, 100, 500, 1000, 2000}.
Experiment Setup: CCVR’s improvements generally increase with stronger data heterogeneity, linking classifier bias to the imbalance of training data.The study varies CIFAR-10 Dirichlet concentration parameters α ∈ {0.5, 0.1, 0.05}, where smaller α is more imbalanced.

6 Limitations

The study focuses on classifier behavior, but finds that classifier calibration remains insufficient under highly heterogeneous data and depends on representation quality. Its evaluation is centered on image classification and may not extend to language tasks or other architectures.

Classifier calibration alone cannot match IID accuracy in highly heterogeneous settings.The authors attribute this limitation to the quality of learned representations.
Classifier-calibration benefits are strongest when methods learn high-quality representations but poor classifiers.The authors describe this as a trade-off between representation and classifier quality.
The experiments mainly study image classification with convolutional neural networks.The Gaussian assumption is reported to work well for visual models such as CNNs.
The reported Gaussian assumption may not hold for language tasks or architectures such as LSTM and Transformer.The authors identify extensions to other tasks and architectures as an open direction.

7 Conclusion

The paper analyzes how non-IID federated learning affects neural-network layers and finds particularly low similarity and significant bias in local classifiers. It then introduces CCVR, which calibrates classifiers using virtual features from an approximated GMM and improves performance across three image datasets.

The classifiers of different local models are less similar than any other layer and exhibit significant bias.The analysis uses representation techniques to compare corresponding layers across clients.
CCVR calibrates the classifier by sampling virtual features from an approximated Gaussian Mixture Model.This avoids uploading raw features to the server.
CCVR steadily improves over several popular federated learning algorithms on three image datasets.

A Derivation of Global Mean and Covariance

This material introduces the global-statistics derivation and explains the CKA-based experimental setup and feature-reconstruction privacy check. It also reports that CCVR alleviates classifier-weight imbalance.

A Derivation of Global Mean and Covariance: The derivation assumes Nc,k ≥ 1 and Nc ≥ 2 for each class c and client k when forming global statistics.The accompanying text states that the mean formula also holds when Nc,k = 1.
B Details of Centered Kernel Alignment (CKA): CKA compares output-feature representations from two neural networks given the same input data.The representations are organized as feature matrices extracted from a selected dataset and then pre-processed for similarity computation.
B Details of Centered Kernel Alignment (CKA): The experiments use linear CKA with 256-dimensional outputs and 50,000 selected CIFAR-10 samples.A local model is the global model optimized on a client’s local data for 10 epochs.
Raw representations can resemble corresponding raw images after feature inversion, whereas reconstruction from mean representations is examined as a privacy comparison.The experiments use ImageNet images and a pre-trained ResNet-50.
CCVR observably alleviates imbalanced classifier-weight L2-norm distributions for FedAvg, FedProx, and MOON.The comparison is conducted on CIFAR-10 with concentration parameter α = 0.1.

D.2 Why not choose regularization during training instead of post-calibration?

The experiments find that regularizing the classifier during training yields only minor improvements in most cases, motivating comparison with post-calibration.

Classifier regularization during training improves performance only a little in most cases.
FedProx and MOON surpass FedAvg while enhancing CKA similarity across layers.The passage uses this comparison to examine why training-time regularization produces minor improvements.

D.3 Comparison of the effectiveness of CCVR on FedAvg, FedProx and MOON.

CCVR improves accuracy across FedAvg, FedProx, and MOON, with gains depending on the underlying method, data heterogeneity, and number of virtual features. MOON benefits less on CINIC-10 because its learned feature space is more discriminative and its classifier is less tied to label distribution early in training.

3.75% versus 9.79% and 9.53%: CCVR improves CINIC-10 less for MOON than for FedAvg and FedProx.The comparison is reported for MOON, FedAvg, and FedProx, respectively.
MOON’s classifier L2-norm distribution is unrelated to label distribution in Round 1, unlike FedAvg and FedProx.By the end of federated training, MOON’s classifier tends to favor different classes than the other two methods.
MOON produces low-entropy, highly intra-class-compact feature clusters and a larger decision margin, making its feature space more discriminative.These representations provide greater tolerance to classifier bias.
Larger Mc yields higher accuracy under highly heterogeneous distributions with α = 0.05 and α = 0.1, while Mc is more sensitive under more balanced distributions.This conclusion holds for CCVR applied to FedProx and MOON as well as the previously studied FedAvg setting.
CCVR steadily improves accuracy across all methods when the number of clients is N ∈ {10, 50, 100}.The experiments use CIFAR-10 with α = 0.1, and varying the client count does not affect CCVR’s effectiveness.

D.6 Full results of classifier calibration with whole data and partial data.

The section reports full CIFAR-10 classifier-calibration results using whole and partial data, alongside feature visualizations comparing classifier predictions before and after CCVR.

Table 5 reports CIFAR-10 Accuracy@1 (%) for classifier calibration with whole data and partial data.
Figure 13 visualizes CINIC-10 features from FedAvg, FedProx, and MOON, colored by ground truth and classifier predictions before and after CCVR.The visualization compares predictions before and after applying CCVR.

E.1 Datasets

The experiments use heterogeneous non-IID client label distributions across CIFAR-10, CIFAR-100, and CINIC-10, with specified architectures, hyperparameters, and client settings.

CIFAR-10, CIFAR-100, and CINIC-10 are visualized with heterogeneous label distributions across 10 clients.Under the non-IID partition, sample counts vary by client and each client may contain only a few categories.
The concentration parameter α is set to 0.5 for CIFAR-100 and CINIC-10.
The CIFAR-10 model is a simple convolutional neural network, while CIFAR-100 and CINIC-10 use MobileNetV2.For CIFAR-100, the classifier output dimension is changed to 100.
The experiments repeat all runs with three different random seeds and summarize hyperparameters in Table 8.
Figure 14 studies CCVR Accuracy@1 (%) on CIFAR-10 with different numbers of virtual features, while Tables 4 and 5 report client-count and calibration-data settings.