The Muon optimiser has become a workhorse for large language model pretraining, delivering state-of-the-art efficiency gains over AdamW. Its defining trick is orthogonalising matrix-valued momentum, replacing every nonzero singular value with one while preserving singular directions. But Muon's global convergence guarantees remain fragile. Stochastic Muon with fixed mini-batches can fail to converge even on smooth, strongly convex objectives. This gap between empirical success and theoretical understanding motivates the question: can we embed Muon's spectral normalisation into a stochastic framework with rigorous long-time behaviour?

Makras and Sabanis answer this by connecting Muon to underdamped Langevin dynamics. They introduce a family of non-quadratic kinetic energies that yield a new underdamped Langevin system with momentum preconditioning. The gradient of this kinetic energy acts as a smooth spectral taming of the momentum. Under assumptions that allow nonconvex potentials and superlinear gradient growth, they prove the dynamics leaves the target Gibbs measure invariant, establish exponential convergence to equilibrium in a weighted total variation distance, and show that the Euler-Maruyama discretisation admits moment bounds uniform in time, without modifying the potential gradient.

Why Standard Langevin Fails on Matrix Spaces

Classical underdamped Langevin dynamics uses a quadratic kinetic energy, giving a position update proportional to momentum. When the potential is neither convex nor globally gradient-Lipschitz, the explicit Euler-Maruyama discretisation becomes unstable. Moments diverge as the stepsize shrinks on a fixed time interval, a phenomenon documented by Hutzenthaler, Jentzen, and Kloeden. Tamed Euler methods control this by rescaling large gradient components, but applying taming directly to the force in overdamped Langevin alters the invariant distribution and leads to overtamed algorithms.

The authors' insight is to tame the momentum entering the position equation instead, while retaining the original gradient force in the momentum equation. This preserves the Gibbs position marginal exactly. The taming mechanism derives from the choice of kinetic energy, not from an ad hoc modification of the dynamics.

The Relativistic Kinetic Energy and Its Taming Map

For a taming parameter δ > 0, define the generalised kinetic energy on matrix space ℝⁿˣᵐ:

K_δ(p) = (1/δ) Tr[(I_n + δ p p^⊤)^(1/2) - I_n]

This is a singular-value generalisation of the pseudo-Huber potential. If p = U Σ V^⊤ is a singular value decomposition with Σ = diag(σ₁(p), …, σ_d(p)) where d = min(n, m), then

K_δ(p) = ∑ᵢ k_δ(σᵢ(p)),  k_δ(s) = (√(1 + δ s²) - 1)/δ

The gradient of this kinetic energy is the singular-value taming map:

T_δ(p) = ∇_p K_δ(p) = (I_n + δ p p^⊤)^(-1/2) p = U diag(σᵢ(p)/√(1 + δ σᵢ(p)²)) V^⊤

The scalar map s ↦ s/√(1 + δ s²) is approximately linear near zero and saturates smoothly at δ^(-1/2). It coincides with the relativistic momentum-to-velocity relation for unit rest mass and light speed δ^(-1/2), hence the term "relativistic kinetic energy." The same smooth spectral potential appeared in the deterministic Hamiltonian formulation of regularised Muon by Mustafi et al., but here it is used to construct stochastic dynamics with a prescribed position marginal.

To preserve linear friction, the authors define M_δ(p) = (I_n + δ p p^⊤)^(1/2) and set the momentum diffusion coefficient to γ M_δ(p). The identity M_δ(p) T_δ(p) = p then ensures the momentum equation has linear damping −γ p plus a divergence correction Γ_δ(p). Using the SDE-to-invariant-measure recipe of Ma, Chen, and Fox, the resulting underdamped Langevin system is:

dX_t = T_δ(P_t) dt
dP_t = [-∇U(X_t) - γ P_t + (γ/β) Γ_δ(P_t)] dt + √(2γ/β) M_δ(P_t)^(1/2) dW_t

where W_t is an n×m matrix of independent standard Brownian motions. The state-dependent noise and its divergence correction are chosen jointly with the linear damping. This diffusion admits the Gibbs invariant law associated with the separable Hamiltonian H_δ(x, p) = U(x) + K_δ(p):

