PaperlayerView exact revision on arXiv

Source-linked AI summary

Tackling Data Heterogeneity in Federated Learning with Class Prototypes

Yutong Dai, Zeyuan Chen, Junnan Li, Shelby Heinecke, Lichao Sun, Ran Xu

arXiv:2212.02758v2cs.LGcs.AI

TL;DR

Federated learning faces data heterogeneity, including class imbalance that can degrade global and personalized models. FedNH combines uniformly distributed class prototypes with prototype semantics, achieving competitive or better classification accuracy with lower computation costs while improving consistency across tested cross-device cases.

  • Problem

    Federated learning clients can have different data distributions, and class imbalance is a data-heterogeneity challenge for global and personalized classification models.

  • Method

    FedNH combines uniformity and semantics in class prototypes, using uniform prototype initialization and a prototype-learning strategy intended to reduce minority representation collapse.

  • Results

    FedNH improves personalized and global classification accuracy on CIFAR10, CIFAR100, and TinyImageNet, matching or exceeding state-of-the-art methods with significantly fewer computation costs.

  • Takeaways & Limitations

    Across tested cross-device cases, FedNH provides consistent personalized-model gains and achieves the best PM(V) accuracy in almost all cases with minimal computation overhead.

  • Takeaways & Limitations

    The convergence analysis assumes bounded-below, smooth local losses and a bounded stochastic-gradient condition.

Abstract

from arXiv · show

Data heterogeneity across clients in federated learning (FL) settings is a widely acknowledged challenge. In response, personalized federated learning (PFL) emerged as a framework to curate local models for clients' tasks. In PFL, a common strategy is to develop local and global models jointly - the global model (for generalization) informs the local models, and the local models (for personalization) are aggregated to update the global model. A key observation is that if we can improve the generalization ability of local models, then we can improve the generalization of global models, which in turn builds better personalized models. In this work, we consider class imbalance, an overlooked type of data heterogeneity, in the classification setting. We propose FedNH, a novel method that improves the local models' performance for both personalization and generalization by combining the uniformity and semantics of class prototypes. FedNH initially distributes class prototypes uniformly in the latent space and smoothly infuses the class semantics into class prototypes. We show that imposing uniformity helps to combat prototype collapse while infusing class semantics improves local models. Extensive experiments were conducted on popular classification datasets under the cross-device setting. Our results demonstrate the effectiveness and stability of our method over recent works.

1 Introduction

Federated learning faces client-level distribution differences, including overlooked class imbalance, motivating personalized models and FedNH’s evaluation on standard benchmarks.

  • Class imbalance and data heterogeneity remain overlooked together in personalized federated learning, where some clients may lack samples from particular classes.
  • Figure 1 visualizes balanced and imbalanced spiral datasets alongside learned representations and class prototypes under free and uniform classification heads.
  • FedNH improves personalized and global classification accuracy on Cifar10, Cifar100, and TinyImageNet, matching or exceeding state-of-the-art methods with fewer computation costs.
  • The paper defines neural-network bodies as representation learners and heads, or prototypes, as the final linear classification layer.

2 Related Work

Prior work addresses personalization and class imbalance through model decomposition, resampling, losses, client selection, auxiliary data, and classifier retraining.

  • Personalized federated learning methods include parameter decoupling, regularization, and model interpolation; FedPer and FedRep decouple shared bodies from personalized heads.
  • Non-federated approaches mitigate class imbalance through over-sampling, under-sampling, class-balanced losses, and decoupled training procedures.
  • Federated approaches include complementary-client selection, server-side auxiliary balanced data, and privacy-preserving classifier retraining based on federated features.

3 Methodology

