Faster high-accuracy log-concave sampling via algorithmic warm starts

Introduction

Yet, despite several decades of progress, many fundamental theoretical questions remain open about the complexity of sampling. Arguably one of the foremost questions in this field is:

This problem of determining the query complexity for sampling has remained open even for the canonical and seemingly simple class of strongly log-concave and log-smooth (in brief, “well-conditioned”) distributions $\pi$, let alone in more complicated settings. It is worth emphasizing that this state of affairs for sampling is in sharp contrast to that for optimization—indeed, the analogous query complexity questions for convex optimization were solved long ago in celebrated results from the 1980s [NY83, Nes18].

Within the literature, it is of central interest to understand this complexity question in the high-accuracy regimeThroughout, we use the standard terminology low accuracy to refer to complexity results which scale polynomially in $1/\varepsilon$, and the term high accuracy for results which scale polylogarithmically in $1/\varepsilon$; here, $\varepsilon$ is the desired target accuracy. These two regimes require different algorithms and analyses, as explained in the sequel. , since classical high-accuracy samplers such as the Metropolis-adjusted Langevin algorithm (MALA) and the Metropolized Hamiltonian Monte Carlo algorithm (MHMC) remain the de facto gold standard in practice. Yet the complexity for this high-accuracy setting has been particularly difficult to pin down, as we explain shortly.

The purpose of this paper is to develop faster high-accuracy samplers, and in doing so move towards a better understanding of the first-order complexity of sampling. For simplicity of exposition, let us presently assume that $\pi$ is well-conditioned, since by the proximal reduction framework [LST21a, Che+22], it is known that improvements to the complexity of well-conditioned sampling lead to improvements in more general settings such as when $\pi$ is (non-strongly) log-concave, or even non-log-concave but satisfies standard isoperimetric assumptions such as the log-Sobolev or Poincaré inequality. (Indeed, our results improve upon the state-of-the-art for all these settings.)

A central motivation of this paper is the large gap between (our current understanding of) the complexity of low-accuracy samplers and high-accuracy samplers. To explain this gap, let us briefly provide relevant background on both classes of algorithms.

Low-accuracy samplers arise as discretizations of stochastic processes with stationary distribution $\pi$, such as the Langevin diffusion [[, the sampling analog of the gradient flow, see]]jordan1998variational, wibisono2018sampling or the underdamped Langevin diffusion [[, the sampling analog of the accelerated gradient flow, see]]ma2021there. Once discretized, however, the resulting discrete-time Markov chain is typically biased, i.e., its stationary distribution is no longer equal to $\pi$. In order to control the size of the bias, the step size of the algorithm is chosen to scale polynomially with $\varepsilon$, and hence the overall running time scales polynomially with $1/\varepsilon$. Despite this drawback, the discretization analysis is by now well-understood, with state-of-the-art results achieving a complexity of $\widetilde{O}(d^{1/3}/\varepsilon^{2/3})$ [SL19, FLO21, BM22]; see [CLW21] for a discussion of tightness.

High-accuracy samplers, in contrast, are typically designed in such a way that there is no bias. This is achieved by, e.g., appending a Metropolis–Hastings filter to each step (see Appendix D.1 for background). Common examples of these algorithms include MALA and MHMC, which are routinely deployed in large-scale applications and are the default implementations of sampling routines in many modern software packages [GLG15, Aba+16a]. However, the filter which debiases the algorithm also greatly complicates the analysis, and thus far the best complexity result for these algorithmsWe discuss the result of [LW22] for the zigzag sampler further in §1.3. is $\widetilde{O}(d\log^{O(1)}(1/\varepsilon))$ [Dwi+18, Che+20, LST20]. Note that the dimension dependence of this result is substantially worse than what is known in the low-accuracy regime and is at odds with the popularity of high-accuracy samplers in practice.

A promising first step towards resolving this gap was put forth in [Che+21] and later refined in [WSC22]: when initialized from a warm start (i.e., a measure $\mu_{0}$ with $\chi^{2}(\mu_{0}\mathbin{\|}\pi)\leq O(1)$), the complexity of MALA improves to $\widetilde{O}(d^{1/2}\log^{2}(1/\varepsilon))$ since it can safely take much larger step sizes (of size $d^{-1/2}$ rather than $d^{-1}$). This raises the natural question: is the warm start condition merely an artefact of the analyses? Rather surprisingly, it was shown in [LST21] that there exist bad initializations for MALA for which the dimension dependence is at least $\widetilde{\Omega}(d)$. Taken together, these results show that the complexity of MALA fundamentally hinges on the warmness of its initialization.

The key question is thus: can such a warm start be obtained algorithmically? Or more precisely:

Is there an algorithm which makes $\widetilde{O}(d^{1/2})$ queries to a first-order oracle for $f$

and outputs a measure $\mu_{0}$ with $\chi^{2}(\mu_{0}\mathbin{\|}\pi)\leq O(1)$?

The requirement that the algorithm makes $\widetilde{O}(d^{1/2})$ queries is essential, else the cost of obtaining the warm start dominates the subsequent cost of running MALA. Yet this was the state of affairs—previously, the fastest algorithms took significantly longer to produce a warm start than to actually use it, defeating the purpose of the warm start. Resolving this discrepancy has been posed as an important question in many papers, e.g., [Che+21, LST21, Che+22a, LW22, WSC22].

The main challenge for answering this warm start question is that the chi-squared divergence is quite a strong performance metric. (We emphasize that it is essential to obtain the warm start in the chi-squared divergence, or more generally in a Rényi divergence $\mathcal{R}_{q}$ of order $q>1$, rather than other common metrics such as total variation, Wasserstein, or KL divergence; see §1.2 for an in-depth discussion.) The aforementioned results in the low-accuracy regime fall short of achieving this goal, since they only hold in the Wasserstein metric (for which standard coupling arguments are readily available). Despite significant effort, the best known guarantee for producing a warm start—achieved by the Langevin Monte Carlo (LMC) algorithm [Che+22a]—is far too costly as it requires $\widetilde{O}(d)$ queries, which defeats the purpose of the warm start.

Towards this hope of algorithmic warm starts, [WSC22] made the promising empirical observation that MALA mixes much faster if it is initialized at the output of the underdamped Langevin Monte Carlo (ULMC) algorithm. However, they left open the question of rigorously proving that this yields a warm start. While it is widely believed that ULMC is substantially faster than LMC, the previous best results for computing a warm start with ULMC had dimension dependence $\widetilde{O}(d^{5/2})$ (implicit from [GT20]) or very recently $\widetilde{O}(d^{2})$ (implicit from [Zha+23]), see the prior work section §1.3 for details. We emphasize that this dimension dependence is not only a far cry from the elusive $\widetilde{O}(d^{1/2})$ goal, but moreover is even worse than known results for the simpler LMC algorithm. Unfortunately, any improvement to these ULMC warm start bounds appears to require overcoming fundamental difficulties with studying hypocoercive differential equations which remain unsolved today, despite being the focus of intensive research activity within the PDE community since the work of Kolmogorov [Kol34]. For a further discussion of these technical obstacles, see §1.2.

1 Contributions

In this paper, we develop techniques which bypass longstanding challenges for analyzing hypocoercive dynamics, thereby establishing the first $\widetilde{O}(d^{1/2})$ Rényi mixing results for ULMC. This resolves the aforementioned warm start conjecture, which has been raised in a number of prior works, e.g., [Che+21, LST21, Che+22a, LW22, WSC22]. As discussed above, this enables us to design significantly faster high-accuracy samplers—both for the log-concave setting and far beyond. Finally, this also closes the long line of work devoted to understanding the complexity of MALA (see Table 1). We present our results in more detail below, and then discuss our new techniques in §1.2.

Our first main result is an improvement of the state-of-the-art Rényi mixing bounds for ULMC from $\widetilde{O}(d^{2})$ to $\widetilde{O}(d^{1/2})$. This resolves the warm start question in the affirmative. We remark that although the warm start problem was stated above for $\chi^{2}$ convergence, our result actually holds more generally for Rényi divergences $\mathcal{R}_{q}$ of any order $q\geq 1$, and thus we state it as such. (For the purpose of warm starts, it suffices to take $q=2$ since $\chi^{2}=\exp(\mathcal{R}_{2})-1$ is of constant size when $\mathcal{R}_{2}$ is.) Below, $\alpha$ and $\beta$ denote the strong log-concavity and log-smoothness bounds; their ratio $\kappa:=\beta/\alpha$ is the condition number.

As we detail in §1.2, the main barrier to obtaining this result is that the underdamped Langevin dynamics falls within a class of PDEs known as hypocoercive equations, for which fundamental questions remain unresolved.Implications of our techniques for the analysis of hypocoercive diffusions are explored in shortly forthcoming work.

Theorem 1.1 provides the first algorithm for computing warm starts that is not significantly slower than the use of the warm start. This enables us to exploit, for the first time, the recent breakthroughs on MALA [Che+21, WSC22] which improve the complexity of MALA from $\widetilde{O}(d)$ to $\widetilde{O}(d^{1/2})$ from a warm start.We remark that all of our results could replace MALA with the zigzag algorithm [LW22]. Indeed, the zigzag sampler has the same key issue as MALA: it requires a warm start in chi-squared divergence for the known $d^{1/2}$ mixing result to apply. However, we focus on MALA because MALA’s robust empirical performance has made it a central focus of study in the MCMC literature for nearly three decades [Bes+95]. By combining this with additional algorithmic tools for improving the dependence on the condition number, we obtain our second main result, which substantially advances the state-of-the-art for high-accuracy log-concave sampling.

The algorithmic warm start result of Theorem 1.1 confirms the aforementioned empirical conjecture of [WSC22] and provides the final missing piece in our understanding of the complexity of MALA, closing the line of work developed in [RR98, Dwi+18, Che+20, LST20, Che+21, LST21, WSC22] (Table 1). Indeed, due to matching lower bounds in [Che+21, LST21, WSC22], the complexities $\widetilde{O}(\kappa d)$ and $\widetilde{O}(\kappa d^{1/2})$ with or without a warm start are known to be tight, and hence the key remaining question was whether a warm start is actually efficiently computable.

The complexity in Theorem 1.2 constitutes a natural barrier for high-accuracy sampling. Indeed, regarding the dimension dependence, any further progress beyond $\widetilde{O}(d^{1/2})$ would seem to require completely different algorithms—both for obtaining a warm start and also for exploiting a warm start. For example, the $\widetilde{O}(d^{1/2})$ complexity of MALA is unimprovable even under arbitrarily warm starts [Che+21, WSC22]. And regarding the condition number dependence, any further progress beyond $\widetilde{O}(\kappa)$ in the high-dimensional regimeAnalogous to classical optimization results, there are sampling algorithms which achieve logarithmic dependence on $\kappa$ at the expense of larger polynomial dependence on $d$. The open question mentioned here is really: can one improve the condition dependence beyond near-linear while also maintaining comparable dimension dependence? would constitute a major breakthrough in the complexity of sampling since it is currently unknown whether an acceleration phenomenon holds in the sampling context.

More broadly, our result provides evidence of the potential for designing faster high-accuracy samplers by combining low-accuracy samplers for computing a warm start, together with improved high-accuracy mixing from the warm start. We believe that this research program may be crucial for future progress in high-accuracy sampling, since faster mixing from a warm start seems likely to hold for other Metropolized algorithms. See §6 for further discussion in this direction.

High-accuracy log-concave sampling is the key to obtaining state-of-the-art complexity results for a wide class of distributions beyond log-concavity. This is achieved by using our faster log-concave sampler in Theorem 1.2 to improve the per-iteration complexity of the proximal sampler [LST21a, Che+22]. This approach is overviewed in the techniques section §1.2, and leads to the following result.