μ_{β,δ}(dx, dp) = Z_{β,δ}^(-1) e^(-β[U(x) + K_δ(p)]) dx dp = π_β^x(dx) π_{β,δ}^p(dp)

Singular-value taming changes the equilibrium momentum distribution while preserving the intended position marginal exactly. Taking δ → 0 recovers the classical underdamped Langevin system with quadratic kinetic energy.

From Continuous Dynamics to the Discrete Algorithm

The Euler-Maruyama discretisation with stepsize λ > 0 gives the explicit scheme:

X_{k+1} = X_k + λ T_δ(P_k)
P_{k+1} = (1 - γλ) P_k - λ ∇U(X_k) + (λγ/β) Γ_δ(P_k) + √(2λγ/β) M_δ(P_k)^(1/2) ξ_{k+1}

where ξ_k are i.i.d. standard Gaussian matrices. For fixed δ, γ, λ, the stochastic increment and divergence correction vanish as β → ∞. Setting m_k = -γ P_k, β₁ = 1 - γλ, η = λ/γ with 0 < γλ < 1, the limiting iterative scheme becomes:

m_{k+1} = β₁ m_k + (1 - β₁) ∇U(X_k)
X_{k+1} = X_k - η T_{δ/γ²}(m_k)

Defining the ε-regularised orthogonalisation map Orth_ε(p) = U diag(σᵢ(p)/√(ε² + σᵢ(p)²)) V^⊤, one finds T_{δ/γ²}(p) = (γ/√δ) Orth_{γ/√δ}(p). This identifies the scheme as regularised Muon with regularisation ε = γ/√δ and effective learning rate λ/√δ. The zero-temperature limit of the sampling algorithm recovers Muon's spectral normalisation.

Assumptions: Beyond Convexity and Gradient-Lipschitz

The analysis works under three assumptions that are strictly weaker than convexity plus global gradient-Lipschitz:

  • Assumption 1. U ∈ C^∞(ℝⁿˣᵐ; ℝ).
  • Assumption 2. The Hessian satisfies an L₀/L₁ growth bound: ||∇²U(x)||_{F→F} ≤ L₀ + L₁ |∇U(x)| for all x, with L₀, L₁ ≥ 0. This allows polynomially Lipschitz gradients.
  • Assumption 3. The gradient is 2-dissipative: ⟨x, ∇U(x)⟩ ≥ w₁ |x|² - w₂ for all x, with w₁ > 0, w₂ ≥ 0. This implies coercivity and |∇U(x)| → ∞ as |x| → ∞.

These assumptions cover semiconvex potentials with polynomial gradient growth and an L₀/L₁ Hessian bound. The authors note this class includes objective functions of LLMs with smooth activation functions, citing Lytras et al. (2026) for further discussion.

Exponential Convergence in Weighted Total Variation

The first main result establishes exponential convergence of the continuous diffusion to equilibrium. Using Harris' theorem, the proof combines a geometric drift condition with a minorisation argument. The Lyapunov function is constructed as Ψ(x, p) = exp(α U(x) + a β K_δ(p) + b ⟨x, T_δ(p)⟩ + c |p|²) for suitable constants α, a, b, c. The geometric drift condition requires RΨ ≤ ϱ Ψ + C_𝒱 for some ϱ ∈ (0, 1). The minorisation condition is verified using Hörmander's condition, controllability, and positive recurrence arguments on a compact sublevel set.

The theorem states: under the standing assumptions and admissible parameter choices, the diffusion admits the unique invariant probability measure μ_{β,δ}. There exist a Lyapunov weight V ≥ 1 and constants C, c > 0 such that for every initial law μ₀ with μ₀(V) < ∞,

||ℒ(X_t, P_t) - μ_{β,δ}||_{TV,V} ≤ C e^{-c t} μ₀(V),  t ≥ 0

where ||·||_{TV,V} denotes weighted total variation distance. The corollary gives uniform-in-time polynomial moment bounds for the continuous process.

Uniform Moment Bounds for the Explicit Discretisation

