Resolving the Mixing Time of the Langevin Algorithm to its Stationary Distribution for Log-Concave Sampling

• •

  Dependence on error $\varepsilon$. As is standard, Theorem 1.2 is stated for mixing to constant error, here $\varepsilon=1/4$. A basic fact about Markov chains is that a mixing bound to constant error can be boosted to arbitrarily small error $\varepsilon>0$ at an exponential rate, see, e.g., [58, §4.5]. Specifically, Theorem 3.1 implies $T_{\mathrm{mix},\,\mathsf{TV}}\left(\varepsilon\right)\leqslant\frac{D^{2}}{\eta}\log\frac{1}{\varepsilon}$.

• •

  Dependence on diameter $D$ and stepsize $\eta$. Note that $\Theta(D^{2}/\eta)$ is the number of iterations it takes for a univariate random walk with i.i.d. $\mathcal{N}(0,2\eta)$ Gaussian increments to move a distance $D$ from its initialization. Thus, $\Theta(D^{2}/\eta)$ is the number of iterations it takes for the Langevin Algorithm to just reach the opposite side of a univariate interval $\mathcal{K}=[-D/2,D/2]$ even in the simple setting of zero potentials (a.k.a., uniform $\pi$). Since reachability is obviously a pre-requisite for mixing, $\Theta(D^{2}/\eta)$ is therefore a lower bound on the mixing time. (This is the intuition behind the lower bound in Theorem 3.2.) The matching upper bound in Theorem 3.1 shows that, modulo a constant factor, the Langevin Algorithm takes the same amount of time to mix as it does to just reach the opposite end of $\mathcal{K}$.

• •

  Dependence on other parameters. Note that the fast mixing in Theorem 1.2 depends on no other parameters, and in particular has no dimension dependence for $\eta$ of constant size (a common regime in differentially private optimization). However, if one seeks to sample from $\pi$, then the discretization bias $\pi_{\eta}\approx\pi$ must be small, which requires $1/\eta$ to scale polynomially in the dimension and inverse accuracy. This is why mixing bounds to $\pi$ inevitably have polynomial dependence on the dimension and inverse accuracy. Theorem 1.2 clarifies that these polynomial dependencies arise solely through the single factor of $1/\eta$.

We prove Theorem 3.1 by introducing a technique from the differential privacy literature to the sampling literature. See §1.2 below for an overview of this technique, called Privacy Amplification by Iteration. This approach gives a short, simple proof of optimal mixing bounds—and moreover has several additional appealing properties, which are a mix of technical and pedagogical in nature:

• •

  Minimal assumptions. Theorem 3.1 requires an essentially minimal set of assumptions. This is in part because of our direct analysis approach, and also in part because our analysis distenagles the study of the mixing and the bias of the Langevin Algorithm. Indeed, assumptions like curvature of the set and Lipschitzness of the potential²²2In the constrained setting, Lipschitzness is implied by smoothness; however, in the unconstrained setting, Lipschitzness is not implied and yet our techniques still apply despite this. appear necessary for discretization bounds (see, e.g,. [14])—but as we show, these are unnecessary for fast mixing. Note also that our result requires no warm start.

  In fact, Theorem 3.1 only makes three assumptions. The first is that $\mathcal{K}$ is bounded; this is easily relaxed, see the following paragraph. The second is that the potential $f$ and the set $\mathcal{K}$ are convex; it is unavoidable to have at least some type of assumption in this direction since arbitrary non-convexity leads to computational hardness due to reductions to non-convex optimization (see §5). The third is that the stepsize $\eta\leqslant 2/M$ is not too large, else $\pi_{\eta}$ is a meaningless approximation of $\pi$ and in fact is even transient in the unconstrained setting.³³3This can be seen even in $1$ dimension with the quadratic potential $f(x)=\tfrac{M}{2}x^{2}$ on $\mathcal{K}=[-D/2,D/2]$. Briefly, if $\eta=(2+\varepsilon)/M$, then $\pi_{\eta}$ concentrates away from $0$, whereas $\pi$ concentrates around $0$, in fact arbitrarily so as $M\to\infty$. Thus $\pi$ and $\pi_{\eta}$ can be arbitrarily different if $\eta>2/M$. Moreover in the unconstrained setting, this Markov chain is transient.