For each of the following classes of distributions, we obtain complexity bounds which improve by a factor of $d^{1/2}$ over the state-of-the-art results in [Che+22]:

$\pi$ is log-smooth and weakly log-concave.

$\pi$ is log-smooth and satisfies a log-Sobolev inequality.

$\pi$ is log-smooth and satisfies a Poincaré inequality.

The latter two assumptions of log-Sobolev and Poincaré—called functional inequalities—capture strictly richer classes of target distributions than strong-log-concavity. There are two major motivations for studying the complexity of sampling in this setting. First, functional inequalities are quite flexible, as they are preserved under common operations such as bounded perturbations and Lipschitz mappings (see §A). Consequently, they often capture the breadth of settings encountered in practice, including non-log-concave settings. Second, these functional inequalities classically imply convergence of diffusions in continuous time, making them natural assumptions under which to study the corresponding discretizations.

Despite the appeal of this program, proving sampling guarantees under functional inequalities introduces a number of additional technical complications and was only accomplished recently, starting with [VW19] and continued in the works [Wib19, LE20, Ma+21, Che+22a]. Our result continues this line of work, and in particular highlights the use of high-accuracy samplers for well-conditioned distributions as a powerful algorithmic tool for the broader problem of sampling under isoperimetry.

2 Challenges and techniques

To explain what properties are needed for a warm start requires first explaining why a warm start helps. Briefly, the complexity of MALA is governed by the largest possible step size for which the algorithm still accepts a reasonable fraction of the proposals (see Appendix D.1 for background on MALA). The basic reason why we might expect to improve the complexity of MALA from $\widetilde{O}(d)$ to $\widetilde{O}(d^{1/2})$ is that at the stationary distribution $\pi$, the step size can be increased significantly from $d^{-1}$ to $d^{-1/2}$ while keeping the acceptance probability high. More precisely, with step size $d^{-1/2}$, the acceptance probability is large from a typical point from $\pi$; however, it can be exponentially small in regions that are atypical (i.e., exponentially rare under $\pi$). The existence of such regions implies that there are “bottlenecks” in the state space which take exponentially long to traverse. The role of a warm start initialization is to avoid such bottlenecks.

In other words, a key property that a warm start $\mu_{0}$ must satisfy is that if $\pi$ assigns exponentially small probability to an event, then so must $\mu_{0}$. Crucially, this property does not hold if $\mu_{0}$ is only known to be close to $\pi$ in common probability metrics such as total variation, Wasserstein, or KL divergence—but this property does hold if $\mu_{0}$ is close to $\pi$ in the chi-squared divergence, or more generally any Rényi divergence $\mathcal{R}_{q}$ of order $q>1$.This is the same reason why differential privacy requires guarantees in Rényi divergences [Mir17].

In the preceding discussion, taking large step sizes from a non-warm initialization was problematic due to the rejections in the Metropolis–Hastings filter step. A natural idea, then, is to remove the filter for the initial stage of the algorithm and later reinstate it when the law of the iterate is closer to the target $\pi$. Since the proposal of MALA is just one step of the LMC algorithm, this amounts to using LMC to procure the warm start. More generally, we can consider using any low-accuracy sampler as our warm start algorithm, and indeed, as we discuss next, it will be crucial to consider ULMC instead of LMC in order to achieve the desired $\widetilde{O}(d^{1/2})$ dimension dependence.

At a high level, if we discretize a diffusion with step size $h$ for continuous time $T$, then the total number of iterations is $N=T/h$. In order to understand the dimension dependence of the algorithm, one must therefore understand both $h$ and $T$. These two terms reflect two distinct aspects of mixing analysis: the discretization bias and the convergence time.

The first part—the discretization bias—is now relatively well-understood (see the prior work discussions in §1.3), even for the chi-squared divergence and more general Rényi divergences. In particular, it is known that the Rényi bias of LMC is controlled for step sizes $h\lesssim 1/(dT)$, and the Rényi bias of ULMC is controlled for step sizes $h\lesssim 1/\sqrt{dT}$. (In fact, we streamline arguments in the literature in order to provide a shorter and simpler proof of this in Appendix C.3.) Since the Langevin diffusion does not reach approximate stationarity until time $T\geq\Omega(\log d)$, it follows that LMC requires at least $N=T/h=\widetilde{\Omega}(d)$ iterations, which is too slow for our purposes.

ULMC is more promising, as the discretization bounds lead to iteration complexity bounds of $N=T/h=d^{1/2}T^{3/2}$. However, in order to reach our warm start goal of $N=\widetilde{O}(d^{1/2})$, this means that the convergence time $T$ must be nearly dimension-free, i.e., of size $\widetilde{O}(1)$.

Since the Rényi divergence to $\pi$ initially scales as $\widetilde{\Theta}(d)$, in order to obtain nearly dimension-free bounds on $T$, we require the diffusion to converge to stationarity in Rényi divergence with an exponential rate. This is a strong property of the diffusion, which we call hyperequilibration.

Hyperequilibration was not even known for the simpler (standard, overdamped) Langevin diffusion (LD) until quite recently [CLL19, VW19]. While a spectral gap for LD (or equivalently, a Poincaré inequality for $\pi$) classically implies exponential decay of the chi-squared divergence, this is far weaker than hyperequilibration. Indeed, hyperequilibration requires exponential decay of $\mathcal{R}_{2}$, which amounts to doubly exponential decay of the chi-squared divergence, since $\mathcal{R}_{2}=\log(\chi^{2}+1)$. Under the stronger assumption of a log-Sobolev inequality for $\pi$, it is well-known that the KL divergence decays exponentially fast, but it was unclear that the same holds for the Rényi divergence which, as discussed above, is crucial for warm starts. It was only through the inspired semigroup calculations of [CLL19, VW19] that we now know this to be true, namely, a log-Sobolev inequality implies hyperequilibration for LD.This explains our choice of the terminology hyperequilibration: it is inspired by the analogy to the classical property of hypercontractivity, which is equivalent to the logarithmic Sobolev inequality (LSI) [Gro75].

Recall, though, that the LD incurs too much discretization bias. To obtain sufficient control over both the discretization bias and the convergence time, we therefore need to establish hyperequilibration for the underdamped Langevin diffusion (ULD). However, this question brings us to longstanding challenges from the theory of hypocoercive PDEs.

To recap: for LD, we have exponential decay of the chi-squared divergence under a Poincaré inequality, exponential decay of the KL divergence under a log-Sobolev inequality, and finally hyperequilibration under a log-Sobolev inequality. What, then, are the analogous results for ULD? Since its introduction in the 1930s by Kolmogorov [Kol34], the regularity and convergence of ULD have been the focus of intensive research. It took nearly half a century to establish mixing [Tro77], and a further 30 years and Villani’s “slightly miraculous-looking computations” [Vil09, pg. 42] to prove exponential decay of the KL divergence under a log-Sobolev inequality. Establishing hyperequilibration for ULD remains out of reach for existing techniques.

The reason for this sudden jump in difficulty from the overdamped to the underdamped diffusions is due to a fundamental issue: the degeneracy of ULD. In brief, whereas LD is driven by a full-dimensional Brownian motion, ULD is driven by a degenerate one which is only added to a subset of the coordinates. For sampling purposes, this degeneracy is a desirable feature as it leads to smoother sample paths and smaller discretization error; however, this same degeneracy is also the source of deep questions in PDE theory which have motivated research in that field for nearly a century. The key challenge here is that the standard tools of Markov semigroup theory—which provide the backbone of the analysis for LD—completely break down for ULD. To address this difficulty, the theory of hypocoercivity, inspired by Hörmander’s groundbreaking work on hypoellipticity [Hör67], was laid down by Villani in the monograph [Vil09] as a principled framework for the study of degenerate diffusions. However, this is still a relatively nascent area of PDE and many important questions remain wide open; see the prior work in §1.3 for further background.

In contrast, we note that it is well-known how to obtain fast rates of convergence in the Wasserstein metric via standard coupling arguments. Consequently, the state-of-the-art $\widetilde{O}(d^{1/2})$ guarantees for the ULMC algorithm hold in the Wasserstein metric or the KL divergence [Che+18, SL19, DR20, Ma+21, Zha+23], whereas for Rényi divergence bounds, it was previously unknown how to obtain rates which are better than even $\widetilde{O}(d^{2})$.

2.2 Settling the warm start conjecture: regularization via privacy

To settle the warm start conjecture, we adopt a fundamentally different perspective. Namely, instead of trying to directly establish hyperequilibration via hypocoercivity techniques, we ask whether it can be deduced from simpler Wasserstein coupling arguments. At the heart of this approach is the fact that diffusions often enjoy strong regularizing properties, which allow for bounding stronger metrics (e.g., Rényi) in terms of weaker ones (e.g., Wasserstein). Such regularization results are typically established for continuous-time diffusions via abstract calculus methods, such as the theory of Markov semigroups [BGL14]. However, as discussed above, these techniques do not extend to ULD due to the fundamental issue of degeneracy.

Our key insight is to prove a regularization result for the discrete-time algorithm directly. This is enabled by the fact that although the noise added to each iteration of ULMC is nearly degenerate—and indeed degenerates as the step size $h\searrow 0$, as it must because ULD is degenerate—this ULMC noise remains non-degenerate for any positive step size $h>0$. Hence, we can expect some mild amount of regularization for ULMC, a fact that we establish for the first time. On a technical level, we accomplish this via a more sophisticated version of techniques from the differential privacy literature—namely, the shifted Rényi analysis—which we describe next.

The regularization result we seek is of the following form: if we initialize two copies of our process of interest at the distributions $\mu_{0}$, $\nu_{0}$, and arrive at distributions $\mu_{n}$, $\nu_{n}$ respectively at iteration $n$, we wish to control $\mathcal{R}_{q}(\mu_{n}\mathbin{\|}\nu_{n})$ in terms of an initial Wasserstein distance $W(\mu_{0},\nu_{0})$. In our application, the process of interest—namely ULMC—is an instance of what is sometimes called a “contractive noisy iteration” (CNI): an algorithm that interleaves Lipschitz mappings with (Gaussian) noise convolution steps. This notion of a Contractive Noisy Iteration is of broad interest as it captures algorithms in differential privacy (e.g., noisy optimization algorithms) and in sampling (e.g., discretizations of diffusions), and we therefore place our results in a framework which encompasses these various use cases.

A generalization of the regularization result we seek is to prove that for a CNI,

where $\mathcal{R}_{q}^{(w)}$ is the shifted Rényi divergence, defined as

see §3 for details. Indeed, if we take $w=W(\mu_{0},\nu_{0})$ in ($\star$ ‣ 1.2.2), then the term $\mathcal{R}_{q}^{(w)}(\mu_{0}\mathbin{\|}\nu_{0})$ vanishes, and we will have controlled $\mathcal{R}_{q}(\mu_{n}\mathbin{\|}\nu_{n})$ in terms of $W(\mu_{0},\nu_{0})$ as desired. However, ($\star$ ‣ 1.2.2) is more general, as it allows for carefully tracking the shift parameter $w$ throughout. This proof technique, called shifted divergence analysis, was first introduced in the context of differential privacy by [Fel+18] for the purpose of establishing Privacy Amplification by Iteration, and was recently honed into a form amenable to sampling analyses in [AT22, AT22a].

