Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices

Abstract

We study the Unadjusted Langevin Algorithm (ULA) for sampling from a probability distribution $\nu = e^{-f}$ on $\mathbb {R}^n$. We prove a convergence guarantee in Kullback-Leibler (KL) divergence assuming $\nu$ satisfies a log-Sobolev inequality and the Hessian of f is bounded. Notably, we do not assume convexity or bounds on higher derivatives. We prove convergence guarantees in Rényi divergence of order $q > 1$ assuming the limit of ULA satisfies isoperimetry, namely either the log-Sobolev or Poincaré inequality. We also prove a bound on the bias of the limiting distribution of ULA assuming third-order smoothness of f, without requiring isoperimetry.

This work was supported in part by NSF awards CCF-1717349, DMS-1839323, CCF-2007443, and CCF-2106644.

Appendix

1.1 Review on Notation and Basic Properties

Throughout, we represent a probability distribution $\rho$ on $\mathbb {R}^n$ via its probability density function with respect to the Lebesgue measure, so $\rho \colon \mathbb {R}^n \to \mathbb {R}$ with $\int _{\mathbb {R}^n} \rho (x) dx = 1$. We typically assume $\rho$ has full support and smooth density, so $\rho (x) > 0$ and $x \mapsto \rho (x)$ is differentiable. Given a function $f \colon \mathbb {R}^n \to \mathbb {R}$, we denote the expected value of f under $\rho$ by

$$\displaystyle \begin{aligned} \mathbb{E}_\rho[f] = \int_{\mathbb{R}^n} f(x) \rho(x) \, dx. \end{aligned} $$

We use the Euclidean inner product $\langle x,y \rangle = \sum _{i=1}^n x_i y_i$ for $x = (x_i)_{1 \le i \le n}, y = (y_i)_{1 \le i \le n} \in \mathbb {R}^n$. For symmetric matrices $A, B \in \mathbb {R}^{n \times n}$, let $A \preceq B$ denote that $B-A$ is positive semidefinite. For $\mu \in \mathbb {R}^n$, $\Sigma \succ 0$, let $\mathcal {N}(\mu ,\Sigma )$ denote the Gaussian distribution on $\mathbb {R}^n$ with mean $\mu$ and covariance matrix $\Sigma$.

Given a smooth function $f \colon \mathbb {R}^n \to \mathbb {R}$, its gradient$\nabla f \colon \mathbb {R}^n \to \mathbb {R}^n$ is the vector of partial derivatives:

$$\displaystyle \begin{aligned} \nabla f(x) = \left(\frac{\partial f(x)}{\partial x_1}, \dots, \frac{\partial f(x)}{\partial x_n} \right). \end{aligned} $$

The Hessian$\nabla ^2 f \colon \mathbb {R}^n \to \mathbb {R}^{n \times n}$ is the matrix of second partial derivatives:

$$\displaystyle \begin{aligned} \nabla^2 f(x) = \left(\frac{\partial ^2f(x)}{\partial x_i x_j} \right)_{1 \le i,j \le n}. \end{aligned} $$

The Laplacian$\Delta f \colon \mathbb {R}^n \to \mathbb {R}$ is the trace of its Hessian:

$$\displaystyle \begin{aligned} \Delta f(x) = \mathrm{Tr}(\nabla^2 f(x)) = \sum_{i=1}^n \frac{\partial ^2 f(x)}{\partial x_i^2}. \end{aligned} $$

Given a smooth vector field $v = (v_1,\dots ,v_n) \colon \mathbb {R}^n \to \mathbb {R}^n$, its divergence$\nabla \cdot v \colon \mathbb {R}^n \to \mathbb {R}$ is

$$\displaystyle \begin{aligned} (\nabla \cdot v)(x) = \sum_{i=1}^n \frac{\partial v_i(x)}{\partial x_i}. \end{aligned} $$

In particular, the divergence of gradient is the Laplacian:

$$\displaystyle \begin{aligned} (\nabla \cdot \nabla f)(x) = \sum_{i=1}^n \frac{\partial ^2 f(x)}{\partial x_i^2} = \Delta f(x). \end{aligned} $$

For any function $f \colon \mathbb {R}^n \to \mathbb {R}$ and vector field $v \colon \mathbb {R}^n \to \mathbb {R}^n$ with sufficiently fast decay at infinity, we have the following integration by parts formula:

$$\displaystyle \begin{aligned} \int_{\mathbb{R}^n} \langle v(x), \nabla f(x) \rangle dx = -\int_{\mathbb{R}^n} f(x) (\nabla \cdot v)(x) dx. \end{aligned} $$

Furthermore, for any two functions $f, g \colon \mathbb {R}^n \to \mathbb {R}$,

$$\displaystyle \begin{aligned} \int_{\mathbb{R}^n} f(x) \Delta g(x) dx = -\int_{\mathbb{R}^n} \langle \nabla f(x), \nabla g(x) \rangle dx = \int_{\mathbb{R}^n} g(x) \Delta f(x) dx. \end{aligned} $$

When the argument is clear, we omit the argument $(x)$ in the formulae for brevity. For example, the last integral above becomes

$$\displaystyle \begin{aligned} {} \int f \, \Delta g \, dx = -\int \langle \nabla f, \nabla g \rangle \, dx = \int g \, \Delta f \, dx. \end{aligned} $$

(B.1)

Derivation of the Fokker-Planck Equation

Consider a stochastic differential equation

$$\displaystyle \begin{aligned} {} dX_t = v(X_t) \, dt + \sqrt{2} \, dW_t \end{aligned} $$

(B.2)

where $v \colon \mathbb {R}^n \to \mathbb {R}^n$ is a smooth vector field and $(W_t)_{t \ge 0}$ is the Brownian motion on $\mathbb {R}^n$ with $W_0 = 0$.

We will show that if $X_t$ evolves following (B.2), then its probability density function $\rho _t(x)$ evolves following the Fokker-Planck equation:

$$\displaystyle \begin{aligned} {} \frac{\partial \rho_t}{\partial t} = -\nabla \cdot (\rho_t v) + \Delta \rho_t. \end{aligned} $$

(B.3)

We can derive this heuristically as follows; we refer to standard textbooks for rigorous derivation [58].

For any smooth test function $\phi \colon \mathbb {R}^n \to \mathbb {R}$, let us compute the time derivative of the expectation

$$\displaystyle \begin{aligned} A(t) = \mathbb{E}_{\rho_t}[\phi] = \mathbb{E}[\phi(X_t)]. \end{aligned} $$

On the one hand, we can compute this as

$$\displaystyle \begin{aligned} {} \dot A(t) = \frac{d}{dt} A(t) = \frac{d}{dt} \int_{\mathbb{R}^n} \rho_t(x) \phi(x) \, dx = \int_{\mathbb{R}^n} \frac{\partial \rho_t(x)}{\partial t} \phi(x) \, dx. \end{aligned} $$

(B.4)

On the other hand, by (B.2), for small $\epsilon > 0$ we have

$$\displaystyle \begin{aligned} X_{t+\epsilon } &= X_t + \int_t^{t+\epsilon } v(X_s) ds + \sqrt{2} (W_{t+\epsilon }-W_t) \\ &= X_t + \epsilon v(X_t) + \sqrt{2} (W_{t+\epsilon }-W_t) + O(\epsilon ^2) \\ &\stackrel{d}{=} X_t + \epsilon v(X_t) + \sqrt{2\epsilon } Z + O(\epsilon ^2) \end{aligned} $$

where $Z \sim \mathcal {N}(0,I)$ is independent of $X_t$, since $W_{t+\epsilon }-W_t \sim \mathcal {N}(0,\epsilon I)$. Then by Taylor expansion,

$$\displaystyle \begin{aligned} \phi(X_{t+\epsilon }) &\stackrel{d}{=} \phi\left(X_t + \epsilon v(X_t) + \sqrt{2\epsilon } Z + O(\epsilon ^2)\right) \\ &= \phi(X_t) + \epsilon \langle \nabla \phi(X_t), v(X_t) \rangle + \sqrt{2\epsilon } \langle \nabla \phi(X_t), Z \rangle \\ &\quad + \frac{1}{2} 2\epsilon \langle Z, \nabla^2 \phi(X_t) Z \rangle + O(\epsilon ^{\frac{3}{2}}). \end{aligned} $$

Now we take expectation on both sides. Since $Z \sim \mathcal {N}(0,I)$ is independent of $X_t$,

