Tight Complexity of Muon under Heavy-Tailed Noise

arXiv:2605.18528 · math.OC, cs.LG · Submitted 2026-05-18 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Tight Complexity of Muon under Heavy-Tailed Noise".

Jane: The gist A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: The title itself, "Tight Complexity of Muon under Heavy-Tailed Noise," really sets the stage for what they're doing here. It sounds very specific and technical, which is exactly where this research lives.

Jane: And it’s about taking an existing optimizer, Muon, and testing it under a kind of noise that we actually see in real deep learning training.

Lu: What they are really investigating is how the problem class defined by input-output norms changes when you introduce these heavy-tailed stochastic gradient noises.

Meng: So we're looking at how the dimension of the model, with m and n, affects how many times you need to call an oracle to find a stable point in this setup.

Lalam: It sounds like they are connecting the geometry of the matrices—the input-output norm—directly to the required computational effort for optimization.

Tom: Exactly. They’re building on this idea that optimizer design needs to respect how the model is parametrized, which is a big lesson in itself.

Jane: And they're showing that even with these complex noise conditions, there’s a predictable dimension dependence we have to deal with when using gradient-based methods.

The paper's summary: Tom: Let’s look at the core of "Tight Complexity of Muon under Heavy-Tailed Noise." They summarize their contributions by focusing on two central questions they want to answer.

Jane: The first question is about what dimension dependence is unavoidable in stochastic nonconvex matrix optimization when you move beyond just Frobenius geometry.

Lu: They establish a sharp dimension dependence for stochastic nonconvex matrix optimization specifically when the problem is defined by the spectral norm and under p th-moment heavy-tailed noise.

Meng: For any gradient-based method trying to find an epsilon-stationary point, they show you need ((m, n) epsilon-3p - 2p - one) oracle calls if the dimension product condition is met.

Tom: And then they show that a batched Scion method can actually achieve that matching upper bound of O((m, n) epsilon-3p - 2p - one).

Jane: That means their method matches the theoretical lower bound exactly in this specific setup, which is quite strong.

Lalam: It’s a concrete result showing that we can predict the necessary complexity based on the model's dimensions and the noise level.

The paper's improvements: Tom: Beyond just matching bounds, they propose an improvement for Scion methods by introducing something called a transported Scion method.

Jane: This transportation technique is designed to leverage higher-order smoothness when the Hessian matrix is Lipschitz.

Lu: When that condition holds, they prove an improved bound of O((m, n) epsilon-5p - three 2p-two). That's a better rate than what we saw before.

Meng: That acceleration means needing fewer oracle calls to find that epsilon-stationary point, which is a practical win for training deep models.

Tom: And they didn't stop there; they incorporated practical heuristics into this transported method and used it on CNNs and transformers.

Jane: The experiments showed that this transportation technique is compatible with real-world training pipelines, even outperforming baseline methods like NorMuon in some tests.

Conclusion: Tom: So to wrap up the "Tight Complexity of Muon under Heavy-Tailed Noise," we’ve seen that they established a sharp lower bound on oracle calls based on the input-output norm geometry and noise level.

Jane: They also showed a way to speed things up with their transported Scion method when smoothness conditions are met, leading to faster convergence rates.

Lu: This work gives us a solid theoretical foundation for understanding how scale-invariant optimizers behave under the heavy-tailed noise typical of modern training data.

Meng: From an engineering standpoint, knowing these complexity bounds helps us set realistic expectations for how much compute we need before we start training massive models.

Lalam: It suggests that designing optimizers to respect the layerwise structure is not just theoretical; it has tangible consequences for how stable and architecture-aware our updates can be in practice.

Tom: Yeah, so this paper really solidifies the idea that optimizer design should deeply consider how the model is parametrized to handle these heavy-tailed conditions.

Columbia University

math.OC, cs.LG

Submitted: 2026-05-18

Updated: 2026-10-11

Comments: 46 pages, 3 figures

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 91/100

