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

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 (Schatten-$\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.

Training massive AI models requires enormous computing power and energy. To speed up this process, systems analyze data in groups called "batches," but finding the optimal batch size is difficult: too small wastes time, while too large degrades the model quality. Existing automatic batch-sizing tools rely on outdated assumptions that clash with the advanced algorithms driving today's AI training. Thus, we developed a technique that continuously measures data "noise"—or uncertainty— tailored for the modern algorithms. Our system then uses this noise signal to automatically fine-tune the batch size on the fly. This adaptive approach trains models just as effectively as highly tuned manual setups but requires up to 66% fewer training steps, making large-scale AI training faster, cheaper, and much more energy-efficient.