Authors: Igor Sokolov, Alexander Tyurin, Peter Richtárik
Organizations: Center of Excellence for Generative AI, King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia. · AXXX, Moscow, Russia. · Applied AI Institute, Moscow, Russia.
The empirical success of matrix-aware optimizers, such as Muon (Jordan et al., 2024) and Shampoo (arXiv:1802.09568), highlights the potential of non-Euclidean geometry for deep learning, including large language model (LLM) training. Despite this progress, the theoretical understanding of stochastic optimization in general normed spaces remains limited, particularly in distributed settings. We analyze minibatch stochastic gradient descent (SGD) for smooth nonconvex optimization with respect to an arbitrary norm on a finite-dimensional space. Using the framework of κ-regular normed spaces (arXiv:0809.0813), we derive high-probability oracle complexity bounds under a light-tail condition stated directly in the dual norm induced by the update norm. The bounds make the dependence on geometry explicit through the regularity constant κ of the dual space. In the Euclidean case, our bound recovers the standard minibatch SGD rate up to logarithmic factors. For some non-Euclidean geometries and noise models, it is sharper than bounds obtained by converting a Euclidean noise assumption, reducing a linear dimension dependence to a logarithmic one. Building on this framework, we introduce Rennala NSGD, a semi-asynchronous distributed algorithm, and establish wall-clock time complexity guarantees under fixed, heterogeneous worker computation times. In the Euclidean case (κ=1), our bound matches the parameter dependence of classical Rennala SGD up to logarithmic factors.
Figures & tables
Space (E∗\mathchar59∥⋅∥∗)
Euclidean/Frobenius a
Geometry-native, Rademacher b
Geometry-native, Gaussian b
(Rd\mathchar59∥⋅∥1)
O(\nicefracd2LΔε2)
O(\nicefracd3LΔε2)
O(\nicefracd3LΔε2)
(Rd\mathchar59∥⋅∥2)
O(\nicefracdLΔε2)
O(\nicefracdLΔε2)
O(\nicefracdLΔε2)
(Rd\mathchar59∥⋅∥∞)
O(\nicefracdLΔε2)
O(\nicefracln(ed)LΔε2)
O(\nicefrac(ln(ed))2LΔε2)
Table 1: Statistical oracle-cost terms, omitting the factor ν2 and the common leading term O(LΔ/ε) . Formulas hold for d≥1 . Within each row, green and red mark smaller and larger asymptotic orders; blue marks agreement. Colours do not imply strict improvement in every finite dimension.
Figure 2
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Symbol
Meaning
Reference / scope
Spaces and norms
E , E∗
Parameter space and its dual; 1≤dim(E)<∞ . Gradients belong to E∗ .
Appendix A
E∗∗
Bidual space, canonically identified with E in finite dimensions.
Theorem 5
x , s
Generic primal vector x∈E and dual vector s∈E∗ .
Appendix A
⟨s\mathchar59x⟩
Dual pairing: the value of the linear functional s at x .
Appendix A
∥⋅∥
Primal norm, called the update norm.
Section 3.1
Appendix
Table 2: Notation reference. Definitions with a stated local scope do not replace the conventions used elsewhere.
Space (E∗\mathchar59∥⋅∥∗)
κ scaling b
Auxiliary norm ∥⋅∥∗\mathchar59+
Upper bound
Lower bound
Any d -dimensional normed space, d≥1 a
Θ(d)
Inner-product norm whose unit ball is the maximum-volume ellipsoid contained in B∗={s∈E∗:∥s∥∗≤1} (the John ellipsoid).
Prop. 10
Prop. 11
(Rd\mathchar59∥⋅∥p) , 1≤p<2
Θ(dp2−1)
dp1−21∥⋅∥2 (scaled Euclidean norm)
Prop. 12
Prop. 13
(Rd\mathchar59∥⋅∥p) , 2≤p≤∞
Θ(min{p\mathchar59ln(ed)})
∥⋅∥ρ∗ , where ρ∗=2 for 1≤d<8 , and ρ∗=min{p\mathchar59ρ^(d)} for d≥8 . c
Recently, several instances of non-Euclidean SGD, including SignSGD, Lion, and Muon, have attracted significant interest from the optimization community due to their practical success in training deep neural networks. Consequently, a number of works have attempted to explain this success by developing theoretical convergence analyses. Unfortunately, these results cannot properly justify the superior performance of these methods, as they could not beat the convergence rate of vanilla Euclidean SGD. We resolve this important open problem by developing a new unified convergence analysis under the structured smoothness and gradient noise assumption. In particular, our results indicate that non-Euclidean SGD (i) can exploit the sparsity or low-rank structure of the upper bounds on the Hessian and gradient noise, (ii) can provably benefit from popular algorithmic tools such as extrapolation or momentum variance reduction, and (iii) can match the state-of-the-art convergence rates of adaptive and more complex optimization algorithms such as AdaGrad and Shampoo.
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 (ℓ∞) and stochastic spectral descent (specSGD) / Muon (S∞). 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.
Hiroki Naganuma, Shagun Gupta, Youssef Briki +4
Mila, Montreal, Canada · Université de Montréal, Montreal, Canada · Meta Platforms, Menlo Park, California, USA
Non-Euclidean optimisation methods with matrix-valued updates, such as Muon and Scion, have recently shown strong empirical performance for training Transformer models, yet their theoretical advantages over Euclidean methods remain poorly understood. We address this gap in the heavy-tailed non-convex regime, where stochastic gradients have bounded p-th central moments, p∈(1,2]. We show that certain non-Euclidean methods achieve optimal sample complexity under stronger stationarity measures, while Euclidean methods incur additional dimension-dependent costs. As a consequence, for m×n matrices, Muon finds an ε-stationary point in nuclear norm within O(min{m,n}ε2Δ1L(εσ)p−1p) samples, absorbing heavy-tailed noise without extra dimension dependence, unlike Euclidean methods. We further prove this sample complexity, including its dimension dependence, is optimal for all first-order methods under nuclear-norm stationarity. Experiments on large language models support our theory. Surprisingly, our results suggest that other Schatten geometries beyond the spectral geometry of Muon can perform competitively in certain settings.
Florian Hübler, Thomas Pethick, Suvrit Sra
Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, Technical University of Munich, Germany · Munich Center for Machine Learning (MCML)