• •

  Unification of different problem settings. Our analysis technique extends nearly unchanged and yields tight bounds regardless of whether the Langevin Algorithm uses mini-batch stochastic gradients, projections, or strongly log-concave targets $\pi$. In the stochastic setting, our mixing result is identical regardless of the batch size. In the unconstrained setting, our mixing result is identical except with $D$ replaced by the diameter of a ball that captures all but $\varepsilon$ mass of $\pi_{\eta}$; such a diameter proxy is finite and leads to tight mixing bounds, details in Appendix A.4. And in the strongly log-concave setting, our mixing result improves exponentially in $D$; details in §4.

• •

  Pedagogical hope towards unifying sampling and optimization. While the problems of optimization and sampling are well-known to exhibit appealing syntactic similarities in terms of problem formulations and algorithms, it is less clear to how these correspondences should extend to analyses and error bounds. One particularly popular approach, initiated by the seminal work [88], exploits the JKO Theorem which states that the Langevin Diffusion is equivalent to a KL gradient flow over the space $\mathcal{P}(\mathbb{R}^{d})$ of probability distributions [52]. See the previous work section §1.3 for details. The present paper offers a different approach towards unifying the analyses and error bounds of sampling and optimization algorithms that does not require lifting from $\mathbb{R}^{d}$ to $\mathcal{P}(\mathbb{R}^{d})$, thereby avoiding the (beautiful yet) high-powered associated mathematical machinery. From an analysis perspective, we exploit convexity through the contractivity of a gradient update (in a different way from [26, 27]). This is reminiscent of how convexity is used in textbook proofs of Gradient Descent (see, e.g., [13, pages 271 and 279] respectively for the convex and strongly convex settings), and leads us to error bounds for sampling that in many ways mirror classical results in optimization, see Table 2.⁴⁴4In fact, it is worth mentioning that our techniques yield analyses and error bounds for sampling which are even simpler than those for optimization in the setting of stochastic gradients—since our sampling bounds are unchanged for any minibatch size $b$, whereas optimization bounds certainly are not.

Due to the many applications of sampling, there is an extensive literature on different sampling algorithms spanning multiple communities and decades. Much of this work was inspired by the seminal paper [38], which gave the first randomized polynomial-time algorithm for computing the volume of a convex body via a reduction to (log-concave) sampling. A beautiful line of work from the theoretical computer science community has improved the (originally large) polynomial runtime factors by developing elegant algorithms based on geometric random walks such as the Metropolized Random Walk [69, 79], the Ball Walk [63, 64], and the Hit-and-Run Walk [66, 65, 62, 53, 8], to name a few; for details see the survey [87]. A key distinguishing aspect of this line of work is that these geometric random walk algorithms only access the target distribution $\pi$ by evaluating the density up to a constant factor—in contrast to the Langevin Algorithm which also evaluates gradient information about the density. This distinction is formalized through the notion of zero-th order and first-order algorithms: the former access $\pi$ only through evaluations of the potential $f$, whereas the latter also evaluate gradients of $f$. Most first-order algorithms are variants, at least intuitively, of the Langevin Algorithm. These variations come in many forms and can be essential for faster runtimes in theory, practice, or both. A few examples of these variations include: dealing with constrained sets by using projected [14], proximal [81, 12, 34], or mirror [1, 93, 49] variants of the Langevin Algorithm; augmenting the state space by using Hamiltonian Monte Carlo [29, 71] or its variants such as Riemannian Hamiltonian Monte Carlo [47, 55, 54], NUTS [48], or the Underdamped Langevin Algorithm [18, 39, 15]; or using different discretization schemes such as the Randomized Midpoint Discretization [84] or the Leap-Frog Integrator [11]. These algorithms can be combined with Metropolis-Hastings filters to correct the bias of the algorithm’s stationary distribution; this typically improves mixing times to $\pi$ from polynomial to polylogarithmic in the inverse accuracy. A notable example is the Metropolis-Adjusted Langevin Algorithm [78], whose mixing time has been the focus of much recent work [22, 36] and was recently characterized in [91]. While there have been many more exciting developments, this literature is far too large to attempt to comprehensively cite here, and we refer the reader to textbooks and surveys such as [4, 76, 61, 77, 20] for further background.

