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

arXiv: https://arxiv.org/abs/2607.01474 · HTML: https://mlsanity.com/paper/2607.01474/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 35 references that are in the map.
- First posted on arXiv 1 Jul 2026.

## Builds on (most cited first, 16 of 16)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [A systematic study of the class imbalance problem in convolutional neural networks](https://mlsanity.com/paper/1710.05381/) (2017)
- [Dynamic Curriculum Learning for Imbalanced Data Classification](https://mlsanity.com/paper/1901.06783/) (2019)
- [Generative Adversarial Minority Oversampling](https://mlsanity.com/paper/1903.09730/) (2019)
- [Gradient Reweighting: Towards Imbalanced Class-Incremental Learning](https://mlsanity.com/paper/2402.18528/) (2024)
- [FedNoRo: Towards Noise-Robust Federated Learning by Addressing Class Imbalance and Label Noise Heterogeneity](https://mlsanity.com/paper/2305.05230/) (2023)
- [Implicit Gradient Alignment in Distributed and Federated Learning](https://mlsanity.com/paper/2106.13897/) (2021)
- [A Theoretical Analysis of the Learning Dynamics under Class Imbalance](https://mlsanity.com/paper/2207.00391/) (2022)
- [Fed-GraB: Federated Long-tailed Learning with Self-Adjusting Gradient Balancer](https://mlsanity.com/paper/2310.07587/) (2023)
- [A Survey on Class Imbalance in Federated Learning](https://mlsanity.com/paper/2303.11673/) (2023)
- [Adaptive Gradient Normalization and Independent Sampling for (Stochastic) Generalized-Smooth Optimization](https://mlsanity.com/paper/2410.14054/) (2024)

## Abstract (arXiv)

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.