Non-asymptotic convergence bounds for modified tamed unadjusted Langevin algorithm in non-convex setting
Ariel Neufeld, Matthew Ng Cheng En, Ying Zhang
TL;DR
The paper develops a non-asymptotic analysis for a modified taming of the Unadjusted Langevin Algorithm (mTULA) to sample from $\pi_\beta \propto e^{-\beta U}$ in high dimensions, even when $U$ is non-convex with super-linearly growing gradients. It introduces a taming factor $h_\lambda(\theta)=\frac{h(\theta)}{\bigl(1+\lambda|\theta|^{2r}\bigr)^{1/2}}$ and proves explicit Wasserstein-1 and Wasserstein-2 bounds with convergence rates $O(\lambda)$ and $O(\lambda^{1/2})$, under a convex-at-infinity condition and polynomial Lipschitz assumptions. The work positions mTULA relative to existing taming approaches by removing the need for global strong convexity while achieving faster discretisation rates, supported by high-dimensional numerical experiments. These results enhance robust, non-asymptotic sampling for non-convex targets and provide guidance on step-size selection in practice.
Abstract
We consider the problem of sampling from a high-dimensional target distribution $π_β$ on $\mathbb{R}^d$ with density proportional to $θ\mapsto e^{-βU(θ)}$ using explicit numerical schemes based on discretising the Langevin stochastic differential equation (SDE). In recent literature, taming has been proposed and studied as a method for ensuring stability of Langevin-based numerical schemes in the case of super-linearly growing drift coefficients for the Langevin SDE. In particular, the Tamed Unadjusted Langevin Algorithm (TULA) was proposed in [Bro+19] to sample from such target distributions with the gradient of the potential $U$ being super-linearly growing. However, theoretical guarantees in Wasserstein distances for Langevin-based algorithms have traditionally been derived assuming strong convexity of the potential $U$. In this paper, we propose a novel taming factor and derive, under a setting with possibly non-convex potential $U$ and super-linearly growing gradient of $U$, non-asymptotic theoretical bounds in Wasserstein-1 and Wasserstein-2 distances between the law of our algorithm, which we name the modified Tamed Unadjusted Langevin Algorithm (mTULA), and the target distribution $π_β$. We obtain respective rates of convergence $\mathcal{O}(λ)$ and $\mathcal{O}(λ^{1/2})$ in Wasserstein-1 and Wasserstein-2 distances for the discretisation error of mTULA in step size $λ$. High-dimensional numerical simulations which support our theoretical findings are presented to showcase the applicability of our algorithm.