Stochastic Runge-Kutta Accelerates Langevin Monte Carlo and Beyond

Introduction

Sampling from a probability distribution is a fundamental problem that arises in machine learning, statistics, and optimization. In many situations, the goal is to obtain samples from a target distribution given only the unnormalized density . A prominent approach to this problem is the method of Markov chain Monte Carlo (MCMC), where an ergodic Markov chain is simulated so that iterates converge exactly or approximately to the distribution of interest .

MCMC samplers based on numerically integrating continuous-time dynamics have proven very useful due to their ability to accommodate a stochastic gradient oracle . Moreover, when used as optimizations algorithms, these methods can deliver strong theoretical guarantees in non-convex settings . A popular example in this regime is the unadjusted Langevin Monte Carlo (LMC) algorithm . Fast mixing of LMC is inherited from exponential Wasserstein decay of the Langevin diffusion, and numerical integration using the Euler-Maruyama scheme with a sufficiently small step size ensures the Markov chain tracks the diffusion. Asymptotic guarantees of this algorithm are well-studied , and non-asymptotic analyses specifying explicit constants in convergence bounds were recently conducted .

Our contributions can be summarized as follows:

We provide a broadly applicable theorem for establishing convergence rates of sampling algorithms based on discretizing Itô diffusions exhibiting exponential Wasserstein-2 contraction to the target invariant measure. The convergence rate is explicitly expressed in terms of the contraction rate of the diffusion and local properties of the numerical scheme, both of which can be easily derived.

We provide examples and numerical studies of sampling from both convex and non-convex potentials with SRK methods and show they lead to better stability and lower asymptotic errors.

The convergence analyses of sampling using the overdamped and underdamped Langevin diffusion were extended to the non-convex setting . For the Langevin diffusion, the most common assumption on the potential is strong convexity outside a ball of finite radius, in addition to Lipschitz smoothness and twice differentiability . More generally, Vempala and Wibisono showed that convergence in the KL divergence of LMC can be derived assuming a log-Sobolev inequality on the target measure. For general Itô diffusions, the notion of distant dissipativity is used to study convergence to target measures with non-convex potentials in the $1$-Wasserstein distance. Different from these works, our non-convex convergence analysis, due to conducted in $W_{2}$, requires the slightly stronger uniform dissipativity condition . In optimization, non-asymptotic results for stochastic gradient Langevin dynamics and its variants have been established for non-convex objectives .

Sampling with Discretized Diffusions

We study the problem of sampling from a target distribution $p(x)$ with the help of a candidate Itô diffusion given as the solution to the following stochastic differential equation (SDE):

It is straightforward to verify (2) for the above diffusion which implies that the target $p(x)$ is its invariant measure. Moreover, strong convexity of $f$ implies uniform dissipativity and ensures that the diffusion achieves fast convergence.

In practice, the Itô diffusion (1) (similarly (3)) cannot be simulated in continuous time and is instead approximated by a discrete-time numerical integration scheme. Owing to its simplicity, a common choice is the Euler-Maruyama (EM) scheme , which relies on the following update rule,

where $\sigma_{i}(x)$ denotes the $i$th column of $\sigma(x)$. Then, applying Itô’s lemma to the integral form of the SDE (1) with the starting point $X_{0}$ yields the following expansion around $X_{0}$ :

The expansion justifies the update rule of the EM scheme, since the discretization is nothing more than taking the first three terms on the right hand side of (6). Similarly, a mean-square order 1.0 SRK scheme for general Itô diffusions – introduced in Section 4.2 – approximates the first four terms. In principle, one may recursively apply Itô’s lemma to terms in the expansion to obtain a more fine-grained approximation. However, the appearance of non-Gaussian terms in the guise of iterated Brownian integrals presents a challenge for simulation. Nevertheless, it is clear that the above SRK scheme will be a more accurate local approximation than the EM scheme, due to accounting more terms in the expansion. As a result, the local deviation between the continuous-time process and Markov chain will be smaller. We characterize this property of a numerical scheme as follows.

Convergence Rates of Numerical Schemes for Sampling

We present a user-friendly and broadly applicable theorem that establishes the convergence rate of a diffusion-based sampling algorithm. We develop our explicit bounds in the $2$-Wasserstein distance based on two crucial steps. We first verify that the candidate diffusion exhibits exponential Wasserstein-2 contraction and thereafter compute the uniform local deviation orders of the scheme.

where $\delta_{x}P_{t}$ denotes the distribution of the diffusion $X_{t}$ starting from $x$. Moreover, if $r(t)=e^{-\alpha t}$ for some $\alpha>0$, then we say the diffusion has exponential $W_{2}$-contraction.

The above condition guarantees fast mixing of the sampling algorithm. For Itô diffusions, uniform dissipativity suffices to ensure exponential $W_{2}$-contraction $r(t)=e^{-\alpha t}$ [24, Prop. 3.3].

A diffusion defined by (1) is $\alpha$-uniformly dissipative if

For Itô diffusions with a constant diffusion coefficient, uniform dissipativity is equivalent to one-sided Lipschitz continuity of the drift with coefficient $-2\alpha$. In particular, for the overdamped Langevin diffusion (3), this reduces to strong convexity of the potential. Moreover, for this special case, exponential $W_{2}$-contraction of the diffusion and strong convexity of the potential are equivalent . We will ultimately verify uniform dissipativity for the candidate diffusions, but we first use $W_{2}$-contraction to derive the convergence rate of a diffusion-based sampling algorithm.