$$\displaystyle \begin{aligned} A(t+\epsilon ) &= \mathbb{E}[\phi(X_{t+\epsilon })]\\ &= \mathbb{E}\Big[\phi(X_t) + \epsilon \langle \nabla \phi(X_t), v(X_t) \rangle + \sqrt{2\epsilon } \langle \nabla \phi(X_t), Z \rangle \\ &\qquad + \epsilon \langle Z, \nabla^2 \phi(X_t) Z \rangle \Big] + O(\epsilon ^{\frac{3}{2}}) \\ &= A(t)+ \epsilon \left(\mathbb{E}[\langle \nabla \phi(X_t), v(X_t) \rangle] + \mathbb{E}[\Delta \phi(X_t)]\right) + O(\epsilon ^{\frac{3}{2}}). \end{aligned} $$

Therefore, by integration by parts, this second approach gives

$$\displaystyle \begin{aligned} \dot A(t) &= \lim_{ \epsilon \to 0} \frac{A(t+\epsilon )-A(t)}{\epsilon } \notag \\ &= \mathbb{E}[\langle \nabla \phi(X_t), v(X_t) \rangle] + \mathbb{E}[\Delta \phi(X_t)] \notag \\ &= \int_{\mathbb{R}^n} \langle \nabla \phi(x), \rho_t(x) v(x) \rangle dx + \int_{\mathbb{R}^n} \rho_t(x) \Delta \phi(x) \, dx \notag \\ &= -\int_{\mathbb{R}^n} \phi(x) \nabla \cdot (\rho_t v)(x) \, dx + \int_{\mathbb{R}^n} \phi(x) \Delta \rho_t(x) \, dx \notag \\ &= \int_{\mathbb{R}^n} \phi(x) \left(-\nabla \cdot (\rho_t v)(x) + \Delta \rho_t(x)\right) \, dx. {} \end{aligned} $$

(B.5)

Comparing (B.4) and (B.5), and since $\phi$ is arbitrary, we conclude that

$$\displaystyle \begin{aligned} \frac{\partial \rho_t(x)}{\partial t} = -\nabla \cdot (\rho_t v)(x) + \Delta \rho_t(x) \end{aligned} $$

as claimed in (B.3).

When $v = -\nabla f$, the stochastic differential equation (B.2) becomes the Langevin dynamics (7) from Sect. 2.3, and the Fokker-Planck equation (B.3) becomes (8).

In the proof of Lemma 3, we also apply the Fokker-Planck equation (B.3) when $v = -\nabla f(x_0)$ is a constant vector field to derive the evolution equation (30) for one step of ULA.

1.2 Remaining Proofs

1.2.1 Proof of Lemma 16

Proof of Lemma 16

Let $g \colon \mathbb {R}^n \to \mathbb {R}$ be a smooth function, and let $\tilde g \colon \mathbb {R}^n \to \mathbb {R}$ be the function $\tilde g(x) = g(T(x))$. Let $X \sim \nu$, so $T(X) \sim \tilde \nu$. Note that

$$\displaystyle \begin{aligned} \mathbb{E}_{\tilde \nu}[g^2] &= \mathbb{E}_{X \sim \nu}[g(T(X))^2] = \mathbb{E}_{\nu}[\tilde g^2], \\ \mathbb{E}_{\tilde \nu}[g^2 \log g^2] &= \mathbb{E}_{X \sim \nu}[g(T(X))^2 \log g(T(X))^2] = \mathbb{E}_{\nu}[\tilde g^2 \log \tilde g^2]. \end{aligned} $$

Furthermore, we have $\nabla \tilde g(x) = \nabla T(x) \, \nabla g(T(x))$. Since T is L-Lipschitz, $\|\nabla T(x)\| \le L$. Then

$$\displaystyle \begin{aligned} \|\nabla \tilde g(x)\| \le \|\nabla T(x)\| \, \|\nabla g(T(x))\| \le L \|\nabla g(T(x))\|. \end{aligned} $$

This implies

$$\displaystyle \begin{aligned} \mathbb{E}_{\tilde \nu}[\|\nabla g\|{}^2] = \mathbb{E}_{X \sim \nu}[\|\nabla g(T(X))\|{}^2] \ge \frac{\mathbb{E}_{\nu}[\|\nabla \tilde g\|{}^2]}{L^2}. \end{aligned} $$