A subtle yet essential technical issue that arises in establishing ($\star$ ‣ 1.2.2) is: which Wasserstein metric $W$ do we use? All previous versions of ($\star$ ‣ 1.2.2) required the $W_{\infty}$ metric, which is problematic for our setting as the $W_{\infty}$ metric is infinite at initialization. Here, our main insight is to use a non-standard Wasserstein metric, called the Orlicz–Wasserstein metric, based on the sub-Gaussian Orlicz norm. As we discuss in Remark 3.9, this is exactly the right metric to use: in fact, ($\star$ ‣ 1.2.2) cannot hold for any weaker metric (e.g., $W_{p}$ for any finite $p$), and the initialization bound cannot be finite for any stronger metric. We then show that for Orlicz–Wasserstein shifts, ($\star$ ‣ 1.2.2) indeed holds, with the caveat that the order of the shifted Rényi divergence on the right-hand side of ($\star$ ‣ 1.2.2) is increased. This increase in the order also means that additional care is required when applying ($\star$ ‣ 1.2.2), as the inequality cannot be iterated too many times, but we bypass this issue by showing that it suffices to only exploit the regularization from a single step.

Finally, we note that our analysis answers the open question raised in [AT22a] of how to use the shifted divergence technique in order to obtain sampling guarantees for discretized diffusions w.r.t. the true target distribution $\pi$, rather than w.r.t. the biased limit of the algorithm.

2.3 From warm starts to faster high-accuracy samplers

In light of the discussion thus far, combining our warm start result with the recent advances on MALA [Che+21, WSC22] immediately improves the dimension dependence of high-accuracy log-concave sampling to $\widetilde{O}(d^{1/2})$. However, two further issues remain. First, thus far we have ignored the dependence on the condition number $\kappa$ for simplicity of exposition, but the combined approach of ULMC and MALA incurs suboptimal dependence on $\kappa$, namely $\kappa^{3/2}$ rather than $\kappa$. Second, the result only holds for strongly log-concave targets. We address both of these issues simultaneously by adding a third algorithmic building block: the proximal sampler. Below, we briefly overview the proximal sampler and the final remaining technical challenges in its application.

If $f$ is $\beta$-smooth, and the step size $h$ is chosen as $h\asymp\frac{1}{\beta}$, one can check that $\pi^{X\mid Y=y}$ is strongly log-concave and log-smooth with condition number $O(1)$. Hence:

The requirement of sampling from the RGO to high accuracy arises to avoid accumulation of the errors from inexact implementation of the RGO.

So far, we have not made use of any assumptions on $\pi$ beyond smoothness of $f$. Additional assumptions on $\pi$, such as log-concavity, can then used to control the number of outer loops. This program was carried out in [Che+22], which carefully studied the outer loop complexity of the proximal sampler under a variety of assumptions on the target $\pi$ which, when combined with the implementation of the RGO via existing high-accuracy samplers, yielded state-of-the-art complexity bounds for sampling under those assumptions. Our faster high-accuracy log-concave sampler provides a better implementation of the RGO, and hence we improve upon these prior results by a factor of roughly $d^{1/2}$ in each setting. Moreover, in the strongly log-concave setting, the number of outer iterations of the proximal sampler is shown to be $\widetilde{O}(\kappa)$ [LST21a], so using ULMC + MALA to implement the RGO boosts the condition number dependence of the overall sampler to near-linear. This resolves the two issues described above, but in doing so we must also develop an inexact error analysis for the proximal sampler.

In order to apply the proximal reduction framework, we must understand how the error from inexact implementation of the RGO propagates into the final sampling error. This was carried out in [LST21a] for the TV distance via a simple coupling argument, which amounts to a union bound over failure events at each iteration. Similarly, it is straightforward to carry out the inexact error analysis in the Wasserstein metric due to the availability of the triangle inequality. However, to establish our guarantees in §5, which hold also in the KL and $\chi^{2}$ metrics, we must perform an error analysis in $\chi^{2}$ (or equivalently, in Rényi). This is also complicated by the fact while the outer loop of the proximal sampler converges exponentially fast in the strongly log-concave setting, which facilitates summing up the geometrically decaying errors from each iteration, the convergence in the weakly log-concave setting does not have an exponential rate and moreover uses a modified Lyapunov functional, changing the nature of the error analysis. We remark that prior works such as [Che+22] did not encounter such issues, since their rejection sampling implementation of the RGO is exact. Therefore, we believe that our inexact error analysis will also be useful for any future applications of the proximal sampler.

We also remark that our application of the proximal sampler, and the ensuing need for careful inexact error analysis, resembles the use of the (accelerated) proximal point method in optimization, e.g., [Fro+15, LMH15].

3 Related work

Rényi guarantees for sampling are relatively recent. Indeed, [VW19] proved fast Rényi mixing for LD and LMC to their respective stationary distributions, and this was translated into Rényi sampling guarantees for LMC in [GT20, Che+21, Che+22a, EHZ22], for the proximal sampler in [Che+22], and for ULMC in [GT20, Zha+23]. These lines of work have led to $\widetilde{O}(d)$ dimension dependence for LMC and the proximal sampler, but for ULMC the rates are much worse, namely $\widetilde{O}(d^{5/2})$ dependence [GT20] and only very recently $\widetilde{O}(d^{2})$ dependence [Zha+23]. In Theorem 4.1, we obtain the first $\widetilde{O}(d^{1/2})$ rate in Rényi.

In contrast, there are many more works which break the $\widetilde{O}(d)$ barrier in the Wasserstein metric: the randomized midpoint discretization of Langevin [HBE20], unadjusted Hamiltonian Monte Carlo (HMC) [CV19], ULMC [Che+18, DR20, Mon21], and more sophisticated discretizations of ULMC and HMC [SL19, FLO21, BM22]. Among these algorithms, at present we only understand how to perform Rényi discretization analysis for ULMC, but for ULMC it is the convergence of the corresponding continuous-time diffusion which remains elusive, as we review next.

The underdamped Langevin diffusion has a rich history, dating back to Kolmogorov [Kol34]. The PDE governing the evolution of its marginal density is referred to as the kinetic Fokker–Planck equation. Unlike the Langevin diffusion, which is driven by a full-dimensional Brownian motion and for which regularity and convergence fall within the purview of classical elliptic and parabolic PDE theory, the underdamped Langevin diffusion is the canonical example of a degenerate diffusion for which these and related questions remain active areas of research within PDE. See §4.1 for background.

The question of regularity for these equations was largely solved by Hörmander [Hör67] in arguably one of the most influential breakthroughs in PDE theory of the last century through the introduction of the theory of hypoellipticity. In turn, it inspired Villani to coin the study of the convergence of such equations hypocoercivity in his seminal monograph [Vil09].

While convergence of this diffusion has been studied for nearly a century, early convergence results were qualitative in nature. It took intensive developments in the PDE community to get to a point where quantitative rates could be extracted, beginning in the 1970s [Tro77]. We do not attempt to comprehensively survey the extensive literature here. We refer to the monograph [Vil09] for history; see also, e.g., the papers [DMS09, Bau17, RS18] for more modern references. We also mention the recent space-time Poincaré approach of [Alb+19, CLW20], which is also directly inspired by Hörmander’s hypoelliptic theory. As we discuss in §1.2, however, all of these approaches fall short of establishing the key property of hyperequilibration.

MALA has been intensely studied over the past three decades since its introduction in [Bes+95], in large part due to its strong practical performance—in fact, it and its variants comprise the default implementations of sampling routines in many modern software packages [GLG15, Aba+16a]. Many classical works studied the geometric ergodicity and asymptotic properties of MALA. With regards to the dimension dependence, particularly influential was the optimal scaling result of [RR98], which showed that taking step size $h\propto d^{-1/3}$ leads to a non-trivial diffusion limit for MALA as $d\to\infty$, at least for product measures $\pi$ satisfying strong regularity assumptions and when initialized at stationarity. Modern analysis techniques have enabled an understanding of the non-asymptotic complexity of MALA [Dwi+18, Che+20, LST20, Che+21, LST21, WSC22], see Table 1 for a summary of the progress in this direction. Our work closes this line of work by showing that the warm start rate of [WSC22], which is tight due to their matching lower bound, is achievable. Moreover, our work provides theoretical justification for the improved empirical performance of MALA after using a low-accuracy algorithm for warm starts, as observed in [WSC22].

The proximal sampler is an algorithmic framework introduced in [TP18, LST21a]. In [LST21a], it was used as a mechanism for boosting the condition number dependence of any high-accuracy log-concave sampler to near-linear, which was then used to design samplers for composite and finite-sum potentials. Then, in [Che+22], it was shown that the proximal sampler reduces the problem of sampling from distributions satisfying weak log-concavity or functional inequalities to the problem of high-accuracy log-concave sampling. In this work, we exploit both these properties of the proximal sampler, and we contribute to its inexact error analysis (see §1.2).

We also mention that in recent work, the proximal sampler has been connected to stochastic localization, leading to recent progress on the KLS conjecture [KP21, CE22, KL22], as well as to diffusion models [Che+23]. There are also applications to sampling from semi-smooth or non-smooth potentials [LC21, LC22, LC22a], and to differential privacy [GLL22, Gop+23, Gop+23a].

The zigzag sampler is an alternative high-accuracy sampler that was recently proposed in [BFR19]. Instead of using a Metropolis–Hastings filter, the zigzag sampler is a piecewise deterministic Markov process which can be implemented without discretization bias. It was recently shown in [LW22] that similarly to MALA, the zigzag sampler has a dimension dependence of $\widetilde{O}(d^{1/2})$ from a warm start. Morever, in [LW22, Corollary 1.4], Lu and Wang show that by using LMC with a large step size to warm start the algorithm, one obtains a high-accuracy log-concave sampler with dimension dependence $\widetilde{O}(d^{4/5})$. Indeed, the same strategy can be used with the warm start results of [Che+21, WSC22] to obtain complexities strictly better than $\widetilde{O}(d)$; however, it is clear that such an approach can never reach the desired complexity of $\widetilde{O}(d^{1/2})$—and in fact there is a fundamental barrier even at $\widetilde{O}(d^{3/4})$ because it is bottlenecked by the discretization bias of LMC. The goal of this paper is achieving $\widetilde{O}(d^{1/2})$ complexity as this is this a natural barrier for high-accuracy samplers given a warm start, and LMC cannot work for this goal.With regards to dimension dependence, running LMC with step size $h$ for $1/h$ steps yields a distribution $\mu$ with $\log\chi^{2}(\mu\mathbin{\|}\pi)\leq\widetilde{O}(dh)$ [Che+22a]. By optimizing the step size $h$ and combining this with the best known complexity $\widetilde{O}(d^{1/2}\log^{3/2}\chi^{2}(\mu\mathbin{\|}\pi))$ of the zigzag sampler, one obtains the final complexity $\widetilde{O}(d^{4/5})$ [LW22]. We point out that even if the complexity of the zigzag sampler were improvable to $\widetilde{O}(d^{1/2}\log\chi^{2}(\mu\mathbin{\|}\pi))$, the total complexity would still be at least $1/h+d^{1/2}\,(dh)\geqslant\widetilde{\Omega}(d^{3/4})$. Thus $d^{3/4}$ is a natural barrier for any warm starting approach using LMC. In direct analogy to our use of MALA, our new complexity result for ULMC (Theorem 4.1) can also be used to warm start the zigzag sampler, leading to the same final complexity bound of $\widetilde{O}(\kappa d^{1/2}\log^{O(1)}(1/\varepsilon))$. This answers the open questions in [LW22] regarding warm starting the zigzag sampler.

