Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning
For federated learning practitioners dealing with class imbalance, this work provides a novel client-side optimizer and hyperparameter tuning method that improve performance on long-tailed data while preserving privacy.
FedCGNM addresses class imbalance in federated learning by partitioning classes into groups, maintaining normalized momentum per group, and using their sum as the update direction, which equalizes gradient magnitudes and reduces noise. It also introduces FedHOO, an X-armed-bandit algorithm for efficient hyperparameter exploration in small-client regimes, achieving consistent improvements over baselines on multiple benchmarks.
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.