For a diffusion with invariant measure $\nu^{*}$, exponentially contracting $W_{2}$-rate $r(t)=e^{-\alpha t}$, and Lipschitz drift and diffusion coefficients, suppose its discretization based on a numerical integration scheme has uniform local deviation orders $(p_{1},p_{2})$ where $p_{1}\geq{1}/{2}$ and $p_{2}\geq p_{1}+{1}/{2}$. Let $\nu_{k}$ be the measure associated with the Markov chain obtained from the discretization after $k$ steps starting from the dirac measure $\nu_{0}=\delta_{x_{0}}$. Then, for constant step size $h$ satisfying

where $C_{h}$ is the step size constraint for obtaining the uniform local deviation orders, we have

Moreover, if $p_{1}>1/2$ and the step size additionally satisfies

Theorem 1 directly translates mean-square order results in the SDE literature to convergence rates of sampling algorithms in $W_{2}$. The proof deferred to Appendix A follows from an inductive argument over the local deviation at each step (see e.g. ), and the convergence is provided by the exponential $W_{2}$-contraction of the diffusion. To invoke the theorem and obtain convergence rates of a sampling algorithm, it suffices to (i) show that the candidate diffusion is uniformly dissipative and (ii) derive the local deviation orders for the underlying discretization. Below, we demonstrate this on both the overdamped Langevin and general Itô diffusions when the EM scheme is used for discretization, as well as the underdamped Langevin diffusion when a linearization is used for discretization . For these schemes, local deviation orders are either well-known or straightforward to derive. Thus, convergence rates for corresponding sampling algorithms can be easily obtained using Theorem 1.

Consider sampling from a target distribution whose potential is strongly convex using the underdamped Langevin diffusion:

Cheng et al. show that the continuous-time process $\{(X_{t},X_{t}+V_{t})\}_{t\geq 0}$ exhibits exponential $W_{2}$-contraction when the coefficients $\gamma$ and $u$ are appropriately chosen [8, Thm. 5]. Moreover, the scheme devised by linearizing the degenerate SDE for the augmented state has uniform local deviation orders $(1.5,2.0)$ Cheng et al. derive the uniform local mean-square deviation order. Jensen’s inequality implies that the local mean deviation is of the same uniform order. This entails uniform local deviation orders are $(2.0,2.0)$ and hence also $(1.5,2.0)$ when step size constraint $C_{h}\leq 1$; note $p_{2}\geq p_{1}+1/2$ is required to invoke Theorem 1. [8, Thm. 9]. Theorem 1 implies that the convergence rate is $\mathcal{O}(d^{1/2}\epsilon^{-1})$, where the dimension dependence is extracted from explicit bounds. This recovers the result by Cheng et al. [8, Thm. 1].

Since the first version of this paper appeared on arXiv, several new schemes were devised in the literature. We include the example of deriving the convergence rate for the recently proposed randomized midpoint method . This example demonstrates that Theorem 1 can also be applied to schemes that include additional randomness which is independent of that of the Brownian motion.

Shen and Lee discretize the underdamped Langevin diffusion with a variant of the midpoint method, where the midpoint is computed with the linearization scheme at a random time $\tau$ uniformly selected in $[0,h]$. It can be shown that the local mean-square deviation order is the same as that of the one-step linearization scheme [56, Lem. 2]. However, one sees that the local mean deviation order improves by inspecting the following bound

While computing the local deviation orders of a numerical scheme for a single step is often straightforward, it is not immediately clear how one might verify them uniformly for each iteration. This requires a uniform bound on moments of the Markov chain defined by the numerical scheme. As our second principal contribution, we explicitly bound the Markov chain moments of SRK schemes which, combined with Theorem 1, leads to improved rates by only accessing the first-order oracle.

Sampling with Stochastic Runge-Kutta and Improved Rates

We show that convergence rates of sampling can be significantly improved if an Itô diffusion with exponential $W_{2}$-contraction is discretized using SRK methods. Compared to the EM scheme, SRK schemes we consider query the same order oracle and improve on the deviation orders.

Theorem 1 hints that one may expect the convergence rate of sampling to improve as more terms of the Itô-Taylor expansion are incorporated in the numerical integration scheme. However, in practice, a challenge for simulation is the appearance of non-Gaussian terms in the form of iterated Itô integrals. Fortunately, since the overdamped Langevin diffusion has a constant diffusion coefficient, efficient SRK methods can still be applied to accelerate convergence.

The proof of this theorem is given in Appendix B where we provide explicit constants. The basic idea of the proof is to match up the terms in the Itô-Taylor expansion to terms in the Taylor expansion of the discretization scheme. However, extreme care is needed to ensure a tight dimension dependence.

For large-scale Bayesian inference, computing the full gradient of the potential can be costly. Fortunately, for SRK-LD, the convergence rate is retained when we replace the first-order oracle with an unbiased stochastic one, provided queries of the latter have a variance not overly large. We provide an informal discussion in Appendix E.

2 Sampling from Non-Convex Potentials with Itô Diffusions

For the Langevin diffusion, the conclusions of Theorem 1 only apply to distributions with strongly convex potentials, as exponential $W_{2}$-contraction of the Langevin diffusion is equivalent to strong convexity of the potential. This shortcoming can be addressed using a non-constant diffusion coefficient which allows us to sample from non-convex potentials using uniformly dissipative candidate diffusions. Below, we use a mean-square order 1.0 SRK scheme for general diffusions and achieve an improved convergence rate compared to sampling with the EM scheme.

