Source-linked AI summary
Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials
Nikolaos Makras, Sotirios Sabanis
TL;DR
The paper asks whether a stable Langevin framework can provide theoretical guarantees for spectral-normalised Muon and its long-time behavior. It introduces momentum taming that preserves the intended Gibbs position marginal, and establishes exponential convergence and uniform-in-iteration moment bounds for the discretization.
Problem
The paper asks whether spectral-normalised Muon can be studied in a stable Langevin framework with theoretical guarantees and characterisable long-time behavior.
Source: §1 · pp. 1–2Method
The method incorporates singular-value taming into the momentum-to-position map, with momentum diffusion and a divergence correction chosen to preserve linear friction and the desired Gibbs position marginal.
Source: §1 · p. 2Results
The diffusion has the Gibbs invariant law, converges exponentially, and its explicit discretization has moment bounds uniform in the number of iterations.
Source: §1 · p. 2Takeaways & Limitations
Singular-value taming changes the equilibrium momentum distribution while preserving the intended position marginal exactly.
Source: §1 · p. 2The guarantees are established for smooth semiconvex potentials with polynomially Lipschitz gradients, an L0/L1 Hessian growth bound, and convexity at infinity.
Source: §1.1 · p. 3
Abstract
from arXiv · showhide
We consider the problem of sampling from Gibbs distributions on matrix spaces whose potential energies are neither convex nor globally gradient-Lipschitz. We introduce a family of non-quadratic kinetic energies that lead to a new underdamped Langevin system with momentum preconditioning, in which the gradient of the kinetic energy acts as a smooth spectral taming of the momentum. We prove that, under these relaxed assumptions on the potential, the resulting dynamics leaves the target Gibbs measure invariant, and we establish exponential convergence to equilibrium in a weighted total variation distance. Finally, we show that the corresponding Euler-Maruyama discretization admits moment bounds that are uniform in time, without any modification of the potential gradient, which ensures the stability of the resulting sampling algorithm.
1. Introduction
The introduction asks whether a stable Langevin algorithm can have theoretical guarantees while recovering regularised Muon at zero temperature. It answers by taming momentum in a kinetic Langevin system, preserving the target Gibbs position marginal, and establishing convergence and discretisation stability under nonconvex, superlinear-growth assumptions.
1. Introduction
- The proposed kinetic Langevin system incorporates singular-value taming into the momentum-to-position map, with momentum diffusion and a divergence correction chosen to preserve linear friction and the desired Gibbs position marginal.The construction retains the original gradient force in the momentum equation rather than replacing it with a tamed approximation.Source: §1 · p. 2
1.1. Contributions
- In the zero-temperature limit, the discrete scheme recovers regularised Muon, connecting the sampling construction to the optimiser’s spectral update.The introduction poses this recovery alongside the goal of obtaining a stable algorithm with theoretical guarantees.Source: §1 · p. 2, §1.1 · p. 3
- The diffusion has a Gibbs invariant law whose position marginal is exactly the prescribed target, independently of the kinetic regularisation parameter.The kinetic regularisation changes the equilibrium momentum distribution while preserving the position marginal.Source: §1 · p. 2, §1.1 · p. 3
- The diffusion converges exponentially in weighted total variation under assumptions that allow nonconvex potentials and gradients that are not globally Lipschitz.The proof combines Lyapunov estimates with hypoellipticity, controllability, and a minorisation argument.Source: §1.1 · p. 3, §1.2 · p. 4
- The explicit discretisation has moment bounds uniform over iterations despite superlinear gradient growth.For sufficiently small fixed stepsizes, the informal theorem states polynomial moments of every order, given an initial state with finite expectation of the discrete analysis’s exponential Lyapunov weight.Source: §1.1 · p. 3
1.2. Related literature
- Prior work studies spectral or force taming and non-quadratic kinetic energies, while this work places singular-value taming in the momentum-to-position map and targets a prescribed Gibbs measure.The authors state that, to their knowledge, prior work had not identified Muon as the zero-temperature limit of a sampling algorithm for a prescribed Gibbs measure.Source: §1.2 · p. 4
Notation
The notation section sets conventions for matrix operations and derivatives, probability measures and Markov kernels, and weighted total variation.
Notation
- The Frobenius inner product and norm, spectral norm, and nuclear norm provide the matrix inner-product and norm conventions used in the paper.The Frobenius norm reduces to the Euclidean norm for vectors; the spectral norm is the largest singular value.Source: Notation · p. 4
- Gradients and Hessians are defined entrywise for matrix-valued variables, with vectorization identifying the gradient with its vector form.The Hessian is treated as a self-adjoint operator on the matrix space.Source: Notation · p. 4
- For a Markov kernel, Rf denotes the expectation of f under the kernel, and the V-total variation distance compares measures using test functions bounded by V.When V=1, this distance reduces to total variation.Source: Notation · pp. 4–5
2. Kinetic energy and derivation of the SDE
The paper builds an underdamped Langevin system from a nonquadratic kinetic energy whose gradient smoothly tames momentum by singular value. The diffusion is chosen to preserve linear friction and the target Gibbs density, and its discretization recovers classical underdamped Langevin as the taming parameter vanishes.
2. Kinetic energy and derivation of the SDE
- The kinetic energy is a singular-value generalization of the pseudo-Huber potential, with gradient Tδ(p) providing the momentum-to-position map.The preconditioning acts only on singular values and enters through the kinetic energy.Source: §2 · p. 5
- The matrix taming map generalizes the relativistic momentum-to-velocity relation, while its massless limit maps positive singular values to one as in Muon orthogonalization.As δ → 0, Kδ(p) → |p|²/2, Tδ(p) → p, Mδ(p) → I_n, and Γδ(p) → 0, recovering classical underdamped Langevin.Source: §2 · pp. 5–6
- The target Hamiltonian is H(x, p) = U(x) + Kδ(p), and the construction seeks Gibbs density proportional to e^−βH(x,p).The Langevin recipe uses a symmetric positive semidefinite diffusion matrix D and a skew-symmetric matrix Q.Source: §2 · p. 5
- Choosing Mδ(p) so Mδ(p)Tδ(p) = p preserves linear momentum damping −γp while retaining the tamed position velocity.A constant diffusion coefficient would instead give bounded friction; the construction adds a divergence correction to the momentum diffusion.Source: §2 · pp. 5–6
- The Euler–Maruyama discretization is defined for a time step λ ∈ (0, δ].The supplied passage identifies the step-size range but does not give the update equations.Source: §2 · p. 6
3. Assumptions & associated Remarks
The assumptions require a smooth potential with controlled Hessian growth and a 2-dissipative gradient. These conditions imply local Lipschitz continuity of the gradient and coercivity of the potential.
3. Assumptions & associated Remarks
- The potential U is assumed to be infinitely differentiable, and its Hessian satisfies an L0/L1 growth bound.The supplied assumption passage names the bound but does not provide its formula.Source: §3 · p. 6
- The smoothness and 2-dissipativity assumptions imply that ∇U is locally Lipschitz and that U is coercive.Coercivity also entails |∇U(x)| → ∞ as |x| → ∞; the paper sets min_x U(x) = 0 without loss of generality.Source: §3 · p. 6
4. Main Results
Under the standing assumptions, the diffusion has a unique invariant law and converges exponentially in weighted total variation. Its explicit discretization also has polynomial moments of every order uniformly over iterations under the stated parameter, step-size, and initialization conditions.
4. Main Results
- The diffusion has a unique invariant probability measure and converges exponentially in weighted total variation for initial laws with finite Lyapunov-weight expectation.The proof invokes Harris’ theorem after establishing geometric drift and minorization conditions.Source: §4 · pp. 6–7
- The SDE has a unique global strong solution, with time-uniform polynomial moment bounds when the initial exponential Lyapunov-weight expectation is finite.The argument uses a Lyapunov estimate and non-explosion; the weight dominates every polynomial in the state variables.Source: §4 · p. 7
- For sufficiently small fixed λ and admissible parameters, the explicit scheme has polynomial moments of every order uniformly over iterations, including for deterministic initial states.The stated initialization condition is finite expectation of the discrete Lyapunov function V.Source: §4 · p. 7
- The discrete Lyapunov drift inequality bounds E[V_n] uniformly in n by iteration.The estimate is E[V_{n+1} | F_n] ≤ e^−ηλ V_n + Bλ.Source: §4 · p. 7
5. Taming map Properties
The taming map is a well-behaved spectral transformation: it is 1-Lipschitz, and its gradient inherits convexity-related inequalities. The associated spectral potential also has a dimension-dependent Laplacian bound.
5. Taming map Properties
- Tλ is 1-Lipschitz, and its gradient satisfies the one-sided Lipschitz and co-coercivity inequalities.The cited properties follow from the scalar derivative bound and convexity of Kλ.Source: §5 · p. 8
- Convexity of the auxiliary spectral function yields I − ∇²Kλ(p) ⪰ 0, while ΔKλ ≤ mn.The Laplacian bound is stated for the matrix kinetic energy Kλ.Source: §5 · p. 9
- The Laplacian of Kλ is a smooth symmetric, even function of the singular values, and its gradient has a singular-value spectral form.The gradient is expressed using U, the diagonal of the singular-value gradient, and V⊤.Source: §5 · p. 9
6. Geometric Drift Condition
The geometric drift argument uses a Lyapunov function built around the tamed kinetic energy; for fixed δ, β, and γ, a sufficiently small choice of b gives the drift condition, with constants and the compact set allowed to depend on δ.
6. Geometric Drift Condition
- The Lyapunov construction uses Qδ(p) to precondition momentum, and the tamed kinetic energy permits analysis despite Qδ lacking uniform positive definiteness.The text defines Tδ(p) = Qδ(p)p and notes that Qδ(p) is not uniformly bounded below by cI.Source: §6 · p. 10
- For 0 < a < 1, the exponential-energy drift contains dissipation proportional to −a(1−a)βγ⟨p, Tδ(p)⟩, which vanishes at a = 1.The paper identifies a = 1 as the upper threshold for expecting uniform-in-time exponential moment bounds.Source: §6 · pp. 11–12
- Under Assumptions 1–3, the drift condition holds for fixed δ, β, and γ when b is sufficiently small; its constants and compact set may depend on δ.The middle restriction is omitted when Bδ = 0.Source: §6 · pp. 12–14
7. Minorization
The section verifies the invariant measure and develops the regularity, reachability, and minorization properties used in the convergence argument. The bounded spectral speed constrains reachable positions, while compact sets admit a minorization condition.
7.1. Invariant measure identification
- The adjoint calculation gives A*π = 0 after the friction and diffusion terms cancel and the Hamiltonian terms cancel each other.This identifies π as an invariant density for the generator calculation.Source: §7.1 · p. 15
7.2. Hormander’s Condition
- The momentum-direction noise fields and their Lie brackets span the full state space, and under smoothness and nonexplosion the transition kernel has a smooth density.The resulting process is strong-Feller.Source: §7.2 · p. 16
7.3. Controllability
- The spectral-speed bound makes the accessibility set an open lightcone, with |x1 − x0|2 < δ^-1/2T; its closure is the transition-kernel support.The support statement is given as supp RT(z0, ·) = {(x, p): |x − x0|2 ≤ δ^-1/2T}.Source: §7.3 · pp. 17–18, §7.4 · p. 20
7.4. Positive Recurrence
- Controlled trajectories lie in the process support, and the transition-kernel support equals the closed lightcone determined by the spectral-speed bound.The proof establishes the controlled-path inclusion and the converse support inclusion.Source: §7.4 · pp. 19–20
7.5. Minorization
- Every nonempty compact set admits a minorization after a finite time; when the set has nonempty interior, the minorizing probability measure can be supported on it.The proposition supplies positive time and minorization constants for each compact set.Source: §7.5 · p. 20
8. Stability & Lyapunov Condition
The discrete stability analysis uses an exponential Lyapunov function combining the potential with tamed kinetic energy at a shifted momentum. Bounds on its one-step change yield a conditional Lyapunov bound under admissible parameter choices.
8. Stability & Lyapunov Condition
- The stability argument defines y = p + γ^-1∇U(x) and uses V(x, p) = exp[ηU(x) + ηwKδ(y)] as its Lyapunov function.The potential is normalized to be nonnegative, and Kδ is nonnegative by construction.Source: §8 · p. 21
8.1. Supporting Lemmas
- The supporting estimates bound the one-step change of F by λAyKδ(Yn) + λApKδ(Pn) + λĈ3.The conditional noise increment is centered Gaussian with covariance 2γλβ^-1Mδ(Pn) ⊕ Id.Source: §8.1 · pp. 22–23
8.2. The Lyapunov bound
- The conditional Lyapunov bound uses the parameter restrictions Ay ≤ 0 and Ap + A3 ≤ 0.These are identified as the required hyperparameter conditions for the preceding drift argument.Source: §8.2 · p. 24