Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent

Abstract

To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, relying on heuristics that are brittle and costly to tune. Existing adaptive strategies based on gradient noise scale (GNS) offer a principled alternative. However, their assumption of SGD's Euclidean geometry creates a fundamental mismatch with popular optimizers based on generalized norms, such as signSGD / Signum ($\ell_\infty$) and stochastic spectral descent (specSGD) / Muon ($\mathcal{S}_\infty$). In this work, we derive gradient noise scales for signSGD and specSGD that naturally emerge from the geometry of their respective dual norms. To practically estimate these non-Euclidean metrics, we propose an efficient variance estimation procedure that leverages the local mini-batch gradients on different ranks in distributed data-parallel systems. Our experiments demonstrate that adaptive batch size strategies using non-Euclidean GNS enable us to match the validation loss of constant-batch baselines while reducing training steps by up to 66\% for Signum and Muon on a 160 million parameter Llama model.

Research summary

Problem

Euclidean gradient-noise-scale rules are geometrically mismatched with Signum and Muon-style spectral optimizers.

Method

The work derives non-Euclidean noise scales from the optimizers' dual norms and estimates them efficiently from per-rank mini-batch gradients in distributed training.

Result

Adaptive batches matched constant-batch validation loss while reducing training steps by up to 66% for Signum and Muon on a 160M-parameter Llama model.

Significance

It makes adaptive batch sizing consistent with the geometry of modern non-Euclidean optimizers and directly usable in distributed systems.

Limitations

The headline result is reported on a 160M-parameter model; scaling behavior beyond the evaluated setup needs further study.

Research links

• arXiv

How to cite

Hiroki Naganuma, Shagun Gupta, Youssef Briki, Ioannis Mitliagkas, Irina Rish, Parameswaran Raman, Hao-Jun Michael Shi. “Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent.” International Conference on Machine Learning, 2026.

@inproceedings{Naganuma2026Adaptive,
  title = {Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent},
  author = {Hiroki Naganuma and Shagun Gupta and Youssef Briki and Ioannis Mitliagkas and Irina Rish and Parameswaran Raman and Hao-Jun Michael Shi},
  year = {2026},
  booktitle = {International Conference on Machine Learning},
  eprint = {2602.03001},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2602.03001}
}