FedNH addresses class imbalance by combining uniformly distributed class prototypes with semantic updates during federated training. The method fixes a shared classification head during local training, aggregates client representations, and updates global prototypes to improve consistency and avoid prototype collapse.

  • 3.1 A Motivating Example: Imbalanced training pulls minority-class prototypes toward majority classes and collapses minority representations, motivating uniformly distributed prototypes.The motivating example places prototypes uniformly on a unit hypersphere and fixes them during training.
  • 3.2 Proposed Method: FedNH combines uniformity and semantic prototype updates to improve local and global model generalization under heterogeneous class distributions.The convergence analysis provides an informal convergence result, with formal details deferred to Theorem 2 in Appendix B.
  • 3.2 Proposed Method: FedNH formulates prototype initialization to maximize pairwise separation while constraining every class prototype to have unit norm.The first constraint enforces separation, while the second ensures equal prototype magnitude.
  • 3.2 Proposed Method: During local training, clients learn representation bodies with a fixed shared classification head, giving selected clients a consistent learning goal.Clients subsequently compute averaged representations for the classes they own and return them with updated body parameters.
  • Server Update: The server averages received body parameters and updates global class prototypes using client-provided class representations.The prototype update is designed to infuse class semantics, while the server normalizes each global prototype to unit norm.

4 Experiments

Experiments evaluate FedNH under class-imbalanced cross-device federated learning using personalized and global accuracy, fairness, semantic prototype analysis, and smoothing sensitivity. FedNH is consistently competitive while adding little computation overhead.

  • Setups: Dirichlet(β) distributions with β ∈ {0.3, 1.0} create clients with different class distributions and sample counts.When β ≤ 1.0, clients are likely to have one or two dominating classes.
  • Evaluation: Table 1 reports global accuracy and personalized accuracy under PM(V) and PM(L), which use different class-weighting assumptions.PM(V) weights appearing classes equally, whereas PM(L) assumes matching training and testing distributions.
  • Results: FedNH delivers the most consistent personalized-model gains across tested cases and achieves the best PM(V) accuracy in almost all cases.It also adds minimum computation overhead compared with methods such as FedROD.
  • Results: PM(V) is less sensitive to class imbalance than PM(L), whose accuracy changes substantially from Dir(0.3) to Dir(1.0).PM(V) remains relatively stable across these imbalance settings.
  • Results: A strong global model often corresponds to strong personalized models when comparing GM and PM metrics.This comparison links global-model performance with personalized-model performance in the experiments.
  • Analysis: FedNH captures fine-grained semantic similarities among Cifar100 class prototypes, whereas FedAvg treats all classes as different.For smoothing sensitivity, ρ = 0.1 gives the worst performance, consistent with the theorem’s recommendation that ρ not be too small.
  • Fairness: FedNH improves personalized-model accuracy more equally across clients, as measured by the standard deviation of client accuracies.The fairness analysis reports mean accuracy ± standard deviation.

5 Conclusion

The paper concludes that FedNH addresses class-imbalance heterogeneity by combining uniformity and semantics in class prototypes. Its current scope is classification, with future work targeting intermediate layers and other vision tasks.

  • Conclusion: FedNH combines uniformity and semantics of class prototypes to learn high-quality representations for classification.The authors report effectiveness and robustness over recent works.
  • Limitations and Future Work: The method currently applies only to classification, and its prototype inductive bias is imposed only on the neural-network head.Future work will explore intermediate layers and different vision tasks.

A.1 Description of Datasets

The experiments use Cifar10, Cifar100, and TinyImageNet, spanning datasets with 10, 100, and 200 image classes and specified train/test counts.

  • Datasets: Cifar10 has 10 classes with 5000 training and 1000 testing images per class.
  • Datasets: Cifar100 has 100 classes with 500 training and 100 testing images per class.
  • Datasets: TinyImageNet has 200 classes with 500 training and 50 testing images per class, using 64 × 64 images.

A.2 Implementation Details

