1 Introduction↩︎

1.1 Motivation and background↩︎

What is the complexity of sampling from a smooth distribution \(\pi \propto \exp(-f)\) over \(\mathbb{R}^d\)? When \(\pi\) is log-concave or satisfies an isoperimetric inequality, then this question is well-studied, with many results providing iteration complexity guarantees which scale polynomially in the dimension and the isoperimetric constants; see [1] for an introduction to some of these developments.

In this work, we are interested in the opposite scenario, in which \(\pi\) is allowed to be highly non-log-concave, with the only assumption being that \(f\) is \(L\)-smooth, in the sense that \(-LI \preceq \nabla^2 f \preceq LI\). In this case, simple worst-case examples demonstrate that generating a sample close in total variation distance to \(\pi\) can require exponentially many queries to \(f\) (e.g., with respect to the ambient dimension \(d\); see [2].). One way to overcome this difficulty is to impose additional problem structure, such as assuming that \(\pi\) is a mixture of strongly log-concave distributions [3] or to assume that an annealing path has small “action” [4], which requires designing tailor-made algorithms for each problem class under consideration. Another approach is to weaken the notion of progress by seeking an approximately “stationary” solution.

The latter idea has been especially fruitful in the sister field of optimization [5], namely to minimize (or find a stationary point of) an objective function \(f \colon \mathbb{R}^d \to \mathbb{R}\), in which the corresponding question becomes: How many queries to \(\nabla f\) are necessary to reach a point \(x\) with \(\norm{\nabla f(x)} \le \varepsilon\)? In the case where \(f\) is \(L\)-smooth but potentially non-convex, gradient descent finds such an approximate stationary point in \(O(1/\varepsilon^2)\) iterations, and this rate is optimal in high dimension [6]. The optimal complexity in low dimension is still open, and it turns out that the results of this paper take on an interesting interpretation in light of this story; see Section 1.3 below. Although finding a stationary point is still a far cry from finding a global optimum, it is often the starting point for deeper investigations, such as the development of algorithms which escape “undesirable” stationary points (e.g., strict saddles). More broadly, this theory has the advantage of having great generality, as it is one of the few frameworks available to quantitatively study general non-convex optimization problems.

Developing an analogous theory for the task of sampling is more subtle, as one must first decide on a suitable definition of an “approximate stationary point”. Such a definition was proposed in [7], and it is based on the celebrated JKO theorem [8] which interprets the Langevin diffusion as a gradient flow, as explained next. The Langevin diffusion is the stochastic process for \(X_t \in \mathbb{R}^d\) governed by the stochastic differential equation (SDE) \[\begin{align} \label{eq:langevin} \mathrm{d}X_t = -\nabla f(X_t)\,\mathrm{d}t+ \sqrt 2\,\mathrm{d}B_t\,, \end{align}\tag{1}\] where \(B_t\) is the standard Brownian motion in \(\mathbb{R}^d\). Langevin diffusion is the basis of some of the most popular sampling algorithms, such as the unadjusted Langevin algorithm (ULA) and the Metropolis-adjusted Langevin algorithm (MALA), as well as the proximal sampler studied in this paper, all of which can be viewed as time discretization of the Langevin diffusion. If \(\rho_t\) denotes the law of \(X_t\) evolving along Langevin diffusion as in 1 , then [8] interprets the curve of probability measures \((\rho_t)_{t\ge 0}\) as a gradient flow of the relative entropy or KL divergence functional \(\mathsf{KL}(\cdot \,\|\,\pi)\) over the space of probability measures equipped with the Wasserstein \(W_2\) metric. See [9] for further discussion on this perspective.

The gradient flow perspective also brings with it a set of calculation rules, usually termed “Otto calculus”, such as the following identity for the dissipation of the KL divergence along Langevin diffusion 1 , classically known as de Bruijn’s identity: \[\begin{align} \label{eq:debrujin} \frac{\mathrm{d}}{\mathrm{d}t}\, \mathsf{KL}(\rho_t \,\|\,\pi) = -\mathsf{FI}(\rho_t\,\|\,\pi)\,, \end{align}\tag{2}\] where \(\mathsf{FI}(\rho_t\,\|\,\pi)\) is the relative Fisher information of \(\rho_t\) relative to \(\pi\): \[\begin{align} \mathsf{FI}(\rho_t \,\|\,\pi) \mathrel{\vcenter{:}}=\int \Bigl\| \nabla \log \frac{\mathrm{d}\rho_t}{\mathrm{d}\pi}\Bigr\|^2\,\mathrm{d}\rho_t\,. \end{align}\] In the context of Otto calculus, the relative Fisher information takes on a new interpretation: it is the squared norm of the Wasserstein gradient of \(\mathsf{KL}(\cdot\,\|\,\pi)\) at \(\rho_t\). This motivated [7] to define an \(\varepsilon\)-stationary point for sampling to be a probability measure \(\rho\) for which \(\sqrt{\mathsf{FI}(\rho \,\|\,\pi)} \le \varepsilon\), and the complexity question for non-log-concave sampling becomes:

In that same paper, [7] showed that the unadjusted Langevin algorithm, i.e., the Euler–Maruyama discretization of 1 , can achieve this goal in \(O(1/\varepsilon^4)\) steps.³ Subsequently, the work of [10] established a query lower bound of \(\Omega(1/\varepsilon^{2-o(1)})\) in high dimension, and [11] developed an improved algorithm using parallel computation which runs for \(\widetilde{O}(1/\varepsilon^2)\) parallel rounds and makes \(\widetilde{O}(1/\varepsilon^6)\) total queries.

Approximate stationarity in sampling admits a compelling interpretation as measuring closeness to metastability. Indeed, it is well-established that in the presence of free energy barriers, e.g., produced by well-separated modes of the target distribution \(\pi\), local Markov chains often get stuck at certain distributions which persist for exponential stretches of time, called metastable states. Although the concept of a metastable state is only loosely defined, in light of 2 , we can view the relative Fisher information as a quantitative measure of metastability, since it implies that a measure \(\rho\) with small relative Fisher information will experience slow decrease of the KL divergence under the Langevin diffusion. However, we also note that even when \(\pi\) is non-log-concave and has multiple modes, the KL divergence \(\rho \mapsto \mathsf{KL}(\rho \,\|\,\pi)\) does not have spurious local minima; the only stationary point is the global minimizer \(\rho = \pi\).

