Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning
Pith reviewed 2026-07-03 21:03 UTC · model grok-4.3
The pith
Class grouping by variance plus per-group momentum normalization equalizes gradients across imbalanced classes in federated learning.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
FedCGNM 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. The approach supplies a convergence analysis that accounts for time-varying resampling rates and introduces FedHOO to search those rates efficiently via federated parallelism.
What carries the argument
Class-grouped normalized momentum: classes are partitioned by minimum within-group variance, momentum is tracked and normalized to unit length per group, and the normalized vectors are summed to form the update direction.
If this is right
- Gradient magnitudes become comparable between majority and minority groups.
- Noise from rare-class gradients is reduced by the per-group normalization.
- Convergence holds when resampling rates change over training.
- FedHOO evaluates pairs of resampling rates at linear cost by exploiting client parallelism.
- The method improves accuracy on four public long-tailed datasets and one proprietary defect-detection set.
Where Pith is reading between the lines
- The same grouping-plus-normalization step could be tested in centralized training on imbalanced data to see whether the privacy constraint is necessary for the benefit.
- Dynamic adjustment of the number of groups during training might handle shifting class distributions without extra tuning.
- The bandit search for rates could be combined with other client optimizers to explore additional hyperparameters at low extra cost.
- Extreme heterogeneity across clients might require the variance-based grouping to be recomputed periodically rather than once.
Load-bearing premise
That partitioning classes by minimum within-group variance and then normalizing per-group momentums will reliably equalize effective gradient magnitudes and reduce noise for minority classes under the heterogeneity and privacy constraints of federated learning.
What would settle it
Running FedCGNM on a standard long-tailed federated benchmark and finding that minority-class accuracy stays no better than with ordinary momentum would show the equalization step does not work as claimed.
Figures
read the original 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces FedCGNM, a client-side optimizer for federated learning under class imbalance. It partitions classes into a small number of groups by minimizing within-group variance, maintains a momentum vector per group, normalizes each to unit length, and sums the normalized vectors as the update direction. This is claimed to equalize effective gradient magnitudes across majority and minority classes while mitigating noise from rare classes. The manuscript provides a convergence analysis that explicitly incorporates time-varying resampling rates, proposes FedHOO (an X-armed-bandit algorithm that evaluates candidate resampling-rate pairs at linear cost by exploiting federated parallelism), and reports empirical gains on four public long-tailed benchmarks plus a proprietary chip-defect dataset.
Significance. If the central claims hold, the work supplies a practical client-local mechanism for class-imbalance mitigation that respects FL privacy and heterogeneity constraints, together with an efficient hyper-parameter search procedure that scales with the number of clients. The explicit accounting for time-varying resampling rates in the convergence statement and the linear-cost parallel evaluation in FedHOO are concrete strengths that distinguish the contribution from purely heuristic re-weighting approaches.
major comments (2)
- [§4, Theorem 2] §4 (Convergence Analysis), Theorem 2 and the surrounding discussion of the normalized update: the stated bound accounts for the time-varying resampling rates but supplies no explicit control on the directional error that arises when a high-variance momentum estimate (from a minority-class group) is normalized to unit length and injected at full magnitude into the summed direction. Under client heterogeneity this term is not averaged away, so the noise-mitigation claim central to FedCGNM is not yet supported by the analysis.
- [§3.2] §3.2 (Group Construction): the minimum-within-group-variance partitioning is performed once on the global class statistics; the manuscript does not analyze how sensitive the subsequent normalized-momentum update is to errors in this partitioning when only local client data are available, which directly affects whether the magnitude-equalization property holds in practice.
minor comments (2)
- [Eq. (3)] Notation for the per-group momentum vectors is introduced without an explicit index over groups in Eq. (3); adding a group index would improve readability.
- [Table 2] Table 2 caption does not state whether the reported accuracies are averaged over the same number of random seeds as the main tables; consistency of experimental protocol should be clarified.
Simulated Author's Rebuttal
We thank the referee for the constructive feedback on our manuscript. The comments identify important aspects of the convergence analysis and group construction that merit clarification and strengthening. We address each point below and indicate the revisions we will make.
read point-by-point responses
-
Referee: [§4, Theorem 2] §4 (Convergence Analysis), Theorem 2 and the surrounding discussion of the normalized update: the stated bound accounts for the time-varying resampling rates but supplies no explicit control on the directional error that arises when a high-variance momentum estimate (from a minority-class group) is normalized to unit length and injected at full magnitude into the summed direction. Under client heterogeneity this term is not averaged away, so the noise-mitigation claim central to FedCGNM is not yet supported by the analysis.
Authors: We appreciate the referee's careful reading of the analysis. Theorem 2 already incorporates the time-varying resampling rates into the convergence rate. The unit-length normalization is intended to equalize magnitudes, and the momentum update itself damps variance over iterations. However, we agree that an explicit bound on the directional deviation induced by noisy minority-group normalization (which does not average out under heterogeneity) is not derived. In the revision we will extend the proof of Theorem 2 to include a term that controls the expected cosine deviation of each normalized group momentum, using standard concentration arguments on the momentum estimator. This will make the noise-mitigation property explicit in the bound. revision: yes
-
Referee: [§3.2] §3.2 (Group Construction): the minimum-within-group-variance partitioning is performed once on the global class statistics; the manuscript does not analyze how sensitive the subsequent normalized-momentum update is to errors in this partitioning when only local client data are available, which directly affects whether the magnitude-equalization property holds in practice.
Authors: The partitioning step uses global class-frequency statistics, which can be obtained in a privacy-preserving manner via secure aggregation of per-client class counts. We acknowledge that the manuscript provides no formal sensitivity analysis of how local estimation errors or client-specific class distributions affect the resulting groups and therefore the equalization property. In the revised version we will add a short theoretical sensitivity result in §3.2 (showing that small perturbations in the frequency vector induce bounded changes in the normalized direction) together with an empirical study that perturbs the global statistics and measures the resulting performance degradation on the long-tailed benchmarks. revision: yes
Circularity Check
No circularity detected; derivation self-contained
full rationale
The provided abstract and description introduce FedCGNM (class partitioning by minimum within-group variance, per-group momentum, unit normalization, summation as update) and FedHOO (X-armed bandit for rate optimization) as novel client-side methods, with a convergence analysis that explicitly accounts for time-varying resampling rates. No equations, fitted parameters presented as predictions, self-citations, or uniqueness theorems are quoted that would reduce any central claim to its inputs by construction. The design choices are motivated directly from the problem of class imbalance under FL constraints and validated empirically on benchmarks, making the derivation independent rather than tautological. This matches the default expectation of no significant circularity.
Axiom & Free-Parameter Ledger
Reference graph
Works this paper leans on
-
[1]
Artificial Intelligence Review , volume=
A survey on imbalanced learning: latest research, applications and future directions , author=. Artificial Intelligence Review , volume=. 2024 , publisher=
work page 2024
-
[2]
Survey on deep learning with class imbalance , author=. Journal of Big Data , volume=. 2019 , publisher=
work page 2019
-
[3]
Neural Processing Letters , volume=
Learning from imbalanced data sets with weighted cross-entropy function , author=. Neural Processing Letters , volume=. 2019 , publisher=
work page 2019
-
[4]
Proceedings of the IEEE International Conference on Computer Vision , pages=
Focal loss for dense object detection , author=. Proceedings of the IEEE International Conference on Computer Vision , pages=
-
[5]
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
Class-balanced loss based on effective number of samples , author=. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
-
[6]
Proceedings of the AAAI Conference on Artificial Intelligence , volume=
Addressing class imbalance in federated learning , author=. Proceedings of the AAAI Conference on Artificial Intelligence , volume=
-
[7]
International Conference on Machine Learning , pages=
A theoretical analysis of the learning dynamics under class imbalance , author=. International Conference on Machine Learning , pages=. 2023 , organization=
work page 2023
-
[8]
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
Gradient reweighting: Towards imbalanced class-incremental learning , author=. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
-
[9]
International Conference on Database Systems for Advanced Applications , pages=
FedGR: Federated learning with gravitation regulation for double imbalance distribution , author=. International Conference on Database Systems for Advanced Applications , pages=. 2023 , organization=
work page 2023
-
[11]
Frontiers in Digital Health , volume=
A review on over-sampling techniques in classification of multi-class imbalanced datasets: Insights for medical problems , author=. Frontiers in Digital Health , volume=. 2024 , publisher=
work page 2024
-
[12]
Journal of Artificial Intelligence Research , volume=
SMOTE: synthetic minority over-sampling technique , author=. Journal of Artificial Intelligence Research , volume=
-
[13]
Proceedings of the IEEE/CVF International Conference on Computer Vision , pages=
Generative adversarial minority oversampling , author=. Proceedings of the IEEE/CVF International Conference on Computer Vision , pages=
-
[14]
Engineering Applications of Artificial Intelligence , volume=
An improved generative adversarial network to oversample imbalanced datasets , author=. Engineering Applications of Artificial Intelligence , volume=. 2024 , publisher=
work page 2024
-
[15]
IEEE Transactions on Parallel and Distributed Systems , volume=
Fast: enhancing federated learning through adaptive data sampling and local training , author=. IEEE Transactions on Parallel and Distributed Systems , volume=. 2023 , publisher=
work page 2023
-
[16]
Resampling approaches to handle class imbalance: a review from a data perspective , author=. Journal of Big Data , volume=. 2025 , publisher=
work page 2025
-
[17]
A systematic study of the class imbalance problem in convolutional neural networks , author=. Neural Networks , volume=. 2018 , publisher=
work page 2018
-
[18]
International Journal of Data Science and Analytics , pages=
Leveraging local data sampling strategies to improve federated learning , author=. International Journal of Data Science and Analytics , pages=. 2024 , publisher=
work page 2024
-
[19]
Proceedings of the IEEE/CVF International Conference on Computer Vision , pages=
Dynamic curriculum learning for imbalanced data classification , author=. Proceedings of the IEEE/CVF International Conference on Computer Vision , pages=
- [20]
-
[21]
World Journal of Methodology , volume=
Challenges and limitations of synthetic minority oversampling techniques in machine learning , author=. World Journal of Methodology , volume=
- [22]
-
[23]
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
Large-scale long-tailed recognition in an open world , author=. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages=
-
[24]
Proceedings of the AAAI Conference on Artificial Intelligence , volume=
Implicit gradient alignment in distributed and federated learning , author=. Proceedings of the AAAI Conference on Artificial Intelligence , volume=
-
[25]
UCI Machine Learning Repository , volume=
Adult , author=. UCI Machine Learning Repository , volume=
-
[26]
2015 Military Communications and Information Systems Conference (MilCIS) , pages=
UNSW-NB15: a comprehensive data set for network intrusion detection systems (UNSW-NB15 network data set) , author=. 2015 Military Communications and Information Systems Conference (MilCIS) , pages=. 2015 , organization=
work page 2015
-
[27]
Enhanced network intrusion detection using a hybrid CNN-LSTM approach on the UNSW-NB15 dataset , author=. 2024 21st International Conference on Electrical Engineering, Computing Science and Automatic Control (CCE) , pages=. 2024 , organization=
work page 2024
-
[28]
Finite-time analysis of the multiarmed bandit problem , author=. Machine Learning , volume=. 2002 , publisher=
work page 2002
-
[29]
Artificial Intelligence and Statistics , pages=
Communication-efficient learning of deep networks from decentralized data , author=. Artificial Intelligence and Statistics , pages=. 2017 , organization=
work page 2017
-
[30]
Learning multiple layers of features from tiny images , author=. 2009 , publisher=
work page 2009
-
[31]
Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages=
Deep residual learning for image recognition , author=. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages=
-
[32]
International Conference on Learning Representations , year=
An agnostic approach to federated learning with class imbalance , author=. International Conference on Learning Representations , year=
-
[33]
Advances in Neural Information Processing Systems , volume=
Fed-grab: Federated long-tailed learning with self-adjusting gradient balancer , author=. Advances in Neural Information Processing Systems , volume=
-
[36]
Proceedings of Machine learning and systems , volume=
Federated optimization in heterogeneous networks , author=. Proceedings of Machine learning and systems , volume=
-
[38]
International Conference on Machine Learning , pages=
The non-iid data quagmire of decentralized machine learning , author=. International Conference on Machine Learning , pages=. 2020 , organization=
work page 2020
-
[39]
M., Albalkhi, I., and Naswhan, A
Alkhawaldeh, I. M., Albalkhi, I., and Naswhan, A. J. Challenges and limitations of synthetic minority oversampling techniques in machine learning. World Journal of Methodology, 13 0 (5): 0 373, 2023
work page 2023
-
[40]
Finite-time analysis of the multiarmed bandit problem
Auer, P., Cesa-Bianchi, N., and Fischer, P. Finite-time analysis of the multiarmed bandit problem. Machine Learning, 47 0 (2): 0 235--256, 2002
work page 2002
-
[41]
Aurelio, Y. S., De Almeida, G. M., de Castro, C. L., and Braga, A. P. Learning from imbalanced data sets with weighted cross-entropy function. Neural Processing Letters, 50: 0 1937--1949, 2019
work page 1937
-
[42]
Becker, B. and Kohavi, R. Adult. UCI Machine Learning Repository, 10: 0 C5XW20, 1996
work page 1996
-
[43]
Bubeck, S., Munos, R., Stoltz, G., and Szepesv \'a ri, C. X-armed bandits. Journal of Machine Learning Research, 12 0 (5), 2011
work page 2011
-
[44]
Buda, M., Maki, A., and Mazurowski, M. A. A systematic study of the class imbalance problem in convolutional neural networks. Neural Networks, 106: 0 249--259, 2018
work page 2018
-
[45]
Carvalho, M., Pinho, A. J., and Br \'a s, S. Resampling approaches to handle class imbalance: a review from a data perspective. Journal of Big Data, 12 0 (1): 0 71, 2025
work page 2025
-
[46]
Chawla, N. V., Bowyer, K. W., Hall, L. O., and Kegelmeyer, W. P. Smote: synthetic minority over-sampling technique. Journal of Artificial Intelligence Research, 16: 0 321--357, 2002
work page 2002
-
[47]
Chen, C.-C., Lin, Y.-S., and Chen, H.-Y. Tackling class-imbalanced learning issues based on local neighborhood information and generative adversarial networks. Sensors & Materials, 36, 2024 a
work page 2024
-
[48]
Chen, W., Yang, K., Yu, Z., Shi, Y., and Chen, C. P. A survey on imbalanced learning: latest research, applications and future directions. Artificial Intelligence Review, 57 0 (6): 0 137, 2024 b
work page 2024
-
[49]
Class-balanced loss based on effective number of samples
Cui, Y., Jia, M., Lin, T.-Y., Song, Y., and Belongie, S. Class-balanced loss based on effective number of samples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.\ 9268--9277, 2019
work page 2019
-
[50]
Implicit gradient alignment in distributed and federated learning
Dandi, Y., Barba, L., and Jaggi, M. Implicit gradient alignment in distributed and federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp.\ 6454--6462, 2022
work page 2022
-
[51]
Leveraging local data sampling strategies to improve federated learning
D \"u sing, C., Cimiano, P., and Paa en, B. Leveraging local data sampling strategies to improve federated learning. International Journal of Data Science and Analytics, pp.\ 1--20, 2024
work page 2024
-
[52]
A theoretical analysis of the learning dynamics under class imbalance
Francazi, E., Baity-Jesi, M., and Lucchi, A. A theoretical analysis of the learning dynamics under class imbalance. In International Conference on Machine Learning, pp.\ 10285--10322. PMLR, 2023
work page 2023
-
[53]
Gradient reweighting: Towards imbalanced class-incremental learning
He, J. Gradient reweighting: Towards imbalanced class-incremental learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.\ 16668--16677, 2024
work page 2024
-
[54]
Deep residual learning for image recognition
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp.\ 770--778, 2016
work page 2016
-
[55]
Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification
Hsu, T.-M. H., Qi, H., and Brown, M. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06335, 2019
work page internal anchor Pith review Pith/arXiv arXiv 1909
-
[56]
Johnson, J. M. and Khoshgoftaar, T. M. Survey on deep learning with class imbalance. Journal of Big Data, 6 0 (1): 0 1--54, 2019
work page 2019
-
[57]
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. 2009
work page 2009
-
[58]
K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V
Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2: 0 429--450, 2020
work page 2020
-
[59]
Focal loss for dense object detection
Lin, T.-Y., Goyal, P., Girshick, R., He, K., and Doll \'a r, P. Focal loss for dense object detection. In Proceedings of the IEEE International Conference on Computer Vision, pp.\ 2980--2988, 2017
work page 2017
-
[60]
Liu, Z., Miao, Z., Zhan, X., Wang, J., Gong, B., and Yu, S. X. Large-scale long-tailed recognition in an open world. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.\ 2537--2546, 2019
work page 2019
-
[61]
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pp.\ 1273--1282. PMLR, 2017
work page 2017
-
[62]
Moustafa, N. and Slay, J. Unsw-nb15: a comprehensive data set for network intrusion detection systems (unsw-nb15 network data set). In 2015 Military Communications and Information Systems Conference (MilCIS), pp.\ 1--6. IEEE, 2015
work page 2015
-
[63]
Mullick, S. S., Datta, S., and Das, S. Generative adversarial minority oversampling. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.\ 1695--1704, 2019
work page 2019
-
[64]
An improved generative adversarial network to oversample imbalanced datasets
Pan, T., Pedrycz, W., Yang, J., and Wang, J. An improved generative adversarial network to oversample imbalanced datasets. Engineering Applications of Artificial Intelligence, 132: 0 107934, 2024
work page 2024
-
[65]
Pear, Z. T. and Kibria, H. B. Enhanced network intrusion detection using a hybrid cnn-lstm approach on the unsw-nb15 dataset. In 2024 21st International Conference on Electrical Engineering, Computing Science and Automatic Control (CCE), pp.\ 1--6. IEEE, 2024
work page 2024
-
[66]
An agnostic approach to federated learning with class imbalance
Shen, Z., Cervino, J., Hassani, H., and Ribeiro, A. An agnostic approach to federated learning with class imbalance. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=Xo0lbDt975
work page 2022
-
[67]
Addressing class imbalance in federated learning
Wang, L., Xu, S., Wang, X., and Zhu, Q. Addressing class imbalance in federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp.\ 10165--10173, 2021
work page 2021
-
[68]
Dynamic curriculum learning for imbalanced data classification
Wang, Y., Gan, W., Yang, J., Wu, W., and Yan, J. Dynamic curriculum learning for imbalanced data classification. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.\ 5017--5026, 2019
work page 2019
-
[69]
Fast: enhancing federated learning through adaptive data sampling and local training
Wang, Z., Xu, H., Xu, Y., Jiang, Z., Liu, J., and Chen, S. Fast: enhancing federated learning through adaptive data sampling and local training. IEEE Transactions on Parallel and Distributed Systems, 35 0 (2): 0 221--236, 2023
work page 2023
-
[70]
Wu, N., Yu, L., Jiang, X., Cheng, K.-T., and Yan, Z. Fednoro: Towards noise-robust federated learning by addressing class imbalance and label noise heterogeneity. arXiv preprint arXiv:2305.05230, 2023
-
[71]
T., Wu, J., Yang, H., and Liu, Z
Xiao, Z., Chen, Z., Liu, S., Wang, H., Feng, Y., Hao, J., Zhou, J. T., Wu, J., Yang, H., and Liu, Z. Fed-grab: Federated long-tailed learning with self-adjusting gradient balancer. Advances in Neural Information Processing Systems, 36: 0 77745--77757, 2023
work page 2023
-
[72]
Yang, Y., Tripp, E., Sun, Y., Zou, S., and Zhou, Y. Adaptive gradient normalization and independent sampling for (stochastic) generalized-smooth optimization. arXiv preprint arXiv:2410.14054, 2024
-
[73]
A survey on class imbalance in federated learning
Zhang, J., Li, C., Qi, J., and He, J. A survey on class imbalance in federated learning. arXiv preprint arXiv:2303.11673, 2023
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.