Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials
The paper tackles sampling from Gibbs distributions on matrix spaces with potentials that are neither convex nor globally gradient-Lipschitz. Its core idea is to replace the quadratic kinetic energy of underdamped Langevin dynamics with a non-quadratic, relativistic-like energy $K_\delta(p)=\delta^{-1}\operatorname{Tr}[(I_n+\delta pp^\top)^{1/2}-I_n]$, whose gradient tames momentum singular values smoothly. By carefully coupling a momentum-dependent diffusion matrix $M_\delta(p)$ and its divergence correction $\Gamma_\delta(p)$, the authors preserve the exact target Gibbs measure, recover regularised Muon in the zero-temperature limit, and prove exponential convergence of the diffusion plus uniform-in-time moment bounds for an explicit Euler-Maruyama scheme.
The paper makes a clean theoretical advance by placing Muon's spectral preconditioning inside a rigorously defined stochastic dynamics whose invariant measure is correct and whose explicit discretisation is stable under weak assumptions. However, the supplied manuscript ends abruptly in the middle of the invariant-measure proof in Section 7.1, leaving the minorization argument and parts of the discrete-time analysis incomplete. Moreover, the zero-temperature connection to regularised Muon is formal rather than quantitative, and there are no experiments or explicit parameter recipes to translate the theory into practice.
The strongest aspect is the principled derivation of the SDE. The taming map $T_\delta(p)=(I_n+\delta pp^\top)^{-1/2}p$ emerges as the gradient of a unitarily invariant convex kinetic energy, which guarantees that the classical underdamped Langevin system is recovered as $\delta\to 0$ and ensures the invariant distribution has the correct position marginal. The authors also carefully verify non-trivial matrix-analytic properties---Lipschitzness, dissipativity, and Laplacian bounds for the taming map---that underpin the stability analysis. The Lyapunov function $\Psi(x,p)=\exp\{a\beta H+b\langle G(x),p\rangle+bd/\sqrt{\delta}\}$ and the resulting geometric drift condition in Lemma 7 are sophisticated and well adapted to the matrix-valued multiplicative-noise setting.
First, the provided text is physically truncated at equation (58) in Section 7.1, so the minorization argument and the complete discrete-time stability proof cannot be evaluated. Second, the main claims in the introduction are explicitly flagged as informal, and the zero-temperature identification with regularised Muon (equations (9)--(10)) rests on a formal $\beta\to\infty$ limit that collapses the stochastic terms but is not accompanied by quantitative optimisation or concentration bounds.
Third, although the assumptions relax global Lipschitz gradients, they still impose $C^\infty$ smoothness, an $L_0/L_1$ Hessian growth bound, and 2-dissipativity---standard for this literature but stronger than what one can verify for typical LLM losses. Finally, the admissible stepsize satisfies $\lambda\in(0,\min\{\delta,\gamma^{-1}\}]$, and the Lyapunov constants depend on $\delta$ in a complicated, implicit fashion, making it hard to extract practical step-size guidance.
The positioning against related work is precise and fair. The authors distinguish their approach from Brosse et al. (2019) and Johnston et al. (2024) by emphasising that prior schemes tame the gradient force, whereas they tame the momentum-to-position map and preserve the Gibbs target via a carefully designed state-dependent diffusion. They also correctly contrast with Mustafi et al. (2026), noting that the latter requires bounded gradients and curvature. The paper's main weakness on this axis is the total absence of numerical evidence: there are no simulations comparing the proposed sampler to standard kinetic Langevin or to practical Muon variants, so the claimed stability benefits remain purely theoretical.
No code, data, or supplementary material is mentioned, blocking independent reproduction of the matrix-valued SDE and its spectral-tamed discretisation. Implementing the scheme requires evaluating $(I_n+\delta pp^\top)^{1/2}$ and its matrix divergence $\Gamma_\delta(p)$ at every step, which is non-trivial to do efficiently and accurately on GPUs despite the SVD-based closed forms. Hyperparameter selection is also involved: the drift condition demands that $b$ satisfy a triple upper bound in equation (47), the discrete stability theorem requires $A_y<0$ and $A_p+A_3<0$, and the stepsize is restricted to $\lambda\in(0,\min\{\delta,\gamma^{-1}\}]$. Without explicit constants, a reference implementation, or the complete text of the proofs, reproducing the claimed uniform moment bounds would be extremely difficult.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.