These interpretations have been further developed in the setting where \(\pi\) is a mixture, \(\pi = \sum_{k=1}^K w_k \pi_k\) for some weights \(w_1,\dots,w_K > 0\) with \(\sum_{k=1}^K w_k = 1\), and some component distributions \(\pi_1,\dots,\pi_K\). For example, when \(\pi\) is a mixture of Gaussians, [12] showed that a FI guarantee \(\mathsf{FI}(\rho\,\|\,\pi) \le \varepsilon^2\) implies mixing for each of the components which have non-trivial mass under \(\rho\). More recently, [13] showed that if each \(\pi_k\) satisfies a functional inequality (e.g., log-Sobolev), then \(\mathsf{FI}(\rho\,\|\,\pi) \le \varepsilon^2\) implies that \(\rho\) is close in KL divergence to a reweighted mixture \(\sum_{k=1}^K w_k' \pi_k\) for some (explicit) weights \(w_1',\dots,w_K'\). Generally, these results indicate that the complexity of stationarity in sampling captures the complexity of local mixing, which is a useful complement to worst-case global mixing results for non-log-concave sampling.

There remains, however, a gap between the upper and lower bounds in the sequential setting, leading to the question of determining an optimal algorithm for local mixing. In particular, analogously to the picture for optimization, is an optimal algorithm given by some discretization of the Langevin diffusion—the sampling equivalent of the gradient flow—and if so, which one?

In this work, we close this gap and answer this question affirmatively. We develop an algorithm, based on the proximal sampler [14], [15], which achieves the optimal complexity of \(O(1/\varepsilon^2)\) up to logarithmic factors. Moreover, as we discuss below, the dimension dependence of our result is the best possible without further advances in the complexity of log-concave sampling.

1.2 Contributions↩︎

Our results are based on the proximal sampler algorithm, which is reviewed in Section 2. Briefly, the proximal sampler first augments the target distribution \(\pi^X(x) \propto \exp(-f(x))\) into a joint density over two variables, \(\pi^{XY}(x,y) \propto \exp(-f(x)-\frac{1}{2h}\,\|y-x\|^2)\), and then applies Gibbs sampling. Concretely, this consists of alternating between the following two steps: given \(X\), sample \(Y \sim \pi^{Y\mid X} = \mathcal{N}(X, hI_d)\); given \(Y\), sample \(X \sim \pi^{X\mid Y}\), where \(\pi^{X\mid Y}\) is known as the restricted Gaussian oracle (RGO). When \(\pi^X\) is log-smooth and satisfies an isoperimetric inequality, such as a log-Sobolev inequality, [15] shows a high-accuracy guarantee for the proximal sampler in Rényi divergence. When \(\pi^X\) is log-smooth and strongly log-concave, [16] shows a high-accuracy guarantee for the proximal sampler in relative Fisher information. In this paper, we work only under the assumption that \(\pi^X\) is log-smooth, without assuming log-concavity or isoperimetry.

Our first contribution is to show that under log-smoothness alone, the relative Fisher information decays in an average sense along the proximal sampler. Namely, we establish that \[\begin{align} \label{eq:ideal95FI95bd} \mathsf{FI}(\bar\rho_K^X \,\|\,\pi^X) \le \frac{1}{K} \sum_{k=1}^K \mathsf{FI}(\rho_k^X \,\|\,\pi^X) \le \frac{\mathsf{KL}_0}{h\,(1-\frac{Lh}{2})\,K}\,, \end{align}\tag{3}\] where \(\rho_k^X\) denotes the law of the \(k\)-th iterate of the proximal sampler algorithm, \(\bar\rho_K^X \mathrel{\vcenter{:}}=\frac{1}{K} \sum_{k=1}^K \rho_k^X\) is the averaged measure, and \(\mathsf{KL}_0 \mathrel{\vcenter{:}}=\mathsf{KL}(\rho_0^X \,\|\,\pi^X)\); see Theorem 1 below. We note such a result is expected in light of the de Bruijn identity 2 , which implies that along the Langevin diffusion, \[\begin{align} \frac{1}{T} \int_0^T \mathsf{FI}(\rho_t\,\|\,\pi^X)\,\mathrm{d}t\le \frac{\mathsf{KL}_0}{T}\,. \end{align}\] Establishing the corresponding guarantee 3 for the proximal sampler is more involved since it requires tracking the relative Fisher information along simultaneous evolution of both arguments along the heat flow and its time reversal, and here we leverage the recent results of [16].

However, 3 is not yet an algorithmic result, as it assumes that we can implement the RGO exactly. Our next contributions address this point. First, by choosing \(h \asymp 1/(dL)\) and by implementing the RGO exactly via rejection sampling following [15], the result 3 implies an \(O(\frac{Ld \, \mathsf{KL}_0}{\varepsilon^2})\) query complexity guarantee for the proximal sampler to achieve \(\mathsf{FI}(\bar\rho_K^X \,\|\,\pi^X) \le \varepsilon^2\); see Corollary 1. We note this already improves upon the \(O(\frac{L^2 d \, \mathsf{KL}_0}{\varepsilon^4})\) complexity obtained in [7] via ULA. Our next results leverage the recent approximate RGO implementations [17]–[19] to further improve the dimension dependence.

By choosing \(h = 1/(2L)\), one can see that the RGO \(\pi^{X \mid Y=y}\) is \(L\)-strongly log-concave and \(3L\)-log-smooth, which means that it is a log-concave distribution with condition number at most \(3\). Thus, it is natural to use existing log-concave samplers to implement the RGO. There are two main obstacles to this approach:

1. Even if we have a log-concave sampler which implements the RGO accurately in the sense of relative Fisher information, i.e., the sampler yields an approximate RGO \(\hat{\pi}^{X\mid Y=y}\) such that \(\mathsf{FI}(\hat{\pi}^{X\mid Y=y} \,\|\,\pi^{X\mid Y=y}) \le \delta^2\), it is non-trivial to track the propagation of this error along the proximal sampler, as the relative Fisher information lacks natural properties such as a data-processing inequality. See Section 4.3 for discussion.

2. Prior works on log-concave sampling usually do not yield relative Fisher information guarantees in the first place.

We address these issues by showing that, given a high-accuracy log-concave sampler with Rényi divergence guarantees—namely, an algorithm which yields a sample from any \(O(1)\)-conditioned log-concave density up to error \(\varepsilon^2\) in \(\mathcal{R}_3(\cdot \,\|\,\pi^X)\) (Rényi divergence of order \(3\)) in \(\widetilde{O}(d^p \mathop{\mathrm{polylog}}(1/\varepsilon))\) iterations—perturbing the output of the sampler by a Gaussian with a small but well-chosen variance yields an RGO implementation whose error is well-controlled along the proximal sampler. This leads to an algorithm which obtains \(\mathsf{FI}(\cdot\,\|\,\pi^X) \le \varepsilon^2\) using \[\begin{align} \label{eq:informal95result} \frac{L\,\mathsf{KL}_0}{\varepsilon^2} \times d^p \times \mathop{\mathrm{polylog}}\bigl(\Delta_0, L/\varepsilon^2, d\bigr)\qquad\text{iterations}\,, \end{align}\tag{4}\] where \(\Delta_0 \mathrel{\vcenter{:}}=\mathcal{R}_2(\rho_0^X \,\|\,\pi^X)\) arises as an artifact of our analysis; see Theorem 2. In particular, although the optimal exponent \(p\) is still unknown, current state-of-the-art log-concave samplers achieve this high-accuracy guarantee with exponent \(p=\frac{1}{2}\) [17]–[19], which we can easily plug into the result above.

Note that, up to logarithmic factors, the dependence on \(L/\varepsilon^2\) is optimal in high dimension, in light of the lower bound of [10]. If we compare against the analogous problem of finding approximate stationary points in non-convex optimization, there the optimal complexity in high dimension is \(L\Omega_0/\varepsilon^2\), where \(\Omega_0 \mathrel{\vcenter{:}}= f(x_0) - \inf f\) is the initial objective gap [6]. Thus, our complexity bound is entirely similar to the one in optimization, except that we incur an extra dimensional factor \(d^p\), which is the complexity of high-accuracy log-concave sampling. Thanks to our general reduction, any improvement to the complexity of the latter task would yield an improved Fisher information complexity result for non-log-concave sampling.

However, one could still ask whether there is a different approach which avoids the reduction to log-concave sampling altogether and could therefore potentially lead to improved dimension dependence. Our final result says, in essence, that the answer is no. More precisely, we show that any Fisher information complexity bound of the form 4 implies the existence of a log-concave sampler with complexity \(\widetilde{O}(d^p \mathop{\mathrm{polylog}}(1/\varepsilon))\), see Theorem 3. This closes the loop and shows that the optimal dimension dependence for the Fisher information complexity of non-log-concave sampling, and for high-accuracy log-concave sampling, coincide.

1.3 An analogy with optimization↩︎

There is an interesting analogy between our work on Fisher information guarantees for sampling and the study of finding approximate stationary points in optimization. As noted in Section 1.1, it is well-established that gradient descent finds an \(\varepsilon\)-stationary point in \(O(1/\varepsilon^2)\) iterations, and the lower bound of [6] shows that this is the best possible, at least in high dimension. Recent investigations, however, have focused on the complexity in low dimension, where the picture is more nuanced.

To motivate this study, consider first the problem of convex, non-smooth minimization, which is a separate problem altogether but shares some similar features with the stationary point problem. In particular, gradient descent finds an \(\varepsilon\)-minimizer in \(O(1/\varepsilon^2)\) steps, and lower bounds show that this is optimal in high dimension. In low dimension, however, there is a better algorithm: cutting plane methods can find an \(\varepsilon\)-minimizer in \(O(d\log(1/\varepsilon))\) steps. At the level of lower bounds, this manifests in the following way: the lower bound construction in [6] showing that \(O(1/\varepsilon^2)\) is tight is embedded in dimension \(d \gg 1/\varepsilon^2\). The existence of the cutting plane method means that no such constructions exist when \(d \log(1/\varepsilon) \ll 1/\varepsilon^2\), i.e., in dimension \(d \ll 1/(\varepsilon^2 \log(1/\varepsilon))\).

Returning to the stationary point question, the interesting departure from the paragraph above is that thus far, we have not discovered an analogue of cutting plane methods for non-convex optimization. Indeed, the best complexities in low dimension scale as \(O(1/\varepsilon^{2-O(1/d)})\) [20]. If this is fundamental, then there should be lower bound constructions witnessing the \(\Omega(1/\varepsilon^2)\) rate which can be embedded in much smaller dimension: say, \(d \asymp \log(1/\varepsilon)\) or so. However, the lower bound construction of [6] requires a much larger dimension: \(d \gg 1/\varepsilon^2\) for deterministic methods, and \(d \gg \widetilde{\Omega}(1/\varepsilon^4)\) for randomized methods. Thus, at present, it is unclear whether or not there could exist a cutting plane method for finding stationary points.

More recently, [21] showed that in dimension one, when we restrict to deterministic algorithms which only use evaluations of \(\nabla f\) (and not to evaluations of \(f\)), then the \(O(1/\varepsilon^2)\) complexity is already tight; moreover, in dimension two, [22] showed that for deterministic algorithms with access to both \(f\) and \(\nabla f\), the optimal complexity is \(\Theta(1/\varepsilon)\). For general algorithms (randomized, with access to both \(f\) and \(\nabla f\)), the optimal complexity is \(\Theta(\log(1/\varepsilon))\) in dimension one [21], \(\Theta(1/\varepsilon^2)\) in large dimension \(\gg \widetilde{\Omega}(1/\varepsilon^4)\) (as discussed above), and remains open in between.

As for the connection with our work: although the low-dimensional complexity of finding stationary points remains elusive in optimization, our work resolves the analogous question in the sampling setting. Indeed, the lower bound \(\Omega(1/\varepsilon^{2-o(1)})\) in [10] already holds in dimension \(\gg \sqrt{\log(1/\varepsilon)/\log\log(1/\varepsilon)}\), and the present work shows that \(\widetilde{O}(1/\varepsilon^2)\) can be achieved via the proximal sampler in any dimension \(d \ge 1\). Although this does not have direct implications for the optimization question, which remains open, we believe that this further strengthens the rich connections between the two fields.

1.4 Additional related work↩︎

First-order stationarity guarantees have also been investigated for the Stein variational gradient descent (SVGD) algorithm, with initial works focusing on the kernelized Stein discrepancy (KSD), a kernelized analogue of the relative Fisher information [23]–[27]. More recently, relative Fisher information guarantees were obtained for a regularized variant of SVGD [28], [29].

As discussed above, [10] established a lower bound on finding approximate stationary points for sampling which scales as \(\Omega((L/\varepsilon^2)^{1-o(1)})\) in high dimension. In addition, that paper established a second lower bound which shows that the complexity of obtaining a moderate bound of the form \(\mathsf{FI}(\cdot\,\|\,\pi) \lesssim Ld\) can be precisely related to the complexity of finding approximate stationary points in non-convex optimization. In particular, it implies that the Fisher information guarantee for averaged LMC in [7] is already optimal in a certain regime. However, this regime does not cover guarantees \(\mathsf{FI}(\cdot\,\|\,\pi)\le \varepsilon^2\) for small \(\varepsilon \ll 1\), which is the regime of interest in this work. Moreover, their reduction is distinct from the ones we present in Theorem 2 and 3, which relate the complexity of obtaining small Fisher information to the complexity of high-accuracy log-concave sampling.

2 A review of the proximal sampler↩︎

Throughout, we are working with standard notations and definitions, which we review in Appendix 6.

We review the proximal sampler algorithm, originally proposed in [14]. Our goal is to draw a sample from the target distribution \(\pi^X \propto \exp(-f)\) supported on \(\mathbb{R}^d\). Let \(h > 0\) be an arbitrary “step size”. We define the joint probability distribution \(\pi\) on \(\mathbb{R}^{2d}\) with density function: \[\pi^{XY}(x,y) \propto \exp\left( -f(x) - \frac{1}{2h} \|x-y\|^2 \right).\] Observe that the \(x\)-marginal of \(\pi\) is precisely the original target distribution: \(\pi^X(x) = \int_{\mathbb{R}^d} \pi^{XY}(x,\mathrm{d}y)\). Therefore, to sample from \(\pi^X\), it suffices to sample \((X,Y) \sim \pi^{XY}\) and only return the \(X\) component.

To sample from \(\pi^{XY}\), [14] proposed to run Gibbs sampling with an alternating update. This results in the ideal proximal sampler algorithm, where in each iteration \(k \ge 0\), from the current iterate \(x_k \in \mathbb{R}^d\), we perform the following two steps:

1. Forward step: Draw \(y_k \mid x_k \sim \pi^{Y \mid X}( \cdot \mid x_k) = \mathcal{N}(x_k, hI)\).

2. Backward step: Draw \(x_{k+1} \mid y_k \sim \pi^{X \mid Y}(\cdot \mid y_k)\).

By the general theory of Gibbs sampling, the Markov chain update from \((x_k, y_k)\) to \((x_{k+1}, y_{k+1})\) is reversible with respect to the joint distribution \(\pi^{XY}\).

We call the above the ideal proximal sampler since we assume we can implement both the forward and backward steps exactly. The forward step is easy to implement, since it only requires drawing a Gaussian distribution around the current iterate. The backward step is called the restricted Gaussian oracle (RGO) in [14], and it is non-trivial since it requires sampling from the following probability distribution over \(x \in \mathbb{R}^d\), for any fixed \(y \in \mathbb{R}^d\): \[\begin{align} \label{Eq:BackwardCond} \pi^{X \mid Y}(x \mid y) \propto_x \exp\left( -f(x) - \frac{1}{2h} \|x-y\|^2 \right). \end{align}\tag{5}\] Note the proportionality is only over \(x \in \mathbb{R}^d\) (with fixed \(y \in \mathbb{R}^d\)), which we denote with \(\propto_x\).

2.1 Description of proximal sampler at the level of distributions↩︎

For \(h > 0\), let \(\mathsf{Q}^h = (\mathsf{Q}^h_y)_{y \in \mathbb{R}^d}\) denote the Gaussian channel with variance \(h\), which acts on a probability distribution \(\rho\) by: \[\rho \mathsf{Q}^h \mathrel{\vcenter{:}}=\rho \ast \mathcal{N}(0, hI).\] Thus, \(\mathsf{Q}^h_y = \mathcal{N}(y, hI)\) for each \(y \in \mathbb{R}^d\).

Let \(\mathsf{R}^h = (\mathsf{R}^h_y)_{y \in \mathbb{R}^d}\) denote the RGO channel with step size \(h > 0\) formed by \(\pi^X \propto \exp(-f)\), so: \[\mathsf{R}^h_y = \pi^{X \mid Y}(\cdot \mid y)\] where \(\pi^{X \mid Y}(\cdot \mid y)\) is the conditional distribution given in 5 . In our subsequent discussion below, when the step size \(h > 0\) is fixed, we write \(\mathsf{R}\equiv \mathsf{R}^h\) for brevity.

We can describe the ideal proximal sampler as follows. In each iteration \(k \ge 0\), from the current iterate \(x_k \sim \rho_k^X\), we perform the forward step to draw \(y_k \sim \rho_k^Y\) given by: \[\rho_k^Y = \rho_k^X \mathsf{Q}^h \,,\] and then we perform the backward step to draw \(x_{k+1} \sim \rho_{k+1}^X\) given by: \[\rho_{k+1}^X = \rho_k^Y \mathsf{R}^h \,.\] Thus, one step of the ideal proximal sampler is a composition of the Gaussian channel with the RGO channel: \(\rho_{k+1}^X = \rho_k^X \mathsf{Q}^h \mathsf{R}^h \,.\)

From the target distribution \(\pi^X\), we define \[\pi^Y \mathrel{\vcenter{:}}=\pi^X \mathsf{Q}^h = \pi^X \ast \mathcal{N}(0, hI) \,.\] Then, since \(\pi^X\) is invariant under one step of the proximal sampler, we have: \(\pi^X = \pi^X \mathsf{Q}^h \mathsf{R}^h = \pi^Y \mathsf{R}^h \,.\)

In Section 3 we study the guarantee of the proximal sampler with an exact implementation of the RGO. In Section 4 we study the guarantee of the proximal sampler when we use an approximate RGO implementation \(\widehat{\mathsf{R}}^h\) in place of the exact RGO \(\mathsf{R}^h\), with the benefit of obtaining an improved dimension dependence.

3 Guarantees for the ideal proximal sampler↩︎

In this section, we provide guarantees for the ideal proximal sampler algorithm, which means we assume an exact implementation of the RGO.

3.1 Guarantees for one step of the proximal sampler↩︎

We have the following guarantee on one step of the proximal sampler. We provide the proof of Lemma 1 in Appendix 8.1.

3.2 Convergence guarantee for the ideal proximal sampler↩︎

By iterating Lemma 1, we show the following convergence guarantee for the ideal proximal sampler, assuming an exact implementation of the RGO. We provide the proof of Theorem 1 in Appendix 8.2.

3.3 Rejection sampling implementation of the RGO↩︎

For \(h \asymp 1/(Ld)\), [15] showed that the RGO can be implemented exactly via rejection sampling using one evaluation of \(\mathop{\mathrm{prox}}_{hf}\) and \(O(1)\) density evaluations in expectation. This immediately leads to the following result.

Denoting \(\mathsf{KL}_0 \mathrel{\vcenter{:}}=\mathsf{KL}(\rho_0^X \,\|\,\pi^X)\), this already substantially improves upon the \(O(L^2 d\,\mathsf{KL}_0/\varepsilon^4)\) iteration complexity obtained in [7]. However, there is further room for improvement: several subsequent works have developed improved RGO implementations with better iteration complexities [17]–[19]. Following these trends, in the next section, we develop improved bounds based on approximate RGO implementation.

4 Guarantees for the proximal sampler with an approximate RGO implementation↩︎

4.1 Main result↩︎

The bound in Theorem 1 assumes that in each iteration, we can implement the RGO in the backward process exactly. Although this can be accomplished via rejection sampling, it does not yield the best dimension dependence. Here, we show that implementation of the RGO can be reduced to the complexity of high-accuracy log-concave sampling.

The key point is that with step size \(h=1/(2L)\), the RGO distribution \(\mathsf{R}_y\) is \(L\)-strongly log-concave, \(3L\)-log-smooth, with mode at \(\mathop{\mathrm{prox}}_{f/(2L)}(y)\). By translating and rescaling, it follows that \(\mathsf{R}_y\) can be implemented up to \(\mathcal{R}_3\) error \(\delta^2\) using one call to \(\mathop{\mathrm{prox}}_{f/(2L)}(y)\) and one call to \(\mathtt{Alg}\) with tolerance parameter \(\delta\).

We remark that the assumption that we can evaluate \(\mathop{\mathrm{prox}}_{f/(2L)}\) is not restrictive: with our choices of parameters, the computation of \(\mathop{\mathrm{prox}}_{f/(2L)}\) is a strongly convex and smooth optimization problem and can be solved to high accuracy via standard convex optimization algorithms.

Our main reduction is stated below. By the “proximal sampler with smoothed RGO implementation”, we mean that the proximal sampler is run where the RGO is implemented by taking the output of \(\mathtt{Alg}\) on the RGO and adding a Gaussian of appropriate variance to the output; finally, output one of the proximal sampler iterates chosen uniformly at random. See Algorithm 1 for a precise description.

We provide a proof sketch of Theorem 2 in Section 4.3, and the full proof in Appendix 9.7.

Current state-of-the-art samplers (e.g., [17]–[19]) satisfy Assumption 1, where each run of \(\mathtt{Alg}\) uses \(\widetilde{O}(\sqrt d\mathop{\mathrm{polylog}}(1/\delta))\) queries to \(\nabla f\) in expectation. We therefore conclude:

Recall that [10] gave a lower bound construction with \(\mathsf{KL}_0 \le 1\) for which the complexity of obtaining \(\sqrt{\mathsf{FI}} \le \varepsilon\) is \(\widetilde{\Omega}((L/\varepsilon^2)^{1-o(1)})\) for all \(d \ge \widetilde{\Omega}(\sqrt{\log(1/\varepsilon)})\). Interestingly, our reduction in Theorem 2 implies that if this lower bound can be strengthened to exhibit polynomial dependence on \(d\), then it would imply a lower bound on the complexity of high-accuracy log-concave sampling which also exhibits polynomial dependence on \(d\). This would be quite significant as the current best lower bound in that setting only exhibits a dimension dependence of \(\Omega(\log d)\) [30]. Thus, our reduction suggests an alternative approach for this problem which could be easier, e.g., via enabling the use of non-convex bump function constructions.

4.2 Converse reduction↩︎

Here, we argue that the reduction to high-accuracy log-concave sampling in Theorem 2 is the right one, by presenting a reduction in the converse direction via a simple restart strategy. We provide the proof of Theorem 3 in Appendix 9.8.

Note that, aside from the technical condition that \(\mathtt{Alg}\) satisfies \(\mathcal{R}_2(\widehat\pi\,\|\,\pi) \le \Delta_0\), the conclusion of Theorem 2 satisfies the requirement of \(\mathtt{Alg}\) with \(\phi(d) = \widetilde{O}(\sqrt d)\), which by Theorem 3 leads to a high-accuracy log-concave sampler with complexity \(\widetilde{O}(\kappa \sqrt d)\). Morally, the conclusion is that the dimension dependence in Theorem 2 cannot be improved without also improving the dimension dependence of high-accuracy log-concave sampling.

However, we note that there is a slight gap: Theorem 2 assumes the existence of a \(\mathcal{R}_3\) sampler, whereas the output of Theorem 3 is only a \(\mathsf{KL}\) sampler. We do not view this issue as fundamental, as there seems to not be any reason to expect the complexity of sampling in \(\mathsf{KL}\) vs.\(\mathcal{R}_q\) to be substantially different: indeed, guarantees which are initially proven in \(\mathsf{KL}\) are often later improved to hold in \(\mathcal{R}_q\). Moreover, a \(\mathsf{KL}\) guarantee can usually be boosted to a \(\mathcal{R}_3\) guarantee, provided that we have a warm start in a slightly stronger Rényi divergence such as \(\mathcal{R}_4\); see the argument of [18]. Therefore, Theorem 3 should also imply the existence of a \(\mathcal{R}_3\) sampler with complexity \(\widetilde{O}(\kappa\phi(d))\) from a suitable warm start.

4.3 Proof outline of Theorem 2↩︎

Recall that for each \(y \in \mathbb{R}^d\) and step size \(h > 0\), the backward conditional distribution is \[\pi^{X \mid Y}(x \mid y) \propto \exp\Bigl(-f(x) - \frac{1}{2h}\, \|x-y\|^2\Bigr)\,.\] Recall the RGO channel \(\mathsf{R}\) is the map via this backward conditional distribution: For any \(\rho^Y\), the distribution \(\rho^Y \mathsf{R}\) is given by \[(\rho^Y \mathsf{R})(x) = \int_{\mathbb{R}^d} \pi^{X \mid Y}(x \mid y) \, \rho^Y(\mathrm{d}y)\,.\]

4.3.1 Key lemma: Guarantee for the proximal sampler with approximate RGO↩︎

The following key lemma proves a bound on the relative Fisher information along one step of the proximal sampler with an inexact implementation of the RGO. We provide the proof of Lemma 2 in Appendix 9.1.

4.3.2 Applying the key lemma↩︎

In order to use Lemma 2, we must control the quantities \(\bar\varepsilon_\mathsf{KL}\) and \(\bar\varepsilon_\mathsf{FI}\). Since the KL divergence does not satisfy a triangle inequality, we cannot simply bound \(\bar\varepsilon_\mathsf{KL}^2\) by \(\mathsf{KL}(\rho^Y\widehat{\mathsf{R}}\,\|\,\rho^Y\mathsf{R})\); however, we can establish the following substitute. We provide the proof of Lemma 3 in Appendix 9.2.

This lemma implies we can choose the quantity \(\bar\varepsilon_\mathsf{KL}^2\) in ?? to be \[\begin{align} \bar\varepsilon_\mathsf{KL}^2 &\stackrel{(i)}{\le} \mathsf{KL}(\rho^Y \widehat{\mathsf{R}}\,\|\,\rho^Y\mathsf{R}) + (2+\mathcal{R}_2(\rho^Y\mathsf{R}\,\|\,\pi^X))\sqrt{\chi^2(\rho^Y \widehat{\mathsf{R}}\,\|\,\rho^Y\mathsf{R})} \notag \\ &\stackrel{(ii)}{\le} \sup_{y\in\mathbb{R}^d} \mathsf{KL}(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y) + (2+\mathcal{R}_2(\rho^Y\,\|\,\pi^Y)) \sup_{y\in\mathbb{R}^d}\sqrt{\chi^2(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)} \notag \\ &\stackrel{(iii)}{\le} \varepsilon_{\chi^2}^2 + (2+\mathcal{R}_2(\rho^Y\,\|\,\pi^Y)) \,\varepsilon_{\chi^2}\,,\label{eq:KL95bd} \end{align}\tag{6}\] with \(\varepsilon_{\chi^2}^2 \mathrel{\vcenter{:}}=\sup_{y\in\mathbb{R}^d}\chi^2(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)\). In the above, the first inequality \((i)\) is by applying Lemma 3; the second inequality \((ii)\) follows by the joint convexity of KL divergence, which implies \(\mathsf{KL}(\rho^Y \widehat{\mathsf{R}}\,\|\,\rho^Y\mathsf{R}) \le \mathbb{E}_{y \sim \rho^Y}[\mathsf{KL}(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)] \le \sup_{y\in\mathbb{R}^d} \mathsf{KL}(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)\), and similarly for chi-square divergence; and the third inequality \((iii)\) follows from the bound \(\mathsf{KL}(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y) \le \chi^2(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)\).

For the quantity \(\bar\varepsilon_{\mathsf{FI}}\) in ?? , again the triangle inequality fails to hold, and moreover there is an additional difficulty because it is not trivial to bound \(\mathsf{FI}(\rho^Y\widehat{\mathsf{R}}\,\|\,\rho^Y \mathsf{R})\) in terms of \(\sup_{y\in\mathbb{R}^d} \mathsf{FI}(\widehat{\mathsf{R}}_y \,\|\,\mathsf{R}_y)\). Here, we are able to show the following lemma. We provide the proof of Lemma 4 in Appendix 9.3.

To summarize, in order to apply Lemma 2, we must control (1) the Rényi-\(2\) divergence (equivalently, the chi-squared divergence) along the iterates of the proximal sampler, and (2) the accuracies \(\varepsilon_\mathsf{KL}\), \(\varepsilon_{\chi^2}\), \(\varepsilon_{\mathsf{FI}}\) in implementing the RGO. We handle (1) in Lemma 7 in Appendix 9.4.

For (2), the main difficulty is that we wish to implement the RGO using high-accuracy samplers (e.g., MALA), but standard guarantees for such samplers are stated in terms of \(\mathsf{KL}\) or \(\chi^2\), rather than \(\mathsf{FI}\). Rather than establishing separate \(\mathsf{FI}\) guarantees for each high-accuracy sampler, we show via a generic smoothing argument that it suffices to implement the RGO in \(\mathcal{R}_3\). We provide the proof of Corollary 3 in Appendix 9.6.

In fact, the new approximate RGO \(\widehat{\mathsf{R}}\) is easy to implement; given an approximate implementation \(\widetilde{\mathsf{R}}_y\) (which we obtain from Assumption 1), the new approximate RGO \(\widehat{\mathsf{R}}\) is given by \(\widehat{\mathsf{R}}_y = \widetilde{\mathsf{R}}_y \mathsf{Q}^t = \widetilde{\mathsf{R}}_y \ast \mathcal{N}(0, tI)\) with an appropriately chosen smoothing parameter \(t > 0\); practically, given a sample \(X \sim \widetilde{\mathsf{R}}_y\), a sample from \(\widehat{\mathsf{R}}_y\) is given by \(X + \sqrt{t}Z\) for an independent \(Z \sim \mathcal{N}(0,I)\).

Combining together all of these ingredients and performing some bookkeeping yields Theorem 2. We provide the full proof in Appendix 9.7.

5 Conclusion↩︎

In this work, we obtained a new complexity bound for sampling from a measure with small relative Fisher information to the target, which can be interpreted as an approximate first-order stationary measure for sampling. Our bound, which is based on the proximal sampler, is near-optimal in its dependence on the precision \(\varepsilon\). Moreover, the dimension dependence of our bound matches the dimension dependence of high-accuracy log-concave sampling, and the former cannot be improved without also improving the latter. This provides a relatively complete picture of the Fisher information complexity of sampling.

Going forward, one interesting open problem is to improve the Fisher information lower bound of [10] to exhibit polynomial dependence on the dimension since, by our results, it would imply a corresponding lower bound with polynomial dimension dependence for high-accuracy log-concave sampling. Another open problem is to determine the optimal Fisher information complexity in dimension one (see [10]).

6 Notation and definitions↩︎

We are working in the Euclidean space \(\mathbb{R}^d\) of dimension \(d \ge 1\). We denote the \(\ell_2\)-inner product by \(\langle u,v \rangle\) and the \(\ell_2\)-norm by \(\|u\| \mathrel{\vcenter{:}}=\sqrt{\langle u,u \rangle}\) for \(u,v \in \mathbb{R}^d\). For \(x \in \mathbb{R}\), let \(\{x\}_+ = \max\{x,0\}\) and \(\{x\}_- = -\min\{x,0\}\). We use the notation \(a \lesssim b\) (resp.\(a \gtrsim b\)) to mean that \(a\le Cb\) (resp.\(a \ge Cb\)) for a universal constant \(C > 0\), and \(a \asymp b\) to mean \(a\lesssim b\) and \(a \gtrsim b\). If a statement holds provided that \(a\ge Cb\) for a sufficiently large universal constant, then we abbreviate this by writing “if \(a \gg b\), then …”.

For a symmetric matrix \(A \in \mathbb{R}^{d \times d}\), we denote its operator norm by \(\|A\|_\mathsf{op}\), which is the largest absolute value of the eigenvalues of \(A\). We denote the Hilbert-Schmidt (or Frobenius) norm of \(A \in \mathbb{R}^{d \times d}\) by \(\|A\|_{\mathsf{HS}} = \sqrt{\mathsf{Tr}(A^\top A)}\). For a symmetric matrix \(A \in \mathbb{R}^{d \times d}\), we say \(A \succeq 0\) is positive semidefinite if \(u^\top Au \ge 0\) for any \(u \in \mathbb{R}^d\). For symmetric matrices \(A,B \in \mathbb{R}^{d \times d}\), \(A \succeq B\) if \(A-B \succeq 0\). Throughout, \(I_d \in \mathbb{R}^{d \times d}\) denotes the identity matrix.

For a function \(f \colon \mathbb{R}^d \to \mathbb{R}\), we say \(f\) is \(L\)-smooth if \(f\) is twice-differentiable and \(\|\nabla^2 f(x)\|_\mathsf{op}\le L\) for all \(x \in \mathbb{R}^d\). We recall the prox operator of \(f\) is the function \(\mathop{\mathrm{prox}}_f \colon \mathbb{R}^d \to \mathbb{R}^d\) given by, for any \(y \in \mathbb{R}^d\): \[\mathop{\mathrm{prox}}_f(y) = \arg\min_{x \in \mathbb{R}^d} \left\{ f(x) + \frac{1}{2}\, \|x-y\|^2 \right\}\] assuming the inner objective function on the right-hand side above has a minimizer (\(\mathop{\mathrm{prox}}_f\) can return any minimizer, if not unique). Note that for any \(h > 0\), we can write \(\mathop{\mathrm{prox}}_{hf}\) as: \[\mathop{\mathrm{prox}}_{hf}(y) = \arg\min_{x \in \mathbb{R}^d} \left\{ f(x) + \frac{1}{2h}\, \|x-y\|^2 \right\}.\] If \(f\) is \(L\)-smooth and \(h < \frac{1}{L}\), then the inner objective above is a strictly convex function and has a unique minimizer, so \(\mathop{\mathrm{prox}}_{hf}\) is well-defined.

For two measures \(\mu, \nu\) on \(\mathbb{R}^d\), we say \(\mu\) is absolutely continuous with respect to \(\nu\), denoted \(\mu \ll \nu\), if \(\nu(A) = 0\) implies \(\mu(A) = 0\) for any measurable subset \(A \subset \mathbb{R}^d\). If a probability measure \(\mu\) is absolutely continuous with respect to the Lebesgue measure \(\mathrm{d}x\) on \(\mathbb{R}^d\), then we can represent \(\mu\) in terms of its probability density function, which we also denote by \(\mu \colon \mathbb{R}^d \to [0,\infty)\) such that \(\int_{\mathbb{R}^d} \mu(x) \, \mathrm{d}x= 1\). For brevity, we hide the domain of integration and write \(\int\) to denote \(\int_{\mathbb{R}^d}\), unless otherwise specified. We also write \(\mu(x) \, \mathrm{d}x\) and \(\mu(\mathrm{d}x)\) interchangeably.

6.1 Divergences between probability distributions↩︎

Let \(\rho, \pi\) be two probability distributions on \(\mathbb{R}^d\) with \(\rho \ll \pi \ll \mathrm{d}x\). We recall the relative entropy or KL (Kullback-Leibler) divergence between them is: \[\mathsf{KL}(\rho \,\|\, \pi) = \mathbb{E}_\rho\left[\log \frac{\rho}{\pi}\right] = \int_{\mathbb{R}^d} \rho(x) \log \frac{\rho(x)}{\pi(x)} \, \mathrm{d}x.\] If \(\rho\) and \(\pi\) both have differentiable density functions, then we recall the relative Fisher information between \(\rho\) and \(\pi\) is: \[\mathsf{FI}(\rho \,\|\, \pi) = \mathbb{E}_\rho\left[\left\|\nabla \log \frac{\rho}{\pi} \right\|^2\right] = \int_{\mathbb{R}^d} \rho(x) \left\|\nabla \log \frac{\rho(x)}{\pi(x)}\right\|^2 \, \mathrm{d}x.\] We recall the total variation (TV) distance between \(\rho\) and \(\pi\) is: \[\mathop{\mathrm{\mathsf{TV}}}(\rho,\pi) = \sup_{A \subseteq \mathbb{R}^d} |\rho(A) - \pi(A)| = \frac{1}{2} \int_{\mathbb{R}^d} |\rho(x) - \pi(x)| \, \mathrm{d}x.\] We recall the chi-square divergence between \(\rho\) and \(\pi\) is: \[\chi^2(\rho\,\|\,\pi) = \mathsf{Var}_\pi\left(\frac{\rho}{\pi}\right) = \int_{\mathbb{R}^d} \left(\frac{\rho(x)}{\pi(x)}-1\right)^2 \pi(x) \, \mathrm{d}x= \int_{\mathbb{R}^d} \frac{\rho(x)^2}{\pi(x)} \, \mathrm{d}x- 1.\] We recall the Rényi divergence of order \(q > 0, q \neq 1\) between \(\rho\) and \(\pi\) is: \[\mathcal{R}_q(\rho\,\|\,\pi) = \frac{1}{q-1} \log \mathbb{E}_\pi\left[\left(\frac{\rho}{\pi}\right)^q\right] = \frac{1}{q-1} \log \int_{\mathbb{R}^d} \frac{\rho(x)^q}{\pi(x)^{q-1}} \, \mathrm{d}x.\] We recall that \(\mathsf{KL}(\rho\,\|\,\pi) = \lim_{q \to 1} \mathcal{R}_q(\rho\,\|\,\pi)\). Furthermore, we note that \(\mathcal{R}_2(\rho\,\|\,\pi) = \log(1+\chi^2(\rho\,\|\,\pi))\), or equivalently, \(\chi^2(\rho\,\|\,\pi) = \exp(\mathcal{R}_2(\rho\,\|\,\pi))-1\).

We recall the following relations for any \(\rho \ll \pi\), that \(2\mathop{\mathrm{\mathsf{TV}}}(\rho,\pi)^2 \le \mathsf{KL}(\rho\,\|\,\pi) \le \chi^2(\rho\,\|\,\pi)\); the first inequality between KL divergence and TV distance is Pinsker’s inequality. Furthermore, Rényi divergence is increasing in the index, so for any \(1 < q \le q'\), \(\mathsf{KL}(\rho\,\|\,\pi) \le \mathcal{R}_q(\rho\,\|\,\pi) \le \mathcal{R}_{q'}(\rho\,\|\,\pi)\).