Sampling algorithms have been widely used in differential privacy ever since the invention of the exponential mechanism [MT07]; for an exposition of the surrounding history and applications, see the textbook [DR14]. Sampling-inspired analyses have also been recently used to prove privacy properties of optimization algorithms [CYS21, RBP22, YS22]. Most related to this paper are connections in the other direction: the use of techniques from differential privacy in order to analyze sampling. There are two lines of work in this direction. One involves the technique of adaptive composition for Rényi divergences and its use for establishing Rényi bias bounds for LMC and ULMC [GT20, EHZ22, Zha+23]. The other involves the technique of Privacy Amplification by Iteration (PABI), which was originally used to bound the privacy loss of differentially private optimization algorithms [Fel+18, Bal+19, ADP20, FKT20, SBD21, AT22], and its recent use for analyzing the mixing time of LMC to its biased stationary distribution [AT22a]. In this paper, we build upon this technique in several key ways: we show how to improve mixing results for the biased distribution to mixing results for the target distribution, we show how to use these ideas for ULMC rather than LMC, and most importantly we overcome the key issue of unboundedness by replacing $W_{\infty}$ shifts by Orlicz–Wasserstein shifts, see §1.2 for an overview.

4 Simultaneous work

While preparing a draft of this paper for submission, it came to our attention that another group was simultaneously working towards the same problem using completely different techniques and algorithms [FYC23]. We are grateful to them for coordinating simultaneous arXiv submissions. We look forward to reading their paper after it is posted, and we will add a detailed comparison about the differences in a future revision.

5 Organization

We recall preliminaries in §2, especially regarding Rényi divergences. We isolate in §3 our key new technique involving Orlicz–Wasserstein shifted Rényi divergences. We use this technique to obtain faster algorithmic warm starts in §4, and then use these warm starts to develop faster high-accuracy samplers in §5. We conclude in §6 by discussing several future research directions that are motivated by our results. For brevity, we defer various proofs and technical details to the appendices.

Preliminaries

For simplicity of exposition, we assume throughout that we have access to an algorithm for generating independent standard Gaussian random variables. We use the standard notation $\widetilde{O}(g)=g\log^{O(1)}(g)$ to suppress low-order terms. Note that since our final results depend polynomially on the dimension $d$ and condition number $\kappa$, the $\widetilde{O}$ hides polylogarithmic factors in these terms—on the other hand, since $\varepsilon$ occurs only polylogarithmically in our high-accuracy results, we do not hide the polylogarithmic factors in $\varepsilon$.

We say that $f$ is $\alpha$-strongly convex if $\nabla^{2}f\succeq\alpha I_{d}$, and that $f$ is $\beta$-smooth if $\lVert{\nabla^{2}f}\rVert_{\rm op}\leq\beta$. If $f$ is convex, then $f$ is $\beta$-smooth if and only if $\nabla^{2}f\preceq\beta I_{d}$. We always denote by $\kappa$ the condition number $\kappa\coloneqq\beta/\alpha$. If $\pi\propto\exp(-f)$ where $f$ is $\alpha$-strongly convex (resp. $\beta$-smooth), we say that $\pi$ is $\alpha$-strongly log-concave (resp. $\beta$-log-smooth). All other notation is introduced in the main text.

2 Rényi divergences

The Rényi divergence of order $q\in(1,\infty)$ between probability measures $\mu$ and $\nu$ is defined as

if $\mu\ll\nu$, and otherwise is $\infty$. The Rényi divergences of order $q\in\{1,\infty\}$ are defined by continuity.

The Rényi divergence of order $q=1$ coincides with the KL divergence, i.e.,

The Rényi divergence of order $q=2$ is related to the $\chi^{2}$ divergence via the formula

The Rényi divergence of order $q=\infty$ is given by

Our analysis makes use of the following elementary properties of the Rényi divergence. Further details about these properties and their proofs can be found in, e.g., the Rényi divergence survey [VH14] as Theorem 1, Theorem 3, Equation 10, and Remark 1, respectively.

For any Rényi order $q\geqslant 1$, any Markov transition kernel $P$, and any probability distributions $\mu,\nu$,

For any Rényi orders $q^{\prime}\geqslant q\geqslant 1$, and any probability distributions $\mu,\nu$,

For any Rényi order $q\in(1,\infty)$, the corresponding function $\exp((q-1)\;\mathcal{R}_{q}(\cdot\mathbin{\|}\cdot))$ is an $f$-divergence, and thus in particular is jointly convex in its arguments.

We end this section with one last property of Rényi divergences: the weak triangle inequality. The name of this property arises from the fact that although Rényi divergences do not satisfy the triangle inequality, they do satisfy a modified version of it in which the Rényi order is increased and the bound is weakened by a multiplicative factor. Since this property does not appear in the aforementioned survey [VH14] on Rényi divergences, we provide a brief proof for completeness. It can also be found in, e.g., [Mir17, Proposition 11].

For any Rényi order $q>1$, any $\lambda\in(0,1)$, and any probability distributions $\mu,\nu,\pi$,

Expand $\mathcal{R}_{q}(\mu\mathbin{\|}\nu)=\frac{1}{q-1}\log\int fg$ where $f=\mu^{q}/\nu^{q-\lambda}$ and $g=\nu^{q-\lambda}/\pi^{q-1}$, and then apply Hölder’s inequality $\int fg\leqslant(\int f^{a})^{1/a}\,(\int g^{b})^{1/b}$ using Hölder exponents $a=1/\lambda$ and $b=1/(1-\lambda)$. ∎

Improved shifted divergence analysis

In this section we isolate from our analysis a key new technique of independent interest. As overviewed in §1.2, this technique is a strengthening of the “shifted divergence” analysis, a.k.a., “privacy amplification by iteration” (PABI), in which we crucially improve the $\infty$-Wasserstein shift to an Orlicz–Wasserstein shift. This enables obtaining Rényi divergence bounds on the mixing of any Markov chain which interleaves Lipschitz mapping steps (e.g., gradient descent steps) and noise convolution steps (e.g., adding a Gaussian). This notion captures a variety of algorithms from the differential privacy and sampling communities, often called “contractive noisy iterations”.

The main result of this section is formally stated as follows. This result makes use of the Wasserstein metric $W_{\psi_{2}}$ that evaluates a coupling’s quality via the sub-Gaussian Orlicz norm; see §3.1 for background on this notion.

Consider two Markov chains $\{\mu_{n}\}_{n\geqslant 0}$ and $\{\mu_{n}^{\prime}\}_{n\geqslant 0}$ with possibly different initialization, but with the same update transitions

where $P_{n}$ is a Markov transition kernel that is $c$-Lipschitz in the $W_{\psi_{2}}$ metric. Then for any Rényi order $q\geqslant 1$,

so long as N\geqslant\log_{1/c}\Bigl{(}\frac{\sqrt{q\,(q-1)}\,W_{\psi_{2}}(\mu_{0},\mu_{0}^{\prime})}{\sigma\sqrt{2}}\Bigr{)}.

We remark that unlike previous versions of the shifted divergence technique, Theorem 3.1 requires a restriction on the number of iterations $N$.This restriction comes from the fact that with this new Orlicz–Wasserstein shifted Rényi divergence, the new shift-reduction lemma (Lemma 3.7) does not apply to arbitrarily large shifts. But this restriction is mild due to the logarithmic dependence. In fact, it is equivalent to requiring the upper bound in (3.1) to be at most $1/(q-1)$.

The rest of this section is devoted to proving Theorem 3.1. In §3.1 we define a new Lyapunov function, and in §3.2 we use it to prove Theorem 3.1.

Key to our proof of Theorem 3.1 is a new Lyapunov function for tracking how indistinguishable the Markov chains become as they evolve. This new Lyapunov function is a shifted Rényi divergence, but unlike the standard shifted divergence technique, here we measure the shift using an “Orlicz–Wasserstein metric” rather than $W_{\infty}$.

We begin by recalling the definition of a sub-Gaussian Orlicz norm. For shorthand, we drop the adjective “sub-Gaussian” as this is the only Orlicz norm considered in this paper. For further background on Orlicz norms, we refer the reader to, e.g., the textbooks [RR91, Ver18], and we mention that the standard significance of this particular (sub-Gaussian) Orlicz norm is that a random variable is sub-Gaussian if and only if this norm is finite [Ver18, Example 2.7.13].

The Orlicz norm of a random variable $X$ is

Our proof of Theorem 3.1 uses the Orlicz norm for defining an optimal transport metric between probability distributions. In what follows, we write $\mathscr{C}(\mu,\nu)$ to denote the set of couplings between $\mu$ and $\nu$; that is, the set of all jointly defined pairs $(X,Y)$ of random variables with first marginal $\operatorname{law}(X)=\mu$ and second marginal $\operatorname{law}(Y)=\nu$.

The Orlicz–Wasserstein metric between distributions $\mu,\nu$ is

Note that since the Orlicz norm satisfies the triangle inequality, the standard gluing lemma from classical optimal transport theory shows that $W_{\psi_{2}}$ is indeed a metric; see, e.g., [Vil09a, Chapter 6]. The Orlicz–Wasserstein metric has also been considered in prior works [Stu11, Kel17, GHN23], but to our knowledge this paper constitutes the first use of this metric for sampling analysis.

For any Rényi order $q\geqslant 1$ and shift $w\geqslant 0$, the $W_{\psi_{2}}$-shifted Rényi divergence between probability distributions $\mu$ and $\nu$ is defined as

2 Proof of Theorem 3.1

Here we describe how the standard shifted divergence analysis is modified when using shifts in $W_{\psi_{2}}$ rather than $W_{\infty}$, and how this modified argument leads to a proof of Theorem 3.1.

At a high level, the shifted divergence technique—in both its original form and the new form here—is built upon two key lemmas. These two lemmas track how the shifted Rényi divergence evolves when both distributions are either (1) pushed forward through a Lipschitz map; or (2) convolved with Gaussian noise. These two lemmas are called the “contraction-reduction lemmaAlthough we use this name to be consistent with the previous literature on the shifted divergence technique, we note that this map need not be a contraction, i.e., the Lipschitz constant can be greater than $1$.” and the “shift-reduction lemma.”

The contraction-reduction lemma is the simpler of these two lemmas, and extends unchanged—in terms of both statement and proof—when the standard $W_{\infty}$ shift is replaced by our proposed $W_{\psi_{2}}$ shift. For completeness, we provide a brief proof.

Let $\mu^{\prime}$ be the surrogate for $\mu$ in $\mathcal{R}_{q}^{(w/c)}\left(\mu\mathbin{\|}\nu\right)$. Then by definition, $W_{\psi_{2}}(\mu,\mu^{\prime})\leqslant w/c$ and $\mathcal{R}_{q}\left(\mu^{\prime}\mathbin{\|}\nu\right)=\mathcal{R}_{q}^{(w/c)}\left(\mu\mathbin{\|}\nu\right)$. Thus

where the first step is because $W_{\psi_{2}}(\mu P,\mu^{\prime}P)\leqslant c\,W_{\psi_{2}}(\mu,\mu^{\prime})\leqslant c\,(w/c)=w$ by Lipschitzness of $P$; the second step is by the data-processing inequality for Rényi divergences (Lemma 2.3); and the third step is by construction of $\mu^{\prime}$. ∎

The shift-reduction lemma, however, requires substantial modification.

Case 1: initial shift $w=0$. For shorthand, let $\gamma$ denote $\mathcal{N}(0,\sigma^{2}I_{d})$. We bound the Rényi divergence between $\operatorname{law}(X+Z)$ and $\operatorname{law}(Y+Z)$, where $X\sim\mu$, $Y\sim\nu$, and $Z\sim\gamma$. Let $\mu^{\prime}$ be the surrogate for $\mathcal{R}_{2q-1}^{(\delta)}(\mu\mathbin{\|}\nu)$, so that $\mathcal{R}_{2q-1}^{(\delta)}(\mu\mathbin{\|}\nu)=\mathcal{R}_{2q-1}(\mu^{\prime}\mathbin{\|}\nu)$ and $W_{\psi_{2}}(\mu,\mu^{\prime})\leq\delta$. Let $X^{\prime}\sim\mu^{\prime}$ be optimally coupled with $X\sim\mu$ with respect to the Orlicz–Wasserstein metric $W_{\psi_{2}}(\mu,\mu^{\prime})$ so that

