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

Class imbalance is an overlooked form of data heterogeneity in federated classification. FedNH combines uniform class prototypes with class semantics, and experiments report improved personalized and global accuracy, on par with or better than state-of-the-art methods.

  • Problem

    Class imbalance is an overlooked form of data heterogeneity in federated classification.

    Source: §5 · p. 8

  • Method

    FedNH distributes class prototypes uniformly in the latent space and smoothly infuses class semantics into them.

    Source: §5 · p. 8

  • Results

    FedNH improves personalized and global classification accuracy on Cifar10, Cifar100, and TinyImageNet, with results on par with or better than the state of the art and significantly fewer computation costs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings support combining prototype uniformity and semantics to address class imbalance in federated classification.

    Source: §5 · p. 8

    FedNH currently applies only to classification, and its convergence analysis assumes client losses are bounded below and smooth in model parameters.

    Source: §3.3 · p. 5, §5 · p. 8

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

FedNH addresses class-imbalanced, heterogeneous federated classification by combining uniformly initialized class prototypes with gradually infused class semantics. Experiments report improved personalized and global accuracy on Cifar10, Cifar100, and TinyImageNet, with results on par or better than the state of the art and significantly fewer computation costs.

1 Introduction

  • Federated learning enables collaborative learning from decentralized data with privacy protection, but heterogeneous client distributions violate the i.i.d. assumption and can degrade the global model; PFL combines global and local learning for client tasks.Source: §1 · p. 1
  • Class imbalance means clients have differing class distributions and may lack classes; testing on locally distribution-matched splits can bias averaged accuracy toward dominant classes.Source: §1 · p. 1
  • Centralized re-sampling or re-weighting may not transfer because clients have different dominant or missing classes, and re-sampling requires all-class knowledge that may violate privacy.Source: §1 · p. 1
  • Prior imbalanced-learning work [4] [5] separates representation learning from classification, while FL works [7] identify the classifier as a performance bottleneck and suggest strong shared representations can help.Source: §1 · p. 1
  • FedNH uniformly distributes class prototypes in latent space and smoothly infuses class semantics, aiming to improve representation quality and local classification under imbalance.The paper uses “head” and “prototype” interchangeably for the last linear classification layer.Source: §1 · pp. 1–2
  • Experiments on Cifar10, Cifar100, and TinyImageNet show FedNH improves personalized and global classification accuracy, with results on par or better than the state of the art and significantly fewer computation costs.Source: §1 · p. 2

2 Related Work

The related work reviews personalized federated learning strategies and approaches to class imbalance, distinguishing general imbalance methods from the limited work addressing it in federated settings.

2.1 Personalized Federated Learning

  • Personalized FL methods include body/head-decoupling designs such as FedPer [8], FedRep, FedBABU, and FedROD [7], which vary training, sharing, and head personalization.These methods assume the same model architecture across clients.Source: §2.1 · p. 2
  • FedProto [10] shares class prototypes, allowing clients to use different model architectures as bodies.Its prototypes are averaged representations for each class.Source: §2.1 · p. 2

2.2 Class Imbalanced Learning

  • Outside FL, imbalance is addressed through over/undersampling, class-balanced losses [14] [17], decoupled representation/classifier training [4] [5], and multi-expert models [18].Over/undersampling can overfit minority classes or lose majority-class information.Source: §2.2 · p. 2
  • In FL, prior methods select clients with complementary class distributions, use balanced server-side auxiliary data [20], or retrain classifiers using federated features, as in CReFF [21].The client-selection approach requires revealing class distributions to the server; the auxiliary-data approach assumes balanced server data.Source: §2.2 · p. 3
  • FedROD [7] and CReFF [21] are identified as the only FL works addressing class-imbalance heterogeneity without potential privacy leakage or server-side auxiliary datasets.Source: §2.2 · p. 3

3 Methodology

FedNH addresses class imbalance by initializing class prototypes uniformly, holding the classification head fixed during local training, and updating global prototypes to incorporate class semantics. Its convergence analysis considers local training under changing global prototypes.

3.1 A Motivating Example

  • With imbalanced data, minority-class prototypes and representations collapse toward majority classes, while fixed uniform prototypes help counter this collapse and keep clients’ learning goals consistent.In the motivating example, balanced-data prototypes are relatively uniform and similar in magnitude; imposing uniformity does not hurt balanced training.Source: §3.1 · p. 3

3.2 Proposed Method

  • FedNH addresses prototype separation and semantic learning through a classification-head initialization scheme and a separate process for infusing class semantics during training.The method frames these as two questions: how to make prototypes maximally separated and equidistant, and how to encourage them to learn class semantics.Source: §3.2 · pp. 3–4

Initialization.

  • The server initializes unit-norm class prototypes by maximizing their minimum pairwise squared distance, then normalizes network representations to unit length.The optimization is performed once at initialization and its solution can be saved for future use.Source: Initialization · p. 4