Implementation uses standard CNN and ResNet18 architectures with fixed optimization settings, 200 communication rounds, and method-specific baseline configurations.

  • Model Architecture: Cifar10 and Cifar100 use a 5-layer CNN, while TinyImageNet uses ResNet18 with batch normalization replaced by two-group normalization.
  • Optimization: All local training uses SGD with learning rate 0.01, momentum 0.9, exponential decay factor 0.99, and batch size 64.Weight decay is 10^-5 for Cifar10/Cifar100 and 10^-3 for TinyImageNet.
  • Training Schedule: All methods train locally for 5 epochs per round and run for 200 communication rounds.The paper reports no significant testing-accuracy gain after 200 rounds.
  • Baseline Configurations: FedPer shares the body while keeping the last linear layer personalized, and local-training baselines match FedAvg’s selected-round counts.
  • Baseline Configurations: CReFF uses 100 federated features per class, 100 feature-optimization epochs, and 300 classifier-retraining epochs per round.These settings follow its official implementation.
  • Spiral Example: Spiral-example experiments use SGD with learning rate 0.1, momentum 0.9, weight decay 10^-5, and Dir(0.3) allocation across 100 clients.

A.3 Discussions on testing accuracy

The discussion compares personalized-model behavior across baselines in the cross-device setting, highlighting why several methods may underperform under client sampling or limited participation.

  • FedAvg local models serve as strong personalized models, consistent with reports from prior work.The local model can be viewed as fine-tuning the global model, benefiting from federation.
  • FedPer may degrade when client sampling prevents availability of all clients.
  • FedProto may produce unstable, less informative prototypes when few clients participate and their data distributions vary across rounds.
  • Ditto’s personalized model performs worse than FedAvg, making the penalty parameter λ consequential.
  • FedRep’s performance drop may be attributed to limited training epochs on the shared body.

A.4 Some missing visualization

The visualizations examine data distributions, local representations, and prototype similarities to show how FedNH captures semantic structure and consistency across clients.

  • Figure 5: Spiral datasets: Figure 5 visualizes the balanced and unbalanced Spiral training datasets together with their testing datasets.
  • Figure 6: Local representations: FedNH representations place semantically related Leopard and Bear classes closer than Oka tree.The comparison uses PCA and visualizes the first two latent dimensions.
  • Figure 6: Local representations: FedNH learns Leopard and Bear representations consistently across clients, whereas FedAvg switches their latent-space locations.
  • Prototype similarity: FedNH prototypes capture fine-grained class semantics through higher similarity within semantically related Cifar100 super-classes.
  • Prototype similarity: Prototype correlations are high for related animal super-classes but low between animal and architectural super-classes.
  • Baseline comparison: Figure 8 presents prototype similarities for different baseline methods.

A.5 Increasing the participate ratio

Increasing the client participation ratio from 0.1 to 0.3 slightly improves reported global and personalized-model accuracies, while FedNH adds little computation overhead.

  • Slight accuracy improvements occur in GM, PM(V), and PM(I) when participation rises from 0.1 to 0.3.
  • Table 3 reports testing accuracy at a participation ratio of 0.3, with means and standard deviations averaged over three independent runs.
  • FedNH adds minimum computation overhead compared with state-of-the-art methods such as FedROD.FedNH requires additional computation only for prototype updates, whereas FedROD trains two classification heads or a generating sub-network.

B Convergence Analysis

The convergence analysis develops bounds for FedNH under stated assumptions, covering local-update progress, prototype and representation changes, and convergence-related results.

  • The analysis studies Algorithm 1 under the full participation assumption.
  • A sufficient-decrease argument characterizes expected progress from E local stochastic-gradient steps.
  • The proof uses expectations over client-data subsampling together with the stated assumptions.
  • The analysis bounds changes between the end of local updates and receipt of new server prototypes and representation parameters.
  • Lipschitz continuity is established for the objective with respect to both model parameters and prototype-related parameters.
  • The derivation combines intermediate inequalities to obtain the desired result and establishes a lower-bounded expected objective under the assumptions.
Loading 2212.02758v2…