where here and henceforth we write $p_{\eta}$ as shorthand for the law of a random variable $\eta$.

Note that $X+Z$ and $Y+Z$ are the result of the same function applied to the tuples $(X^{\prime},X-X^{\prime}+Z)$ and $(Y,Z)$ respectively. Thus, by the data-processing inequality for Rényi divergences (Lemma 2.3),

By expanding the definition of Rényi divergence and applying Hölder’s inequality, we bound this by

By definition of Rényi divergence and then the construction of $\mu^{\prime}$, the first term 1 simplifies to

Case 2: initial shift $w>0$. Let $\mu^{\prime}$ denote the surrogate for $\mathcal{R}_{2q-1}^{(w+\delta)}(\mu\mathbin{\|}\nu)$, so that $\mathcal{R}_{2q-1}^{(w+\delta)}(\mu\mathbin{\|}\nu)=\mathcal{R}_{2q-1}(\mu^{\prime}\mathbin{\|}\nu)$ and $W_{\psi_{2}}(\mu,\mu^{\prime})\leq w+\delta$. Let $X^{\prime}\sim\mu^{\prime}$ be optimally coupled with $X\sim\mu$ with respect to the Orlicz–Wasserstein metric $W_{\psi_{2}}(\mu,\mu^{\prime})$ so that $\|X-X^{\prime}\|_{\psi_{2}}=W_{\psi_{2}}(\mu,\mu^{\prime})\leqslant w+\delta$. Decompose

where $\tau\coloneqq\delta/(w+\delta)$. Then

Above, the first step is by using $p_{X_{1}^{\prime}}\ast\gamma$ as a surrogate for $\mu\ast\gamma$, which is allowed since

The second step is by using the result from case 1; the third step is by using $\mu^{\prime}$ as a surrogate for $p_{X_{1}^{\prime}}$, which is allowed since $W_{\psi_{2}}(\mu^{\prime},p_{X_{1}^{\prime}})\leqslant\|X^{\prime}-X_{1}^{\prime}\|_{\psi_{2}}=\|X_{2}^{\prime}\|_{\psi_{2}}=\tau\,\|X-X^{\prime}\|_{\psi_{2}}\leqslant\tau\,(w+\delta)=\delta$; and the final step is by construction of $\mu^{\prime}$. ∎

for any $\lambda\in$ and any $\delta\leqslant(1-\lambda)\,\sigma\sqrt{\tfrac{2}{(q-1)\,(q-\lambda)}}$. Choosing $\lambda=1/2$ recovers Lemma 3.7. Different choices of $\lambda$ enable trading off the increase in the Rényi order in the first term against the penalty in the second term. The proof of this generalized bound is identical except for one change: replace the Cauchy–Schwarz inequality in the proof with Hölder’s inequality $\int fg\leqslant(\int\left\lvert f\right\rvert^{1/\lambda})^{\lambda}\,(\int\left\lvert g\right\rvert^{1/(1-\lambda)})^{1-\lambda}$.

We conjecture that (3.3) is tight in that for any input parameters $q$, $\sigma^{2}$, $w$, there exist distributions $\mu,\nu$ such that this bound holds with equality when it is optimized over the knobs $\delta,\lambda$. Details on this conjecture are provided in Appendix B.2.

The Orlicz–Wasserstein metric is the “right” metric to use for the shifted Rényi analysis in the sense that (1) Lemma 3.7 holds with this Orlicz–Wasserstein shifted Rényi divergence; and (2) the Orlicz–Wasserstein distance at initialization is bounded for sampling algorithms (shown in Lemma 4.7).

In contrast, (2) fails for all previous versions of the shifted Rényi divergence analysis, since they use $W_{\infty}$ shifts. And (1) fails for other natural candidates of the Wasserstein metric for which the initialization distance is bounded. This includes the $W_{p}$ metric for any finite $p$, as well as the Orlicz–Wasserstein metric for any Orlicz norm that is weaker than sub-Gaussian. See Appendix B.3 for details. Finally, we remark that this discussion is tailored to the fact that we are analyzing Markov chains with Gaussian noise; if for example, this were replaced by Laplacian noise, then the right notion of shift would be the Orlicz–Wasserstein metric with the sub-exponential Orlicz norm, and our techniques would extend in a straightforward way.

We now use Lemmas 3.6 and 3.7 to prove Theorem 3.1.

Let $\delta=c^{N}\,W_{\psi_{2}}(\mu_{0},\mu_{0}^{\prime})$. By using, in order: the definition of the Markov chain update, Lemma 3.7 (in the form of Remark 3.8 with $\lambda=0$), and then Lemma 3.6, we obtain

Note that the use of Lemma 3.7 is valid since $\delta\leqslant\sigma\sqrt{\frac{2}{q\,(q-1)}}$ by the assumption on $N$.

It suffices to show that the Rényi term in the above display vanishes. To this end, let $Q_{n}$ denote the transition kernel for the $n$-th step of the Markov chain, i.e., $\rho Q_{n}=(\rho P_{n})\ast\mathcal{N}(0,\sigma^{2}I_{d})$. Clearly $Q_{n}$ is $W_{\psi_{2}}$-Lipschitz with parameter $c$ since $W_{\psi_{2}}(\rho Q_{n},\rho^{\prime}Q_{n})\leqslant W_{\psi_{2}}(\rho P_{n},\rho^{\prime}P_{n})\leqslant c\,W_{\psi_{2}}(\rho,\rho^{\prime})$ for any distributions $\rho,\rho^{\prime}$. Thus we may apply Lemma 3.7 $N-1$ times to argue that

Here, the third step is by the choice of $\delta$, and the fourth step is by definition of the shifted Rényi divergence. The proof is complete by combining the above displays. ∎

Low-accuracy sampling with O​(d)𝑂𝑑O(\sqrt{d}) complexity

The main result of the section is the first Rényi convergence guarantee for log-concave sampling that requires a number of first-order queries that scales in the dimension $d$ only as $\sqrt{d}$. This improves over the state-of-the-art which has $d^{2}$ scaling. This result is formally stated as follows.

Suppose that $\pi\propto\exp(-f)$ where $f$ is $\alpha$-strongly-convex and $\beta$-smooth, and let $0<\varepsilon\lesssim\frac{1}{\sqrt{q}}$. There is a randomized algorithm that, given knowledge of the minimizer of $f$ and access to

The algorithm in Theorem 4.1 initializes at the Dirac distribution $\delta_{x^{*}}$. This is reasonable because the cost of using gradient descent to compute $x^{*}$ approximately, using the same first-order oracle access, is dominated by the cost of subsequently running the sampling algorithm. If the algorithm is initialized at some other point $x$, the runtime only increases by a logarithmic factor of $\log W_{\psi_{2}}(\delta_{x},\pi)$, which is lower order unless $x$ is exponentially far from $x^{*}$, since $W_{\psi_{2}}(\delta_{x},\pi)\lesssim\sqrt{d/\alpha}+\|x-x^{*}\|$ by Lemma 4.7 and the triangle inequality.

While the results of [Che+18, DR20, Ma+21, Zha+23] have also shown iteration complexities that scale in the dimension $d$ as $\widetilde{O}(d^{1/2})$, a key difference is that these results do not hold for Rényi divergences. In particular, past work has only proven weaker mixing results in the Wasserstein metric or the KL divergence. As discussed in the introduction, Wasserstein and KL guarantees are insufficient for the purpose of warm-starting high-accuracy sampling algorithms—for this, it is essential to have guarantees in the more stringent Rényi divergence. (Note that Wasserstein bounds are weaker than KL bounds by Talagrand’s $T_{2}$ inequality, and moreover KL bounds are weaker than Rényi bounds by monotonicity of Rényi divergences.) See §1.2 for a detailed discussion of this and of the many longstanding technical difficulties involved with establishing Rényi guarantees.

The algorithm we use is Underdamped Langevin Monte Carlo (ULMC) with certain parameters (stated explicitly in the proof in §4.3). Background on this algorithm is recalled in §4.1. At a high level, the proof of Theorem 4.1 uses the weak triangle inequality for Rényi divergences to decompose the sampling error of ULMC into the following two terms, both measured in Rényi divergence:

The “bias” error between the stationary distribution of ULMC and the target distribution $\pi$.

The “discrete mixing” error of ULMC to its biased stationary distribution.

The bias error (1) is readily handled by recent results such as [GT20, Zha+23]. Bounding the discrete mixing error (2) is the key technical challenge; see §1.2 for a detailed discussion of the technical obstacles related to this, and the connections to open problems about hypocoercivity in the PDE literature. The key contribution of this section is to bound this quantity in Theorem 4.4 below. To do this, we use the new shifted divergence technique developed in §3.

The exposition in this subsection is based on the corresponding section in [Che23]; we refer the interested reader there for further details.

Studied since the work of Kolmogorov in the 1930s [Kol34], the underdamped Langevin diffusion—sometimes also called the kinetic Langevin diffusion—is the solution to the stochastic differential equation

where ${(B_{t})}_{t\geq 0}$ is a standard $d$-dimensional Brownian motion. Analogous to the classical theory of convex optimization, here the auxiliary state variable $Y_{t}$ has the physical interpretation of momentum, and the linking parameter $\gamma$ has the physical interpretation of friction. A related interpretation of the underdamped Langevin diffusion is as a variant of the idealized Hamiltonian Monte Carlo algorithm, in which the momentum is refreshed continuously rather than periodically. The stationary distribution for this SDE is the joint distribution

A major obstacle for analyzing the convergence of the underdamped Langevin diffusion is that this process exhibits hypocoercive dynamics, i.e., the standard Markov semigroup approach based on isoperimetric inequalities does not work. It is a longstanding question in PDE theory to develop general tools for establishing fast convergence of hypocoercive dynamics. See the discussion in §1.2. We bypass these issues by instead developing tools for analyzing a discrete-time version of this diffusion. (In forthcoming work, we detail the extent to which these discrete-time analyses enable analysis of continuous-time hypocoercive dynamics.)

There are several ways to discretize the underdamped Langevin diffusion. Perhaps the simplest way is the Euler–Maruyama discretization, as is standard for defining (unadjusted) Langevin Monte Carlo. However, for these underdamped Langevin dynamics, there is a better discretization which dates back at least to 1980 [EB80], namely:

for $t\in[nh,(n+1)h]$. This process is called Underdamped Langevin Monte Carlo (ULMC) or Kinetic Langevin Monte Carlo; we use the former term in this paper. The point of this discretization is that since the gradient is refreshed periodically rather than continuously, the SDE is linear within these periods, and thus can be integrated exactly in closed form (see, e.g., [Che+18, Appendix A]). This is called an “exponential integrator” in the lingo of numerical analysis, and is formalized as follows.

Conditioned on $(X_{nh},Y_{nh})$, the law of $(X_{(n+1)h},Y_{(n+1)h})$ is the Gaussian distribution $\mathcal{N}(F(X_{nh},Y_{nh}),\Sigma\otimes I_{d})$ where

Above, we use the notational shorthand $a\coloneqq\exp(-\gamma h)$.

2 Discrete mixing of Underdamped Langevin Monte Carlo

where $\mathcal{M}$ is the linear map defined in (4.3).