6.1.1 Basic properties of divergences↩︎

All the divergences above (KL, chi-square, and Rényi) are not symmetric, and they do not satisfy the triangle inequality. However, the Rényi divergence satisfies a generalized triangle inequality. For any probability distributions \(\mu, \rho, \pi\) with \(\mu \ll \rho \ll \pi\), and for any \(q > 1\) and \(\lambda \in (0,1)\), we have: \[\mathcal{R}_q(\mu \,\|\, \pi) \le \frac{q-\lambda}{q-1} \, \mathcal{R}_{q/\lambda}(\mu \,\|\, \rho) + \mathcal{R}_{(q-\lambda)/(1-\lambda)}(\rho\,\|\, \pi)\,.\] We apply this in the proof of Lemma 8 with \(q=2\) and \(\lambda = \frac{2}{3}\).

We recall that relative Fisher information is convex (with respect to the usual linear structure of probability distributions) in the first argument. For any probability distributions \(\mu, \rho, \pi\) and any \(\lambda \in [0,1]\), we have: \[\mathsf{FI}(\lambda \mu + (1-\lambda) \rho \,\|\, \pi) \le \lambda \mathsf{FI}(\mu \,\|\, \pi) + (1-\lambda) \mathsf{FI}(\rho \,\|\, \pi)\,.\]