The gist: The gist A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized.

Key concepts

Nonconvex Smooth Stochastic Optimization
This refers to the mathematical problem of finding the best set of parameters (weights) for a neural network when the loss function is not perfectly smooth. It involves using noisy, random samples (stochastic) to estimate gradients, aiming to find a point where improvement stops (stationary point).
Heavy-Tailed Noise
This describes the type of random errors or noise present in the data used during training. Unlike simple Gaussian noise, heavy-tailed noise means that very large errors occur more frequently than expected. This type of noise is common in real-world machine learning scenarios.
Oracle Calls
In optimization, an 'oracle call' represents one evaluation of the gradient (or some related information) with the model. The paper quantifies how many of these evaluations are necessary to reach a desired level of accuracy ($\epsilon$) for solving the optimization problem.

Summary

The gist A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized.

How it works

The study investigates nonconvex smooth stochastic optimization over R m×n equipped with general norms and under p th-moment heavy-tailed noise, aiming to achieve an ϵ-stationary point measured in the dual norm. The first contribution establishes a dimension-dependent lower bound: when max(m,n)(min(m,n))2 is large enough, any gradientbased method requires Ω(min(m, n)ϵ−3p−2p−1) oracle calls for the problem class defined by the spectral norm, which is a common input-output matrix norm. This lower bound matches the upper bound achieved by a scale-invariant batched Scion method with the spectral norm, achieving O(min(m, n)ϵ−3p−2p−1)

Optimization methods and improvements

The paper proposes two main results concerning optimization methods:

  1. A batched Scion method achieves the matching upper bound of O(min(m, n)ϵ−3p−2p−1)

  2. A transported Scion method is proposed to leverage higher-order smoothness and proves an improved bound of O(min(m, n)ϵ−5p−32p−2) when the norm is spectral and the Hessian is Lipschitz

Experimental validation

Experiments on CNNs and transformer models demonstrate the practical utility of these techniques. The results show that transportation techniques are compatible with practical training pipelines, with NorMuon with transportation improving over NorMuon

Sample complexity bounds

The paper derives sample complexity bounds for the algorithms. For Algorithm 1, the required number of stochastic gradient oracles is bounded by O(min(m, n)σp−10ϵ−3p−2p−1) and for Algorithm 2 under higher-order smoothness, the oracle complexity is bounded by O(min(m, n)74σ7p4p−40ϵ−7p4p−4)

This work provides a theoretical foundation for understanding the convergence rates of scale-invariant optimizers under heavy-tailed noise in the context of neural network training.

Improvements for AI systems

  1. Scale-Invariant Optimization for Layerwise Structure: Implement Muon or Scion as primary optimizers to respect layerwise input-output structure of neural networks, leading to updates also support hyperparameter transfer. This allows for more stable and architecture-aware scaling of updates across heterogeneous layers compared to Euclidean methods.

  2. Robustness to Heavy-Tailed Gradients: Utilize the theoretical framework that heavy-tailed stochastic gradient noise necessitates optimization under pth-moment heavy-tailed noise, allowing the system to achieve convergence rates like O(min(m, n)ϵ−3p−2p−1) when defined by spectral and nuclear norms. This provides superior training stability for deep models where gradients are far from sub-Gaussian.

  3. Accelerated Training with Higher-Order Smoothness: Deploy the transported Scion method to leverage higher-order smoothness, improving the convergence bound to O(min(m, n)ϵ−5p−32p−2) when the Hessian is Lipschitz. This acceleration reduces the required number of oracle calls for finding an epsilon-stationary point.

  4. Enhanced Model Performance via Transportation: Integrate a transported Scion method with practical heuristics to achieve superior performance, as demonstrated by its ability to outperform baseline methods in experiments, such as achieving best CIFAR-10 test loss and the best CIFAR-100 test accuracy. This technique demonstrates a flexibility and compatibility with neural network training.

Sources

Related papers