Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials

2026-10-01

AI Summary

This study addresses the challenge of sampling from Gibbs distributions over matrix spaces characterized by non-convex potentials and non-globally Lipschitz gradients. To this end, it proposes a damped Langevin dynamics algorithm incorporating momentum preconditioning. The method introduces a non-quadratic kinetic energy to achieve spectral taming, constraining momentum via the smooth gradient of the kinetic energy without modifying the potential gradient. Theoretical analysis is conducted based on Euler–Maruyama discretization and weighted total variation distances. This work establishes the invariance of the target measure and exponential convergence for the continuous dynamics, while deriving time-uniform moment bounds for the discrete algorithm. By overcoming conventional convexity and global Lipschitz assumptions, these results provide rigorous theoretical guarantees for stable sampling under complex potentials.

Abstract

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.

Problem

• Gibbs distribution sampling
• matrix spaces
• non-convex potentials
• non-gradient-Lipschitz
• underdamped Langevin dynamics

Innovation

• Momentum Preconditioning
• Underdamped Langevin Dynamics
• Non-convex Sampling
• Spectral Taming
• Euler-Maruyama Discretization