The present paper fits into this line of work by analyzing the (Projected Stochastic) Langevin Algorithm. We focus on this algorithm because it is arguably the simplest first-order sampling algorithm and also is the direct sampling analog of (Projected Stochastic) Gradient Descent.

While the Langevin Diffusion converges asymptotically under mild tail assumptions [78], non-asymptotic convergence of the Langevin Algorithm require stronger assumptions. This is in part because sampling is in general computationally hard if the potential is non-convex, e.g., due to reductions from non-convex optimization. Non-asymptotic mixing bounds for the Langevin Algorithm were first derived by the seminal paper [26] in the unconstrained setting where $f$ is strongly convex and smooth. This led to a rapidly growing line of work on mixing times (primarily to $\pi$) in many settings—for example log-concave/strongly log-concave settings and constrained/unconstrained settings as already mentioned above, but also settings where the potential has less stuctural assumptions, for example when it is non-smooth [35, 60, 57, 73, 16], or is non-convex and only satisfies tail-growth or isoperimetric conditions [70, 75, 42, 40, 73, 41, 17, 6, 88]. Another line of work has focused on sampling from non-convex manifolds [59, 19, 46].

We highlight a recent line of work which relaxes (strong) convexity of the potential to functional inequalities of the target distribution $\pi$. This idea was pioneered by the paper [88], which pointed out that smoothness and functional inequalities are sufficient to establish rapid mixing—essentially replacing convexity and strong convexity by Poincaré and Log-Sobolev inequality assumptions, respectively. A rapidly growing line of recent work has established even faster mixing bounds under these assumptions (e.g., [21, 6, 42, 40]) as well as found appealing ways to interpolate between these functional inequalities [21]. It should be noted that for the purpose of mixing times, convexity and functional inequality assumptions are incomparable in the sense that on one hand, functional inequalities allow for (mild) non-convexity; and on the other hand, while convexity implies a Poincaré inequality with a finite constant, there is no uniform bound on this Poincaré constant.

In this paper we focus on the foundational setting of smooth convex potentials $f$ (possibly also with strong convexity and/or bounded domains and/or stochastic gradients) for several reasons: (i) this is the direct sampling analog of the most classical setup in convex optimization; (ii) this setup has many applications in machine learning, statistics, and scientific computing; and (iii) optimal mixing times for $\pi_{\eta}$ were previously unknown even in this seemingly simple setting.

As mentioned earlier, most previous work focuses on mixing to $\pi$. Existing bounds vary depending on the precise set of assumptions, the precise variant of the algorithm, and the notion of mixing error. For several common combinations thereof, state-of-the-art mixing bounds have changed in essentially each of the past few years—and may continue to change in the years to come, as there is no⁶⁶6 At the time of writing, there was no setting in which current mixing bounds for the (Unadjusted) Langevin Algorithm were known to be optimal. The concurrent work [23] also shows an optimality result. Their result concerns an entirely different setting (non-convex potentials where error is measured in Fisher information) and applies to a variant of the Langevin Algorithm (that outputs the averaged iterate) in the specific regime of very large mixing error (namely $\varepsilon\approx\sqrt{d}$, where $d$ is the dimension). We also note that our lower bound (Theorem 4.3) establishes that an upper bound in [88] for mixing to $\pi_{\eta}$ is tight in certain metrics in the unconstrained, non-stochastic, strongly log-concave, and smooth setting. See the main text for details. setting in which current mixing bounds for the (Unadjusted) Langevin Algorithm are known to be optimal.

For example, for the unconstrained, smooth, strongly log-concave setting originally investigated by [26, 31, 32, 27], a line of recent work [88, 45, 42, 21, 35] has led to state-of-the-art mixing results to $\pi$ in Rényi divergence $\mathcal{D}_{\alpha}$ of order roughly $\tilde{O}(\alpha d\kappa^{2}/\varepsilon)$ for $\alpha$ of moderate size [21], and even faster $\tilde{O}(d\kappa/\varepsilon)$ for the KL divergence using averaged iterates [35]. Here $d$ denotes the dimension, $\varepsilon$ denotes the mixing error to $\pi$, $m$ denotes the strong convexity of the potential $f$, $M$ denotes its smoothness, and $\kappa=M/m$ denotes the condition number. In this setting, these Rényi and KL guarantees recover state-of-the-art Wasserstein guarantees [26, 27, 33] by Talagrand’s inequality.