We remark that in a typical use case of this discrete mixing result, $\boldsymbol{\mu_{0}}$ is initialized at a product distribution of the form $\mu_{0}\otimes\mathcal{N}(0,I_{d})$, and is compared to the target distribution $\boldsymbol{\mu_{0}^{\prime}}=\boldsymbol{\pi}=\pi\otimes\mathcal{N}(0,I_{d})$. In this setting, the Orlicz–Wasserstein metric in the upper bound can be simplified to $W_{\psi_{2}}(\mathcal{M}_{\#}\boldsymbol{\mu_{0}},\mathcal{M}_{\#}\boldsymbol{\mu_{0}^{\prime}})\leqslant 2W_{\psi_{2}}(\mu_{0},\pi)$.

To prove Theorem 4.4, we appeal to our new shifted divergence technique developed in §3. This requires analyzing the ULMC iterates in a twisted norm, since an iteration of the ULMC algorithm (or more precisely, the mean-shifting function $F$ defined in Lemma 4.3) is not contractive with respect to the standard Euclidean norm. This twisted norm is the Euclidean norm after the change of coordinates

In these new coordinates, the mean of the next iterate of ULMC started at $(u,v)$ is $\bar{F}(u,v)$, where $\bar{F}=\mathcal{M}\circ F\circ\mathcal{M}^{-1}$. Since $\mathcal{M}^{-1}(u,v)=(u,\frac{\gamma}{2}\,(v-u))$, we can explicitly write

By Lemma 4.3, it follows that conditioned on $(U_{nh},V_{nh})$, the law of $(U_{(n+1)h},V_{(n+1)h})$ is the Gaussian distribution $\mathcal{N}(\bar{F}(U_{nh},V_{nh}),\bar{\Sigma}\otimes I_{d})$ where

We make use of the following two helper lemmas about the dynamics of ULMC in this twisted norm. The first helper lemma shows that the ULMC algorithm sends two iterates to Gaussians with means that are closer in the twisted norm than the original iterates. Since the two Gaussians have the same covariance $\bar{\Sigma}\otimes I_{d}$, this implies that the ULMC Markov transition kernel $P$ is contractive in $W_{\psi_{2}}$, which will allow us to use our new shifted divergence technique from §3. This lemma first appeared in the recent paper [Zha+23]; for completeness, a proof is provided in Appendix C.1.

The second helper lemma estimates the noise of ULMC in this twisted norm. The proof is an explicit computation and is provided in Appendix C.2.

Suppose that $h\lesssim 1/\gamma$. Then the matrix $\bar{\Sigma}$ defined in (4.5) satisfies

Armed with Lemmas 4.5 and 4.6, we are now ready to prove Theorem 4.4.

Let $\{\boldsymbol{\mu_{n}}\}_{n\geqslant 0}$ and $\{\boldsymbol{\mu_{n}^{\prime}}\}_{n\geqslant 0}$ denote the two processes $\boldsymbol{\mu_{n}}=\boldsymbol{\mu_{0}}\boldsymbol{P}^{n}$ and $\boldsymbol{\mu_{n}^{\prime}}=\boldsymbol{\mu_{0}^{\prime}}\boldsymbol{P}^{n}$ obtained by running ULMC from initialization distributions $\boldsymbol{\mu_{0}}$ and $\boldsymbol{\mu_{0}^{\prime}}$, respectively. Define twisted processes $\{\boldsymbol{\nu_{n}}\}_{n\geqslant 0}$ and $\{\boldsymbol{\nu_{n}^{\prime}}\}_{n\geqslant 0}$ by $\boldsymbol{\nu_{n}}=\mathcal{M}_{\#}\boldsymbol{\mu_{n}}$ and $\boldsymbol{\nu_{n}^{\prime}}=\mathcal{M}_{\#}\boldsymbol{\mu_{n}^{\prime}}$, where $\mathcal{M}$ is the change-of-coordinates matrix defined in (4.3). Since $\mathcal{M}$ is invertible, applying the data-processing inequality for Rényi divergences (Lemma 2.3) in both directions implies

We now show that the latter term is at most $\varepsilon^{2}$. For shorthand, define $\lambda\coloneqq\lambda_{\min}(\bar{\Sigma})$ and define $\boldsymbol{Q}$ to be the Markov operator given by $\boldsymbol{\nu}\boldsymbol{Q}=\bar{F}_{\#}\boldsymbol{\nu}\ast\mathcal{N}(0,\bar{\Sigma}\otimes I_{d}-\lambda I_{2d})$. Then by Lemma 4.3 and a change of measure, the law of $\boldsymbol{\nu_{n+1}}$ is

Similarly, the law of $\boldsymbol{\nu_{n+1}^{\prime}}$ is

Thus, letting $c$ denote the Lipschitz constant of the Markov operator $\boldsymbol{Q}$ w.r.t. the $W_{\psi_{2}}$ metric, we may invokeThe application of this result requires the number of iterations $N$ to be large enough that the right hand side of (4.6), later set to $\varepsilon^{2}$, is at most $(\log 2)/(q-1)$. But this holds by assumption. the new shifted divergence result (Theorem 3.1) to bound

We now use the two helper lemmas to quantify the various terms in (4.6). First, by a simple coupling argument and then an application of Lemma 4.5, the Markov operator $\boldsymbol{Q}$ is $W_{\psi_{2}}$-contractive with parameter $c$, where

Second, because ULMC is run with step size $h\lesssim 1/\gamma$, Lemma 4.6 implies

Therefore, by combining the above displays, we conclude that

Setting this bound to $\varepsilon^{2}$ and solving for $N$ completes the proof. ∎

3 Warm start with Underdamped Langevin Monte Carlo

We begin by bounding the distance to the target at initialization. We emphasize that this initialization is not a warm start. Indeed, this initialization is even weaker than what is typically called a “feasible start” in the literature (namely $\mathcal{N}(x^{*},\beta^{-1}I_{d})$), and moreover can be further relaxed to an arbitrary initialization $x_{0}$ so long as the distance between $x_{0}$ and the mode $x^{*}$ of the target distribution is sub-exponentially large (since our final bound depends only logarithmically on this distance).

Suppose that $\pi\propto\exp(-f)$ where $f$ is $\alpha$-strongly convex. Let $x^{*}$ denote the minimizer of $f$. Then

Let $X\sim\pi$ and define $Y\coloneqq\|X-x^{*}\|$. By definition of the Orlicz–Wasserstein metric,

for any $\lambda^{2}\geqslant 16/\alpha$. By combining the above displays and setting $\lambda=\sqrt{32d/\alpha}$, we conclude that

Therefore this choice of $\lambda$ is an upper bound on the Orlicz–Wasserstein norm $W_{\psi_{2}}(\delta_{x^{*}},\pi)$. ∎

The second lemma uses Girsanov’s theorem to bound the bias of ULMC. Here, we build upon recent advances in the literature on Rényi discretization of stochastic processes. Beginning with the works [GT20, Che+21, EHZ22] and culminating in the paper [Che+22a], it is now understood that the Girsanov discretization technique leads to bias bounds for LMC in Rényi divergence matching prior results which only held for weaker divergences, and yet remains flexible enough to cover varying assumptions. The recent paper [Zha+23] extends this technique for ULMC. Since the results of [Zha+23] hold under more general assumptions at the expense of a more involved analysis, and in the interest of keeping our derivations more self-contained, in Appendix C.3 we simplify and streamline the Girsanov argument of [Zha+23] for our setting of interest. In order to clarify where the $d^{1/2}$ comes from, we write the final bound in terms of the total elapsed continuous time $T=Nh$ rather than the number of iterations $N$.

Suppose that $f$ is $\alpha$-strongly-convex and $\beta$-smooth. Let $\boldsymbol{\pi}(x,y)\propto\exp(-f(x)-\operatorname{\frac{1}{2}}\|y\|^{2})$, and let $\boldsymbol{P}$ denote the Markov transition kernel corresponding to an iteration of ULMC with friction parameter $\gamma\asymp\sqrt{\beta}$ and step size $h\lesssim\frac{1}{\beta^{3/4}d^{1/2}q\,{(T\log N)}^{1/2}}$, where $N$ is the total number of iterations and $T=Nh$ is the total elapsed time. Then,

By the weak triangle inequality for Rényi divergence (Lemma 2.7), we can further bound this by

The coefficient $(q-1/2)/(q-1)$ can be crudely bounded by $2$, say, since it suffices to bound the Rényi divergence error for $q\geqslant 3/2$ (indeed, monotonicity of Rényi divergences in the order $q$ then implies the same bound for $q<3/2$).

Now, by combining our discrete mixing result for ULMC (Theorem 4.4), our initialization bound (Lemma 4.7), and the ULMC bias bound (Lemma 4.8), we conclude that

if ULMC is run with friction parameter $\gamma$, step size $h$, and iteration complexity $N$ that satisfy:

By recalling that $T\coloneqq Nh$, solving for these choices of parameters, and omitting logarithmic factors, we conclude that it suffices to run ULMC with the following choices of parameters:

High-accuracy sampling with O​(d)𝑂𝑑O(\sqrt{d}) complexity

Establishing fast mixing results for MALA is a longstanding problem. As detailed in §1, recent breakthroughs have made it clear that the key barrier for fast mixing of MALA is the question of warm starts. In this section, we use the faster low-accuracy sampling result developed in §4 to efficiently warm start MALA. This leads to the fastest known high-accuracy sampling algorithms not only in strongly log-concave settings (details in §5.1), but also in weakly-log-concave and isoperimetric, non-log-concave settings (details in §5.2), for which we improve over state-of-the-art query complexity results by a factor of $\sqrt{d}$.

Moreover, for any $\delta\in(0,1)$ with probability at least $1-\delta$, the number of queries made satisfies

In analogy to the familiar concept from algorithm design for deterministic problems [Cor+22], the algorithm in Theorem 5.1 may be called a “Las Vegas” algorithm because it has a randomized runtime which is small with high probability. The fact that this runtime is randomized is not an issue in practice because the iteration complexity depends on a quantity that is efficiently estimable during the execution of the algorithm.

In the rest of this subsection, we overview the algorithm in Theorem 5.1 and its analysis; see Appendix D for full technical details. This algorithm combines three algorithms as building blocks: ULMC, MALA, and the proximal sampler algorithm. Let us explain this by building up to the full complexity in two steps—both because this will motivate why all three algorithmic components are needed, and also because this is how our analysis actually proceeds.

However, while this simple combination of ULMC and MALA achieves the desired dependence on the dimension $d$, it leads to a suboptimal dependence on the condition number $\kappa$, namely $\widetilde{O}(\kappa^{3/2})$ rather than $\widetilde{O}(\kappa)$. This worsened dependence in $\kappa$ arises from the state-of-the-art bounds on the discretization of ULMC [Zha+23]. For full details on this weak version of Theorem 5.1, see Appendix D.3.

While the main conceptual innovation here is the high-level strategy of combining these three algorithmic building blocks, we remark that an additional technical obstacle for proving Theorem 5.1 is showing mixing in more stringent notions of distance than TV. See the discussion in §1.2. Indeed, while our new ULMC result proves fast mixing in Rényi divergence, existing results on MALA and its combination with the proximal sampler are limited to TV. We boost this mixing in TV to Rényi divergences (and thus all the other desired metrics by standard comparison inequalities) using two additional ideas.

The first improves mixing bounds for the proximal sampler from TV to Rényi divergence. To do this, we control the propagation of error when each step of the proximal sampler algorithm is performed approximately in Rényi divergence. As we show, this is readily accomplished by appealing to the “strong composition rule” of Rényi divergences from the differential privacy literature.

The second improves mixing bounds for MALA from TV to Rényi divergence. We accomplish this by further exploiting the fact that MALA is warm started in Rényi divergence. Note that this means we use the Rényi warm start in two ways: first to show that MALA mixes fast in TV, which is what we can conclude from the above argument and appealing to [Che+21, WSC22]; and second, to boost the TV bound at the final iterate to a more stringent bound. We isolate this TV-to-Rényi boosting technique in the following simple lemma as it may be of independent interest: indeed, since it uses the TV mixing bound in an entirely black-box way, this lemma may be useful for establishing Rényi mixing of other warm-started algorithms. This lemma improves the previous result of [Che+21, Lemma 28] because that result required a warm-start in $\mathcal{R}_{\infty}$ which is currently unavailable algorithmically, whereas this lemma here only requires the weaker condition of a warm start in a Rényi divergence of finite order (stated here with $q=3$ for simplicity).

Let $P$ be a Markov transition kerrnel which has stationary distribution $\pi$. Consider running $P$ from any initialization distribution for $N$ steps to obtain a distribution $\mu_{N}\coloneqq\mu_{0}P^{N}$. Then

See Appendices D.1 and D.2 for background on MALA and the proximal sampler algorithm, respectively; and see Appendices D.3 and D.4 for proofs of the weak version and full version of Theorem 5.1, respectively.

2 Extensions to weakly-log-concave and non-log-concave settings

Our faster algorithm for sampling from well-conditioned targets (Theorem 5.1) yields faster samplers for a variety of other settings, due essentially to the reductions in [Che+22]. We present here several such extensions that concern target distributions which satisfy isoperimetric inequalities, which is quite flexible in the sense that this allows for non-log-concavity and also is preserved under, e.g., bounded perturbations and Lipschitz mappings. See §A for background on these isoperimetric inequalities.

A comment on notation for these isoperimetric settings: we still use the condition number $\kappa$ to denote the ratio $\kappa=\beta/\alpha$, but now $\alpha$ denotes the (inverse) parameter of an isoperimetric bound, rather than the parameter for strong convexity. The motivation behind this notation is that $\alpha$-strong-convexity implies the log-Sobolev inequality with parameter $1/\alpha$, which in turn implies the Poincaré inequality with parameter $1/\alpha$ (see Lemma A.3).

Just as in [Che+22], we could also obtain a result for distributions satisfying a Latała–Oleszkiewicz inequality, which interpolates between PI and LSI. In this setting, we again improve over the previous state-of-the-art bounds by a factor of $d^{1/2}$. However, for the sake of brevity, we omit this extension as it is conceptually similar but requires more involved technical details.

These results are the direct analogs of the state-of-the-art complexity results in [Che+22, Corollary 7], but here with a dimension dependence that is improved by a factor of $d^{1/2}$.

We mention another consequence of our improved high-accuracy sampler for the strongly-log-concave setting. Namely, via the same proximal reduction framework, this gives the following alternative complexity bound for target distributions which are (non-strongly) log-concave, sometimes called weakly-log-concave. This bound is a direct analog of [Che+22, Corollary 6], but here with a dimension dependence that is also improved by a factor of $d^{1/2}$. Note that this theorem is a low-accuracy guarantee; one can also obtain high-accuracy samplers from our results in this log-concave setting by using the fact that log-concavity implies a Poincaré inequality, albeit with a function-dependent constant [KLS95], and then appealing to Theorem 5.4. The resulting low-accuracy and high-accuracy results are incomparable in the sense that each can dominate in different settings—but in any case, our theorems for both settings yield improvements by a factor of $d^{1/2}$.

Proofs for the results in this section are provided in Appendix D. At a high level, the proof of all these results use the same reduction to the problem of sampling from well-conditioned distributions. This reduction is based on the proximal sampler (described in Appendix D.2) and lets us apply our improved sampler for the well-conditioned case (Theorem 5.1). In each case, however, we must track the propagation of error due to the inexact implementation of the backwards step of the proximal sampler, which was not previously done in any work except for in the TV distance.

In Appendix D.7, we provide more explicit, albeit more complicated, statements of these results to address the following two points. (1) The above results depend on the initialization (through $W_{2}(\mu_{0},\pi)$, $\mathcal{R}_{2}(\mu_{0}\mathbin{\|}\pi)$, or $\chi^{2}(\mu_{0}\mathbin{\|}\pi)$) and it may be unclear how large these quantities are in a given application. (2) We assumed that the algorithm has access to a stronger oracle than just a first-order oracle for $f$, namely, we also assumed access to a prox oracle for $hf$ with $h=\frac{1}{2\beta}$. We address (1) by explicitly bounding these initialization quantities in terms of other, more easily computable problem parameters, and we address (2) by removing the assumption of a prox oracle.

Discussion

Here we mention several interesting questions for future research that are inspired by our results.

Are warm starts essential for future progress in high-accuracy sampling? Our work is the first to show the achievability of the faster rates proven for high-accuracy samplers under a warm start assumption. We do this by exhibiting an efficient algorithm for producing the warm start. We believe that this general strategy may be important for future progress in high-accuracy sampling. Indeed, the natural next candidate for improving upon MALA is Metropolized Hamiltonian Monte Carlo [Nea+11], or related variants. For such Metropolized algorithms, we suspect that much of the intuition from §1.2 remains true; namely, that the algorithm can benefit from a more aggressive step size near stationarity. Hence, to extract the full potential of these algorithms, it seems likely that we must again pursue the dual plan of improving the rates under a warm start, and efficiently computing that warm start. Insofar as warm starts continue to play an important role in sampling analysis, the Rényi analysis techniques that we developed in §3 and §4 could prove useful for future progress in this direction.

Can we leverage shifted divergence techniques for further advances in differential privacy and beyond? Core to our results is an improved version of the shifted Rényi divergence technique that uses Orlicz–Wasserstein shifts rather than $W_{\infty}$ shifts. Since their introduction, shifted divergences have been instrumental for advances in differentially private optimization (see the prior work discussion in §1.3), and also very recently in the context of sampling ([AT22a] and Theorem 4.1). We believe that we are only scratching the surface of potential applications, extensions, and refinements of this technique, and we are optimistic that a deeper understanding of our Rényi analysis toolbox will have implications far beyond.

We thank Mufan (Bill) Li and Matthew Zhang for many insightful conversations regarding hypocoercivity. JMA acknowledges funding from an NYU Faculty Fellowship. SC acknowledges funding from NSF TRIPODS program (award DMS-2022448).

Appendix A Background on functional inequalities

In this section, we collect relevant background material for the convenience of the reader.

To go beyond the strongly log-concave case, we can instead assume that the target distribution $\pi$ satisfies a functional inequality, which encodes geometric information (e.g., isoperimetric properties) about $\pi$. We focus primarily on the two most well-studied functional inequalities in this context, namely the log-Sobolev inequality and the Poincaré inequality. For each of these two functional inequalities, the class of distributions satisfying this assumption not only includes all strongly log-concave distributions (see Lemma A.3 below), but also includes many more examples because of closedness properties of these functional inequalities under operations such as bounded perturbations of the potential, pushforwards via Lipschitz mappings, or taking suitable mixtures. We refer to [BGL14] for many of these properties, and to [CCN21] for the closure under taking mixtures.

We next collect together key facts about these functional inequalities. The following results show that the class of distributions satisfying these inequalities is larger than the class of strongly log-concave distributions; see [BGL14, Proposition 5.1.3 and Corollary 5.7.2].

(Bakry–Émery theorem) If $\pi$ is $\alpha$-strongly log-concave, then it satisfies an LSI with constant at most $1/\alpha$.

If $\pi$ satisfies an LSI with constant $C_{\mathsf{LSI}}$, then it also satisfies a PI with constant at most $C_{\mathsf{LSI}}$.

A useful consequence of the LSI is the following sub-Gaussian concentration inequality, typically established via the Herbst argument; see [BGL14, Proposition 5.4.1].

Similarly, the PI implies subexponential concentration, see [BGL14, §4.4.3].

Next we recall two comparison inequalities which enable proving sampling guarantees in Wasserstein distance as an immediate corollary of proving sampling guarantees in other metrics—namely KL divergence in the LSI setting, and chi-squared divergence in the PI setting. Such comparison inequalities are often called transport inequalities. Specifically, the first result, attributed to Otto and Villani [OV00], shows that under an LSI, a transportation inequality between Wasserstein and KL divergence holds (this inequality is often referred to as Talagrand’s $T_{2}$ inequality).

The second result shows a similar transport inequality in the PI setting [Liu20]. Under a PI, Talagrand’s $T_{2}$ inequality does not necessarily hold anymore. Nevertheless, a useful transport inequality still holds if one replaces the KL divergence by the chi-squared divergence.

Finally, we record the following standard second-moment-type bound for strongly log-concave measures; see, e.g., [DKR22, Proposition 2]. We give a short proof sketch for the convenience of the reader.

Suppose that the distribution $\pi\propto\exp(-f)$ is $\alpha$-strongly log-concave, and that $x^{*}$ is the minimizer of $f$. Then,

Appendix B Deferred details for §3

Here, we prove the inequality in Remark 3.4, repeated here for convenience:

By [Ver18, Proposition 2.5.2], if $(X,Y)$ is an optimal coupling of $\mu$ and $\nu$ for the $W_{\psi_{2}}$ distance, then

Observe that for any random variable $Z$, if we denote $\|Z\|_{\infty}\coloneqq\operatorname{ess\,sup}\|Z\|$, then we can bound

and therefore $\|Z\|_{\psi_{2}}\leqslant\|Z\|_{\infty}/\sqrt{\log 2}$ by the definition of the Orlicz norm. Now, applying this bound to the random variable $Z=X-Y$, we conclude the desired inequality

B.2 Remarks on tightness of Lemma 3.7

Here we remark that, conditional on the following plausible conjecture, the generalized version of Lemma 3.7 (as stated in Remark 3.8) is tight. This conjecture states that the shifted Rényi divergence between two isotropic Gaussians with same covariance is achieved by a deterministic shift. Understanding this simple case could be more broadly helpful for understanding tightness of other inequalities and analyses using the shifted Rényi divergence.

where $c\coloneqq\max(0,1-w\sqrt{\log 2}/\|x\|)$.

Of course, the conjecture here is the first equality (the second equality is just the closed-form expression for the Rényi divergence between Gaussians in Lemma 2.5). The direction “$\leqslant$” is clear because $\mathcal{N}(cx,\sigma^{2}I_{d})$ satisfies

and therefore is feasible for the optimization problem defining the shifted Rényi divergence. The direction “$\geqslant$” is the one requiring justification.

In the rest of this subsection, we show the claimed tightness assuming Conjecture B.1. Fix any Rényi order $q\geqslant 1$, noise variance $\sigma^{2}>0$, and initial shift $w\geqslant 0$. Consider distributions $\mu=\delta_{ae_{1}}$ and $\nu=\delta_{0}$, where $a>w\sqrt{\log 2}$. We claim that the bound in Remark 3.8 holds with equality when its parameters $\delta,\lambda$ are optimized; that is,

To this end, supposing Conjecture B.1 holds, the left hand side of (B.2) is equal to

On the other hand, note that $\mathcal{R}_{(q+\lambda-1)/\lambda}^{(w+\delta)}(\mu\mathbin{\|}\nu)$ is equal to if $a\leqslant(w+\delta)\sqrt{\log 2}$, and otherwise is equal to $\infty$. This means that the optimal value of $\delta$ is $a/\sqrt{\log 2}-w$. Thus the right hand side of (B.2) simplifies to

where the final step is because the optimal value of $\lambda$ is at $\lambda=0$. We conclude that the left- and right-hand sides of (B.2) indeed match, as desired.

B.3 Proof for Remark 3.9

$\mathcal{R}_{q}(\mu\ast\mathcal{N}(0,\sigma^{2})\mathbin{\|}\nu\ast\mathcal{N}(0,\sigma^{2}))=\infty$.

This comprises a counterexample to Lemma 3.7 by taking $w=0$ and $\delta=W(\mu,\nu)$ and noting that if $\delta>\sigma/\sqrt{(2q-1)(q-1)}$, then we can simply dilate the space (i.e., replace $\mu(x)$ by $\mu_{b}(x)\propto\mu(bx)$ for a sufficiently large $b>0$) and repeat the same argument.

In the case that $W=W_{p}$, then $W_{p}(\mu,\nu)$ is equal to the $p$-th norm of $\mu$, which is finite for any $p<\infty$. In the case that $W$ is an Orlicz–Wasserstein norm with Orlicz norm weaker than sub-Gaussian, then $W(\mu,\nu)$ is equal to the Orlicz norm of $\nu$, which is finite if we choose $a=b$.

Note that $\mu\ast\mathcal{N}(0,\sigma^{2})$ is not sub-Gaussian, yet $\nu\ast\mathcal{N}(0,\sigma^{2})=\mathcal{N}(0,\sigma^{2})$ is sub-Gaussian. We may therefore appeal to the fact that the Rényi divergence is infinite whenever the first argument is not sub-Gaussian, but the second argument is. For a proof of this fact in the case that $q=2$, see [Che+22a, Lemma 21]; this proof readily extends to any finite $q\in(1,\infty)$ by replacing the Cauchy–Schwarz inequality by Hölder’s inequality.

Appendix C Deferred details for §4

Since $\tfrac{1}{\gamma}(h-\gamma^{-1}\,(1-a))=O(h^{2})$, we have

where we use the notational shorthand $b\coloneqq\frac{2}{\gamma}\,(h+\gamma^{-1}\,(1-a))$.

Since $1-a=O(\gamma h)$ and $b=O(h/\gamma)$, we can approximate $AA^{\mathsf{T}}$ by the following matrix $B$ with error

By a direct computation, the eigenvalues of $B$ are $1+a^{2}\pm(1-a^{2}-b\lambda)$, where $\lambda$ ranges over the eigenvalues of $\nabla^{2}f(u)$. The strong-convexity and smoothness of $f$ implies that $\lambda\in[\alpha,\beta]$. Thus

In order for this to be strictly smaller than $2$, we must take $\gamma>\sqrt{\beta}$. We choose $\gamma=\sqrt{2\beta}$, whereby

C.2 Proof of Lemma 4.6

By definition of $\mathcal{M}$ and $\Sigma$,

The smallest eigenvalue of this matrix is

Above, the first step is by the explicit formula for the eigenvalues of a $2\times 2$ matrix; the second step is by plugging in the entries of $\bar{\Sigma}$ and simplifying; and the third step is by performing a Taylor expansion in the variable $\gamma h$.

C.3 Proof of Lemma 4.8

We invoke the following result, which appears as Lemma 26 in [Zha+23].

Let ${(X_{t},Y_{t})}_{t\geq 0}$ denote the continuous-time underdamped Langevin diffusion (4.1) with potential $f$ that is $\beta$-smooth and minimized at $x^{*}$. Assume that $0<h\lesssim\frac{1}{\sqrt{\beta}\vee\gamma}$ and $0\leq\lambda\lesssim\frac{1}{\gamma dh^{3}}$. Then, conditioned on $(X_{0},Y_{0})$,

This result can be easily adapted from the proof of [Zha+23, Lemma 26], noting that in our situation the bound simplifies as we are assuming $\nabla f$ is $\beta$-Lipschitz rather than merely Hölder continuous. ∎

provided that Novikov’s condition (see [Le ̵16, Theorem 5.23]) holds:

Assuming for the moment that (C.1) is indeed verified, Itô’s formula yields

For $t\in[kh,(k+1)h]$, let $\Delta_{t}\coloneqq\nabla f(X_{t})-\nabla f(X_{kh})$. By the Cauchy–Schwarz inequality,

We claim that the term marked $(\dagger)$ equals $1$; this would follow if the quantity inside the expectation is a martingale. In general, it is only a local martingale, but it is a bona fide martingale provided that Novikov’s condition holds: it suffices to have

Note that this condition is stronger than (C.1).

Towards this end, we bound, for a parameter $p\geq 1$ to be chosen later,This argument avoids the use of the “conditioning lemma” from [GT20] (Lemma 23 in [Zha+23]).

By conditioning on $(X_{kh},Y_{kh})$ and applying Proposition C.1, this is bounded by

provided that $h\lesssim\frac{1}{\beta^{2/3}d^{1/3}p^{1/3}q^{2/3}T^{1/3}}$. We choose $p\asymp\log N$ so that $N^{1/p}\asymp 1$. We now need tail bounds for $\lVert{X_{kh}-x^{*}}\rVert$ and $\lVert{Y_{kh}}\rVert$.

Using the argument in the proof of Lemma 4.7, for $c>0$,

provided that $c\leq\alpha/36$. Therefore,

provided that $h\lesssim\frac{\alpha^{1/4}\gamma^{1/4}}{\beta q^{1/2}\,{(T\log N)}^{1/4}}$.

The same argument applied to $\lVert{Y_{kh}}\rVert$ yields

provided that $h\lesssim\frac{\gamma^{1/2}}{\beta q\,{(T\log N)}^{1/2}}$.

If we put these bounds together and take $\gamma\asymp\sqrt{\beta}$, we deduce that if $h\lesssim\frac{\gamma^{1/2}}{\beta d^{1/3}q\,{(T\log N)}^{1/2}}$ and $T\gtrsim\frac{\sqrt{\beta}}{\alpha}$, then it holds that

This verifies (C.1) and (C.2), and moreover shows that for $h\lesssim\frac{1}{\beta^{3/4}d^{1/2}q\,{(T\log N)}^{1/2}}$,

Applying Proposition C.1 and the above tail estimates,

Hence, by definition of the Rényi divergence, we conclude that

Appendix D Deferred details for §5

Propose a new point $Y_{n}\sim Q(X_{n},\cdot)$;

Accept the new point with probability $1\wedge\frac{\pi(Y_{n})\,Q(Y_{n},X_{n})}{\pi(X_{n})\,Q(X_{n},Y_{n})}$. That is, if the proposal is accepted, we set $X_{n+1}\coloneqq Y_{n}$, otherwise we set $X_{n+1}\coloneqq X_{n}$.

It can be shown that under mild conditions on the proposal $Q$, the resulting Metropolis-adjusted algorithm is a reversible Markov chain with stationary distribution $\pi$. Since Markov chains often exhibit geometrically ergodicity, the family of Metropolis-adjusted chains is often used to design high-accuracy samplers.

In this work, we are primarily considered with the Metropolis-adjusted Langevin algorithm, in which the proposal kernel is taken to be one step of the discretized Langevin algorithm, i.e.,

For this choice of proposal kernel, generating a proposal and computing the acceptance ratio can be carried out with a constant number of evaluations of $f$ (zeroth-order queries) and $\nabla f$ (first-order queries). Actually, in order for the mixing time results we invoke to be valid, we must instead consider the $\frac{1}{2}$-lazy version of the chain, in which each proposal is discarded with probability $\frac{1}{2}$. Since this only affects the mixing time bounds by a factor of $2$, we henceforth ignore this distinction.

D.2 Background on the proximal sampler

Next, we perform Gibbs sampling on the augmented target; namely, we alternate between the following steps.

Forwards step (convolution): Sample $Y_{k}\mid X_{k}\sim\pi^{Y\mid X}(\cdot\mid X_{k})=\mathcal{N}(X_{k},hI_{d})$.

Backwards step (RGO): Sample $X_{k+1}\mid Y_{k}\sim\pi^{X\mid Y}(\cdot\mid Y_{k})$.

The backwards step is known as the restricted Gaussian oracle (RGO), and any successful use of the proximal sampler hinges upon an efficient implementation of the RGO.

Just as the Langevin diffusion is recognized as the Wasserstein gradient flow of the KL divergence w.r.t. $\pi$, thanks to the seminal result of [JKO98], the proximal sampler admits an appealing interpretation as a proximal discretization of that same gradient flow. Indeed, the RGO can be interpreted as the proximal operator for the KL divergence over the Wasserstein space, evaluated at a Dirac measure; moreover, the convergence rates established for the proximal sampler exactly match the classical rates for the proximal point method (PPM) in Euclidean space, a fact which led to a new sharp analysis of the PPM under a Polyak–Łojasiewicz inequality. For this and further discussion, we refer to [Che+22].

From an algorithmic standpoint, we are interested in two key features of the proximal sampler: (1) the proximal reduction framework which allows us to boost the condition number dependence of any high-accuracy sampler to near-linear [LST21a], and (2) its uses for sampling beyond strong log-concavity [Che+22].

The other virtue of the proximal sampler exploited in this work is its application to sampling beyond the strongly log-concave setting, as was put forth in [Che+22]. In this paper, the authors analyzed the outer loop complexity of the proximal sampler under weak log-concavity and under a variety of standard isoperimetric assumptions (e.g., the log-Sobolev or Poincaré inequality). In all of these cases, however, as soon as $f$ is $\beta$-smooth, the approximate implementation of the RGO can be handled, as before, by a high-accuracy sampler for strongly log-concave and log-smooth distributions. Hence, the results of [Che+22] effectively reduce the problem of sampling from each of the aforementioned classes of distributions to the problem of high-accuracy strongly log-concave and log-smooth sampling. They used this reduction to provide new state-of-the-art guarantees for sampling from these classes, and in §5.2 we leverage our faster implementation of the RGO to improve each of their complexity bounds by a factor of $\sqrt{d}$.

D.3 Weak version of Theorem 5.1

Here, we show the following weaker version of Theorem 5.1 as it provides the key building block to prove it. Like Theorem 5.1, this result here shows that from a feasible start, the query complexity of high-accuracy sampling from a strongly-log-concave distribution scales in the dimension $d$ as $\widetilde{O}(d^{1/2})$ rather than $\widetilde{O}(d)$. The difference from Theorem 5.1 is that the weaker result here has a suboptimal dependence on the condition number $\kappa$, namely $\widetilde{O}(\kappa^{3/2})$ rather than $\widetilde{O}(\kappa)$. This dependence will later be boosted using the proximal sampler in §D.4, allowing us to prove the full Theorem 5.1.

To prove convergence in the various metrics, due to standard comparison inequalities it usually suffices to prove a convergence result in the strongest metric, namely, the chi-squared divergence. Indeed, convergence in the KL divergence follows from the monotonicity of Rényi divergences (Lemma 2.4) and convergence in the TV distance follows from Pinsker’s inequality. If $\pi$ satisfies an LSI, then convergence in $W_{2}$ follows from Talagrand’s $T_{2}$ inequality (Lemma A.6); otherwise, if $\pi$ only satisfies a PI, then convergence in $W_{2}$ follows from the quadratic transport-variance inequality (Lemma A.7).

Run ULMC from this arbitrary initialization $\delta_{x_{0}}$ to obtain a $\mathcal{R}_{3}$ warm start (which also implies a $\chi^{2}$ warm start).

Use the $\chi^{2}$ warm start to argue that MALA mixes rapidly in TV.

Use the $\mathcal{R}_{3}$ warm start to argue that the TV mixing guarantee implies mixing guarantees in $\chi^{2}$ (and therefore also the other desired metrics by Remark D.2).

Our main technical contribution here is the ability to implement step 1 in $\widetilde{O}(d^{1/2})$ queries—this is an immediate application of our result on ULMC (Theorem 4.1). Step 2 follows from known results about MALA [Che+22, WSC22], with only minor modification as described below, and step 3 follows from the helper Lemma 5.2.