We refer to the sampling algorithm as SRK-ID, which has the following update rule:

In practice, accurately simulating both the iterated Itô integrals $I_{(j,i)}$ and the Brownian motion increments $I_{(i)}$ simultaneously is difficult. We comment on two possible approximations based on truncating an infinite series in Appendix H.2.

Examples and Numerical Studies

We provide examples of our theory and numerical studies showing SRK methods achieve lower asymptotic errors, are stable under large step sizes, and hence converge faster to a prescribed tolerance. We sample from strongly convex potentials with SRK-LD and non-convex potentials with SRK-ID. Since our theory is in $W_{2}$, we compare with EM on $W_{2}$ and mean squared error (MSE) between iterates of the Markov chain and the target. We do not compare to schemes that require computing derivatives of the drift and diffusion coefficients. Since directly computing $W_{2}$ is infeasible, we estimate it using samples instead. However, sample-based estimators have a bias of order $\Omega(n^{-1/d})$ , so we perform a heuristic correction whose description is in Appendix G.

We consider sampling from a multivariate Gaussian mixture with density

The potential is strongly convex and has Lipschitz gradient and Hessian . One can also verify that it has a Lipschitz third derivative.

To obtain the potential, we generate data from the model with the parameter $\theta_{*}=\mathbf{1}_{d}$ following . To obtain each $x_{i}$, we sample a vector whose components are independently drawn from the Rademacher distribution and normalize it by the Frobenius norm of the sample matrix $X$ times $d^{-1/2}$. Note that our normalization scheme is different from that adopted in , where each $x_{i}$ is normalized by its Euclidean norm. We sample the corresponding $y_{i}$ from the model and fix the regularizer $\alpha=0.3d/\pi^{2}$.

To characterize the true posterior, we sample 50k particles driven by EM with a step size of $0.001$ until convergence. We subsample from these particles 5k examples to represent the true posterior each time we intend to estimate squared $W_{2}$. We monitor the kernel Stein discrepancy Unfortunately, there appear to be two definitions for KSD and the energy distance in the literature, differing in whether a square root is taken or not. We adopt the version with the square root taken. (KSD) using the inverse multiquadratic kernel with hyperparameters $\beta=-1/2$ and $c=1$ to measure the distance between the 100k particles and the true posterior. We confirm that these particles faithfully approximate the true posterior with the squared KSD being less than $0.002$ in all settings.

When sampling from a Gaussian mixture and the posterior of BLR, we observe that SRK-LD leads to a consistent improvement in the asymptotic error compared to the EM scheme when the same step size is used. In particular, Figure 1 (a) plots the estimated asymptotic error in squared $W_{2}$ of different step sizes for 2D and 20D Gaussian mixture problems and shows that SRK-LD is surprisingly stable for exceptionally large step sizes. Figure 1 (b) plots the estimated error in squared $W_{2}$ as the number of iterations increases for 2D BLR. We include additional results on problems in 2D and 20D with error estimates in squared $W_{2}$ and the energy distance along with a wall time analysis in Appendix H.

2 Non-Convex Potentials

We consider sampling from the non-convex potential

where $\beta,\gamma>0$ are scalar parameters of the distribution. The corresponding density is a simplified abstraction for the posterior distribution of Student’s t regression with a pseudo-Huber prior . One can verify that when $\beta+\left\lVert x\right\rVert_{2}^{2}<1$ and $(4\gamma+1)\left\lVert x\right\rVert_{2}^{2}<(2\gamma+1)\sqrt{\beta+\left\lVert x\right\rVert_{2}^{2}}$ , the Hessian has a negative eigenvalue. The candidate diffusion, where the drift coefficient is given by (2) and diffusion coefficient $\sigma(x)=g(x)^{1/2}I_{d}$ with g(x)=\bigl{(}\beta+\|x\|_{2}^{2}\bigr{)}^{1/2}, is uniformly dissipative if $\tfrac{1}{2}-|\gamma-\tfrac{1}{2}|\tfrac{2}{\beta^{1/2}}-\tfrac{d}{8\beta^{1/2}}>0$. Indeed, one can verify that $\mu_{1}(g)\leq 1$, $\mu_{2}(g)\leq\tfrac{2}{\beta^{1/2}}$, and $\mu_{1}(\sigma)\leq\tfrac{1}{2\beta^{1/4}}$. Therefore,

Moreover, $b$ and $\sigma$ have Lipschitz first two derivatives, and the latter satisfies the sublinear growth condition in Theorem 3.

To study the behavior of SRK-ID, we simulate using both SRK-ID and EM. For both schemes, we simulate with a step size of $10^{-3}$ initiated from the same 50k particles approximating the stationary distribution obtained by simulating EM with a step size of $10^{-6}$ until convergence. We compute the MSE between the continuous-time process and the Markov chain with the same Brownian motion for $300$ iterations when we observe the MSE curve plateaus. We approximate the continuous-time process by simulating using the EM scheme with a step size of $10^{-6}$ similar to the setting in . To obtain final results, we average across ten independent runs. We note that the MSE upper bounds $W_{2}$ due to the latter being an infimum over all couplings. Hence, the MSE value serves as an indication of the convergence performance in $W_{2}$.

Figure 1 (c) shows that for $\beta=0.33$, $\gamma=0.5$ and $d=1$, when simulating from a good approximation to the target distribution with the same step size, the MSE of SRK-ID remains small, whereas the MSE of EM converges to a larger value. However, this improvement diminishes as the dimensionality of the sampling problem increases. We report additional results with other parameter settings in Appendix H.2.2. Notably, we did not observe significant differences in the estimated squared $W_{2}$ values. We suspect this is due to the discrepancy being dominated by the bias of our estimator.