For this unconstrained, smooth, strongly-log concave setting, the recent revision⁷⁷7We thank Andre Wibisono for making us aware of this revision. of [88, Lemmas 4,8] provided an upper bound of $\tilde{O}(\frac{\alpha}{m\eta})$ on the mixing time of the Langevin Algorithm to $\pi_{\eta}$. Our lower bound (Theorem 4.3) establishes that their upper bound is tight for Rényi divergences of constant size $\alpha=\Theta(1)$. Our upper bound $\tilde{O}(\frac{1}{\eta m})$ for this setting (Theorem 4.1) has exponentially better (and optimal) dependence on the Rényi parameter $\alpha$, does not require as onerous a warm start, and extends to other settings (e.g., constrained, stochastic, all non-trivial stepsizes, and log-concave target distributions).

For the unconstrained, smooth, (non-strongly) log-concave setting, state-of-the-art mixing results to $\pi$ in $\mathsf{TV}$ are of order roughly $\tilde{O}(d^{2}M^{2}c_{\textrm{PI}}^{2}/\varepsilon^{4})$ [6] and $\tilde{O}(d^{3}M^{2}c_{\textrm{PI}}^{2}/\varepsilon^{2})$ [21]. Here, $c_{\textrm{PI}}$ is the Poincaré constant, which is guaranteed finite but cannot be uniformly bounded. The lack of strong log-concavity poses technical challenges, leading to some results requiring averaged iterates and/or having unstable guarantees as $T$ grows. One approach for bypassing these challenges is to consider “modified” algorithms which are run on a quadratically regularized (and therefore strongly convex) potential [26]; this has led to state-of-the-art guarantees for Wasserstein mixing [28] and MALA [36]. For further details and history, see, e.g., [36, Table 1] and [21, Table 2].

For constrained settings, much less is known. Indeed, the only mixing time result we are aware of for the Projected Langevin Algorithm is [14], which in a tour-de-force analysis shows a polynomial mixing time to $\pi$, albeit with large polynomial factors like $d^{12}$. Faster upper bounds for mixing to $\pi$ were recently shown for different algorithms, notably the Proximal Langevin Algorithm [34, 12] and the Mirror Langevin Algorithm [1]. The paper [14] states two open questions for the Projected Langevin Algorithm: extend their mixing results to mini-batch gradients, and obtain optimal dependence on the dimension and accuracy. We accomplish both goals for $\pi_{\eta}$, and hope that this will help pave the way towards the answers for $\pi$.

The present paper fits into this line of work by establishing tight mixing times to $\pi_{\eta}$ for the (Projected Stochastic) Langevin Algorithm in a variety of settings: constrained/unconstrained, log-concave/strongly log-concave, and exact/stochastic gradients. An appealing aspect of our analysis technique is that it extends to all these settings with essentially no change. An appealing aspect of our results on mixing to $\pi_{\eta}$ is that they clarify that the dimension and error dependencies for mixing to $\pi$ arise only through a single factor of $1/\eta$, see the discussion following Theorem 1.2.

Whereas the study of lower bounds is well developed in the field of optimization, much less is known for sampling. We refer the reader to the recent paper [24] (which settles this question in dimension $1$ for unconstrained, smooth, strongly log-concave sampling) for a discussion of this gap (in dimensions greater than $1$) and the technical challenges for overcoming it. Because of these challenges, existing sampling lower bounds are primarily algorithm-specific, e.g., for the Underdamped Langevin Algorithm [15], or the Unadjusted Langevin Algorithm [23], or the Metropolis Adjusted Langevin Algorithm [91, 56, 22].

The present paper fits into this line of work on algorithm-specific lower bounds by establishing tight lower bounds for the Langevin Algorithm to mix to $\pi_{\eta}$ in the log-concave setting. These are the first tight lower bounds for the mixing time of the (Unadjusted) Langevin Algorithm (to either $\pi_{\eta}$ or $\pi$). See Footnote 6.