The second main result proves stability of the discrete scheme. For fixed δ ∈ (0, 1] and λ ∈ (0, min{δ, γ⁻¹}], with hyperparameters chosen so that certain coefficients A_y, A_p + A₃ are negative, the discrete Lyapunov function 𝒱 satisfies a one-step contraction outside a compact set. Iterating this inequality bounds 𝔼[𝒱_n] uniformly in n. The coercivity of U controls polynomial powers of |x|, while the L₀/L₁ condition bounds the growth of ∇U(x) exponentially in |x|, so every power of |∇U(x)| is dominated by e^{η U(x)}. The linear coercivity of K_δ and the identity p = y - γ⁻¹ ∇U(x) then show 𝒱(x, p) dominates |x|^r + |p|^r for every r > 0.

The result: for every r > 0 there exists C_r < ∞, independent of the number of iterations and initial conditions, such that

sup_{n≥0} 𝔼[|X_n|^r + |P_n|^r] ≤ C_r (1 + 𝔼[𝒱(X₀, P₀)]) < ∞

This holds for every deterministic initial state. Crucially, the moment bounds are uniform in time without any modification of the potential gradient. The explicit scheme remains stable despite superlinear gradient growth because the taming is built into the kinetic energy, not applied as a post-hoc correction to the force.

Properties of the Taming Map

The taming map T_δ(p) = (I_n + δ p p^⊤)^(-1/2) p inherits strong analytical properties from the scalar function τ_δ(s) = s/√(1 + δ s²). The map is 1-Lipschitz in Frobenius norm: |T_δ(A) - T_δ(B)| ≤ |A - B| for all A, B ∈ ℝⁿˣᵐ. Its gradient satisfies |∇T_δ| ≤ 1. The map is co-coercive (since K_δ is convex): ⟨A - B, T_δ(A) - T_δ(B)⟩ ≥ |T_δ(A) - T_δ(B)|². The Frobenius norm of T_δ(p) is bounded by √(d/δ), where d = min(n, m). The Laplacian Δ_p K_δ(p) is bounded by m n. These properties are essential for the Lyapunov analysis and for controlling the divergence correction Γ_δ(p).

Limitations and Trade-offs

The theory requires C^∞ smoothness of the potential, which excludes non-smooth regularisation terms common in practice (e.g., L₁ penalties). The L₀/L₁ Hessian growth bound, while weaker than global Lipschitz, still imposes a structural constraint on how the Hessian can grow relative to the gradient. The 2-dissipativity condition requires the potential to grow at least quadratically at infinity, ruling out flat or slowly growing tails. The exponential convergence rate depends on the Lyapunov weight V, which can be large; the constants C, c are not explicitly quantified. The discrete scheme requires a sufficiently small stepsize λ, and the hyperparameter conditions (A_y < 0, A_p + A₃ < 0) restrict the admissible parameter regime. The connection to Muon only holds in the zero-temperature limit β → ∞, which is a conceptual bridge rather than a practical recipe for finite-temperature sampling.

What This Means for Practitioners

For developers building LLM training systems, this work provides a rigorous foundation for Muon's spectral normalisation. The relativistic kinetic energy gives a differentiable, tunable alternative to hard orthogonalisation: the taming parameter δ controls the trade-off between directional preservation and norm saturation. Small δ approximates classical momentum; large δ approaches Muon's orthogonalisation. The Euler-Maruyama scheme with this kinetic energy yields a sampling algorithm that remains stable even when the loss landscape is nonconvex and gradients grow superlinearly, a common scenario in deep learning.

The zero-temperature connection suggests a principled annealing schedule: start with moderate β for exploration, then increase β to concentrate around minimisers, recovering regularised Muon as β → ∞. The uniform moment bounds mean the explicit discretisation won't suffer from exploding moments as training progresses, a known issue with naive Langevin discretisations on non-Lipschitz potentials.

The paper also clarifies why taming the position update (via kinetic energy) is fundamentally different from taming the gradient directly. The former preserves the target distribution; the latter distorts it. This distinction matters for any practitioner implementing tempered or annealed sampling in high-dimensional matrix spaces.

Open questions remain: how to choose δ and γ adaptively, whether the theory extends to stochastic gradients with mini-batch noise, and how the spectral taming interacts with the Newton-Schulz iterations used in practical Muon implementations. The framework here provides the mathematical tools to investigate these questions rigorously.