Discussion

We established convergence rates of samplings algorithm obtained by discretizing Itô diffusions with exponential $W_{2}$-contraction based on local properties of numerical schemes. The user-friendly conditions promote one to derive rates based on the uniform orders of the local deviation. In addition, we showed that discretizing diffusions with SRK schemes leads to improved rates in $W_{2}$ for both strongly convex potentials and a certain class of non-convex potentials.

Despite focusing on SRK methods, Theorem 1 can be used to obtain convergence rates for other classes of schemes. For the underdamped Langevin diffusion, quasi-symplectic schemes that rely on Runge-Kutta-type updates can achieve mean-square order 2.0 and beyond . For general Itô diffusions, there exist schemes of mean-square order 1.5 and beyond, using the Fourier-Legendre series to approximate the Lévy area .

Compared to some existing proofs for convergence rates in $W_{2}$ (e.g. [8, Thm. 1]), our Theorem 1 requires two conditions (the uniform local mean and mean-square deviation bounds), neither of which can be eliminated in order to obtain a tight convergence bound. The uniform local mean deviation appears in our proof due to a direct expansion of the squared $2$-norm. This can be thought of a natural consequence of our convergence bounds being based on $W_{2}$. An avenue of interest is to see whether additional conditions can be identified to obtain refined bounds in $W_{p}$ for even integer $p>2$.

Another direction of interest is to relax the $W_{2}$-contraction condition on the diffusion to $W_{1}$-contraction or $W_{1}$-decay. This would enable us to leverage results based on distant dissipativity, and consequently allow us to sample from a wider class of non-convex potentials . Orthogonally, for the overdamped Langevin diffusion, the $W_{2}$-contraction condition may be relaxed to a log-Sobolev inequality condition on the target measure, if the discretization analysis is adapted to be based on the KL divergence . This would also broaden the class of non-convex potentials from which we can sample with theoretical guarantees.

Parallel to studying sampling from a mean-square convergence aspect, works in numerical analysis have established convergence results in the weak sense for SRK schemes applied to ergodic SDEs with techniques as aromatic trees and B-series . However, moment bounds in these works are proven by generic arguments (see e.g. [46, Lem. 2.2.2]), and reasoning about the rate’s dimension dependence becomes less obvious. Refined non-asymptotic convergence bounds would provide more insight for these algorithms’ performance on practical problems.

Lastly, the convergence results in $W_{2}$ for SRK-LD and SRK-ID can be augmented to yield generalization bounds for optimization when the excess risk is characterized using the Gibbs distribution .

Acknowledgments

We thank Mufan Li for insightful discussions, and Jimmy Ba, David Duvenaud and Taiji Suzuki for helpful comments on an early draft of this work. We also thank anonymous reviewers for helpful suggestions. MAE is partially funded by NSERC and CIFAR AI Chairs program at Vector Institute.

References

Appendix A Proof of Theorem 1

Let $\{X_{t}\}_{t\geq 0}$ denote the continuous-time process defined by the SDE (1) initiated from the target stationary distribution, driven by the Brownian motion $\{B_{t}\}_{t\geq 0}$. Since the continuous-time transition kernel preserves the stationary distribution, the marginal distribution of $\{X_{t}\}_{t\geq 0}$ remains to be the stationary distribution for all $t\geq 0$.

This we can achieve due to exponential $W_{2}$-contraction. We define the process $\{Z_{s}\}_{s\geq t_{k}}$ as follows

For $s\in[t_{k},t_{k}+h]$, by Young’s inequality, Jensen’s inequality, and Itô isometry,

By the integral form of Grönwall’s inequality for continuous functions,

Combining (13) (17) and (18), by the Cauchy-Schwarz inequality,

Let $\lambda_{3}=8\lambda_{1}^{1/2}\mu_{1}(b)^{1/2}+2\lambda_{2}^{1/2}$. Then, by the Cauchy–Schwarz inequality, we obtain a recursion

where the third to last inequality follows from $e^{-2\alpha h}<1-\alpha h$ when $\alpha h<1/2$, and the second to last inequality follows from the elementary relation below with the choice of $\kappa=\alpha/2$

Let $\eta=1-\alpha h/2\leq e^{-\alpha h/2}\leq 1$. By unrolling the recursion,

Let $\nu_{k}$ and $\nu^{*}$ be the measures associated with the $k$th iterate of the Markov chain and the target distribution, respectively. Since $W_{2}$ is defined as an infimum over all couplings,

To ensure $W_{2}$ is less than some small positive tolerance $\epsilon$, we need only ensure the two terms in the above inequality are each less than $\epsilon/2$. Some simple calculations show that it suffices that

Note that for small enough positive tolerance $\epsilon$, when the step size satisfies (34), it suffices that

Appendix B Proof of Theorem 2

Verifying the order conditions in Theorem 1 for SRK-LD requires bounding the second, fourth, and sixth moments of the Markov chain. In principle, one may employ an exponential moment bound argument using a Lyapunov function. However, in this case, the tightness of the final convergence bound may depend on the selection of the Lyapunov function, and reasoning about the dimension dependence can become less obvious. Here, we directly bound all the even moments by expanding the expression. Intuitively, one expects the $2n$th moments of the Markov chain iterates to be $\mathcal{O}(d^{n})$. The following proofs assume Lipschitz smoothness of the potential to a certain order and dissipativity.

