Muon meets Tamed Langevin: Momentum Preconditioning beyond Convex and gradient-Lipschitz Potentials
A momentum-preconditioned Langevin sampler stays stable for nonconvex, non-Lipschitz matrix potentials.
The paper studies sampling from Gibbs distributions on matrix spaces whose potentials are neither convex nor globally gradient-Lipschitz. Non-quadratic kinetic energies produce an underdamped Langevin system in which the kinetic-energy gradient smoothly tames momentum. Under the relaxed assumptions, the dynamics leave the target measure invariant and converge exponentially in a weighted total variation distance. The Euler-Maruyama discretization has time-uniform moment bounds without modifying the potential gradient.
- Targets nonconvex, non-gradient-Lipschitz matrix potentials.
- Kinetic-energy gradients smoothly tame momentum.
- Dynamics preserve the Gibbs measure and converge exponentially.
- Euler-Maruyama moments stay bounded uniformly in time.
Full article121 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2610.02158