Much of the literature on sampling is inspired by similarities to optimization. These parallels are particularly strong between the Langevin Algorithm and Gradient Descent—the simplest and most canonical first-order algorithms for sampling and optimization, respectively. Broadly speaking, there are two approaches for extending this connection from algorithm definition to analysis. One approach views the Langevin Algorithm as Gradient Descent plus noise, and borrows classical techniques from optimization over $\mathbb{R}^{d}$ to prove sampling bounds—indeed this was the motivation of the original non-asymptotic bounds in [26], see the discussion in [27, §3]. The other approach links sampling on $\mathbb{R}^{d}$ not with optimization on $\mathbb{R}^{d}$, but rather with optimization over the lifted space of probability measures $\mathcal{P}(\mathbb{R}^{d})$ endowed with the geometry of the $2$-Wasserstein metric. This connection is based on the celebrated JKO Theorem [52], which states that the Langevin Diffusion is the gradient flow on $\mathcal{P}(\mathbb{R}^{d})$ with respect to the objective functional $\mathsf{KL}(\cdot\;\|\;\pi)$. This connection is strengthened by the fact that properties of $\pi$ such as log-concavity, strong log-concavity, and log-Sobolev inequalities imply properties of the functional $\mathsf{KL}(\cdot\;\|\;\pi)$ such as convexity, strong convexity, and gradient domination inequalities, respectively. A fruitful line of recent work has exploited this connection by analyzing the Langevin Diffusion (and its variants) via optimization techniques on Wasserstein space, see e.g., [89, 9, 35, 6, 67].

The present paper fits into this interface of optimization and sampling by providing a new way to extract quantitative bounds from the first approach (linking sampling and optimization over $\mathbb{R}^{d}$). Specifically, on the analysis side, we show how the Privacy Amplification by Iteration technique implies tight mixing bounds for the Langevin Algorithm by exploiting the foundational fact from optimization that a Gradient Descent step is contractive (but in a different way from [26, 27]). This leads to error bounds for sampling that in many ways mirror classical results in optimization, see Table 2.

The Langevin Algorithm is identical to Noisy-SGD—the standard algorithm for large-scale machine learning with differential privacy (DP)—modulo a different noise scaling and possibly gradient clipping. What it means for Noisy-SGD to satisfy DP is similar, at least intuitively, to what it means for the Langevin Algorithm to mix. Essentially, DP means that the algorithm yields indistinguishable outputs when run from identical initializations on different objectives, whereas mixing means that the algorithm yields indistinguishable outputs when run from different initializations on identical objectives. These syntactic similarities suggest that analyses might carry over if appropriately modified, and indeed this cross-pollination has proven fruitful in both directions. In one direction, ideas from sampling have recently been used to achieve tight DP guarantees for Noisy-SGD in certain regimes [25, 80, 92]. In the other direction, the technique of adaptive composition from DP has recently been used to prove the first mixing bounds for the Langevin Algorithm in the Rényi divergence [45, 42].

The present paper fits into this interface of sampling and DP in the sense that we introduce to the samping literature a new analysis technique from the DP literature (namely, Privacy Amplification by Iteration), and we show that this gives optimal mixing bounds for the Langevin Algorithm.

As mentioned earlier, the technique of Privacy Amplification by Iteration was first introduced in [43]. It has been used to derive fast algorithms with optimal privacy-utility tradeoffs for stochastic convex optimization [44]. [7] considered generalizations of this approach, and in particular showed a strengthening for strongly convex losses. [5] prove a version of Privacy Amplification by Iteration that holds for bounded sets by using contraction coefficients for hockey-stick divergence, and [85] extend these results to shuffled SGD. Recent work by the authors [3] provides a diameter-aware version of Privacy Amplification by Iteration that also incorporates Privacy Amplification by Sampling for the Sampled Gaussian Noise Mechanism. This is the version of Privacy Amplification by Iteration exploited in this paper.

The present paper fits into this line of work by showing how this technique, previously only used in the privacy literature, also enables understanding questions in the sampling literature.