For constants $\alpha,\beta>0$, the diffusion satisfies the following

For the Langevin diffusion, dissipativity directly follows from strong convexity of the potential . Here, $\alpha$ can be chosen as the strong convexity parameter, provided $\beta$ is an appropriate constant of order $\mathcal{O}(d)$.

Additionally, we assume the discretization has a constant step size $h$ and the timestamp of the $k$th iterate is $t_{k}$ as per the proof of Theorem 1. To simplify notation, we define the following

If the second moment of the initial iterate is finite, then the second moments of Markov chain iterates defined in (11) are uniformly bounded by a constant of order $\mathcal{O}(d)$, i.e.

where $N_{1}=2\pi_{2,2}(f)\left(2+4\pi_{2,1}(f)\right)$ and $N_{2}=2\pi_{2,2}(f)\left(6d+2\right)$.

Additionally, by the Cauchy-Schwarz inequality,

Combining (35) and (36), we obtain the following using AM–GM,

where $N_{3}=2\sqrt{2d\pi_{2,2}(f)}\left(1+\sqrt{\pi_{2,1}(f)}\right)$ and $N_{4}=2\sqrt{2d\pi_{2,2}(f)}\left(1+\sqrt{3d}\right)$.

Additionally, by Stein’s lemma for multivariate Gaussians,

Therefore, by dissipativity and the lower bound (37),

Notice the second equality above also implies

By Taylor’s Theorem with the remainder in integral form,

Since $\nabla f$ is Lipschitz, $\nabla^{2}f$ is bounded, and

Therefore, for $h<1\wedge\alpha/(4\mu_{2}(f)\pi_{2,2}(f))$,

where $N_{5}=\left(\alpha\beta+\mu_{2}(f)\pi_{2,2}(f)+d+5d\mu_{3}(f)\right)+2\beta$.

Putting things together, for $h<3\alpha/(2N_{1}+4)$, we obtain

For $h<1$, by unrolling the recursion, we obtain the following

B.1.2 2​n2𝑛2nth Moment Bound

and constants $N_{7,n}$ to $N_{12,n}$ are given in the proof, if the step size

Our proof is by induction. The base case is given in Lemma 4. For the inductive case, we prove that the $2n$th moment is uniformly bounded by a constant of order $\mathcal{O}(d^{n})$, assuming the $2(n\text{-}1)$th moment is uniformly bounded by a constant of order $\mathcal{O}(d^{n-1})$.

Now, we bound the expectation of $A$ using (42),

where $N_{7,n}=\left(2dn+nN_{5}+4dn(n-1)\right)\mathcal{U}_{2(n-1)}=\mathcal{O}(d^{n})$.

Next, we bound the expectation of $B$. By the Cauchy–Schwarz inequality,

where the second to last inequality follows from Young’s inequality for products with three variables.

Since $N_{8,n}(p)$ does not depend on the dimension, let

In the following, we bound the expectations of $B_{1}$ and $B_{2}$ separately. By Young’s inequality for products and the function $x\mapsto x^{\scriptstyle 1/(2k_{3}+k_{5}+k_{6})}$ being concave on the positive domain,

Thus, when $h<(3n\alpha/8N_{11,n})^{2}$, by (43) and (48),

B.2 Local Deviation Orders

We first provide two lemmas on bounding the second and fourth moments of the change in the continuous-time process. These will be used later when we verify the order conditions.

Suppose $X_{t}$ is the continuous-time process defined by (3) initiated from some iterate of the Markov chain $X_{0}$ defined by (11), then the second moment of $X_{t}$ is uniformly bounded by a constant of order $\mathcal{O}(d)$, i.e.

where $\mathcal{U}_{2}^{\prime}=\mathcal{U}_{2}+2(\beta+d)/\alpha$.

Suppose $X_{t}$ is the continuous-time process defined by (3) initiated from some iterate of the Markov chain $X_{0}$ defined by (11), then

where $C_{0}=2\pi_{2,2}(f)\left(1+\mathcal{U}_{2}^{\prime}\right)+4d$.

Suppose $X_{t}$ is the continuous-time process defined by (3) initiated from some iterate of the Markov chain $X_{0}$ defined by (11), then the fourth moment of $X_{t}$ is uniformly bounded by a constant of order $\mathcal{O}(d^{2})$, i.e.

where $\mathcal{U}_{4}^{\prime}=\mathcal{U}_{4}+(2\beta+6)\mathcal{U}_{2}^{\prime}/\alpha$.

By Itô’s lemma, dissipativity, and Lemma 6,

Suppose $X_{t}$ is the continuous-time process defined by (3) initiated from some iterate of the Markov chain $X_{0}$ defined by (11), then

where $C_{1}=8\pi_{2,4}(f)\left(1+\mathcal{U}_{4}^{\prime}\right)+32d(d+2)$.

Since the two processes share the same Brownian motion,

We bound the second moment of $R(s)$ by bounding those of $R_{1}(s)$, $R_{2}(s)$, and $R_{3}(s)$ separately. For $R_{1}(s)$, by the Cauchy–Schwarz inequality,

Next, we characterize the terms in the Markov chain update. By Taylor’s theorem,

Using the above information, we bound the second moments of $\rho_{1}(t)$ and $\rho_{2}(t)$,

B.2.2 Local Mean Deviation

The proof is similiar to that of Lemma 10 with slight variations on truncating the expansions. Recall since the two processes share the same Brownian motion,

By Taylor’s theorem with the remainder in integral form,

