Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

Abstract: Class imbalance poses a critical challenge in federated learning (FL), where underrepresented classes suffer from poor predictive performance yet cannot be addressed by standard centralized techniques due to privacy and heterogeneity constraints. We propose FedCGNM (Federated Class-Grouped Normalized Momentum), a client-side optimizer in FL that partitions classes into a small number of groups based on minimum within-group variance, maintains a momentum per group, normalizes each group momentum to unit length, and uses the summation of the normalized group momentums as an update direction. This design both equalizes gradient magnitude across majority and minority groups and mitigates the noise inherent in rare-class gradients. We further provide a theoretical convergence analysis explicitly accounting for time-varying resampling-rates. Additionally, to efficiently optimize these rates in small-client regimes, we introduce FedHOO, an X-armed-bandit (XAB) based algorithm that exploits federated parallelism that evaluates many combinations of two candidate rates per client at linear cost. Empirical evaluation on four public long-tailed benchmarks and a proprietary chip-defect dataset demonstrates that FedCGNM consistently outperforms baselines, with FedHOO yielding further gains in small-scale federations.

• The paper proposes FedCGNM, which groups classes and applies unit-norm momentum per group to reduce gradient noise and balance update contributions.
• It introduces FedHOO, a federated bandit algorithm that efficiently explores client resampling rates, accelerating convergence in federated settings.
• Empirical results on industrial and public benchmarks show a 16% performance improvement and enhanced stability in non-IID, imbalanced data scenarios.

Class-Grouped Normalized Momentum and Bandit-Based Hyperparameter Exploration for Class Imbalance in Federated Learning

Introduction

This paper addresses the problem of class imbalance in federated learning (FL), where typical centralized solutions such as reweighting, resampling, or generative augmentation are challenging due to privacy constraints and client/data heterogeneity. Existing optimization-based approaches like Per-Class Normalization (PCN) struggle with multi-class settings, exhibiting instability due to noisy gradient directions and scaling mismatches when the number of classes is large. To overcome these issues, the authors propose Federated Class-Grouped Normalized Momentum (FedCGNM)—an optimization algorithm that partitions classes into a small number of groups (typically "majority" and "minority"), applies unit-norm momentum per group, and combines these to yield stable, balanced updates. Furthermore, they introduce FedHOO, a federated X-armed bandit algorithm exploiting FL's parallelism to efficiently search over client resampling rates, enabling rapid and stable hyperparameter identification in small-client regimes.

Methodological Advancements

Federated Class-Grouped Normalized Momentum (FedCGNM)

PCN equalizes per-class gradient contributions by normalizing each class's gradient. However, when $C$ (number of classes) is large, this introduces two severe problems: aggregation of noisy minority gradients and singularities in update norms. FedCGNM groups classes into $H \ll C$ groups based on minimizing intra-group variance of class sample proportions, then maintains and normalizes a separate momentum per group. Individual group momenta are normalized to unit length, and the sum of these serves as the local update direction.

This approach attenuates gradient noise (by aggregating within groups), prevents scale explosion, and ensures each group—regardless of majority/minority label frequency—exerts equal update influence. The grouping rule is derived as a variance minimization over contiguous class blocks sorted by relative frequency, such that selecting group boundaries via thresholds minimizes expected within-group imbalance. Empirically, the variance-based grouping threshold matches the peak in test accuracy, confirming the optimality of the strategy.

Figure 1: Test accuracy on CIFAR-100-LT ($\xi=20$, $K=5$) versus the number of classes assigned to the minority group, with the red line denoting the threshold found by the variance-based grouping rule.

Theoretical Convergence Analysis

The analysis extends traditional non-i.i.d. FL convergence guarantees by directly incorporating dynamic, per-client sampling rates as part of the global objective evolution. Specifically, the stationary-point convergence rate of $O(T^{-1/2})$ is preserved, provided that the cumulative path variation of sampling rates $V_T$ grows sublinearly relative to the square root of training rounds $T$. Excessive oscillation or rapid resampling schedule changes can thus degrade convergence, underscoring the need for stable yet adaptive rate selection.

FedHOO: Federated Hyperparameter Exploration

