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

Technical Summary: FedCGNM and FedHOO for Class Imbalance in Federated Learning

Problem Statement
Class imbalance remains a critical challenge in Federated Learning (FL), particularly when the global label distribution is long-tailed. In such scenarios, minority classes suffer from poor predictive performance due to underrepresentation in training data. Standard centralized techniques, such as loss reweighting, advanced sampling, or generative augmentation, are difficult to deploy in FL. Privacy constraints prevent data exchange, and synthetic generators often fail to produce realistic samples for sparse regimes or domain-specific tasks (e.g., semiconductor defect detection). Furthermore, existing optimization-level approaches like Per-Class Normalization (PCN) face limitations in multi-class settings: they are susceptible to directional noise from minority classes and introduce scaling mismatches when summing many unit vectors, which destabilizes convergence as the number of classes increases.

Methodology
The paper proposes a two-pronged approach to address these issues: a client-side optimizer named FedCGNM (Federated Class-Grouped Normalized Momentum) and an auxiliary hyperparameter tuner named FedHOO.

1. FedCGNM (Optimizer):

   • Class Grouping: Instead of normalizing gradients for every individual class ($C$), FedCGNM partitions classes into a small number of groups ($H \ll C$, typically majority vs. minority). The grouping rule is derived by minimizing the within-group variance of class proportions, a strategy shown to be a reliable proxy for exhaustive threshold search.
   • Group-Wise Normalized Momentum: For each group, the algorithm maintains a momentum vector. During local training, stochastic gradients are computed per group, and the momentum is updated. Crucially, each group's momentum vector is normalized to unit length (with a stabilizing constant $\delta$) before summation.
   • Update Direction: The final update direction is the sum of these normalized group momentums. This design equalizes the gradient magnitude across majority and minority groups while mitigating the directional noise inherent in rare-class gradients. By operating on a handful of groups rather than individual classes, it avoids the scaling mismatch and instability associated with PCN.
   • Client Alignment: By forcing each group update to have a unit norm, the aggregated update directions across heterogeneous clients become more aligned, addressing a primary cause of performance degradation in FL.

2. FedHOO (Hyperparameter Tuner):

   • Context: Selecting appropriate resampling rates is nontrivial in FL due to the combinatorial search space across clients and the risk of destabilizing optimization with fluctuating schedules.
   • Mechanism: FedHOO is an X-armed bandit (XAB) based algorithm designed for small-scale federations ($K$ is small). It exploits the linear nature of FL aggregation. In each communication round, the server instructs each client to train with two candidate resampling rates.
   • Efficiency: By linearly combining the returned client updates, the server can infer the validation performance of all $2^K$ possible combinations of local rates without requiring $2^K$ separate training rounds. This allows for rapid identification of effective sampling rates early in training while maintaining privacy (clients only return scalar validation metrics).

Theoretical Analysis
The paper provides a convergence analysis for FedCGNM that explicitly accounts for time-varying resampling rates. Under standard assumptions of smoothness, bounded gradients, and unbiasedness, the authors prove that FedCGNM attains a standard $O(T^{-1/2})$ stationary-point rate. The analysis highlights that for convergence to hold, the cumulative variation of the sampling rates across rounds ($V_T$) must satisfy specific conditions, implying that aggressive but unstable resampling can hinder convergence.

Key Contributions

1. FedCGNM Optimizer: The first client-side optimizer that groups classes and applies unit-norm momentum per group. It balances majority/minority influence, reduces noise, and mitigates scale mismatch.
2. Variance-Aware Grouping: A principled, lightweight strategy that frames class partitioning as the minimization of within-group variance, producing optimal splits comparable to exhaustive search.
3. Convergence Analysis: A novel theoretical framework combining FL, group normalization, and dynamic sampling rates, demonstrating the impact of rate schedules on convergence.
4. FedHOO: The first optimization-based algorithm to determine combinatorial local sampling rates in FL using an X-armed bandit framework that leverages federated parallelism for efficient exploration.

Experimental Results
The methods were evaluated on four public long-tailed benchmarks (CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15) and a proprietary semiconductor chip-defect dataset (CDD).

• Performance: FedCGNM consistently outperformed baselines (including FedAvg, FedProx, CLIMB, FedGraB, and PCN variants) across all settings. In small-client regimes ($K=5$), combining FedCGNM with FedHOO yielded the best results, with improvements up to 29% over FedAvg with resampling.
• Industrial Application: On the CDD dataset (1.7% defect rate), the proposed method achieved a 16% improvement over the best baseline, demonstrating effectiveness in real-world, privacy-sensitive industrial settings.
• Robustness: The approach showed resilience in large-client regimes with partial participation and under varying degrees of data heterogeneity (Non-IID).
• Ablation: Experiments confirmed that two groups ($H=2$) offer the best bias-variance trade-off and that moderate momentum factors ($\beta \approx 0.5$) yield optimal performance.

Significance and Claims
The paper claims that FedCGNM and FedHOO form a practical framework for mitigating global class imbalance in privacy-preserving federated settings without relying on data exchange or synthetic generation. The significance lies in addressing class imbalance at the optimization level, which is more generally applicable than data-level tricks restricted by privacy or dataset characteristics. The authors emphasize that their approach is particularly effective in small-scale federations (common in industrial deployments) where joint hyperparameter tuning is feasible and critical, while the grouping strategy ensures stability and alignment in multi-class scenarios where previous normalization techniques fail. The work advances the field by providing both a robust optimizer and a theoretically grounded, efficient method for tuning resampling strategies in the presence of class imbalance.