Now, we show the following equality in a component-wise manner,

To see this, recall that odd moments of the Brownian motion is zero. So, for each $\partial_{i}f$,

Adding the previous two equations together, we obtain the desired equality (53).

Next, we bound the second moments of $\bar{R}_{1}(s)$ and $\bar{R}_{2}(s)$. For $\bar{R}_{1}(s)$, recall from the proof of Lemma 10,

We bound the sixth moments of $v_{1}$, $v_{1}^{\prime}$ and $v_{2}$ using $\nicefrac{{1}}{{2}}+\nicefrac{{1}}{{\sqrt{6}}}<1$, $\nicefrac{{1}}{{2}}-\nicefrac{{1}}{{\sqrt{6}}}<\nicefrac{{1}}{{2}}$ and the closed form moments of a chi-squared random variable with $d$ degrees of freedom $\chi(d)^{2}$ ,

Now, we bound the second moments of $\bar{\rho}_{1}(t)$ and $\bar{\rho}_{2}(t)$ using the derived sixth-moment bounds,

B.3 Invoking Theorem 1

Now, we invoke Theorem 1 with our derived constants. We obtain that if the constant step size

Appendix C Proof of Theorem 3

Verifying the order conditions in Theorem 1 for SRK-ID requires bounding the second and fourth moments of the Markov chain.

The following proofs only assume Lipschitz smoothness of the drift coefficient $b$ and diffusion coefficient $\sigma$ to a certain order and a generalized notion of dissipativity for Itô diffusions.

For constants $\alpha,\beta>0$, the diffusion satisfies the following

For general Itô diffusions, dissipativity directly follows from uniform dissipativity, where $\beta$ is an appropriate constant of order $\mathcal{O}(d)$. Additionally, we assume the discretization has a constant step size $h$ and the timestamp of the $k$th iterate is $t_{k}$ as per the proof of Theorem 1. To simplify notation, we rewrite the update as

Let $p\geq 2$. If $\{G_{s}\}_{s\geq 0}$ is a $d\times m$ matrix-valued process, and $\{B_{t}\}_{t\geq 0}$ is a $d$-dimensional Brownian motion, both of which are adapted to the filtration $\{\mathcal{F}_{s}\}_{s\geq 0}$ such that for some fixed $t>0$, the following relation holds

In particular, equality holds when $p=2$.

The above theorem can be proved directly using Itô’s lemma and Itô isometry, with the help of Hölder’s inequality. The theorem can also be seen as a natural consequence of the Burkholder-Davis-Gundy Inequality .

Let even integer $p\geq 2$. Then, the following relation holds

By Taylor’s Theorem with the remainder in integral form,

To prove the following moment bound lemmas for SRK-ID, we recall a standard quadratic moment bound result whose proof we omit and provide a reference of.

If the second moment of the initial iterate is finite, then the second moments of Markov chain iterates defined in (12) are uniformly bounded, i.e.

and constants $M_{1}$ and $M_{2}$ are given in the proof, if the constant step size

We bound the remaining terms by direct computation. By linear growth,

Putting things together, for $h<1\wedge{\alpha^{2}}/(4M_{1}^{2})$,

Unrolling the recursion gives the following for $h<1\wedge 1/m^{2}$

C.1.2 2​n2𝑛2nth Moment Bound

Before bounding the $2n$th moments, we first generalize Lemma 14 to arbitrary even moments.

The remaining follows easily from Lemma 31. ∎

Our proof is by induction. The base case is given in Lemma 16. For the inductive case, we prove that the $2n$th moment is uniformly bounded by a constant, assuming the $2(n\text{-}1)$th moment is uniformly bounded by a constant.

and with slight abuse of notation, we hide the explicit dependence on $k_{1},\dots,k_{10}$ for the exponents

Note that $p_{1}+p_{2}+p_{3}+p_{4}=2n$. Since $h<1\wedge 1/m^{2}$, we may cancel out the $m$ factor in some of the terms. One can verify that the only remaining term that is $m$-dependent is

Using this information, Lemma 17, Lemma 15, the Cauchy–Schwarz inequality, and $p_{3}+p_{4}\leq 2n$,

By the inductive hypothesis, (55) and (56), and $h<1\wedge n^{2}\alpha^{2}/(4M_{3,n}^{2})$, we obtain the recursion

For $h<1\wedge 1/m^{2}$, by unrolling the recursion, we obtain

C.2 Local Deviation Orders

In this section, we verify the local deviation orders for SRK-ID. The proofs are again by matching up terms in the Itô-Taylor expansion of the continuous-time process to terms in the Taylor expansion of the numerical integration scheme. Extra care needs to be taken for a tight dimension dependence.

Suppose $X_{t}$ is the continuous-time process defined by (1) initiated from some iterate of the Markov chain $X_{0}$ defined by (12), then the second moment of $X_{t}$ is uniformly bounded, i.e.

where $\mathcal{V}_{2}^{\prime}=\mathcal{V}_{2}+\beta/\alpha$.

Suppose $X_{t}$ is the continuous-time process defined by (1) initiated from some iterate of the Markov chain $X_{0}$ defined by (12), then

To bound the fourth moment of change in continuous-time, we use the following lemma.

Assuming $\{X_{t}\}_{t\geq 0}$ is the solution to the SDE (1), under the condition that the drift coefficient $b$ and diffusion coefficient $\sigma$ are Lipschitz. If $\sigma$ satisfies the following sublinear growth condition

and the diffusion is dissipative, then for $n\geq 2$, we have the following relation