We recall the chain rule for KL divergence. For any joint probability distributions \(\rho^{XY}\) and \(\pi^{XY}\), with \(X\)-marginals \(\rho^X\) and \(\pi^X\), and conditional distributions \(\rho^{Y \mid X}\) and \(\pi^{Y \mid X}\), respectively, we have: \[\mathsf{KL}(\rho^{XY} \,\|\, \pi^{XY}) = \mathsf{KL}(\rho^X \,\|\, \pi^X) + \mathbb{E}_{x \sim \rho^X}\left[\mathsf{KL}(\rho^{Y \mid X=x} \,\|\, \pi^{Y \mid X=x}) \right].\]

We also recall the Donsker–Varadhan variational formula for KL divergence: For any probability distributions \(\rho, \pi\), and any function \(g \colon \mathbb{R}^d \to \mathbb{R}\), we have: \[\mathsf{KL}(\rho \,\|\, \pi) \ge \mathbb{E}_\rho[g] - \log \mathbb{E}_\pi[\exp(g)]\,.\]

6.1.2 Channel and data processing inequality↩︎

In this paper, we use the term channel from information theory to describe a process that adds randomness. Formally, a channel \(P\) from \(\mathbb{R}^d\) to \(\mathbb{R}^d\) is a collection of probability distributions, \(P = (P_y)_{y \in \mathbb{R}^d}\), where \(P_y\) is the output distribution of the channel when the input is a point mass at \(y \in \mathbb{R}^d\). Given an input probability distribution \(\rho\) over \(\mathbb{R}^d\), we denote the output of the channel by \(\rho P\), given by: \[\rho P = \int_{\mathbb{R}^d} \rho(y) \, P_y \, \mathrm{d}y.\]

