math.OCMay 18, 2026

Scale-Invariant Neural Network Optimization: Norm Geometry and Heavy-Tailed Noise

Authors: Jiayu ZhangTianyi Lin

Organizations: Department of Industrial Engineering and Operations Research Columbia University

Abstract

A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized. The layerwise input-output structure of neural networks motivates scale-invariant optimizers, such as Muon and Scion, whose updates also support hyperparameter transfer. At the same time, stochastic gradient noise in deep learning is often far from sub-Gaussian and may exhibit heavy tails. These observations have shaped recent algorithmic principles for training neural networks, yet their joint theoretical consequences are underexplored. In particular, it remains unclear what dimension dependence is unavoidable for gradient-based methods given the problem class is defined by input-output norm and under heavy-tailed noise, and whether higher-order smoothness can accelerate training. We study these questions through nonconvex smooth stochastic optimization over Rm×n\mathbb R^{m\times n} equipped with general norms and under pthp^\mathrm{th}-moment heavy-tailed noise, where the goal is to achieve an εε-stationary point in the dual norm. Our first contribution is a dimension-dependent lower bound: when max{m,n}(min{m,n})2\frac{\max\{m,n\}}{(\min\{m,n\})^2} is large enough, any gradient-based method requires Ω(min{m,n}ε3p2p1)Ω(\min\{m, n\}ε^{-\frac{3p-2}{p-1}}) oracles for the problem class defined by the spectral norm, which is a common input-output norm. We prove that a scale-invariant Scion method with the spectral norm can achieve the matching upper bound of O(min{m,n}ε3p2p1)O(\min\{m, n\}ε^{-\frac{3p-2}{p-1}}). To exploit higher-order smoothness, we propose a transported Scion method and improve the bound to O(min{m,n}ε5p32p2)O(\min\{m, n\}ε^{-\frac{5p-3}{2p-2}}) when the Hessian is Lipschitz. Finally, we incorporate heuristics into our transported method and evaluate it across multiple architectures and model sizes, demonstrating its flexibility and compatibility with neural network training.

Explore similar work

Jun 12, 2026math.OC

Free Heavy-Tailed Lunch for Muon: A Theoretical Justification of Empirical Success

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 pp-th central moments, p(1,2]p \in (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×nm \times n matrices, Muon finds an ε\varepsilon-stationary point in nuclear norm within O(min{m,n}Δ1Lε2(σε)pp1)\mathcal{O}\left(\min\{m, n\} \frac{Δ_1 L}{\varepsilon^2} \left(\frac σ\varepsilon \right)^{\frac p {p-1}}\right) 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
Jun 13, 2026cs.LG

When to use what Schatten-p norm in deep learning?

Schatten-\infty based optimizers such as Muon have shown promising empirical performance, but there remains seemingly conflicting observations regarding whether they are beneficial. We resolve this conflict by showing that the conclusion is regime dependent. Even when the objective is smooth in the Schatten-\infty geometry, smaller Schatten-pp geometries can be optimal, specifically in the low-dimensional regime, which we show includes Chinchilla scaling. This conclusion follows from a new noise-robust acceleration result for the SODA framework for p>2p>2. The same analysis explains why Muon-like methods do not require warmup, why they naturally favor large batches, and yields a batch size scaling rule for arbitrary pp.
Thomas Pethick
Sep 8, 2026cs.LG

When Does Scale-Invariant Optimization Become Unstable? An Exact Schedule Law with Weight Decay

Normalization renders large parts of neural networks effectively scale invariant, inducing a hidden feedback loop in which learning-rate schedules and weight decay interact through the parameter norm to control the effective step taken by the optimizer. We show that this interaction is governed by an exact discrete-time law: a single scalar quantity captures all schedule and decay forcing, while norm growth induces an opposing geometric self-quenching effect. This yields a sharp boundary that cleanly separates contraction- and expansion-dominated effective learning rate regimes. To understand the underlying mechanism, we provide exact analysis of a fully solved normalized regression model where the dynamics reduce to two dimensions and show that the balance point is intrinsically unstable, implying that constant learning rate with weight decay cannot stably maintain an interior equilibrium and instead produces recurrent behavior driven by discrete-time Jacobian structure. We further extend this perspective across optimizers through unified homogeneous-optimizer framework that reveals a structural dichotomy in self-quenching strength, providing a first-principles explanation for why adaptive methods exhibit systematically weaker stabilization under normalization. Across dynamical systems and neural networks (MLP, CNN, GPT2 / MNIST, CIFAR, wikiText, OpenWebText), the predicted law holds with high precision and enables direct control of training via the identified scalar, with performance peaking sharply at the predicted boundary. Together, these results isolate a single governing quantity for scale-invariant optimization, providing a precise and actionable lens on training dynamics, optimizer behavior, and schedule design in modern deep learning. Code is available in https://github.com/shasanamin/normalized-optimization-dynamics.
Hasan Amin, Wei-Kai Chang, Rajiv Khanna