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

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.

Many real-world machine learning datasets are highly imbalanced: some types of examples appear often, while rare but important cases appear only a few times. This is especially difficult in federated learning, where multiple clients train a shared model without directly sharing their private data. For example, in industrial defect detection, normal images are common but defect images are rare, even though detecting them is crucial. This paper proposes FedCGNM, a training method that helps federated models learn more fairly from both common and rare classes. Instead of treating every class separately, FedCGNM groups classes with similar frequencies and balances the training signal from each group. This reduces the dominance of frequent classes while making rare-class updates less noisy. We also introduce FedHOO, an optional tool for small federated systems that efficiently selects how strongly each client should rebalance its local data. Experiments on public benchmarks and a real-world chip-defect dataset show that our methods improve performance over existing federated learning baselines. These results suggest that the proposed approach can make federated learning more reliable for imbalanced classification problems in privacy-sensitive and industrial settings.