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

TL;DR

This paper addresses the inefficiency and brittleness of fixed batch-size schedules by introducing geometry-aware gradient-noise scales (GNS) for non-Euclidean optimizers. It derives non-Euclidean GNS metrics for signSGD/Signum (using the $\ell_1$ dual to $\ell_\infty$) and for specSGD/Muon (using the nuclear norm dual to the spectral norm), and formulates CBS-like batch-size rules based on these metrics with a tunable parameter $\theta$. A scalable variance estimator that exploits local mini-batch gradients across data-parallel ranks enables real-time estimation of these non-Euclidean GNS metrics during distributed training. Empirical results on language (Llama 3) and vision (Imagewoof, CIFAR-10) show that adaptive batching guided by non-Euclidean GNS can match baseline validation losses while reducing training steps by up to around 66%, demonstrating practical efficiency gains and scalability. The work offers a principled diagnostic and scheduling tool for modern training stacks, with potential extensions to preconditioned/stateful optimizers and more complex Hessian structures.

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.

Figures (3)

• Figure 1: In \ref{['fig:GNS_intro_plots']} (left), we plot the validation perplexity, batch size, and the exponential moving average of the GNS over steps when training the 160M parameter Llama 3 model for 3.2B tokens on C4 data using signSGD. We compare small-batch ($B = 64$), large-batch ($B = 512$), and adaptive batch size strategies based on Euclidean and non-Euclidean GNS. The plots highlight the improved efficiency of using non-Euclidean GNS measurements for signSGD. In \ref{['fig:GNS_geometry']} (right), we depict how enforcing different GNS ($\ell_2$ versus $\ell_1$) changes the geometry of the allowable region for stochastic gradients (with high probability).
• Figure 2: Comparison of constant batch sizes ($B=64$ and $B=256$) and our adaptive batch size method for the 160M Llama 3 model. For both Signum (Fig \ref{['fig:llama160M_signum_validation_perplexity']}) and Muon (Fig \ref{['fig:llama160M_MUON_validation_perplexity']}), our adaptive strategy matches the final perplexity of the smaller batch size baseline while significantly reducing the total steps.
• Figure 3: Validation perplexity, exponential moving average of GNS, batch size and learning rate for 10 seeds for the adaptive strategy starting with an initial batch size of 64 using signSGD for 160M Llama 3 trained for 3.2B tokens over the C4 dataset. The stability of the validation perplexity across multiple seeds underscores the reliability of our proposed method.

Theorems & Definitions (14)

• Lemma 3.1
• Lemma 1.1
• proof
• Lemma 1.2
• proof
• Lemma 1.3
• proof
• Theorem 1.4
• proof
• Theorem 1.5