Consider running the Langevin Algorithm for $T$ iterations from two different initializations $X_{0}\sim\mu_{0}$ and $X_{0}^{\prime}\sim\pi_{\eta}$, where $\mu_{0}$ is an arbitrary probability distribution supported on $\mathcal{K}$, and $\pi_{\eta}$ is the stationary distribution of the algorithm. Then by definition of the algorithm’s updates, we can couple the random batches $\{B_{t}\}_{t=0}^{T-1}$ and the noise $\{Z_{t}\}_{t=0}^{T-1}$ so that

	$\displaystyle X_{t+1}$	$\displaystyle=\Pi_{\mathcal{K}}\left[X_{t}-\frac{\eta}{b}\sum_{i\in B_{t}}\nabla f_{i}(X_{t})+Z_{t}\right]$
	$\displaystyle X_{t+1}^{\prime}$	$\displaystyle=\Pi_{\mathcal{K}}\left[X_{t}^{\prime}-\frac{\eta}{b}\sum_{i\in B_{t}}\nabla f_{i}(X_{t}^{\prime})+Z_{t}\right]$

where $Z_{t}\sim\mathcal{N}(0,2\eta I_{d})$ are drawn independently.

Define auxiliary sequences $\{Y_{t}\}_{t=0}^{T}$ and $\{Y^{\prime}_{t}\}_{t=0}^{T}$ as follows. Initialize $Y_{0}=X_{0}$ and $Y^{\prime}_{0}=X^{\prime}_{0}$. Then update via

	$\displaystyle Y_{t+1}$	$\displaystyle=\phi_{t+1}(Y_{t})+Z_{t}$
	$\displaystyle Y^{\prime}_{t+1}$	$\displaystyle=\phi_{t+1}(Y^{\prime}_{t})+Z_{t}$

using the same update function, which is defined as

$\displaystyle\phi_{t+1}(y)$

$\displaystyle=\Pi_{\mathcal{K}}\left[y\right]-\frac{\eta}{b}\sum_{i\in B_{t}}\nabla f_{i}(\Pi_{\mathcal{K}}\left[y\right]).$

Observe that by induction, these auxiliary processes satisfy $X_{t}=\Pi_{\mathcal{K}}\left[Y_{t}\right]$ and $X_{t}^{\prime}=\Pi_{\mathcal{K}}\left[Y_{t}^{\prime}\right]$.

In order to show that these two auxiliary processes are $1$-CNI (henceforth simply called CNI), it suffices to check that the stochastic update functions $\phi_{t}$ are contractive, as shown next.

Since we have established that the two auxiliary sequences $\{Y_{t}\}_{t=0}^{T}$ and $\{Y_{t}^{\prime}\}_{t=0}^{T}$ are CNI with identical update functions, Proposition 2.10 with $\alpha=1$ implies

$$\mathsf{KL}\left(Y_{T}\;\|\;Y_{T}^{\prime}\right)\leqslant\frac{D^{2}}{4\eta T}\,.$$

Now make two observations. First, by the data-processing property of the KL divergence (Lemma 2.6) and the fact that $X_{T}=\Pi_{\mathcal{K}}\left[Y_{T}\right]$ and $X^{\prime}_{T}=\Pi_{\mathcal{K}}\left[Y^{\prime}_{T}\right]$, we have $\mathsf{KL}\left(X_{T}\;\|\;X_{T}^{\prime}\right)\leqslant\mathsf{KL}\left(Y_{T}\;\|\;Y_{T}^{\prime}\right)$. Second, by stationarity of $X_{0}^{\prime}\sim\pi_{\eta}$, we have by induction that $X_{T}^{\prime}\sim\pi_{\eta}$. It follows that

$\displaystyle\mathsf{KL}\left(X_{T}\;\|\;\pi_{\eta}\right)=\mathsf{KL}\left(X_{T}\;\|\;X_{T}^{\prime}\right)\leqslant\mathsf{KL}\left(Y_{T}\;\|\;Y_{T}^{\prime}\right)\leqslant\frac{D^{2}}{4\eta T}\,.$

(3.2)

Therefore the KL divergence is at most $1/8$ after $T=2D^{2}/\eta$ iterations. Since a KL divergence bound of $1/8$ implies a TV bound of $1/4$ by Pinkser’s inequality (Lemma 2.5), we conclude that

$\displaystyle T_{\mathrm{mix},\,\mathsf{TV}}(1/4)\leqslant\frac{2D^{2}}{\eta}\,.$

(3.3)