where the (infinitesimal) generator $\mathcal{A}$ is defined as

and the constant $\beta_{n}=\mathcal{O}(d^{\frac{n}{2}})$.

By definition of the generator and dissipativity,

We define the following shorthand notation

Putting things together, we obtain the following bound

where $\beta_{n}=\beta_{n}^{(1)}+\beta_{n}^{(2)}=\mathcal{O}(d^{\frac{n}{2}})$. ∎

Suppose $X_{t}$ is the continuous-time process defined by (1) initiated from some iterate of the Markov chain $X_{0}$ defined by (12), then the fourth moment of $X_{t}$ is uniformly bounded, i.e.

where $\mathcal{V}_{4}^{\prime}=\mathcal{V}_{4}+\beta_{4}/\alpha$.

By Dynkin’s formula applied to the function $(t,x)\mapsto e^{\alpha t}\left\lVert x\right\rVert_{2}^{4}$ and Lemma 21,

Suppose $X_{t}$ is the continuous-time process defined by (1) from some iterate of the Markov chain $X_{0}$ defined by (12), then

Recall the operators $L$ and $\Lambda_{i}$ ($i=1,\dots,m$) defined in (5). By Itô’s lemma,

By Taylor’s theorem with the remainder in integral form,

By Itô isometry and the Cauchy-Schwarz inequality,

Now, we bound the second moments of $\phi_{1}^{(i)}(t)$ and $\phi_{2}^{(i)}(t)$,

Combining (71), (72), (73), (74), and (80),

C.2.2 Local Mean Deviation

Recall the operators $L$ and $\Lambda_{i}$ ($i=1,\dots,m$) defined in (5). By Itô’s lemma,

Now, we bound the second moments of $\bar{S}_{1}(t)$ and $\bar{S}_{2}(t)$,

Now, we bound the second moment of the difference between $\phi_{1}^{(i)}(t)$ and $\phi_{2}^{(i)}(t)$,

C.3 Invoking Theorem 1

Now, we invoke Theorem 1 with our derived constants. We obtain that if the constant step size

Appendix D Convergence Rate for Example 2

Verifying the order conditions in Theorem 1 of the EM scheme for uniformly dissipative diffusions requires bounding the second moments of the Markov chain. Recall, dissipativity (Definition C.1) follows from uniform dissipativity of the Itô diffusion.

If the second moment of the initial iterate is finite, then the second moments of Markov chain iterates defined in (4) are uniformly bounded, i.e.

By odd moments of Gaussian variables being zero and the step size condition,

D.2 Local Deviation Orders

Before verifying the local deviation orders, we first state two auxiliary lemmas. We omit the proofs, since they are almost identical to that of Lemma 6 and Lemma 7, respectively.

Suppose $X_{t}$ is the continuous-time process defined by (1) initiated from some iterate of the Markov chain $X_{0}$ defined by (4), then the second moment of $X_{t}$ is uniformly bounded, i.e.

Suppose $X_{t}$ is the continuous-time process defined by (1) initiated from some iterate of the Markov chain $X_{0}$ defined by (4), then

By Itô isometry and Lipschitz of the drift and diffusion coefficients,

D.2.2 Local Mean Deviation

Since the last two terms in the above inequality are Martingales,

D.3 Invoking Theorem 1

Now, we invoke Theorem 1 with our derived constants. We obtain that if the constant step size

Appendix E Convergence of SRK-LD Under an Unbiased Stochastic Oracle

Similarly, one can derive the new local mean deviation,

One can replace the corresponding terms in (33) and obtain a recursion. Note however, to ensure unrolling the recursion gives a convergence bound, one would need that $\sigma^{2}<\mathcal{O}(\alpha h)$.

Appendix F Auxiliary Lemmas

We list standard results used to develop our theorems and include their proofs for completeness.

Multiplying both sides of the inequality by $m^{n}$ completes the proof. ∎

For the $d$-dimensional Brownian motion $\{B_{t}\}_{t\geq 0}$,

We consider the case where $d=1$. The multi-dimensional case follows naturally, since we assume different dimensions of the Brownian motion vector are independent. Let $t_{k}=\delta k$, we define

Since $S_{m}$ is a sum of Gaussian random variables, it is also Gaussian. By linearity of expectation and independence of Brownian motion increments,

Since $S_{m}\overset{\text{a.s.}}{\to}Z_{t}$ as $m\to\infty$ by the strong law of large numbers, we conclude that $Z_{t}\sim\mathcal{N}\left(0,t^{3}/3\right)$. ∎

Note $\left\lVert B_{t}\right\rVert_{2}^{2}$ may be expressed as the sum of squared Gaussian random variables, i.e.

Observe that this is also a multiple of the chi-squared random variable with $d$ degrees of freedom $\chi(d)^{2}$. Its $n$th moment has the following closed form ,

Then, the vector Laplacian of its gradient is bounded, i.e.

Then, the vector Laplacian of its gradient is $d\mu_{4}$-Lipschitz, i.e.

Let $g(x)=\Delta(f)(x)$. Since $f\in C^{4}$, we may switch the order of partial derivatives,

By Taylor’s theorem with the remainder in integral form,

Note that $\nabla^{2}g(x)$ can be written as a sum of $d$ matrices, each being a sub-tensor of $\nabla^{4}f(x)$, due to the the trace operator, i.e.

Since the operator norm of $\nabla^{4}f(x)$ upper bounds the operator norm of each of its sub-tensor,