A fundamental property in information theory is the data processing inequality (DPI), which states that statistical divergences cannot increase along a channel. Namely, given any channel \(P\), any input distributions \(\rho, \pi\), and statistical divergence \(D\) (including the TV distance, KL divergence, chi-square divergence, and Rényi divergence), we have: \[D(\rho P \,\|\, \pi P) \le D(\rho \,\|\, \pi).\] Notably, the data processing inequality does not hold for the relative Fisher information; see [16] for a counterexample even for the Gaussian channel.

6.2 Properties of probability distributions↩︎

We say a probability distribution \(\pi \propto \exp(-f)\) on \(\mathbb{R}^d\) is \(L\)-log-smooth for some \(L \in (0,\infty)\) if \(f \colon \mathbb{R}^d \to \mathbb{R}\) is twice-differentiable and \(\|\nabla^2 f(x)\|_{\mathsf{op}} \le L\) for all \(x \in \mathbb{R}^d\). We say \(\pi \propto \exp(-f)\) is \(\alpha\)-strongly log concave (SLC) for some \(\alpha > 0\) if \(f \colon \mathbb{R}^d \to \mathbb{R}\) is an \(\alpha\)-strongly convex function. If \(\alpha = 0\), then we say \(\pi\) is log-concave.

We say \(\pi\) satisfies a log-Sobolev inequality (LSI) with constant \(C_\mathsf{LSI}\in (0,\infty)\) if for any probability distribution \(\rho \ll \pi\), we have \[\mathsf{FI}(\rho \,\|\, \pi) \ge \frac{2}{C_\mathsf{LSI}} \mathsf{KL}(\rho \,\|\, \pi).\] We recall that if \(\pi\) is \(\alpha\)-SLC, then \(\pi\) satisfies a LSI with constant \(C_\mathsf{LSI}= 1/\alpha\).

