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

2026-07-01

AI Summary

This work addresses the performance degradation of minority classes in federated learning caused by class imbalance and the infeasibility of centralized resampling. To this end, the authors propose FedCGNM, an optimizer that, for the first time in federated settings, groups classes according to the principle of minimizing intra-group variance and normalizes and fuses momentum per group to balance gradient magnitudes between majority and minority classes. Additionally, they introduce FedHOO, an algorithm based on an X-armed bandit framework that efficiently explores time-varying resampling rate hyperparameters. Theoretical analysis establishes convergence guarantees, and extensive experiments on four public long-tailed datasets and a private chip defect dataset demonstrate significant improvements over existing methods, with FedHOO providing further gains particularly in small-scale federated settings.

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.

Problem

• class imbalance
• federated learning
• privacy constraints
• data heterogeneity
• minority classes

Innovation

• Class-Grouped Normalized Momentum
• Federated Learning
• Class Imbalance
• X-armed Bandit
• Hyperparameter Exploration