Client Update.

  • Each client trains its body using the fixed broadcast classification head, computes averaged local representations for its classes, and returns them with its body parameters for aggregation.The server averages the received body parameters as in FedAvg.Source: Client Update · p. 5, Server Update · p. 5

Server Update.

  • The server updates global prototypes non-parametrically using client-provided class representations and smoothing, with rho suggested near 1 to temper less-informative early updates and client-sampling variation.The method discounts local contributions by 1 − rho so server prototypes change gradually.Source: Server Update · p. 5

3.3 Convergence Analysis

  • The convergence analysis targets small per-client objective-gradient measures despite round-varying global prototypes, under smoothness and stochastic-gradient assumptions.Its informal theorem specifies conditions on rho and the step size eta; formal statements and proofs are in Appendix B.Source: §3.3 · pp. 5–6

4 Experiments

FedNH delivers consistent personalized-model gains, with the best PM(V) accuracy in almost all tested cases; analyses also examine class semantics, smoothing sensitivity, and fairness.

4.1 Setups, Evaluation, and Baselines

  • Experiments simulate class imbalance across 100 clients using Dirichlet β values of 0.3 and 1.0, with 10% client participation.CNNs are used for CIFAR-10 and CIFAR-100, and ResNet18 for TinyImageNet; evaluation includes global and personalized testing accuracy, with several personalized baselines.Source: §4.1 · p. 7

4.2 Results

  • FedNH’s personalized-model gains are the most consistent across tested cases, and it achieves the best PM(V) accuracy in almost all cases.FedAvg personalized models are a strong baseline and can outperform many specialized personalized methods; FedNH adds minimal computation overhead versus methods such as FedROD.Source: §4.2 · p. 7
  • The PM(V) metric is less sensitive to class imbalance than PM(L), whose accuracy changes significantly between Dir(0.3) and Dir(1.0).PM(V) weights all classes appearing in each client’s training data equally, whereas PM(L) assumes training and testing distributions match.Source: §4.2 · pp. 7–8
  • Comparing global-model and personalized-model metrics, the paper observes that a strong global model often leads to strong personalized models.Source: §4.2 · p. 8

4.3 Analysis

  • FedNH captures fine-grained class semantics in Cifar100 prototypes, unlike FedAvg, and similar classes’ learned representations tend to be closer.The analysis groups Cifar100’s 100 classes into 20 super-classes when visualizing prototype similarities.Source: §4.3 · p. 8
  • With smoothing parameter ρ selected from {0.1, 0.3, 0.5, 0.7, 0.9}, ρ = 0.1 gives the worst performance, consistent with Theorem 1’s warning against too-small ρ.Source: §4.3 · p. 8, Figure 4 · p. 8

5 Conclusion

FedNH addresses class-imbalance heterogeneity in federated learning by combining prototype uniformity and semantics. The conclusion also notes its current scope and directions for future work.

5 Conclusion

  • FedNH combines class-prototype uniformity and semantics to learn high-quality classification representations, and experiments show effectiveness and robustness over recent works.Source: §5 · p. 8

Limitation and Future Work

  • FedNH currently applies only to classification and imposes its uniformity and semantics inductive bias on the network head; future work will explore intermediate layers and other vision tasks.Source: §5 · p. 8

A Experiments · A.1 Description of Datasets

The dataset description specifies class counts, per-class training and testing image counts, and image resolution for CIFAR-10, CIFAR-100, and TinyImageNet.

A.1 Description of Datasets

  • Cifar10, Cifar100, and TinyImageNet contain 10, 100, and 200 image classes, respectively.Each class has 5000/1000 training/testing images in Cifar10, 500/100 in Cifar100, and 500/50 in TinyImageNet; TinyImageNet images are 64 × 64.Source: §A.1 · p. 11

A.2 Implementation Details

The implementation standardizes training conditions across algorithms and documents model architectures and baseline-specific choices.

A.2 Implementation Details

  • All algorithms used a unified PyTorch codebase, with SGD training for 200 rounds and fixed seeds to match client configurations across baselines.Local training used a 0.01 initial learning rate, 0.9 momentum, exponential decay by 0.99 per round, five epochs per round, and batch size 64. Weight decay was 10^-5 for CIFAR-10/100 and 10^-3 for TinyImageNet.Source: §A.2 · p. 11
  • The CIFAR-10/100 CNN has five layers, while ResNet18 replaces batch normalization with group normalization using two groups.The CNN consists of two convolutional and three linear layers; the ResNet18 otherwise remains unchanged.Source: §A.2 · p. 11
  • Baseline protocols matched each local-training client’s number of rounds to FedAvg and specified personalized layers and training schedules for FedPer, Ditto, FedRep, and FedProto [10].FedPer keeps the last linear layer private; Ditto uses λ=0.75 and FedProto λ=0.1. FedRep trains the head for 10 epochs, then the body for five.Source: §A.2 · p. 11
  • For CReFF, both learning rates were set to 0.01 because the official 0.1/0.01 choice made cross-device training unstable; FedNH used ρ=0.9 and a trainable scaling parameter.CReFF uses 100 federated features per class, 100 feature-optimization epochs, and 300 classifier retraining epochs per round. FedBABU fine-tunes clients for five epochs; FedROD costs more computation and uses a CNN hyper-network only under the stated GPU-memory constraint.Source: §A.2 · p. 12