Recall the third derivative is $\mu_{4}$-Lipschitz, we obtain

Appendix G Estimating the Wasserstein Distance

For a Borel measure $\mu$ defined on a compact and separable topological space $\mathcal{X}$, a sample-based empirical measure $\mu_{n}$ may asymptotically serve as a proxy to $\mu$ in the $W_{p}$ sense for $p\in[1,\infty)$, i.e.

This is a consequence of the Wasserstein distance metrizing weak convergence and that the empirical measure converges weakly to $\mu$ almost surely .

However, in the finite-sample setting, this distance is typically non-negligible and worsens as the dimensionality increases. Specifically, generalizing previous results based on the $1$-Wasserstein distance , Weed and Bach showed that for $p\in[1,\infty)$,

where $t$ is less than the lower Wasserstein dimension $d_{*}(\mu)$. This presents a severe challenge in estimating the $2$-Wasserstein distance between probability measures using samples.

To better detect convergence, we zero center a simple sample-based estimator by subtracting the null responses and obtain the following new estimator:

where $\hat{\nu}_{n}$ and $\hat{\nu}_{n}^{\prime}$ are based on two independent samples of size $n$ from $\mu$, and similarly for $\hat{\nu}_{n}$ and $\hat{\nu}_{n}^{\prime}$ from $\nu$. This estimator is inspired by the contruction of distances in the maximum mean discrepancy family and the Sinkhorn divergence . Note that the $2$-Wasserstein distance between finite samples can be computed conveniently with existing packages that solves a linear program. Although the new estimator is not guaranteed to be unbiased across all settings, it is unbiased when the two distributions are the same.

where the expectations are approximated via averaging 100 independent draws. Figure 2 reports the deviation across different sample sizes and dimensionalities, where $\mu$ and $\nu$ differ only in either mean or covariance. While the corrected estimator is not unbiased, it is relatively more accurate.

In addition, Figure 3 demonstrates that our bias-corrected estimator becomes more accurate as the two distributions are closer. This indicates that our proposed estimator may provide a more reliable estimate of the 2-Wasserstein distance when the sampling algorithm is close to convergence.

Appendix H Additional Numerical Studies

In this section, we include additional numerical studies complementing Section 5.

We first include additional plots of error estimates in $W_{2}$ and the energy distance for sampling from a Gaussian mixture and the posterior of BLR. The results indicate that the reduction in asymptotic error is consistent across problems with varying dimensionalities that we consider. In the end, we conduct a wall time analysis and show that SRK-LD is competitive in practice.

Figure 4 shows the estimated $W_{2}$ error as the number of iterations increase for the 2D and 20D Gaussian mixture and BLR problems with the parameter settings described in Section 5. We observe consistent improvement in the asymptotic error across different settings in which we experimented.

H.1.2 Asymptotic Error vs Dimensionality and Step Size

Figure 6 (a) and (b) respectively show the asymptotic error against dimensionality and step size for Gaussian mixture sampling. We perform least squares regression in both plots. Plot (a) shows results when a step size of $0.5$ is used. Plot (b) is on semi-log scale, where the quantities are estimated for a 10D problem.

H.1.3 Wall Time

Figure 7 shows the wall time against the estimated $W_{2}^{2}$ of SRK-LD compared to the EM scheme for a 20D Gaussian mixture sampling problem. On a 6-core CPU with 2 threads per core, we observe that SRK-LD is roughly $\times$ 2.5 times as costly as EM per iteration. However, since SRK-LD is more stable for large step sizes, we may choose a step size much larger for SRK-LD compared to EM, in which case its iterates converge to a lower error within less time.

H.2 Non-Convex Potentials

We first discuss how we approximate the iterated Itô integrals, after which we include additional numerical studies varying the dimensionality of the sampling problem.

Simulating both the iterated Itô integrals $I_{(l,i)}$ and the Brownian motion increments $I_{(i)}$ exactly is difficult. We adopt the Kloeden-Platen-Wright approximation, which has an MSE of order $h^{2}/n$, where $n$ is the number of terms in the truncation . The infinite series can be written as follows:

where $\xi_{l,k},\xi_{i,k},\eta_{i,k},\eta_{l,k}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(0,1)$. $A_{(l,i)}$ is known as the Lévy area and is notoriously hard to simulate .

For SDE simulation, in order for the scheme to obtain the same strong convergence order under the approximation, the MSE in the approximation of the Itô integrals must be negligible compared to the local mean-square deviation of the numerical integration scheme. For our experiments, we use $n=3000$, following the rule of thumb that $n\propto h^{-1}$ . Although simulating the extra terms can become costly, the computation may be vectorized, branched off from the main update, and parallelized on an additional thread, since it does not require any information of the current iterate.

Wiktorsson et al. proposed to add a correction term to the truncated series, which results in an approximation that has an MSE of order $h^{2}/n^{2}$. In this case, $n\propto h^{-1/2}$ terms are effectively required. We note that analyzing and comparing between different Lévy area approximations is beyond the scope of this paper.

H.2.2 Additional Results

Figure 8 shows the MSE of simulations starting from a faithful approximation to the target. We adopt the same simulation settings as described in Section 5.2. We observe diminishing gains as the dimensionality increases across all settings with differing $\beta$ and $\gamma$ parameters in which we experimented. These empirical findings corroborate our theoretical results. Note that the corresponding diffusion in all settings are still uniformly dissipativity, yet the potential may become convex when $\beta$ is large. Nevertheless, the potential is never strongly convex when $\beta$ is positive due to the linear growth term.