We say that a probability distribution \(\pi\) on \(\mathbb{R}^d\) is \(\beta\)-sub-Gaussian for some \(\beta \in (0, \infty)\) if for any \(v \in \mathbb{R}^d\), \[\mathbb{E}_{X \sim \pi}[\exp(\langle v, X - \mathbb{E}_\pi[X] \rangle)] \le \exp\left(\frac{1}{2} \beta^2 \|v\|^2\right).\] For a function \(G \colon \mathbb{R}^d \to \mathbb{R}^m\) for \(m \ge 1\), we say “\(G\) is \(\beta\)-sub-Gaussian under \(\pi\)” to mean that the law of \(G(X)\) is \(\beta\)-sub-Gaussian when \(X \sim \pi\). We recall that if \(\pi\) is \(L\)-log-smooth, so \(\nabla \log \pi\) is \(L\)-Lipschitz, then \(\nabla \log \pi\) is \(\sqrt{L}\)-sub-Gaussian under \(\pi\); see [31].

7 Helper results↩︎

7.1 Sub-Gaussian score lemma↩︎

7.2 Rényi divergence along the heat flow↩︎

8 Proofs for the ideal proximal sampler↩︎

8.1 Proof of Lemma 1↩︎

8.2 Proof of Theorem 1↩︎

9 Proofs for the proximal sampler with an approximate RGO implementation↩︎

9.1 Proof of Lemma 2↩︎

9.2 Proof of Lemma 3↩︎

9.3 Proof of Lemma 4↩︎

9.4 Chi-squared divergence along the iterates↩︎

We also need the following lemma, which controls the chi-squared divergence along the iterates.

9.5 Lemma on Fisher information bound from chi-square for approximate RGO↩︎

To prove Corollary 3, we establish the following lemma.

9.6 Proof of Corollary 3↩︎

Corollary 3 follows by combining Lemma 4 and Lemma 8.

9.7 Proof of Theorem 2↩︎

9.8 Proof of Theorem 3↩︎

References↩︎

---

1. Yale University, Department of Statistics and Data Science. Email: sinho.chewi@yale.edu.↩︎

2. Yale University, Department of Computer Science. Email: andre.wibisono@yale.edu. This work was supported by NSF awards CCF–2403391 and CAREER CCF–2443097.↩︎

3. Here, the big-O hides dependence on the dimension \(d\), the initial KL divergence, and the smoothness \(L\).↩︎