A.3 Discussions on testing accuracy

The discussion presents FedAvg local models as strong personalized models and attributes weaker baseline performance to assumptions or training choices that are problematic under cross-device participation. It also describes CReFF’s instability and its weaker performance after tuning in the reported setting.

A.3 Discussions on testing accuracy

  • FedAvg local models are strong personalized models, consistent with reports in [7], and can be viewed as globally fine-tuned models that benefit from federation.The passage characterizes local models as the result of fine-tuning the global model.Source: §A.3 · p. 12
  • FedProto’s full-participation assumption is problematic in cross-device settings, where small, changing client samples yield unstable prototypes that may not reflect the global data or inform newly selected clients.Successive rounds may select clients with drastically different data distributions, making the previous round’s prototypes less informative.Source: §A.3 · p. 12
  • FedPer’s all-clients assumption conflicts with cross-device client sampling, Ditto’s personalized model underperforms FedAvg when its penalty parameter λ may be poorly chosen, and FedRep’s performance drop is attributed to low body-training epochs.The passage says choosing a proper Ditto penalty parameter could make a big difference.Source: §A.3 · p. 12
  • CReFF’s hyperparameters in [21] make cross-device training unstable; tuning learning rates for federated features and the balanced classifier enables convergence but yields weaker performance in the 20-client, 40%-participation setting [21].The passage reports the tuned model’s performance as less strong than that indicated in [21].Source: §A.3 · p. 12

A.4 Some missing visualization

The appendix adds visualizations of the datasets, client representations, and prototype similarities. They show FedNH representations and prototypes capturing class semantics, with client representations more consistent than FedAvg’s.

Client representations

  • FedNH’s Leopard and Bear representations are closer to each other than to Oka tree and occupy consistent latent-space locations across clients, unlike FedAvg’s representations, which switch locations.The comparison uses PCA to visualize the first two dimensions for three classes across clients 0 and 1.Source: §A.4 · p. 12, Figure 6 · p. 13

Dataset visualization

  • Figure 5 visualizes the balanced and unbalanced training datasets alongside the testing datasets used in Section 3.1.Source: §A.4 · p. 12, Figure 5 · p. 13

Prototype similarities

  • On Cifar100, FedNH prototype similarities capture fine-grained class semantics, with higher correlations for some related super-classes and lower correlations for less similar ones.Classes within the same super-class are grouped along the similarity matrix’s main diagonal; off-diagonal examples include aquatic-animal and insect-versus-building super-classes.Source: §A.4 · p. 14, Figure 7 · p. 14

A.5 Increasing the participate ratio · A.6 Discussion on computation costs

Increasing client participation from 0.1 to 0.3 slightly improves accuracy in GM, PM(V), and PM(I). FedNH adds minimal computation overhead, requiring prototype updates rather than additional classification heads.

A.5 Increasing the participate ratio

  • With participation increased from 0.1 to 0.3, accuracy in GM, PM(V), and PM(I) slightly improves.The comparison at participation ratio 0.3 is reported in Table 3.Source: §A.5 · p. 16, Table 3 · p. 16
  • GM accuracy slightly improves when the client participation ratio rises from 0.1 to 0.3.Source: §A.5 · p. 16
  • PM(V) accuracy slightly improves when the client participation ratio rises from 0.1 to 0.3.Source: §A.5 · p. 16
  • PM(I) accuracy slightly improves when the client participation ratio rises from 0.1 to 0.3.Source: §A.5 · p. 16
  • Table 3 reports testing accuracy at participation ratio 0.3 as mean ± standard deviation averaged over three independent runs.Source: Table 3 · p. 16

A.6 Discussion on computation costs

  • FedNH adds minimum computation overhead compared with state-of-the-art methods such as FedROD[7].Source: §A.6 · p. 16
  • FedROD[7] trains two classification heads or a sub-network to generate one, whereas FedNH requires only additional computation for prototype updates (3) and (4).Source: §A.6 · p. 16

B Convergence Analysis

The section analyzes convergence of Algorithm 1 under full participation, deriving one-round progress and a theorem-level guarantee for a uniformly sampled client iterate.

B Convergence Analysis

  • Lemma 3 establishes progress over one communication round for Algorithm 1 under its stated algorithmic choices.The proof combines a sufficient-decrease result for E local stochastic-gradient steps with bounds on the changes between client and server stages.Source: B Convergence Analysis · pp. 17–19
  • Theorem 2 gives a convergence guarantee for a uniformly sampled client iterate under the stated conditions, with κ>0 and the initial objective gap appearing in the bound.The supplied theorem passages do not expose the full bound, so no convergence rate is stated.Source: B Convergence Analysis · p. 19

Related papers

Loading 2212.02758v2…