FedHOO operationalizes a hierarchical optimistic optimization (HOO) strategy tailored for FL, capitalizing on the linearity of update aggregation. Each client runs two local training instances with different resampling rates, then the central server can synthetically produce $2^K$ combinations (where $K$ is the number of clients) during aggregation. By evaluating all combinations in a single communication round, FedHOO rapidly narrows the search interval for optimal per-client rates, balancing exploration and exploitation. This is particularly effective in small-client federations, where the exponential validation overhead is manageable, enabling quick stabilization of rates and improved test accuracy.

Empirical Evaluation

Main Real-World Evaluation: Semiconductor Chip-Defect Detection

On a proprietary chip-defect detection dataset with extreme imbalance (1.7% defect samples), FedCGNM significantly outperformed classical FL optimizers and previous gradient-normalization baselines. The method produced a 16% test performance improvement over the leading baseline, with FedHOO further increasing defect detection accuracy by accelerating convergence to optimal sampling rates.

Figure 2: (a) Sample heterogeneity in chip-defect dataset across factories. (b) Test performance of algorithms on the chip-defect detection task.

Public Benchmarks

Experiments on CIFAR-10/100-LT, Adult Income, and UNSW-NB15 benchmarks under both IID and highly non-IID federated splits confirm consistent outperformance of FedCGNM across all imbalance severities and client configurations. Notably, in the five-client regime, the combination FedCGNM+FedHOO demonstrated the strongest F1 scores, and as the number of clients grew (increasing heterogeneity), FedCGNM’s relative performance penalty was markedly lower than that of competing group-normalized or reweighting techniques.

Sensitivity and Ablation Analyses

The configuration with $H=2$ groups (majority/minority) was empirically optimal, balancing bias-variance and generalization. More granular grouping led to rapid validation loss drift and overfitting, confirming that data-driven grouping is superior to fixed or naive splits. Similarly, moderate group momentum factors ($H \ll C$0 around 0.5) consistently provided the highest accuracy across datasets, with under- or over-regularization leading to subpar generalization.

Figure 3: (Left) F1 scores for different numbers of groups $H \ll C$1. (Right) Training and validation loss curves for CIFAR-100-LT ($H \ll C$2, $H \ll C$3) with varying $H \ll C$4; $H \ll C$5 yields lowest validation loss and best generalization.

Stability and Robustness

FedCGNM exhibits superior stability under heavy heterogeneity, large-client, and partial-participation settings. Qualitative diagnostics of the “PCN” baseline reveal large loss spikes and poor generalization in multi-class settings, while per-class test accuracy histograms under FedCGNM shift mass away from near-zero performance, indicating improved balance and practical utility.

Figure 4: Loss spikes and validation accuracy drops with PCN in multi-class centralized setting, confirming instability and poor generalization.

Implications and Outlook

The combination of optimization-level imbalance remedies, rigorous grouping based on empirical variance, and federated-hyperparameter search forms a comprehensive toolkit for FL scenarios characterized by heterogeneity and severe imbalance. The normalization-by-group approach enhances both update direction alignment and stability, providing theoretical and empirically validated guarantees absent in per-class normalization schemes. Efficient hyperparameter exploration via FedHOO leverages “combinatorial aggregation” properties unique to FL, leading to rapid adaption and improved accuracy without violating privacy or communication constraints.

Practically, this framework supports application in privacy-critical or industrial federations (e.g., distributed defect detection, medical imaging) where global dataset access, synthetic sample generation, or extensive rate sweeps are infeasible. Future research could address scaling hyperparameter search to large $H \ll C$6 regimes, potentially via client clustering or hierarchical exploration, further theoretical analyses of interplay between grouping granularity and non-IID drift, and extension to tasks beyond classification, such as sequence prediction or regression under skew.

Conclusion

FedCGNM, a class-grouped, normalized momentum optimizer, combined with FedHOO for fast, federated hyperparameter exploration, provides a scalable, robust solution to class imbalance in federated learning. Through variance-aware grouping and per-group normalization, it mitigates noise, stabilizes update norms, and enhances alignment. Theoretical convergence guarantees, comprehensive benchmarking, and demonstrated real-world utility on a challenging industrial dataset confirm its efficacy and generalizability. The proposed techniques represent a substantive step toward rigorously principled, practically effective federated training under non-IID, imbalanced regimes.