kTULA: A Langevin sampling algorithm with improved KL bounds under super-linear log-gradients

TL;DR

This paper introduces kTULA, a novel Langevin sampling algorithm designed for distributions with super-linear log-gradients, providing improved convergence guarantees and applicability to high-dimensional problems in deep learning.

Contribution

The paper proposes kTULA, a tamed Langevin algorithm with theoretical convergence guarantees under super-linear log-gradient conditions, advancing sampling methods in challenging settings.

Findings

Achieves a convergence rate of $2-ar{ ext{epsilon}}$ in KL divergence.

Provides improved non-asymptotic bounds in Wasserstein-2 distance.

Demonstrates effectiveness on high-dimensional double-well and neural network optimization problems.

Abstract

Motivated by applications in deep learning, where the global Lipschitz continuity condition is often not satisfied, we examine the problem of sampling from distributions with super-linearly growing log-gradients. We propose a novel tamed Langevin dynamics-based algorithm, called kTULA, to solve the aforementioned sampling problem, and provide a theoretical guarantee for its performance. More precisely, we establish a non-asymptotic convergence bound in Kullback-Leibler (KL) divergence with the best-known rate of convergence equal to $2-\overline{\epsilon}$, $\overline{\epsilon}>0$, which significantly improves relevant results in existing literature. This enables us to obtain an improved non-asymptotic error bound in Wasserstein-2 distance, which can be used to further derive a non-asymptotic guarantee for kTULA to solve the associated optimization problems. To illustrate the applicability of kTULA, we apply the proposed algorithm to the problem of sampling from a high-dimensional double-well potential distribution and to an optimization problem involving a neural network. We show that our main results can be used to provide theoretical guarantees for the performance of kTULA.