Therefore,

$$\displaystyle \begin{aligned} \frac{\mathbb{E}_{\tilde \nu}[\|\nabla g\|{}^2]}{\mathbb{E}_{\tilde \nu}[g^2 \log g^2] \,{-}\, \mathbb{E}_{\tilde \nu}[g^2] \log \mathbb{E}_{\tilde \nu}[g^2]} &\,{\ge}\, \frac{1}{L^2} \, \frac{\mathbb{E}_{\nu}[\|\nabla \tilde g\|{}^2]}{\big(\mathbb{E}_{\nu}[\tilde g^2 \log \tilde g^2] \,{-}\, \mathbb{E}_{\nu}[\tilde g^2] \log \mathbb{E}_{\nu}[\tilde g^2]\big)}\\ &\,{\ge}\, \frac{\alpha}{2L^2} \end{aligned} $$

where the last inequality follows from the assumption that $\nu$ satisfies LSI with constant $\alpha$. This shows that $\tilde \nu$ satisfies LSI with constant $\alpha /L^2$, as desired. □

Proof of Lemma 17

Proof of Lemma 17

We recall the following convolution property of LSI [15]: If $\nu , \tilde \nu$ satisfy LSI with constants $\alpha , \tilde \alpha > 0$, respectively, then $\nu \ast \tilde \nu$ satisfies LSI with constant $\left (\frac {1}{\alpha }+\frac {1}{\tilde \alpha }\right )^{-1}$. Since $\mathcal {N}(0,2tI)$ satisfies LSI with constant $\frac {1}{2t}$, the claim follows. □

Proof of Lemma 19

Proof of Lemma 19

Let $g \colon \mathbb {R}^n \to \mathbb {R}$ be a smooth function, and let $\tilde g \colon \mathbb {R}^n \to \mathbb {R}$ be the function $\tilde g(x) = g(T(x))$. Let $X \sim \nu$, so $T(X) \sim \tilde \nu$. Note that

$$\displaystyle \begin{aligned} \mathrm{Var}_{\tilde \nu}(g) &= \mathrm{Var}_{X \sim \nu}(g(T(X))) = \mathrm{Var}_{\nu}(\tilde g). \end{aligned} $$

Furthermore, we have $\nabla \tilde g(x) = \nabla T(x) \, \nabla g(T(x))$. Since T is L-Lipschitz, $\|\nabla T(x)\| \le L$. Then

$$\displaystyle \begin{aligned} \|\nabla \tilde g(x)\| \le \|\nabla T(x)\| \, \|\nabla g(T(x))\| \le L \|\nabla g(T(x))\|. \end{aligned} $$

This implies

$$\displaystyle \begin{aligned} \mathbb{E}_{\tilde \nu}[\|\nabla g\|{}^2] = \mathbb{E}_{X \sim \nu}[\|\nabla g(T(X))\|{}^2] \ge \frac{\mathbb{E}_{\nu}[\|\nabla \tilde g\|{}^2]}{L^2}. \end{aligned} $$

Therefore,

$$\displaystyle \begin{aligned} \frac{\mathbb{E}_{\tilde \nu}[\|\nabla g\|{}^2]}{\mathrm{Var}_{\tilde \nu}(g)} &\ge \frac{1}{L^2} \, \frac{\mathbb{E}_{\nu}[\|\nabla \tilde g\|{}^2]}{\mathrm{Var}_{\nu}(\tilde g)} \ge \frac{\alpha}{L^2} \end{aligned} $$

where the last inequality follows from the assumption that $\nu$ satisfies Poincaré inequality with constant $\alpha$. This shows that $\tilde \nu$ satisfies Poincaré inequality with constant $\alpha /L^2$, as desired. □

Proof of Lemma 20

Proof of Lemma 20

We recall the following convolution property of Poincaré inequality [23]: If $\nu , \tilde \nu$ satisfy Poincaré inequality with constants $\alpha , \tilde \alpha > 0$, respectively, then $\nu \ast \tilde \nu$ satisfies Poincaré inequality with constant $\left (\frac {1}{\alpha }+\frac {1}{\tilde \alpha }\right )^{-1}$. Since $\mathcal {N}(0,2tI)$ satisfies Poincaré inequality with constant $\frac {1}{2t}$, the claim follows. □