When is the Convergence Time of Langevin Algorithms Dimension Independent? A Composite Optimization Viewpoint

Yoav Freund^⋄ University of California, San Diego Yi-An Ma^⋄ University of California, San Diego Tong Zhang^⋄ Google Research The Hong Kong University of Science and Technology

Abstract

There has been a surge of works bridging MCMC sampling and optimization, with a specific focus on translating non-asymptotic convergence guarantees for optimization problems into the analysis of Langevin algorithms in MCMC sampling. A conspicuous distinction between the convergence analysis of Langevin sampling and that of optimization is that all known convergence rates for Langevin algorithms depend on the dimensionality of the problem, whereas the convergence rates for optimization are dimension-free for convex problems. Whether a dimension independent convergence rate can be achieved by Langevin algorithm is thus a long-standing open problem. This paper provides an affirmative answer to this problem for large classes of either Lipschitz or smooth convex problems with normal priors. By viewing Langevin algorithm as composite optimization, we develop a new analysis technique that leads to dimension independent convergence rates for such problems.

1 Introduction

Two of the major themes in machine learning are point prediction and uncertainty quantification. Computationally, they manifest in two types of algorithms: optimization and Markov chain Monte Carlo (MCMC). While both strategies have developed relatively separately for decades, there is a recent trend in relating both strands of research and translating nonasymptotic convergence guarantees in gradient based optimization methods to those in MCMC Dalalyan (2017); Durmus and Moulines (2019); Dalalyan and Karagulyan (2017); Cheng et al. (2018b); Cheng and Bartlett (2018); Dwivedi et al. (2018); Mangoubi and Smith (2017); Mangoubi and Vishnoi (2018); Bou-Rabee et al. (2018); Chatterji et al. (2018); Vempala and Wibisono (2019); Wibisono (2018). In particular, the Langevin sampling algorithm Rossky et al. (1978); Roberts and Stramer (2002); Durmus and Moulines (2017) has been shown to be a form of gradient descent on the space of probabilities Jordan et al. (1998); Wibisono (2018); Bernton (2018); Durmus et al. (2019). Many convergence rates on Langevin algorithm have emerged thenceforward, based on different assumptions on the posterior distribution (e.g., Dalalyan and Riou-Durand, 2018; Ma et al., 2019; Lee et al., 2018; Cheng et al., 2018a; Shen and Lee, 2019; Ma et al., 2021; Mou et al., 2021; Zou and Gu, 2019, to list a few).

A conspicuous distinction between the convergence analysis of Langevin sampling and that of gradient descent is that all known convergence rates for Langevin algorithms depend on the dimensionality of the problem, whereas the convergence rates for gradient descent are dimension-free for convex problems. This prompts us to ask:

Can Langevin algorithm achieve dimension independent convergence rate under the usual convex assumptions?

In order to answer this question formally, we make two assumptions on the negative log-likelihood function. One is that the negative log-likelihood is convex. Another is that the negative log-likelihood is either Lipschitz continuous or Lipschitz smooth. We also employ a known and tractable prior distribution that is strongly log-concave—often times taken to be a normal distribution—to serve as a parallel to the $L_{2}$ regularizer in gradient descent.

Under such assumptions, we answer the above highlighted question in the affirmative. In particular, we prove that a Langevin algorithm converges similarly as convex optimization for this class of problems. In the analysis, we observe that the number of gradient queries required for the algorithm to converge does not depend on the dimensionality of the problem for either the Lipschitz continuous log-likelihood or the Lipschitz smooth log-likelihood equipped with a ridge separable structure.

To obtain this result, we first follow recent works (reference Durmus et al. (2019) in particular) and formulate the posterior sampling problem as optimizing over the Kullback-Leibler (KL) divergence, which is composed of two terms: (regularized) entropy and cross entropy. We then decompose the Langevin algorithm into two steps, each optimizing one part of the objective function. With a strongly convex and tractable prior, we explicitly integrate the diffusion along the prior distribution, optimizing the regularized entropy; whereas gradient descent over the convex negative log-likelihood optimizes the cross entropy. Via analyzing an intermediate quantity in this composite optimization procedure, we achieve a tight convergence bound that corresponds to the gradient descent’s convergence for convex optimization on the Euclidean space. This dimension independent convergence time for Lipschitz continuous log-likelihood and Lipschitz smooth log-likelihood endowed with a ridge separable structure carries over to the stochastic versions of the Langevin algorithm.

2 Preliminaries

2.1 Two Problem Classes

We consider sampling from a posterior distribution over parameter $w\in{\mathbb{R}}^{d}$ , given the data set ${\mathbf{z}}$ :

p(w|{\mathbf{z}})\propto p({\mathbf{z}}|w)\pi(w)\propto\exp\left(-U(w)\right),

where the potential function $U$ decomposes into two parts: $U(w)=\beta^{-1}\left(f(w)+g(w)\right)$ .

While the formulation is general, in the machine learning setting, $f(w)$ usually corresponds to the negative log-likelihood, and $g(w)$ corresponds to the negative log-prior. The parameter $\beta$ is the temperature, which often takes the value of $1/n$ in machine learning, where $n$ is the number of training data. The key motivation to consider this decomposition is that we assume that $g$ is “simple” so that an SDE involving $g$ can be solved to high precision. We will take advantage of this assumption in our algorithm design.

Assumption on function $g$

A0

We assume that function $g$ is $m$ -strongly convex ( $g(w)-\frac{m}{2}\left\|w\right\|^{2}$ is convex) ¹¹1We also say that the density proportional to $\exp\left(-\beta^{-1}g(w)\right)$ is $\beta^{-1}m$ -strongly log-concave in this case. and can be explicitly integrated.

Assumption on function $f$

We assume that function $f$ is convex (Assumption A1) and consider two cases regarding its regularity.

•

In the first case, we assume that function $f$ is $G$ -Lipschitz continuous (Assumption $\rm{A2_{L}}$ ).
•

In the second case, we assume that function $f$ is $L$ -Lipschitz smooth (Assumption $\rm{A2_{S}}$ ). We then instantiate the result by endowing it with a ridge separable structure (Assumptions R1 and R2).

The first case stems from Bayesian classification problems, where one has a simple strongly log-concave prior and a log-concave and log-Lipschitz likelihood that encodes the complexity of the data. Examples include Bayesian neural networks for classification tasks Neal (1996), Bayesian logistic regression Gelman et al. (2004), as well as other Bayesian classification problems Sollich (2002) with Gaussian or Bayesian elastic net priors. The second case corresponds to the regression type problems, where the entire posterior is strongly log-concave and log-Lipschitz smooth. In this case, one can separate the negative log-posterior into two parts: $\beta^{-1}g(w)=\frac{\beta^{-1}m}{2}\left\|w\right\|^{2}$ and $\beta^{-1}f(w)=\left(-\log p(w|{\mathbf{z}})-\frac{\beta^{-1}m}{2}\left\|w\right\|^{2}\right)$ , which is convex and $\beta^{-1}L$ -Lipschitz smooth. We therefore directly let $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ in Section 6.

2.2 Objective Functional and Convergence Criteria

We take the KL divergence $\beta^{-1}Q(p)$ to be our objective functional and solve the following optimization problem:

	$\displaystyle p_{*}=$	$\displaystyle\arg\min_{p}Q(p),$		(1)
		$\displaystyle Q\left(p\right)=\int p(w)\ln\frac{p}{p(w\|{\mathbf{z}})}dw={\mathbb{E}}_{w\sim p}\left[f(w)+g(w)+\beta\ln p(w)\right].$

The minimizer that solves the optimization problem (1) is the posterior distribution:

\displaystyle p_{*}(w)\propto\exp\left(-\beta^{-1}(f(w)+g(w))\right).

(2)

We further define the entropy functional as

H(p)=\beta\;{\mathbb{E}}_{w\sim p}\ln p(w),

so that the objective functional decomposes into the regularized entropy plus cross entropy:

Q(p)=\big{(}H(p)+{\mathbb{E}}_{w\sim p}\left[g(w)\right]\big{)}+{\mathbb{E}}_{w\sim p}\left[f(w)\right].

With this definition of the objective function, we state that the difference in $Q$ leads to the KL divergence.

Proposition 1.

Let $p$ be the solution of (1), and $p^{\prime}$ be another distribution on $w$ . We have

{\mathrm{KL}}(p^{\prime}\|p)=\beta^{-1}[Q(p^{\prime})-Q(p)].

This result establishes that the convergence in the objective $\beta^{-1}Q(p^{\prime})$ is equivalent to the convergence in ${\mathrm{KL}}$ -divergence. Therefore our analysis will focus on the convergence of $\beta^{-1}Q(p^{\prime})$ .

We also define the $2$ -Wasserstein distance between two distributions that will become useful in our analysis.

Definition 1.

Given two probability distributions $p(x)$ and $p^{\prime}(y)$ on ${\mathbb{R}}^{d}$ , and let $\Pi(p,p^{\prime})$ be the class of distributions $q(x,y)$ on ${\mathbb{R}}^{d}\times{\mathbb{R}}^{d}$ so that the marginals $q(x)=p(x)$ and $q(y)=p^{\prime}(y)$ . The $W_{2}$ Wasserstein distance of $p$ and $p^{\prime}$ is defined as

W_{2}(p,p^{\prime})^{2}=\min_{q\in\Pi(p,p^{\prime})}{\mathbb{E}}_{(x,y)\sim q}\|x-y\|_{2}^{2}.

A celebrated relationship between the ${\mathrm{KL}}$ -divergence and the $2$ -Wasserstein distance is known as the Talagrand inequality Otto and Villani (2000).

Proposition 2.

Assume that probability density $p_{*}$ is $\widehat{m}$ -strongly log-concave, and $p^{\prime}$ defines another distribution on ${\mathbb{R}}^{d}$ . Then $p_{*}$ satisfies the log-Sobolev inequality with constant $\widehat{m}/2$ ²²2This fact follows from the Bakry-Emery criterion Bakry and Emery (1985)., and yields the following Talagrand inequality:

W_{2}^{2}(p_{*},p^{\prime})\leq\widehat{m}^{-1}{\mathrm{KL}}(p_{*}\|p^{\prime}).

3 Related Works

Some previous works have aimed to sample from posteriors of the similar kind and obtain convergence in the KL divergence or the squared $2$ -Wasserstein distance.

In the Lipschitz continuous case,

where the negative log-likelihood is convex and $\widehat{G}$ -Lipschitz continuous, composed with an $\widehat{m}$ -strongly convex and $M$ -Lipschitz smooth negative log-prior, the convergence time to achieve $W_{2}^{2}(\widetilde{p}_{T},p_{*})\leq\epsilon$ is $\widetilde{\Omega}\left(\frac{dM+\widehat{G}^{2}}{\widehat{m}^{2}\epsilon^{2}}\right)$ (Corollary 22 of Durmus et al., 2019). Similarly, Chatterji et al. (2020) uses Gaussian smoothing to obtain a convergence time of $\widetilde{\Omega}\left(\frac{d(M+\widehat{G}^{2})}{\widehat{m}^{2}\epsilon}\right)$ (in Theorem 3.4), which improves the dependence on accuracy $\epsilon$ . In Mou et al. (2019), the Metropolis-adjusted Langevin algorithm is levaraged with a proximal sampling oracle to remove the polynomial dependence on the accuracy $\epsilon$ (in total variation distance) and achieve a $\widetilde{\Omega}\left(d\log(\frac{1}{\epsilon})\right)$ convergence time for a related composite posterior distribution. Unfortunately, an additional dimension dependent factor is always introduced into the overall convergence rate. This work demonstrates that if the $m$ -strongly convex regularizer is explicitly integrable, then the convergence time for the Langevin algorithm to achieve $\mathrm{KL}(\widetilde{p}_{T}\|p_{*})\leq\epsilon$ is dimension independent: $T=\Omega\left(\frac{\widehat{G}^{2}}{\widehat{m}\epsilon}\right)$ . This is proven in Theorem 1 for the full gradient Langevin algorithm, and in Theorem 2 for the stochastic gradient Langevin algorithm. Using Proposition 2, the result implies a bound of $T=\Omega\left(\frac{\widehat{G}^{2}}{\widehat{m}^{2}\epsilon}\right)$ to achieve $W_{2}^{2}(\widetilde{p}_{T},p_{*})\leq\epsilon$ .

In the Lipschitz smooth case,

where the negative log-posterior $U$ is $\widehat{m}$ -strongly convex and $\widehat{L}$ -Lipschitz smooth, the overdamped Langevin algorithm has been shown to converge in $\widetilde{\Omega}\left(\frac{\widehat{L}^{2}}{\widehat{m}^{2}}\frac{d}{\epsilon}\right)$ number of gradient queries Dalalyan (2017); Dalalyan and Karagulyan (2017); Cheng and Bartlett (2018); Durmus and Moulines (2017, 2019); Durmus et al. (2019), while the underdamped Langevin algorithm converges in $\widetilde{\Omega}\left(\frac{\widehat{L}^{3/2}}{\widehat{m}^{2}}\sqrt{\frac{d}{\epsilon}}\right)$ gradient queries Cheng et al. (2018b); Ma et al. (2021); Dalalyan and Riou-Durand (2018), to ensure that ${\mathrm{KL}}(\widetilde{p}_{T}\|p_{*})\leq\epsilon$ and $W_{2}^{2}(\widetilde{p}_{T},p_{*})\leq\epsilon$ . Using a randomized midpoint integration method for the underdamped Langevin dynamics, this convergence time can be reduced to $\widetilde{\Omega}\left(\frac{\widehat{L}}{\widehat{m}^{4/3}}\left(\frac{d}{\epsilon}\right)^{1/3}\right)$ for convergence in squared $2$ -Wasserstein distance Shen and Lee (2019). This paper establishes that for overdamped Langevin algorithm, the convergence time can be sharpened to $\Omega\left(\frac{{\mathrm{trace}}(\widehat{H}^{2})}{\widehat{m}^{2}\epsilon}\right)$ to achieve ${\mathrm{KL}}(\widetilde{p}_{T}\|p_{*})\leq\epsilon$ , where matrix $\widehat{H}$ is an upper bound for the Hessian of function $U$ .

Previous works have also focused on the ridge separable potential functions studied in this work. There is a literature that requires incoherence conditions on the data vectors and/or high-order smoothness conditions on the component functions to achieve a $\widetilde{\Omega}\left(\left(\frac{d}{\epsilon}\right)^{1/4}\right)$ convergence time for $W_{2}^{2}(\widetilde{p}_{T},p_{*})\leq\epsilon$ using Hamiltonian Monte Carlo methods Mangoubi and Smith (2017); Mangoubi and Vishnoi (2018). Making further assumptions that the differential equation of the Hamiltonian dynamics is close to the span of a small number of basis functions, this bound can be improved to polynomial in $\log(d)$ Lee et al. (2018). Another thread of work alleviates these assumptions and achieves the $\widetilde{\Omega}\left(\left(\frac{d}{\epsilon}\right)^{1/4}\right)$ convergence time for the general ridge separable potential functions via higher order Langevin dynamics and integration schemes Mou et al. (2021). We follow this general ridge separable setting and assume that each individual log-likelihood is Lipschitz smooth. Under this assumption, we demonstrate in this paper, by instantiating the bound for the general Lipschitz smooth case, that the Langevin algorithm converges in $\Omega\left(\frac{1}{\epsilon}\right)$ number of gradient queries to achieve ${\mathrm{KL}}(\widetilde{p}_{T}\|p_{*})\leq\epsilon$ (see Corollary 1 and Corollary 2).

4 Langevin Algorithms

We consider the following variant of the Langevin Algorithm 1.

Input: Initial distribution

p_{0}

{\mathbb{R}}^{d}

, stepsize

\eta_{t}

\beta=1

1 Draw

w_{0}

from

p_{0}

2 for $t=1,2,\ldots,T$ do

3 Sample

\widetilde{w}_{t}

from

\widetilde{w}_{t}(\eta_{t})

with the following SDE on

{\mathbb{R}}^{d}

and initial value

\widetilde{w}_{t}(0)=w_{t-1}

\displaystyle\widetilde{w}_{t}(\eta_{t})=w_{t-1}-\int_{0}^{\eta_{t}}\nabla g(\widetilde{w}_{t}(s))ds+\sqrt{2\beta}\int_{0}^{\eta_{t}}dB_{s},

(3)

where

dB_{s}

is the standard Brownian motion on

{\mathbb{R}}^{d}

4 Let

{w}_{t}=\widetilde{w}_{t}-\widetilde{\eta}_{t}\nabla f(\widetilde{w}_{t})

(4)

5 end for

return $\widetilde{w}_{T}$

Algorithm 1 Langevin Algorithm with Prior Diffusion

In this method, we assume that the prior diffusion equation (3) can be solved efficiently. When the prior distribution is a standard normal distribution where $g(w)=\frac{m}{2}\left\|w\right\|_{2}^{2}$ on ${\mathbb{R}}^{d}$ , we can instantiate equation (3) to be:

\displaystyle\text{Sample}\quad\widetilde{w}_{t}\sim\mathcal{N}\left(e^{-m\eta_{t}}w_{t-1},\frac{1-e^{-2m\eta_{t}}}{m}\beta{\mathrm{I}}\right).

(5)

In general, the diffusion equation (3) can also be solved numerically for separable $g(w)$ of the form

g(w)=\sum_{j=1}^{d}g_{j}(w_{j}),

where $w=[w_{1},\ldots,w_{d}]$ . In this case, we only need to solve $d$ one-dimensional problems, which are relatively simple. For example, this includes the $L_{1}-L_{2}$ regularization arising from the Bayesian elastic net Li and Lin (2010),

g(w)=\frac{m}{2}\|w\|_{2}^{2}+\alpha\|w\|_{1},

among other priors that decompose coordinate-wise.

We will also consider the stochastic version of Algorithm 1, the stochastic gradient Langevin dynamics (SGLD) method, with a strongly convex function $g(w)$ . Assume that function $f$ decomposes into $f(w)=\frac{1}{n}\sum_{i=1}^{n}\ell(w,z_{i})$ . Let $D$ be the distribution over the dataset $\Omega$ such that expectation over it provides the unbiased estimate of the full gradient: ${\mathbb{E}}_{z\sim D}\nabla_{w}\ell(w,z)=\nabla f(w)$ . Then the new algorithm takes the following form and can be instantiated in the same way as Algorithm 1.

Input: Initial distribution

p_{0}

{\mathbb{R}}^{d}

, stepsize

\eta_{t}

\beta=1/n

1 Draw

w_{0}

from

p_{0}

2 for $t=1,2,\ldots,T$ do

3 Sample

\widetilde{w}_{t}

from

\widetilde{w}_{t}(\eta_{t})

with the following SDE on

{\mathbb{R}}^{d}

and initial value

\widetilde{w}_{t}(0)=w_{t-1}

\displaystyle\widetilde{w}_{t}(\eta_{t})=w_{t-1}-\int_{0}^{\eta_{t}}\nabla g(\widetilde{w}_{t}(s))ds+\sqrt{2\beta}\int_{0}^{\eta_{t}}dB_{s},

(6)

where

dB_{s}

is the standard Bronwian motion on

{\mathbb{R}}^{d}

4 Draw minibatch

\mathcal{S}

where each

z_{i}\in\mathcal{S}

are i.i.d. draws:

z_{i}\sim D

. Let

{w}_{t}=\widetilde{w}_{t}-\widetilde{\eta}_{t}\frac{1}{|\mathcal{S}|}\sum_{z_{i}\in\mathcal{S}}\nabla_{w}\ell(\widetilde{w}_{t},z_{i}).

(7)

5 end for

return $\widetilde{w}_{T}$

Algorithm 2 Stochastic Gradient Langevin Algorithm with Prior Diffusion

This algorithm becomes the streaming SGLD method where in each iteration we take one data point $z\sim D$ .

In the analysis of Algorithm 1, we will use $p_{t-1}$ to denote the distribution of $w_{t-1}$ , and $\widetilde{p}_{t}$ to denote the distribution of $\widetilde{w}_{t}$ , where the randomness include all random sampling in the algorithm. We also denote $\mu_{t}$ and $\widetilde{\mu}_{t}$ as the measures associated with random variables $w_{t}$ and $\widetilde{w}_{t}$ , respectively. When using samples along the Markov chain to estimate expectations over function $\phi(\cdot)$ , we take a weighted average, so that

\hat{\phi}(p)=\frac{1}{\sum_{s=1}^{T}\eta_{s}}\sum_{t=1}^{T}\eta_{t}\phi(\widetilde{w}_{t}),

which is equivalent to the expectation with respect to the weighted averaged distribution:

\bar{\mu}_{T}=\frac{1}{\sum_{s=1}^{T}\eta_{s}}\sum_{t=1}^{T}\eta_{t}\widetilde{\mu}_{t}.

We prove in what follows the convergence of the distribution $\widetilde{p}_{t}$ along the updates of (3) and (4) towards the posterior distribution (2).

5 Langevin Algorithms in Lipschitz Convex Case

For the posterior $p(w|{\mathbf{z}})\propto\left(-\beta^{-1}(f(w)+g(w))\right)$ , we assume that function $f$ satisfies the following two conditions common to convex analysis.

Assumptions for the Lipschitz Convex Case:

A1

Function $f:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is convex.

$\rm{A2_{L}}$

Function $f$ is $G$ -Lipschitz continuous on ${\mathbb{R}}^{d}$ : $\|\nabla f(w)\|_{2}\leq G$ .

We also assume that function $g:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is $m$ -strongly convex. Note that we have assumed that the gradient of function $f$ exists but have not assumed that function $f$ is smooth.

5.1 Full Gradient Langevin Algorithm Convergence in Lipschitz Convex Case

Our main result for Full Gradient Langevin Algorithm in the case that $f$ is Lipshitz can be stated as follows.

Theorem 1.

Assume that function $f$ satisfies the convex and Lipschitz continuous Assumptions A1 and $\rm{A2_{L}}$ . Further assume that function $g(w)$ satisfies Assumption 1. Then for $\widetilde{p}_{T}$ following the Langevin Algorithm 1, it satisfies ( for $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2/\left(m(t+2)\right)$ ):

\sum_{t=1}^{T}\frac{1+0.5t}{T+0.25T(T+1)}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{5G^{2}}{\beta mT}.

By the convexity of the KL divergence, this leads to the convergence time of

T=\frac{5G^{2}}{\beta m\epsilon},

for the averaged distribution to convergence to $\epsilon$ accuracy in the KL-divergence $\beta^{-1}Q$ .

We devote the rest of this section to prove Theorem 1.

Proof of Theorem 1.

We take a composite optimization approach and analyze the convergence of the Langevin algorithm in two steps. First we characterize the decrease of the regularized entropy ${\mathbb{E}}_{w\sim p}\left[g(w)+H(p)\right]$ along the diffusion step (3).

Lemma 1 (For Regularized Entropy).

We generalize Lemma 5 of Durmus et al. (2019) and have for $\widetilde{p}_{t}$ being the density of $\widetilde{w}_{t}$ following equation (3) and $p$ being another probability density,

\frac{2}{m}(1-e^{-m\eta_{t}})\left({\mathbb{E}}_{w\sim\widetilde{p}_{t}}[g(w)+H(\widetilde{p}_{t})]-{\mathbb{E}}_{w\sim p}[g(w)+H(p)]\right)\leq e^{-m\eta_{t}}W_{2}^{2}({p}_{t-1},p)-W_{2}^{2}(\widetilde{p}_{t},p),

where $m$ is the strong convexity of $g(w)$ .

We then capture the decrease of the cross entropy ${\mathbb{E}}_{w\sim p}\left[f(w)\right]$ along the gradient descent step (4). This result follows and parallels the standard convergence analysis of gradient descent (see Zinkevich, 2003; Zhang, 2004, for example).

Lemma 2.

Given probability density $p$ on ${\mathbb{R}}^{d}$ . Define

f(p)={\mathbb{E}}_{w\sim p}f(w),

then we have for $p_{t}$ being the density of $w_{t}$ following equation (4):

2\widetilde{\eta}_{t}[f(\widetilde{p}_{t})-f(p)]\leq W_{2}^{2}(\widetilde{p}_{t},p)-W_{2}^{2}(p_{t},p)+\widetilde{\eta}_{t}^{2}G^{2}.

We then combine the two steps to prove the overall convergence rate for the Langevin algorithm. It is worth noting that by aligning the diffusion step (3) and the gradient descent step (4) at $\widetilde{p}_{t}$ , we combine ${\mathbb{E}}_{w\sim\widetilde{p}_{t}}[g(w)+H(\widetilde{p}_{t})]$ with $f(\widetilde{p}_{t})$ and cancel out $W_{2}^{2}(\widetilde{p}_{t},p)$ perfectly and achieve the same convergence rate as that of stochastic gradient descent in optimization.

Proposition 3.

Set $\widetilde{\eta}_{t}=(1-e^{-m\eta_{t}})/m=\tau\cdot(\tau/\widetilde{\eta}_{0}+mt)^{-1}$ for some $\tau\geq 1$ and $\widetilde{\eta}_{0}>0$ . Then

\sum_{t=1}^{T}\widetilde{\eta}_{t}^{1-\tau}[Q(\widetilde{p}_{t})-Q(p)]\leq\widetilde{\eta}_{0}^{-\tau}W_{2}^{2}(p_{0},p)+G^{2}\sum_{t=1}^{T}\widetilde{\eta}_{t}^{2-\tau}.

Choosing $\tau=2$ and $p=p_{*}$ , we have

\sum_{t=1}^{T}\frac{1+0.5t}{T+0.25T(T+1)}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{4}{\beta m\widetilde{\eta}_{0}^{2}T(T+1)}W_{2}^{2}(p_{0},p_{*})+\frac{4G^{2}}{\beta m(T+1)}.

(8)

The learning rate schedule of $\eta_{t}=1/mt$ (with $\tau=1$ ) was introduced to SGD analysis for strongly convex objectives in Shalev-Shwartz et al. (2011), which leads to a similar rate as that of Proposition 3, but with an extra $\log(T)$ term than (9). The use of $\tau>1$ has been adopted in more recent literature of SGD analysis, as an effort to avoid the $\log(T)$ term (for example, see Lacoste-Julien et al. (2012)). The resulting bound in the SGD analysis becomes identical to that of Proposition 3, and this rate is optimal for nonsmooth strongly convex optimization Rakhlin et al. (2012). In addition, it is possible to implement for Langevin algorithm a similar scheme using moving averaging, as discussed in Shamir and Zhang (2013).

It can be observed that taking a large step size $\widetilde{\eta}_{0}$ will grant rapid convergence. The largest one can take is to choose $\eta_{0}=+\infty$ and consequently $\widetilde{\eta}_{0}=1/m$ , leading to a learning rate schedule of $\widetilde{\eta}_{t}=2/(m\cdot(t+2))$ . In this case, we are effectively initializing from $\widetilde{p}_{1}\propto\exp\left(-\beta^{-1}g(w)\right)$ . Choosing the same $p_{0}\propto\exp\left(-\beta^{-1}g(w)\right)$ , we can bound the initial error $W_{2}^{2}(p_{0},p_{*})$ via the Talagrand inequality in Proposition 1 and the log-Sobolev inequality Bakry and Emery (1985); Ledoux (2000) for the $\beta^{-1}m$ -strongly log-concave distribution $p_{*}$ :

W_{2}^{2}(p_{0},p_{*})\leq\frac{\beta}{m}{\mathrm{KL}}(p_{*}\|p_{0})\leq\frac{\beta^{2}}{2m^{2}}{\mathbb{E}}_{p_{*}}\left[\left\|\nabla\log\frac{p_{*}}{p_{0}}\right\|^{2}\right]\leq\frac{G^{2}}{2m^{2}},

since $\left\|\nabla\log\frac{p_{*}}{p_{0}}(w)\right\|=\left\|\beta^{-1}\nabla f(w)\right\|\leq\beta^{-1}G$ . Plugging this bound and $\widetilde{\eta}_{0}=1/m$ into equation (9), and noting that $T\geq 1$ , we arrive at our result that

\sum_{t=1}^{T}\frac{1+0.5t}{T+0.25T(T+1)}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{5G^{2}}{\beta mT}.

(9)

∎

Proof of Proposition 3.

We can add the inequalities in Lemma 1 and Lemma 2 to obtain:

\displaystyle\widetilde{\eta}_{t}[Q(\widetilde{p}_{t})-Q(p)]\leq e^{-m\eta_{t}}W_{2}({p}_{t-1},p)^{2}-W_{2}(p_{t},p)^{2}+\widetilde{\eta}_{t}^{2}G^{2}.

This is equivalent to

\widetilde{\eta}_{t}^{1-\tau}[Q(\widetilde{p}_{t})-Q(p)]\leq(1-m\widetilde{\eta}_{t})\widetilde{\eta}_{t}^{-\tau}W_{2}({p}_{t-1},p)^{2}-\widetilde{\eta}_{t}^{-\tau}W_{2}(p_{t},p)^{2}+\widetilde{\eta}_{t}^{2-\tau}G^{2}.

(10)

We first show that

(1-m\widetilde{\eta}_{t})\widetilde{\eta}_{t}^{-\tau}\leq\widetilde{\eta}_{t-1}^{-\tau}.

(11)

Let $s=t+\tau/(m\eta_{0})\geq 1$ for $t\geq 1$ , $\tilde{\eta}_{t}=\tau/(ms)$ and $\tilde{\eta}_{t}=\tau/(m(s-1))$ . Therefore (11) is equivalent to

(1-\tau/s)s^{\tau}\leq(s-1)^{\tau}.

This inequality follows from the fact that for $z=1/s\in[0,1]$ and $\tau\geq 1$ : $\psi(z)=(1-z)^{\tau}$ is convex in $z$ , and thus $(1-\tau z)=\psi(0)+\psi^{\prime}(0)z\leq\psi(z)=(1-z)^{\tau}$ .

By combining (10) and (11), we obtain

\displaystyle\widetilde{\eta}_{t}^{1-\tau}[Q(\widetilde{p}_{t})-Q(p)]\leq\widetilde{\eta}_{t-1}^{-\tau}W_{2}({p}_{t-1},p)^{2}-\widetilde{\eta}_{t}^{-\tau}W_{2}(p_{t},p)^{2}+\widetilde{\eta}_{t}^{2-\tau}G^{2}.

By summing over $t=1$ to $t=T$ , we obtain the bound.

∎

5.2 Streaming SGLD Convergence in Lipschitz Convex Case

To analyze the streaming stochastic gradient Langevin algorithm, we assume that function $f$ decomposes:

f(w)=\frac{1}{n}\sum_{i=1}^{n}\ell(w,z_{i})={\mathbb{E}}_{z\sim D}[\ell(w,z)],

where $D$ is the distribution over the data samples. In this case, we modify Assumption $\rm{A2_{L}}$ and assume that the individual log-likelihood satisfies the Lipschitz condition.

Assumptions on individual loss $\ell$

$\rm{A2_{L}^{SG}}$

Function $\ell$ is $G_{\ell}$ -Lipschitz continuous on ${\mathbb{R}}^{d}$ : $\|\nabla\ell(w,z)\|_{2}\leq G_{\ell}$ , $\forall z\in\Omega$ .

In the case that $\ell(w,z)$ is Lipschitz, our main result for SGLD is the following counterpart of Theorem 1.

Theorem 2.

Assume that function $f$ satisfies the convex assumption A1 and the Lipschitz continuous assumption for the individual log-likelihood $\rm{A2_{L}^{SG}}$ . Further assume that function $g(w)$ satisfies Assumption 1. Then for $\widetilde{p}_{T}$ following the streaming SGLD Algorithm 2, it satisfies ( for $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2/\left(m(t+2)\right)$ ):

\sum_{t=1}^{T}\frac{1+0.5t}{T+0.25T(T+1)}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{5G_{\ell}^{2}}{\beta mT}.

leading to the convergence time of

T=\frac{5G_{\ell}^{2}}{\beta m\epsilon},

for the averaged distribution to convergence to $\epsilon$ accuracy in the KL-divergence $\beta^{-1}Q$ .

We devote the rest of this section to prove Theorem 2.

Proof of Theorem 2.

Same as in the previous section, convergence of the regularized entropy ${\mathbb{E}}_{w\sim p}\left[g(w)\right]+H(p)$ along equation (6) follows Lemma 1.

For the convergence of the cross entropy ${\mathbb{E}}_{w\sim p}\left[f(w)\right]$ along equation (7), the following Lemma follows the standard analysis of SGD.

Lemma 3.

Adopt Assumption $\rm{A2_{L}^{SG}}$ that $\ell(w,z)$ is $G_{\ell}$ -Lipschitz for all $z\in\Omega$ . Also adopt Assumption A1 that $f(w)={\mathbb{E}}_{z\sim D}\ell(w,z)$ is convex. We have for all $w\in{\mathbb{R}}^{d}$ :

2\widetilde{\eta}_{t}{\mathbb{E}}_{z\sim D}[\ell(\widetilde{w}_{t},z)-\ell(w,z)]\leq\|\widetilde{w}_{t}-w\|_{2}^{2}-{\mathbb{E}}_{w_{t}|\widetilde{w}_{t}}\|w_{t}-w\|_{2}^{2}+\widetilde{\eta}_{t}^{2}G_{\ell}^{2}.

(12)

It implies the following bound, which modifies Lemma 2.

Lemma 4.

Given any probability density $q$ on ${\mathbb{R}}^{d}$ . Define

\ell(q)={\mathbb{E}}_{w\sim q}{\mathbb{E}}_{z\sim D}\ell(w,z),

then we have

2\widetilde{\eta}_{t}[\ell(\widetilde{p}_{t})-\ell(p)]\leq W_{2}(\widetilde{p}_{t},p)^{2}-W_{2}(p_{t},p)^{2}+\widetilde{\eta}_{t}^{2}G_{\ell}^{2}.

Initializing from the prior distribution, we can follow the same proof as in Proposition 3 and obtain a similar convergence rate as in the non-stochastic case.

Proposition 4.

Set $\widetilde{\eta}_{t}=(1-e^{-m\eta_{t}})/m=\tau\cdot(\tau/\widetilde{\eta}_{0}+mt)^{-1}$ for some $\tau\geq 1$ and $\widetilde{\eta}_{0}>0$ . Then

\sum_{t=1}^{T}\widetilde{\eta}_{t}^{1-\tau}[Q(\widetilde{p}_{t})-Q(p)]\leq\eta_{0}^{-\tau}W_{2}(p_{0},p)^{2}+G_{\ell}^{2}\sum_{t=1}^{T}\widetilde{\eta}_{t}^{2-\tau}.

We can choose $\tau=2$ , and then for $p=p_{*}$ , we have

\sum_{t=1}^{T}\frac{1+0.5t}{T+0.25T(T+1)}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{4}{\beta m\widetilde{\eta}_{0}^{2}T(T+1)}W_{2}(p_{0},p_{*})^{2}+\frac{4G_{\ell}^{2}}{\beta m(T+1)}.

(13)

Following the same steps as in the full gradient case, we arrive at the result. ∎

6 Langevin Algorithms in Smooth Convex Case

For the posterior $p(w|{\mathbf{z}})\propto\left(-\beta^{-1}(f(w)+g(w))\right)$ , we make the following assumptions on function $f$ .

Assumptions for the smooth convex case:

A1

Function $f:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is convex and positive.

$\rm{A2_{S}}$

Function $f$ is $L$ -Smooth on ${\mathbb{R}}^{d}$ : $\|\nabla f(w)-\nabla f(w^{\prime})\|_{2}\leq L\|w-w^{\prime}\|_{2}$ .

We also assume that function $g:{\mathbb{R}}^{d}\rightarrow{\mathbb{R}}$ is $m$ -strongly convex. Note that this is equivalent to the cases where we simply assume the entire negative log-posterior to be $\beta^{-1}m$ -strongly convex and $\left(\beta^{-1}(L+m)\right)$ -Lipschitz smooth: one can separate the negative log-posterior into two parts: $\frac{\beta^{-1}m}{2}\left\|w\right\|^{2}$ and $\left(-\log p(w|{\mathbf{z}})-\frac{\beta^{-1}m}{2}\left\|w\right\|^{2}\right)$ , which is convex and $\beta^{-1}L$ -Lipschitz smooth. We therefore directly let $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ in what follows.

6.1 Full Gradient Langevin Algorithm Convergence in Smooth Convex Case

Our main result for Full Gradient Langevin Algorithm in the case that $f$ is smooth can be stated as follows. Compared to Theorem 1, the result of Theorem 3 is useful for loss functions such as least squares loss that are smooth but not Lipschitz continuous.

Theorem 3.

Assume that function $f$ satisfies the convex and Lipschitz smooth Assumptions A1 and $\rm{A2_{S}}$ . Also assume that $\nabla^{2}f(w)\preceq H$ . Further let function $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ . Then for $\widetilde{p}_{T}$ following Algorithm 1 and initializing from $p_{0}\propto\exp(-\beta^{-1}g)$ , it satisfies ( for $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2\cdot\left(8L+mt\right)^{-1}$ ):

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\\ \leq\frac{64L^{2}}{m^{2}T(T+1)}\cdot\left(\frac{8}{m^{2}}{\mathrm{trace}}\left(H^{2}\right)+2U(0)\right)+\frac{16}{T+1}\cdot\left(\frac{8}{m^{2}}{\mathrm{trace}}\left(H^{2}\right)+2U(0)\right).

leading to the convergence time of

T=64\cdot\max\left\{\frac{8{\mathrm{trace}}\left(H^{2}\right)}{m^{2}\epsilon},\frac{2U(0)}{\epsilon}\right\},

for the averaged distribution to convergence to $\epsilon\leq 1$ accuracy in the KL-divergence $\beta^{-1}Q$ .

Ridge Separable Case

Assume that function $f$ decomposes into the following ridge-separable form:

\displaystyle f(w)=\frac{1}{n}\sum_{i=1}^{n}s_{i}(w^{\top}z_{i}),

(14)

We make some assumptions on the activation function $s_{i}$ and the data points $z_{i}$ .

Assumptions in ridge separable case

R1

$\forall i\in\{1,\dots,n\}$ , the one dimensional activation function $s_{i}(\cdot)$ has a bounded second derivative: $\left|s_{i}^{\prime\prime}(x)\right|\leq L_{s}$ , for any $x\in{\mathbb{R}}$ .
R2

$\forall i\in\{1,\dots,n\}$ , data point $z_{i}$ has a bounded norm: $\left\|z_{i}\right\|^{2}\leq R_{z}$ .

Assumptions R1 and R2 combines to give a Lipschitz smoothness constant of $L_{s}R_{z}$ for the individual log-likelihood.

Corollary 1.

We make the convexity Assumption A1 on function $f$ and let it take the ridge-separable form (14) (also let function $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ ). Further adopt Assumptions R1 and R2 on the activation functions and the data points, respectively. Then the convergence time of Algorithm 1 initializing from $p_{0}\propto\exp(-\beta^{-1}g)$ (with step size $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2\left(8L_{s}R_{z}+mt\right)^{-1}$ ) is

T=64\cdot\max\left\{\frac{8L_{s}^{2}R_{z}^{2}}{m^{2}\epsilon},\frac{2U(0)}{\epsilon}\right\},

for the averaged distribution to convergence to $\epsilon$ accuracy in the KL-divergence $\beta^{-1}Q$ .

Proof.

From Assumptions R1 and R2, we know that $\nabla^{2}f(w)\preceq\frac{1}{n}L_{s}ZZ^{\top}=H$ , where

{\mathrm{trace}}\left(H^{2}\right)=L_{s}^{2}\frac{1}{n^{2}}{\mathrm{trace}}\left(Z^{\top}ZZ^{\top}Z\right)=L_{s}^{2}\frac{1}{n^{2}}\left\|Z^{\top}Z\right\|_{F}^{2}=L_{s}^{2}\frac{1}{n^{2}}\sum_{i,j=1}^{n}\left(z_{i}^{\top}z_{j}\right)^{2}\leq L_{s}^{2}\frac{1}{n}\sum_{i=1}^{n}\left\|z_{i}\right\|^{4}\leq L_{s}^{2}R_{z}^{2}.

Plugging this into Theorem 3 leads to the convergence time of

T=64\cdot\max\left\{\frac{8L_{s}^{2}R_{z}^{2}}{m^{2}\epsilon},\frac{2U(0)}{\epsilon}\right\}.

∎

We devote the rest of this section to the proof of Theorem 3.

Proof of Theorem 3.

Same as in Section 5.1, convergence of the regularized entropy ${\mathbb{E}}_{w\sim p}\left[g(w)\right]+H(p)$ along equation (6) follows Lemma 1.

For the decrease of the cross entropy ${\mathbb{E}}_{w\sim p}\left[f(w)\right]$ along the gradient descent step (4), we use the following derivation for $L$ -Lipschitz smooth $f$ . For $p_{t}$ being the density of $w_{t}$ following equation (4) and for $p$ being another probability density,

		$\displaystyle 2\widetilde{\eta}_{t}[{\mathbb{E}}_{w\sim\widetilde{p}_{t}}f(w)-{\mathbb{E}}_{w^{\prime}\sim p}f(w^{\prime})]$
	$\displaystyle\leq$	$\displaystyle[W_{2}(\widetilde{p}_{t},p)^{2}-W_{2}({p}_{t},p)^{2}]+\eta_{t}^{2}{\mathbb{E}}_{w\sim\widetilde{p}_{t}}\\|\nabla f(w)\\|_{2}^{2}$
	$\displaystyle\leq$	$\displaystyle[W_{2}(\widetilde{p}_{t},p)^{2}-W_{2}({p}_{t},p)^{2}]+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{(w,w^{\prime})\sim\gamma_{t}}\\|\nabla f(w)-\nabla f(w^{\prime})\\|_{2}^{2}+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{w^{\prime}\sim p}\\|\nabla f(w^{\prime})\\|_{2}^{2},$

where $\gamma_{t}\in\Gamma_{\text{opt}}(\widetilde{p}_{t},p)$ is the optimal coupling between distributions with densities $\widetilde{p}_{t}$ and $p$ .

With $\widetilde{\eta}_{t}=(1-e^{-m\eta_{t}})/m$ , we have

	$\displaystyle 2\widetilde{\eta}_{t}[Q(\widetilde{p}_{t})-Q(p)]$	$\displaystyle\leq(1-m\widetilde{\eta}_{t})W_{2}(\widetilde{p}_{t-1},p)^{2}-W_{2}({p}_{t},p)^{2}$
		$\displaystyle+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{(w,w^{\prime})\sim\gamma_{t}}\\|\nabla f(w)-\nabla f(w^{\prime})\\|_{2}^{2}+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{w^{\prime}\sim p}\\|\nabla f(w^{\prime})\\|_{2}^{2}.$		(15)

We also have the following lemma.

Lemma 5.

Let $\gamma_{t}\in\Gamma_{\text{opt}}(\widetilde{p}_{t},p)$ be the optimal coupling of $\widetilde{p}_{t}$ and $p$ , and let $p$ the solution of (1). Then we have

{\mathbb{E}}_{(w,w^{\prime})\sim\gamma_{t}}\|\nabla f(w)-\nabla f(w^{\prime})\|_{2}^{2}\leq 2L[Q(\widetilde{p}_{t})-Q(p)].

We note that Lemma 5 and equation (15) imply that

\displaystyle 2\widetilde{\eta}_{t}(1-2L\widetilde{\eta}_{t})[Q(\widetilde{p}_{t})-Q(p)]\leq(1-m\widetilde{\eta}_{t})W_{2}(p_{t-1},p)^{2}-W_{2}(p_{t},p)^{2}+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{w\sim p}\|\nabla f(w)\|_{2}^{2},

(16)

where $p$ satisfies (2).

Next we bound the last term of equation (16) at $p_{*}$ : ${\mathbb{E}}_{w\sim p_{*}}\|\nabla f(w)\|_{2}^{2}$ .

Lemma 6.

Assume that

\nabla^{2}f(w)\preceq H,\quad g(w)=\frac{m}{2}\|w\|_{2}^{2}.

Let

w_{*}=\arg\min_{w}\left[f(w)+g(w)\right],

and $p_{*}\propto\exp(-\beta^{-1}(f(w)+g(w)))$ be the solution of (2). Then

{\mathbb{E}}_{w\sim p_{*}}\|\nabla f(w)\|_{2}^{2}\leq 16\frac{\beta}{m}{\mathrm{trace}}\left(H^{2}\right)+2m^{2}\|w_{*}\|_{2}^{2}.

With these lemmas, we are ready to prove the convergence time of the Langevin algorithm 1.

We note that the shrinking step size scheduling of $\widetilde{\eta}_{t}=\tau\cdot\left(\frac{\tau}{\widetilde{\tau}_{0}}+mt\right)^{-1}$ satisfies:

\frac{\widetilde{\eta}_{t}^{\tau}}{\widetilde{\eta}_{t-1}^{\tau}}\leq\left(1-m\widetilde{\eta}_{t}\right).

Using this inequality and combining Lemma 6 and equation (16) at $p=p_{*}$ , we obtain that

		$\displaystyle 2\widetilde{\eta}_{t}^{1-\tau}\left(1-2L\widetilde{\eta}_{t}\right)[Q(\widetilde{p}_{t})-Q(p_{*})]$
	$\displaystyle\leq$	$\displaystyle\widetilde{\eta}_{t-1}^{-\tau}W_{2}(p_{t-1},p_{})^{2}-\widetilde{\eta}_{t}^{-\tau}W_{2}(p_{t},p_{})^{2}+4\widetilde{\eta}_{t}^{2-\tau}\left[8\left(\beta/m\right){\mathrm{trace}}\left(H^{2}\right)+m^{2}\left\\|w_{*}\right\\|^{2}\right].$

Summing over $t=1,\dots,T$ ,

		$\displaystyle 2\sum_{t=1}^{T}\widetilde{\eta}_{t}^{1-\tau}\left(1-2L\widetilde{\eta}_{t}\right)[Q(\widetilde{p}_{t})-Q(p_{*})]$
	$\displaystyle\leq$	$\displaystyle\widetilde{\eta}_{0}^{-\tau}W_{2}(p_{0},p_{})^{2}+4\left[8\left(\beta/m\right){\mathrm{trace}}\left(H^{2}\right)+m^{2}\left\\|w_{}\right\\|^{2}\right]\sum_{t=1}^{T}\widetilde{\eta}_{t}^{2-\tau}.$

Denote $\Delta=4\left[8\left(\beta/m\right){\mathrm{trace}}\left(H^{2}\right)+m^{2}\left\|w_{*}\right\|^{2}\right]$ and take $\widetilde{\eta}_{t}=\tau\cdot\left(\frac{\tau}{\widetilde{\eta}_{0}}+mt\right)^{-1}$ . Then for $\widetilde{\eta}_{t}\leq\frac{1}{4}\frac{1}{L}$ and for $\tau=2$ ,

m\sum_{t=1}^{T}\left(1/\left(m\widetilde{\eta}_{0}\right)+t/2\right)[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{1}{\widetilde{\eta}_{0}^{2}}W_{2}(p_{0},p_{*})^{2}+\Delta\cdot T,

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\frac{4}{m\widetilde{\eta}_{0}^{2}T(T+1)}W_{2}(p_{0},p_{*})^{2}+\frac{4\Delta}{m(T+1)}.

(17)

Inspired by the Lipschitz continuous case, we take $p_{0}(w)\propto\exp\left(-\beta^{-1}g(w)\right)$ . Then by the Talagrand and log-Sobolev inequalities,

W_{2}(p_{0},p_{*})^{2}\leq\frac{\beta}{m}\mathrm{KL}\left(p_{*}\|p_{0}\right)\leq\frac{\beta^{2}}{2m^{2}}{\mathbb{E}}_{p_{*}}\left[\left\|\nabla\log\frac{p_{*}}{p_{0}}\right\|^{2}\right]=\frac{\beta^{2}}{2m^{2}}{\mathbb{E}}_{p_{*}}\left[\left\|\beta^{-1}\nabla f(w)\right\|^{2}\right].

Applying Lemma 6 to the above inequality, we obtain that $W_{2}(p_{0},p_{*})^{2}\leq\frac{1}{m^{2}}\left(8\left(\beta/m\right){\mathrm{trace}}\left(H^{2}\right)+m^{2}\left\|w_{*}\right\|^{2}\right)$ . Then taking $\widetilde{\eta}_{0}=\frac{1}{4}\frac{1}{L}$ , we obtain that the weighted-averaged KL divergence is upper bounded:

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\\ \leq\frac{64L^{2}}{m^{2}T(T+1)}\cdot\left(\frac{8}{m^{2}}{\mathrm{trace}}\left(H^{2}\right)+\beta^{-1}m\left\|w_{*}\right\|^{2}\right)+\frac{16}{T+1}\cdot\left(\frac{8}{m^{2}}{\mathrm{trace}}\left(H^{2}\right)+\beta^{-1}m\left\|w_{*}\right\|^{2}\right).

(18)

Since $L^{2}\leq{\mathrm{trace}}\left(H^{2}\right)$ , $\forall\epsilon\leq 1$ , the weighted-averaged KL divergence $\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\epsilon$ if

T\geq 64\cdot\max\left\{\frac{8{\mathrm{trace}}\left(H^{2}\right)}{m^{2}\epsilon},\frac{\beta^{-1}m\left\|w_{*}\right\|^{2}}{\epsilon}\right\}.

Plugging in the bound that $m\left\|w_{*}\right\|^{2}\leq 2f(0)=2\beta U(0)$ gives the final result.

∎

6.2 SGLD Convergence in Smooth Convex Case

Similar to the Lipschitz continuous case, we assume that function $f$ decomposes:

f(w)=\frac{1}{n}\sum_{i=1}^{n}\ell(w,z_{i})={\mathbb{E}}_{z\sim D}[\ell(w,z)],

where $D$ is the distribution over the data samples. Making the following assumption, which modifies Assumption $\rm{A2_{S}}$ , that the individual log-likelihood satisfies the Lipschitz smooth condition yields the convergence rate for the SGLD method.

Assumptions on individual loss $\ell$

$\rm{A{{2}}_{S}^{SG}}$

Function $\ell$ is $L_{\ell}$ -Lipschitz smooth on ${\mathbb{R}}^{d}$ : $\ell(y,z)\geq\ell(x,z)+\nabla\ell(x,z)^{\top}(y-x)+\frac{1}{2L_{\ell}}\left\|\nabla\ell(y,z)-\nabla\ell(x,z)\right\|^{2}$ , $\forall z\in\Omega$ .
$\rm{A{{3}}_{S}^{SG}}$

The stochastic gradient variance at the mode $w_{*}$ is bounded: ${\mathbb{E}}_{z\sim D}\left[\left\|\nabla\ell(w_{*},z)-\nabla f(w_{*})\right\|^{2}\right]\leq b^{2}$ .

Assumption $\rm{A{{2}}_{S}^{SG}}$ ensures that the stochastic estimates of $f$ are $L_{\ell}$ Lipschitz smooth.

Under the above assumptions, we obtain in what follows the convergence rate for the SGLD method with minibatch size $|\mathcal{S}|$ . This result is the counterpart of its full gradient version in Theorem 3.

Theorem 4.

We make the convexity Assumptions A1 on function $f$ and the regularity Assumptions $\rm{A{{2}}_{S}^{SG}}$ and $\rm{A{{3}}_{S}^{SG}}$ on its components $\ell$ . Also assume that $\nabla^{2}f(w)\preceq H$ . Let function $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ . Then taking $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2\cdot\left(8L_{s}R_{z}+mt\right)^{-1}$ , the convergence time of the SGLD Algorithm 2 initializing from $p_{0}\propto\exp(-\beta^{-1}g)$ is

\displaystyle T=\Omega\left(\max\left\{\frac{L_{\ell}{\mathrm{trace}}(H)}{m^{2}\epsilon},\frac{U(0)}{\epsilon},\frac{n}{|\mathcal{S}|}\frac{b^{2}}{m\epsilon}\right\}\right),

to achieve an accuracy of

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\epsilon.

Ridge Separable Case

Assume that the individual component $\ell$ take the following form so that function $f$ becomes ridge-separable:

\displaystyle\ell(w,z_{i})=s_{i}(w^{\top}z_{i}).

(19)

To ensure bounded stochastic gradient variance at the mode of the posterior, we additionally assume that at the mode $w_{*}$ , the derivatives of the activation functions are bounded.

Assumption in ridge separable case on bounded variance

$\rm{R3^{SG}}$

$\exists b_{s}>0$ , so that $\left|s_{i}^{\prime}(w_{*}^{\top}z_{i})\right|\leq b_{s}$ , $\forall i\in\{1,\dots,n\}$ , where $w_{*}=\arg\min_{w}\left[f(w)+g(w)\right]$ .

Assumption $\rm{R3^{SG}}$ ensures that the stochastic gradient variance is bounded at the mode. Then we have the following corollary instantiating Theorem 4.

Corollary 2.

We make the convexity Assumption A1 on function $f$ and let it take the ridge-separable form (14) (also let function $g(w)=\frac{m}{2}\left\|w\right\|^{2}$ ). Further adopt Assumptions R1, R2, and $\rm{R3^{SG}}$ . Then taking $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m=2\cdot\left(8L_{s}R_{z}+mt\right)^{-1}$ , the convergence time of Algorithm 2 initializing from $p_{0}\propto\exp(-\beta^{-1}g)$ is

T=\Omega\left(\max\left\{\frac{L_{s}^{2}R_{z}^{2}}{m^{2}\epsilon},\frac{U(0)}{\epsilon},\frac{n}{|\mathcal{S}|}\frac{R_{z}b_{s}^{2}}{m\epsilon}\right\}\right),

to achieve an accuracy of

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\epsilon.

Proof of Corollary 2.

Since function $\ell$ takes form (19), we can compute that $\nabla^{2}\ell(w,z_{i})=s_{i}(w^{\top}z_{i})z_{i}z_{i}^{\top}$ . Using Assumptions R1 and R2, the Lipschitz smoothness $L_{\ell}=L_{s}R_{z}$ . Same as in Corollary 1, we know that $\nabla^{2}f(w)\preceq\frac{1}{n}L_{s}ZZ^{\top}=H$ . Therefore,

{\mathrm{trace}}\left(H\right)=L_{s}\cdot\frac{1}{n}{\mathrm{trace}}\left(ZZ^{\top}\right)\leq L_{s}R_{z},

leading to the fact that $L_{\ell}\cdot{\mathrm{trace}}\left(H\right)\leq L_{s}^{2}R_{z}^{2}$ .

For the stochastic gradient bound $b$ at $w_{*}$ , we apply Assumptions R2 and $\rm{R3^{SG}}$ to obtain

\left\|\nabla\ell(w_{*},z_{i})-\nabla\ell(w_{*},z_{j})\right\|=\left\|s^{\prime}_{i}(w_{*}^{\top}z_{i})z_{i}-s^{\prime}_{j}(w_{*}^{\top}z_{j})z_{j}\right\|\leq 2\sqrt{R_{z}}b_{s}.

We thus have

\left\|\nabla\ell(w_{*},z)-\nabla f(w_{*})\right\|=\Big{\|}\nabla\ell(w,z_{i})-\frac{1}{n}\sum_{j=1}^{n}\nabla\ell(w,z_{j})\Big{\|}\leq 2\sqrt{R_{z}}b_{s},

leading to the fact that

{\mathbb{E}}_{z\sim D}\left[\left\|\nabla\ell(w_{*},z)-\nabla f(w_{*})\right\|^{2}\right]\leq 4R_{z}b_{s}^{2}.

Therefore, the stochastic gradient variance bound in Assumption $\rm{A{{3}}_{S}^{SG}}$ , $b=2\sqrt{R_{z}}b_{s}$ . Plugging these bounds into Theorem 4 proves the corollary. ∎

We devote the rest of this section to the proof of Theorem 4.

Proof of Theorem 4.

We first note that because each $\ell(\cdot,z_{i})$ is $L_{\ell}$ -Lipschitz smooth, the stochastic estimate of function $f$ , $\widetilde{f}=\frac{1}{|\mathcal{S}|}\sum_{z_{i}\in\mathcal{S}}\ell(\widetilde{w}_{t},z_{i})$ is $L_{\ell}$ -Lipschitz smooth:

	$\displaystyle\left\\|\frac{1}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|^{2}$	$\displaystyle\leq\frac{1}{\|\mathcal{S}\|^{2}}\left(\sum_{z_{i}\in\mathcal{S}}\left\\|\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|\right)^{2}$
		$\displaystyle\leq\frac{1}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\left\\|\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|^{2}$
		$\displaystyle\leq\frac{2L_{\ell}}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\left(\ell(y,z_{i})-\ell(x,z_{i})-\nabla\ell(x,z_{i})^{\top}(y-x)\right)$
		$\displaystyle=2L_{\ell}\left(\widetilde{f}(y)-\widetilde{f}(x)+\nabla\widetilde{f}(x)^{\top}(y-x)\right).$

We thereby invoke the next lemma.

Lemma 7.

Assume that function $f$ is convex, and that its stochastic estimate $\widetilde{f}$ is $L_{\ell}$ -Lipschitz smooth. Then

	$\displaystyle W_{2}^{2}(p_{t},p)$	$\displaystyle\leq W_{2}^{2}(\widetilde{p}_{t},p)-2\widetilde{\eta}_{t}\left(f(\widetilde{p}_{t})-f(p)\right)$
		$\displaystyle+\widetilde{\eta}_{t}^{2}\left(4L_{\ell}[Q(\widetilde{p}_{t})-Q(p)]+2{\mathbb{E}}_{(\widetilde{w}_{t},w^{\prime})\sim\gamma_{t}}\left[{\mathbb{E}}_{\mathcal{S}\|\widetilde{w}_{t}}\left\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\\|^{2}\right]\right),$

where $f(q)={\mathbb{E}}_{w\sim q}f(w)$ , and $\gamma_{t}\in\Gamma_{\mathrm{opt}}(\widetilde{p}_{t},p)$ is the optimal coupling between $\widetilde{p}_{t}$ and $p$ .

Taking $\widetilde{\eta}_{t}=\left(1-e^{-m\eta_{t}}\right)/m$ and combining Lemma 1 and Lemma 7, we obtain that

\displaystyle 2\widetilde{\eta}_{t}\left(1-2\widetilde{\eta}_{t}L_{\ell}\right)\left(Q(\widetilde{p}_{t})-Q(p)\right)\leq e^{-m\eta_{t}}W_{2}^{2}(p_{t-1},p)-W_{2}^{2}(p_{t},p)+2\widetilde{\eta}_{t}^{2}{\mathbb{E}}_{(\widetilde{w}_{t},w^{\prime})\sim\gamma_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\widetilde{w}_{t}}\left\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}\right].

(20)

We then adapt Lemma 6 to the stochastic gradient method.

Lemma 8 (Stochastic Gradient Counterpart of Lemma 6).

Assume that

\nabla^{2}f(w)\preceq H,\quad g(w)=\frac{m}{2}\|w\|_{2}^{2}.

Let

w_{*}=\arg\min_{w}\left[f(w)+g(w)\right],

and $p$ be the solution of (2). Then for $L_{\ell}$ -Lipschitz smooth function $\widetilde{f}$ , at $p=p_{*}$ and consequently $\gamma_{t}\in\Gamma_{\mathrm{opt}}(\widetilde{p}_{t},p_{*})$ ,

\displaystyle{\mathbb{E}}_{(\widetilde{w}_{t},w^{\prime})\sim\gamma_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\widetilde{w}_{t}}\left\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}\right]\leq\frac{2\beta L_{\ell}}{m}{\mathrm{trace}}(H)+4m^{2}\left\|w_{*}\right\|^{2}+4{\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})-\nabla f(w_{*})\right\|^{2}.

(21)

For the last piece of information, we establish the variance of the stochastic gradient at the mode, ${\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})-\nabla f(w_{*})\right\|_{2}^{2}$ . For samples $z_{i}$ that are i.i.d. draws from the data set and are unbiased estimators of $\nabla f(w_{*})=\frac{1}{n}\sum_{j=1}^{n}\nabla\ell(w_{*},z_{j})$ , we have

\displaystyle{\mathbb{E}}_{\mathcal{S}}\left[\left\|\sum_{z_{i}\in\mathcal{S}}\left(\nabla\ell(w_{*},z_{i})-\frac{1}{n}\sum_{j=1}^{n}\nabla\ell(w_{*},z_{j})\right)\right\|^{2}\right]=|\mathcal{S}|\cdot{\mathbb{E}}_{z\sim D}\left[\left\|\nabla\ell(w_{*},z)-\frac{1}{n}\sum_{j=1}^{n}\nabla\ell(w_{*},z_{j})\right\|^{2}\right]\leq|\mathcal{S}|\cdot b^{2},

Leading to the bound that

\displaystyle{\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})-\nabla f(w_{*})\right\|_{2}^{2}=\frac{1}{|\mathcal{S}|^{2}}{\mathbb{E}}_{\mathcal{S}}\left[\left\|\sum_{z_{i}\in\mathcal{S}}\left(\nabla\ell(w_{*},z_{i})-\frac{1}{n}\sum_{j=1}^{n}\nabla\ell(w_{*},z_{j})\right)\right\|^{2}\right]\leq\frac{b^{2}}{|\mathcal{S}|}.

Plugging this result and Lemma 8 into equation (6.2) at $p=p_{*}$ , we obtain the final bound that

	$\displaystyle 2\widetilde{\eta}_{t}\left(1-2\widetilde{\eta}_{t}L_{s}R_{z}\right)\left(Q(\widetilde{p}_{t})-Q(p_{*})\right)$	$\displaystyle\leq\left(1-m\widetilde{\eta}_{t}\right)W_{2}^{2}(p_{t-1},p_{})-W_{2}^{2}(p_{t},p_{})$
		$\displaystyle+4\widetilde{\eta}_{t}^{2}\left(\beta\frac{L_{\ell}}{m}{\mathrm{trace}}(H)+2m^{2}\left\\|w_{*}\right\\|^{2}+2\frac{b^{2}}{\|\mathcal{S}\|}\right).$

This leads to a convergence rate, similar to the full gradient case, of

\displaystyle T=\Omega\left(\max\left(\frac{L_{\ell}{\mathrm{trace}}(H)}{m^{2}\epsilon},\frac{\beta^{-1}m\left\|w_{*}\right\|^{2}}{\epsilon},\frac{\beta^{-1}}{|\mathcal{S}|}\frac{b^{2}}{m\epsilon}\right)\right),

so that the weighted-averaged KL divergence:

\sum_{t=1}^{T}\frac{1/\left(m\widetilde{\eta}_{0}\right)+t/2}{T/\left(m\widetilde{\eta}_{0}\right)+T(T+1)/4}\beta^{-1}[Q(\widetilde{p}_{t})-Q(p_{*})]\leq\epsilon.

Since $\beta=1/n$ , and that $m\left\|w_{*}\right\|^{2}\leq 2f(0)=2\beta U(0)$ ,

\displaystyle T=\Omega\left(\max\left(\frac{L_{\ell}{\mathrm{trace}}(H)}{m^{2}\epsilon},\frac{U(0)}{\epsilon},\frac{n}{|\mathcal{S}|}\frac{b^{2}}{m\epsilon}\right)\right).

∎

7 Proofs of the Supporting Lemmas

7.1 Proofs of Lemmas in the Lipschitz Continuous Case

7.1.1 Proof of Lemma 1

Before proving Lemma 1, we first state a result in (Theorem 23.9 of Villani, 2009) that establishes the strong subdifferential of the Wasserstein-2 distance.

Lemma 9.

Assume that $\mu_{t},\hat{\mu}_{t}$ solve the following continuity equations

\frac{\partial\mu_{t}}{\partial t}+\nabla\cdot(\xi_{t}\mu_{t})=0,\qquad\frac{\partial\hat{\mu}_{t}}{\partial t}+\nabla\cdot(\hat{\xi}_{t}\hat{\mu}_{t})=0.

Then

\frac{1}{2}\frac{d}{dt}W_{2}^{2}(\mu_{t},\hat{\mu}_{t})=-\int\left\langle\tilde{\nabla}\psi_{t},\xi_{t}\right\rangle d\mu_{t}-\int\left\langle\tilde{\nabla}\hat{\psi}_{t},\hat{\xi}_{t}\right\rangle d\hat{\mu}_{t},

where $\psi_{t}$ and $\hat{\psi}_{t}$ are the optimal transport vector fields:

\exp(\tilde{\nabla}\psi_{t})_{\#}\mu_{t}=\hat{\mu_{t}},\qquad\exp(\tilde{\nabla}\hat{\psi}_{t})_{\#}\hat{\mu}_{t}=\mu_{t}.

Writing $p_{t}$ and $\hat{p}_{t}$ as the density functions of $\mu_{t}$ and $\hat{\mu}_{t}$ , we take $\xi_{t}=-\beta\nabla\log p_{t}-\nabla g$ and $\hat{\xi}_{t}=0$ so that $\mu_{t}$ follows the Fokker-Planck equation equation associated with process (3) and $\hat{\mu}_{t}=\nu$ is a constant measure. This leads to the following equation

\frac{1}{2}\frac{d}{ds}W_{2}^{2}(\mu_{s},\nu)=\int\left\langle\beta\nabla\log p_{s}+\nabla g,(\tilde{\nabla}\psi)_{\mu_{s}}^{\nu}\right\rangle d\mu_{s}.

For $\mu$ being the probability measure associated with its density $p$ , define relative entropy $\beta^{-1}F(\mu)$ , where $F(\mu)=\mathbb{E}_{w\sim p}\left[g(w)\right]+H(p)$ . We can then use the fact that the relative entropy $\beta^{-1}F$ is $\beta^{-1}m$ -geodesically strongly convex (see Proposition 9.3.2 of Ambrosio et al. (2008)) to prove the following Lemma.

Lemma 10.

For $p$ being the density of $\mu$ ,

F(\nu)-F(\mu)-\frac{m}{2}W_{2}^{2}(\mu,\nu)\geq\int\left\langle\beta\nabla\log p+\nabla g,(\tilde{\nabla}\psi)_{\mu}^{\nu}\right\rangle d\mu.

Proof.

Let $\mu_{t}$ be the geodesic between $\mu$ and $\nu$ . $\beta^{-1}m$ -geodesic strong convexity of $\beta^{-1}F$ states that (see Proposition 9.3.2 of Ambrosio et al. (2008)):

\beta^{-1}F(\mu_{t})\leq t\beta^{-1}F(\nu)+(1-t)\beta^{-1}F(\mu)-\frac{\beta^{-1}m}{2}t(1-t)W_{2}^{2}(\mu,\nu),

and consequently

\displaystyle\frac{F(\mu_{t})-F(\mu)}{t}\leq F(\nu)-F(\mu)-\frac{m}{2}(1-t)W_{2}^{2}(\mu,\nu).

By the definition of subdifferential (c.f. Villani, 2009, Theorem 23.14) we also have along the diffusion process defined by equation (3):

\lim\inf_{t\downarrow 0}\frac{F(\mu_{t})-F(\mu)}{t}\geq\int\left\langle\beta\nabla\log p+\nabla g,(\tilde{\nabla}\psi)_{\mu}^{\nu}\right\rangle d\mu.

Taking the limit of $t\rightarrow 0$ , we obtain the result. ∎

Proof of Lemma 1.

Combining Lemma 9 and 10, we obtain that

\displaystyle\frac{1}{2}\frac{d}{ds}W_{2}^{2}(\mu_{s},\nu)=\int\left\langle\beta\nabla\log p_{s}+\nabla g,(\tilde{\nabla}\psi)_{\mu_{s}}^{\nu}\right\rangle d\mu_{s}\leq F(\nu)-F(\mu_{s})-\frac{m}{2}W_{2}^{2}(\mu_{s},\nu).

(22)

Along the Fokker-Planck equation associated with process (3), $\frac{d}{ds}F(\mu_{s})=-\mathbb{E}_{p_{s}}\left[\left\|\beta\nabla\log p_{s}+\nabla g\right\|_{2}^{2}\right]\leq 0$ , meaning that $F(\mu_{s})$ is monotonically decreasing. We obtain from equation (22) for $s\in[0,t]$ ,

\displaystyle\frac{1}{2}\frac{d}{ds}W_{2}^{2}(\mu_{s},\nu)\leq\sup_{s\in[0,t]}\left[F(\nu)-F(\mu_{s})\right]-\frac{m}{2}W_{2}^{2}(\mu_{s},\nu)=F(\nu)-F(\mu_{t})-\frac{m}{2}W_{2}^{2}(\mu_{s},\nu).

Applying the Gronwall’s inequality, we arrive at the conclusion that

\displaystyle\frac{2}{m}\left(1-e^{-m\Delta t}\right)\left(F(\mu_{t})-F(\nu)\right)\leq e^{-m\Delta t}W_{2}^{2}(\mu_{0},\nu)-W_{2}^{2}(\mu_{t},\nu).

Taking $d\mu_{t}=\tilde{p}_{t}dx$ , $d\mu_{0}=p_{t-1}dx$ , $d\nu=pdx$ , and ${\Delta t}=\eta_{t}$ finishes the proof. ∎

7.1.2 Proof of Lemma 2

Proof of Lemma 2.

We first state a point-wise result along the gradient descent step (4):

\displaystyle 2\eta_{t}\left(f(\widetilde{w}_{t})-f(w)\right)\leq\left\|\widetilde{w}_{t}-w\right\|_{2}^{2}-\left\|w_{t}-w\right\|_{2}^{2}+\eta_{t}^{2}G^{2}.

(23)

This is because

	$\displaystyle\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-\eta_{t}\nabla f(\tilde{w}_{t})-w\right\\|_{2}^{2}$
		$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla f(\tilde{w}_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}\left\\|\nabla f(\tilde{w}_{t})\right\\|_{2}^{2}$
		$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left(f(\tilde{w}_{t})-f(w)\right)+\eta_{t}^{2}G^{2},$

where the last step follows from the convexity and Lipschitz smoothness of $f$ .

We then denote the measures corresponding to random variables $w_{t}$ and $\tilde{w}_{t}$ to be: $w_{t}\sim\mu_{t}$ and $\tilde{w}_{t}\sim\tilde{\mu}_{t}$ . From the definitions, we know that they have densities $p_{t}$ and $\tilde{p}_{t}$ .

Denote an optimal coupling between $\tilde{\mu}_{t}$ and $\mu$ (where measure $\mu$ has density $p$ , which is the stationary distribution) to be $\gamma\in\Gamma_{opt}(\tilde{\mu}_{t},\mu)$ . We then take expectations over $\gamma(\tilde{w}_{t},w)$ on both sides of equation (23):

	$\displaystyle 2\eta_{t}\left(f(\tilde{p}_{t})-f(p)\right)$	$\displaystyle=2\eta_{t}\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[f(\tilde{w}_{t})-f(w)\right]$
		$\displaystyle\leq\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\\|\tilde{w}_{t}-w\\|_{2}^{2}\right]-\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\\|w_{t}-w\\|_{2}^{2}\right]+\eta_{t}^{2}G^{2}$
		$\displaystyle=W_{2}^{2}(\tilde{p}_{t},p)-\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\\|w_{t}-w\\|_{2}^{2}\right]+\eta_{t}^{2}G^{2}.$

From the relationship $w_{t}=\tilde{w}_{t}-\eta_{t}\nabla f(\tilde{w}_{t})$ , we know that the joint distribution of $(w_{t},w)$ is $\left(\mathrm{id}-\eta_{t}\nabla f,\,\mathrm{id}\right)_{\#}\gamma$ . Note that $\tilde{\gamma}=\left(\mathrm{id}-\eta_{t}\nabla f,\,\mathrm{id}\right)_{\#}\gamma$ also defines a coupling, and therefore

	$\displaystyle\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\\|w_{t}-w\\|_{2}^{2}\right]$	$\displaystyle=\mathbb{E}_{(w_{t},w)\sim\tilde{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\right]$
		$\displaystyle\geq\inf_{\hat{\gamma}\in\Gamma(\mu_{t},\mu)}\mathbb{E}_{(w_{t},w)\sim\hat{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\right]=W_{2}^{2}(p_{t},p).$

Therefore,

\displaystyle 2\eta_{t}\left(f(\tilde{p}_{t})-f(p)\right)\leq W_{2}^{2}(\tilde{p}_{t},p)-W_{2}^{2}(p_{t},p)+\eta_{t}^{2}G^{2}.

∎

7.1.3 Proofs of Lemma 12 and 4 for the streaming SGLD algorithm 2

Proof of Lemma 12.

By the definitions of $w_{t}$ and $\tilde{w}_{t}$ ,

	$\displaystyle\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-\eta_{t}\nabla\ell(\tilde{w}_{t},z_{t})-w\right\\|_{2}^{2}$
		$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla\ell(\tilde{w}_{t},z_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}\left\\|\nabla\ell(\tilde{w}_{t},z_{t})\right\\|_{2}^{2}.$

We now take expectation with respect to $z_{t}$ , conditioned on $\tilde{w}_{t}$ , to obtain

	$\displaystyle{\mathbb{E}}_{z_{t}\|\tilde{w}_{t}}\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla f(\tilde{w}_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}G_{\ell}^{2}$
		$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left(f(\tilde{w}_{t})-f(w)\right)+\eta_{t}^{2}G_{\ell}^{2}.$

The last step follows from the convexity of $f$ . Therefore, the desired bound follows. ∎

Proof of Lemma 4.

We first denote the measures corresponding to random variables $w_{t}$ and $\tilde{w}_{t}$ to be: $w_{t}\sim\mu_{t}$ and $\tilde{w}_{t}\sim\tilde{\mu}_{t}$ . From the definitions, we know that they have densities $p_{t}$ and $\tilde{p}_{t}$ .

	$\displaystyle 2\eta_{t}\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\mathbb{E}_{z}\left[\ell(\tilde{w}_{t},z)-\ell(w,z)\right]\right]$
	$\displaystyle\leq\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[\\|\tilde{w}_{t}-w\\|_{2}^{2}\right]-\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t}\right]\right]+\eta_{t}^{2}G_{\ell}^{2}$
	$\displaystyle=W_{2}^{2}(\tilde{p}_{t},p)-\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t}\right]\right]+\eta_{t}^{2}G_{\ell}^{2}.$

From the relationship $w_{t}=\tilde{w}_{t}-\eta_{t}\nabla\ell(\tilde{w}_{t},z_{t})$ , we know that conditional on $z_{t}$ , the joint distribution of $(w_{t},w)$ is $\left(\mathrm{id}-\eta_{t}\nabla\ell,\,\mathrm{id}\right)_{\#}\gamma$ . Note that $\tilde{\gamma}=\left(\mathrm{id}-\eta_{t}\nabla\ell,\,\mathrm{id}\right)_{\#}\gamma$ also defines a coupling, and therefore

	$\displaystyle\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t}\right]\right]$
	$\displaystyle=\mathbb{E}_{z_{t}}\left[\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t},z_{t}\right]\right]\right]$
	$\displaystyle=\mathbb{E}_{z_{t}}\left[\mathbb{E}_{(w_{t},w)\sim\tilde{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\|z_{t}\right]\right]$
	$\displaystyle\geq\inf_{\hat{\gamma}\in\Gamma(\mu_{t},\mu)}\mathbb{E}_{(w_{t},w)\sim\hat{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\right]=W_{2}^{2}(p_{t},p).$

Plugging this result and the Lipschitz assumption on $\ell$ in, we obtain that

2\widetilde{\eta}_{t}[\ell(\widetilde{p}_{t})-\ell(p)]\leq W_{2}(\widetilde{p}_{t},p)^{2}-W_{2}(p_{t},p)^{2}+\widetilde{\eta}_{t}^{2}G_{\ell}^{2}.

∎

7.2 Proofs of Lemmas in the Lipschitz Smooth Case

7.2.1 Proofs of Lemmas 5 and 6 for the full gradient Langevin algorithm 1

Proof of Lemma 5.

By the geodesic convexity of the entropy function $H(p)=\beta{\mathbb{E}}_{w\sim p}\left[\ln p(w)\right]$ ,

H(\widetilde{p_{t}})-H(p)\geq\beta\int\left\langle\nabla\ln p(w^{\prime}),\left(T_{p^{\prime}}^{p}-\mathrm{id}\right)(w^{\prime})\right\rangle p(w^{\prime})\ dw^{\prime}.

where $T_{p}^{\widetilde{p}_{t}}$ is the optimal transport from $p$ to $\widetilde{p}_{t}$ . Using optimal coupling $\mu_{t}\in\Pi(\widetilde{p},p)$ ,

H(\widetilde{p_{t}})-H(p)\geq\beta{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}\left[\left\langle\nabla\ln p(w^{\prime}),w-w^{\prime}\right\rangle\right].

In addition, convexity of $f$ and $g$ implies that

{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}[g(w)-g(w^{\prime})]\geq{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}\left[\left\langle\nabla g(w^{\prime}),w-w^{\prime}\right\rangle\right]

and

{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}[f(w)-f(w^{\prime})]\geq{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}\left[\left\langle\nabla f(w^{\prime}),w-w^{\prime}\right\rangle\right].

Adding the above three inequalities, and note that the following holds point-wise

\beta\nabla\ln p(w^{\prime})+\nabla g(w^{\prime})+\nabla f(w^{\prime})=0,

we obtain that

H(\widetilde{p_{t}})-H(p)+{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}[g(w)-g(w^{\prime})]\geq-{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}\left[\left\langle\nabla f(w^{\prime}),w-w^{\prime}\right\rangle\right],

and that

\displaystyle Q(\widetilde{p_{t}})-Q(p)\geq{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}[f(w)-f(w^{\prime})]-{\mathbb{E}}_{w,w^{\prime}\sim\mu_{t}}\left[\left\langle\nabla f(w^{\prime}),w-w^{\prime}\right\rangle\right].

(24)

Since the potential function $f(w)$ is convex and $L$ -smooth,

f(w)\geq f(w^{\prime})+\nabla f(w^{\prime})^{\top}(w-w^{\prime})+\frac{1}{2L}\|\nabla f(w^{\prime})-\nabla f(w)\|_{2}^{2}.

(25)

Combining equations (24) and (25), we obtain the desired bound. ∎

Proof of Lemma 6.

We have

	$\displaystyle{\mathbb{E}}_{w\sim p}\\|\nabla f(w)\\|_{2}^{2}$	$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\\|\nabla f(w)-\nabla f(w_{})\\|_{2}^{2}+2\\|\nabla f(w_{})\\|_{2}^{2}$
		$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\left\\|\left(\int_{0}^{1}\nabla^{2}f\left(\tau w+(1-\tau)w_{}\right)d\tau\right)\left(w-w_{}\right)\right\\|^{2}+2m^{2}\\|w_{*}\\|_{2}^{2}$
		$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\left\\|H(w-w_{})\right\\|^{2}+2m^{2}\\|w_{}\\|_{2}^{2}.$

In what follows we bound ${\mathbb{E}}_{w\sim p}\left\|H(w-w_{*})\right\|^{2}$ . We first note that the posterior density:

p\propto\exp(-\beta^{-1}(f(w)+g(w)))

is $\frac{m}{\beta}$ strongly log-concave. By Hargé (2004), for any convex function $h$ ,

{\mathbb{E}}_{w\sim p}\left[h\left(w-{\mathbb{E}}_{w\sim p}\left[w\right]\right)\right]\leq{\mathbb{E}}_{\tilde{w}\sim\mathcal{N}\left(0,\beta/m\cdot{\mathrm{I}}\right)}\left[h(\tilde{w})\right].

Taking $h(w)=\left\|Hw\right\|^{2}$ , we obtain that

{\mathbb{E}}_{w\sim p}\left\|H\left(w-{\mathbb{E}}_{w\sim p}\left[w\right]\right)\right\|^{2}\leq{\mathbb{E}}_{\tilde{w}\sim\mathcal{N}\left(0,\beta/m\cdot{\mathrm{I}}\right)}\left\|H\tilde{w}\right\|^{2}=\frac{\beta}{m}{\mathrm{trace}}\left(H^{2}\right).

On the other hand, by the celebrated relation between mean and mode for 1-unimodal distributions (see, e.g., Theorem $7$ of Basu and DasGupta, 1996),

\left(w_{*}-{\mathbb{E}}_{w\sim p}[w]\right)^{\top}\Sigma^{-1}\left(w_{*}-{\mathbb{E}}_{w\sim p}[w]\right)\leq 3,

where $\Sigma$ is the covariance of $p$ . This results in the following bound

\left\|H\left(w_{*}-{\mathbb{E}}_{w\sim p}[w]\right)\right\|^{2}\leq 3\frac{\beta}{m}\left\|H\right\|_{2}^{2}.

Combining the two bounds, we obtain that

	$\displaystyle{\mathbb{E}}_{w\sim p}\left\\|H(w-w_{*})\right\\|^{2}$	$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\left\\|H\left(w-{\mathbb{E}}_{w\sim p}\left[w\right]\right)\right\\|^{2}+2\left\\|H\left(w_{*}-{\mathbb{E}}_{w\sim p}[w]\right)\right\\|^{2}$
		$\displaystyle\leq 2\frac{\beta}{m}{\mathrm{trace}}\left(H^{2}\right)+6\frac{\beta}{m}\left\\|H\right\\|_{2}^{2}\leq 8\frac{\beta}{m}{\mathrm{trace}}\left(H^{2}\right).$

∎

7.2.2 Proofs of Lemma 7 and 8 for the SGLD algorithm 2

Proof of Lemma 7.

By the definitions of $w_{t}$ and $\tilde{w}_{t}$ ,

	$\displaystyle\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-\eta_{t}\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})-w\right\\|_{2}^{2}$
		$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}\left\\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\\|_{2}^{2}.$

We now take expectation with respect to $\mathcal{S}$ , conditioned on $\tilde{w}_{t}$ , to obtain

	$\displaystyle{\mathbb{E}}_{w_{t}\|\tilde{w}_{t}}\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla f(\tilde{w}_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\\|_{2}^{2}$
		$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left(f(\tilde{w}_{t})-f(w)\right)+\eta_{t}^{2}{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\\|_{2}^{2}.$		(26)

We then upper bound ${\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\|_{2}^{2}$ by introducing variable $w^{\prime}$ that is distributed according to $p$ and couples optimally with the law of $\tilde{w}_{t}$ :

\displaystyle{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\|_{2}^{2}\leq 2{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})-\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}+2{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}.

(27)

For function $\widetilde{f}$ being $L_{\ell}$ -Lipschitz smooth,

\widetilde{f}(\tilde{w}_{t},\mathcal{S})\geq\widetilde{f}(w^{\prime},\mathcal{S})+\nabla\widetilde{f}(w^{\prime},\mathcal{S})^{\top}(w-w^{\prime})+\frac{1}{2L_{\ell}}\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})-\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\|_{2}^{2}.

Taking expectation over the randomness of minibatch assignment $\mathcal{S}$ on both sides leads to the fact that

f(\tilde{w}_{t})\geq f(w^{\prime})+\nabla f(w^{\prime})^{\top}(w-w^{\prime})+\frac{1}{2L_{\ell}}{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})-\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\|_{2}^{2}.

Combining this equation with equation (24), we adapt Lemma 5 to the stochastic gradient method:

{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})-\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\|_{2}^{2}\right]\leq 2L_{\ell}[Q(\widetilde{p}_{t})-Q(p)].

Applying this result to equation (27) and taking expectation of $(\tilde{w}_{t},w^{\prime})\sim\mu_{t}$ on both sides, we obtain:

{\mathbb{E}}_{\tilde{w}_{t}\sim\widetilde{p}_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\|_{2}^{2}\right]\leq 4L_{\ell}[Q(\widetilde{p}_{t})-Q(p)]+2{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}\right].

Therefore,

	$\displaystyle{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{w_{t}\|\tilde{w}_{t}}\left\\|w_{t}-w\right\\|_{2}^{2}\right]$	$\displaystyle\leq{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}\right]-2\eta_{t}\left(f(\widetilde{p}_{t})-f(p)\right)$
		$\displaystyle+\eta_{t}^{2}\left(4L_{\ell}[Q(\widetilde{p}_{t})-Q(p)]+2{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\\|^{2}\right]\right),$

leading to the final result that

	$\displaystyle W_{2}^{2}(p_{t},p)$	$\displaystyle\leq{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{w_{t}\|\tilde{w}_{t}}\left\\|w_{t}-w\right\\|_{2}^{2}\right]$
		$\displaystyle\leq W_{2}^{2}(\widetilde{p}_{t},p)-2\eta_{t}\left(f(\widetilde{p}_{t})-f(p)\right)$
		$\displaystyle+\eta_{t}^{2}\left(4L_{\ell}[Q(\widetilde{p}_{t})-Q(p)]+2{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\\|^{2}\right]\right).$

∎

Proof of Lemma 8.

We have

	$\displaystyle{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\\|^{2}\leq$	$\displaystyle 2{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})-\nabla\widetilde{f}(w_{},\mathcal{S})\\|_{2}^{2}+2{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\\|\nabla\widetilde{f}(w_{},\mathcal{S})\\|_{2}^{2}$
	$\displaystyle\leq$	$\displaystyle 4L_{\ell}\left[f(w)-f(w_{})-\nabla f(w_{})^{\top}(w-w_{})\right]+2{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\\|\nabla\widetilde{f}(w_{},\mathcal{S})\\|_{2}^{2}.$

Taking expectation on both sides, we obtain that

	$\displaystyle{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\\|^{2}\right]$	$\displaystyle\leq 4L_{\ell}{\mathbb{E}}_{w\sim p}\left[f(w)-f(w_{})-\nabla f(w_{})^{\top}(w-w_{*})\right]$
		$\displaystyle+2{\mathbb{E}}_{\mathcal{S}}\left\\|\nabla\widetilde{f}(w_{*},\mathcal{S})\right\\|^{2}.$		(28)

We now upper bound ${\mathbb{E}}_{w\sim p}\left[f(w)-f(w_{*})-\nabla f(w_{*})^{\top}(w-w_{*})\right]$ . Let $p_{0}$ be the normal distribution $N(w_{*},(\beta/m)I)$ , and define

	$\displaystyle\Delta f(w)$	$\displaystyle=f(w)-f(w_{})-\nabla f(w_{})^{\top}(w-w_{*}),$		(29)
	$\displaystyle\Delta g(w)$	$\displaystyle=g(w)-g(w_{})-\nabla g(w_{})^{\top}(w-w_{*}),$		(30)

then $p$ can be expressed as

p\propto\exp(-\beta^{-1}(\Delta f(w)+\Delta g(w))),

which is the solution of

p=\arg\min_{p}{\mathbb{E}}_{w\sim p}\left[\Delta f(w)+\beta\ln\frac{p(w)}{p_{0}(w)}\right].

Therefore

$\displaystyle{\mathbb{E}}_{w\sim p}\Delta f(w)\leq$	$\displaystyle{\mathbb{E}}_{w\sim p}\left[\Delta f(w)+\beta\ln\frac{p(w)}{p_{0}(w)}\right]$
$\displaystyle=$	$\displaystyle-\beta\ln{\mathbb{E}}_{w\sim p_{0}}\exp(-\beta^{-1}\Delta f(w))$
$\displaystyle\leq$	$\displaystyle-\beta\ln{\mathbb{E}}_{w\sim p_{0}}\exp(-0.5\beta^{-1}(w-w_{})^{\top}H(w-w_{}))$
$\displaystyle=$	$\displaystyle 0.5\beta\ln\|I+H/m\|\leq 0.5\beta{\mathrm{trace}}(H/m).$	(31)

We then decompose ${\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}$ :

{\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}\leq 2\left\|\nabla f(w_{*})\right\|^{2}+2{\mathbb{E}}_{\mathcal{S}}\left\|\nabla f(w_{*})-\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}.

Since $w_{*}$ is the minimum of $f+g$ , $\nabla f(w_{*})=-\nabla g(w_{*})=-mw_{*}$ . Hence

\displaystyle{\mathbb{E}}_{\mathcal{S}}\left\|\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}\leq 2m^{2}\left\|w_{*}\right\|^{2}+2{\mathbb{E}}_{\mathcal{S}}\left\|\nabla f(w_{*})-\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}.

(32)

Plugging inequalities (31) and (32) into inequality (28) proves the desired result that

\displaystyle{\mathbb{E}}_{(\tilde{w}_{t},w^{\prime})\sim\mu_{t}}\left[{\mathbb{E}}_{\mathcal{S}|\tilde{w}_{t}}\left\|\nabla\widetilde{f}(w^{\prime},\mathcal{S})\right\|^{2}\right]\leq\frac{2\beta L_{\ell}}{m}{\mathrm{trace}}(H)+4m^{2}\left\|w_{*}\right\|^{2}+4{\mathbb{E}}_{\mathcal{S}}\left\|\nabla f(w_{*})-\nabla\widetilde{f}(w_{*},\mathcal{S})\right\|^{2}.

(33)

∎

8 Conclusion

This paper investigated the convergence of Langevin algorithms with strongly log-concave posteriors. We assume that the strongly log-concave posterior can be decomposed into two parts, with one part being simple and explicitly integrable with respect to the underlying SDE. This is analogous to the situation of proximal gradient methods in convex optimization. Using a new analysis technique which mimics the corresponding analysis of convex optimization, we obtain convergence results for Langenvin algorithms that are independent of dimension, both for Lipschitz and for a large class of smooth convex problems in machine learning. Our result addresses a long-standing puzzle with respect to the convergence of the Langevin algorithms. We note that the current work focused on the standard Langevin algorithm, and the resulting convergence rate in terms of $\epsilon$ dependency is inferior to the best known results leveraging underdamped or even higher order Langevin dynamics such as Cheng et al. (2018b); Dalalyan and Riou-Durand (2018); Shen and Lee (2019); Ma et al. (2021); Mou et al. (2021), which corresponds to accelerated methods in optimization. It thus remains open to investigate whether dimension independent bounds can be combined with these accelerated methods to improve $\epsilon$ dependence as well as condition number dependence.

References

Ambrosio et al. (2008) L. Ambrosio, N. Gigli, and G. Savaré. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2nd edition, 2008.
Bakry and Emery (1985) D. Bakry and M. Emery. Diffusions hypercontractives. In Séminaire de Probabilités XIX 1983/84, pages 177–206. 1985.
Basu and DasGupta (1996) S. Basu and A. DasGupta. The mean, median, and mode of unimodal distributions: a characterization. Teor. Veroyatnost. i Primenen., 41:336–352, 1996.
Bernton (2018) E. Bernton. Langevin Monte Carlo and JKO splitting. In Proceedings of the 31st Conference on Learning Theory (COLT), pages 1777–1798, 2018.
Bou-Rabee et al. (2018) N. Bou-Rabee, A. Eberle, and R. Zimmer. Coupling and convergence for Hamiltonian Monte Carlo. arXiv:1805.00452, 2018.
Chatterji et al. (2018) N. Chatterji, N. Flammarion, Y.-A. Ma, P. Bartlett, and M. Jordan. On the theory of variance reduction for stochastic gradient Monte Carlo. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pages 764–773, 2018.
Chatterji et al. (2020) Niladri S. Chatterji, Jelena Diakonikolas, Michael I. Jordan, and Peter L. Bartlett. Langevin monte carlo without smoothness. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS 2020), volume 108, 2020.
Cheng and Bartlett (2018) X. Cheng and P. L. Bartlett. Convergence of Langevin MCMC in KL-divergence. In Proceedings of the 29th International Conference on Algorithmic Learning Theory (ALT), pages 186–211, 2018.
Cheng et al. (2018a) X. Cheng, N. S. Chatterji, Y. Abbasi-Yadkori, P. L. Bartlett, and M. I. Jordan. Sharp convergence rates for Langevin dynamics in the nonconvex setting. arXiv:1805.01648, 2018a.
Cheng et al. (2018b) X. Cheng, N. S. Chatterji, P. L. Bartlett, and M. I. Jordan. Underdamped Langevin MCMC: A non-asymptotic analysis. In Proceedings of the 31st Conference on Learning Theory (COLT), pages 300–323, 2018b.
Dalalyan (2017) A. S. Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. J. Royal Stat. Soc. B, 79(3):651–676, 2017.
Dalalyan and Karagulyan (2017) A. S. Dalalyan and A. G. Karagulyan. User-friendly guarantees for the Langevin Monte Carlo with inaccurate gradient. arXiv:1710.00095, 2017.
Dalalyan and Riou-Durand (2018) A. S. Dalalyan and L. Riou-Durand. On sampling from a log-concave density using kinetic Langevin diffusions. arXiv:1807.09382, 2018.
Durmus and Moulines (2017) A. Durmus and E. Moulines. Nonasymptotic convergence analysis for the unadjusted Langevin algorithm. Ann. Appl. Probab., 27(3):1551–1587, 06 2017.
Durmus and Moulines (2019) A. Durmus and E. Moulines. High-dimensional Bayesian inference via the unadjusted Langevin algorithm. Bernoulli, 25(4A):2854–2882, 2019.
Durmus et al. (2019) Alain Durmus, Szymon Majewski, and Błażej Miasojedow. Analysis of Langevin Monte Carlo via convex optimization. The Journal of Machine Learning Research, 20(1):2666–2711, 2019.
Dwivedi et al. (2018) R. Dwivedi, Y. Chen, M. J. Wainwright, and B. Yu. Log-concave sampling: Metropolis-Hastings algorithms are fast! arXiv:1801.02309, 2018.
Gelman et al. (2004) J. B. Gelman, A.and Carlin, H. S. Stern, and D. B. Rubin. Bayesian Data Analysis. Chapman and Hall, New York, 2004.
Hargé (2004) G. Hargé. A convex/log-concave correlation inequality for Gaussian measure and an application to abstract Wiener spaces. Probab. Theory Related Fields, 130:415–440, 2004.
Jordan et al. (1998) R. Jordan, D. Kinderlehrer, and F. Otto. The variational formulation of the Fokker-Planck equation. SIAM J. Math. Anal., 29(1):1–17, January 1998.
Lacoste-Julien et al. (2012) Simon Lacoste-Julien, Mark Schmidt, and Francis R. Bach. A simpler approach to obtaining an o(1/t) convergence rate for the projected stochastic subgradient method. CoRR, abs/1212.2002, 2012.
Ledoux (2000) M. Ledoux. The geometry of Markov diffusion generators. Ann Fac Sci Toulouse Math, 9(6):305–366, 2000.
Lee et al. (2018) Y.-T. Lee, Z. Song, and S. S. Vempala. Algorithmic theory of ODEs and sampling from well-conditioned logconcave densities. arXiv:1812.06243, 2018.
Li and Lin (2010) Q. Li and N. Lin. The Bayesian elastic net. Bayesian Analysis, 5(1):151–170, 2010.
Ma et al. (2019) Y.-A. Ma, Y. Chen, C. Jin, N. Flammarion, and M. I. Jordan. Sampling can be faster than optimization. Proc. Natl. Acad. Sci. U.S.A., 116:20881–20885, 2019.
Ma et al. (2021) Y.-A. Ma, N. S. Chatterji, X. Cheng, N. Flammarion, P. L. Bartlett, and M. I. Jordan. Is there an analog of Nesterov acceleration for MCMC? Bernoulli, 27(3):1942–1992, 2021.
Mangoubi and Smith (2017) O. Mangoubi and A. Smith. Rapid mixing of Hamiltonian Monte Carlo on strongly log-concave distributions. arXiv:1708.07114, 2017.
Mangoubi and Vishnoi (2018) O. Mangoubi and N. K. Vishnoi. Dimensionally tight running time bounds for second-order Hamiltonian Monte Carlo. arXiv:1802.08898, 2018.
Mou et al. (2021) W. Mou, Y.-A. Ma, M. J. Wainwright, P. L. Bartlett, and M. I. Jordan. High-order Langevin diffusion yields an accelerated MCMC algorithm. Journal of Machine Learning Research, 22:1–41, 2021.
Mou et al. (2019) Wenlong Mou, Nicolas Flammarion, Martin J. Wainwright, and Peter L. Bartlett. An efficient sampling algorithm for non-smooth composite potentials. arXiv:1910.00551, 2019.
Neal (1996) R. M. Neal. Bayesian Learning for Neural Networks. Number 118 in Lecture Notes in Statistics. Springer-Verlag, New York, 1996.
Otto and Villani (2000) F. Otto and C. Villani. Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality. J. Funct. Anal., 173(2):361–400, 2000.
Rakhlin et al. (2012) Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. In International conference on machine learning, pages 1571–1578, 2012.
Roberts and Stramer (2002) G. O. Roberts and O. Stramer. Langevin diffusions and Metropolis-Hastings algorithms. Methodol. Comput. Appl. Probab., 4:337–357, 2002.
Rossky et al. (1978) P. J. Rossky, J. D. Doll, and H. L. Friedman. Brownian dynamics as smart Monte Carlo simulation. J. Chem. Phys., 69(10):4628, 1978.
Shalev-Shwartz et al. (2011) Shai Shalev-Shwartz, Yoram Singer, Nathan Srebro, and Andrew Cotter. Pegasos: Primal estimated sub-gradient solver for svm. Mathematical programming, 127(1):3–30, 2011.
Shamir and Zhang (2013) Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning, pages 71–79. PMLR, 2013.
Shen and Lee (2019) R. Shen and Y. T. Lee. The randomized midpoint method for log-concave sampling. In Advances in Neural Information Processing Systems, pages 2098–2109, 2019.
Sollich (2002) P. Sollich. Bayesian methods for support vector machines: Evidence and predictive class probabilities. Machine Learning, 46:21–52, 2002.
Vempala and Wibisono (2019) S. S. Vempala and A. Wibisono. Rapid convergence of the unadjusted Langevin algorithm: Log-Sobolev suffices. arXiv:1903.08568, 2019.
Villani (2009) C. Villani. Optimal Transport, Old and New, volume 338 of Grundlehren der Mathematischen Wissenschaften. Springer, Berlin, Heidelberg, 2009.
Wibisono (2018) A. Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Proceedings of the 31st Conference on Learning Theory (COLT), pages 2093–3027, 2018.
Zhang (2004) Tong Zhang. Solving large scale linear prediction problems using stochastic gradient descent algorithms. In Proceedings of the twenty-first international conference on Machine learning, page 116, 2004.
Zinkevich (2003) Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference on machine learning (icml-03), pages 928–936, 2003.
Zou and Gu (2019) Difan Zou and Quanquan Gu. Stochastic gradient hamiltonian monte carlo methods with recursive variance reduction. In Advances in Neural Information Processing Systems (NeurIPS) 32. 2019.

	$\displaystyle\left\\|\frac{1}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|^{2}$	$\displaystyle\leq\frac{1}{\|\mathcal{S}\|^{2}}\left(\sum_{z_{i}\in\mathcal{S}}\left\\|\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|\right)^{2}$
		$\displaystyle\leq\frac{1}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\left\\|\nabla\ell(y,z_{i})-\nabla\ell(x,z_{i})\right\\|^{2}$
		$\displaystyle\leq\frac{2L_{\ell}}{\|\mathcal{S}\|}\sum_{z_{i}\in\mathcal{S}}\left(\ell(y,z_{i})-\ell(x,z_{i})-\nabla\ell(x,z_{i})^{\top}(y-x)\right)$
		$\displaystyle=2L_{\ell}\left(\widetilde{f}(y)-\widetilde{f}(x)+\nabla\widetilde{f}(x)^{\top}(y-x)\right).$

	$\displaystyle\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-\eta_{t}\nabla f(\tilde{w}_{t})-w\right\\|_{2}^{2}$
		$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla f(\tilde{w}_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}\left\\|\nabla f(\tilde{w}_{t})\right\\|_{2}^{2}$
		$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left(f(\tilde{w}_{t})-f(w)\right)+\eta_{t}^{2}G^{2},$

	$\displaystyle\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t}\right]\right]$
	$\displaystyle=\mathbb{E}_{z_{t}}\left[\mathbb{E}_{(\tilde{w}_{t},w)\sim\gamma}\left[{\mathbb{E}}_{w_{t}}\left[\\|w_{t}-w\\|_{2}^{2}\|\widetilde{w}_{t},z_{t}\right]\right]\right]$
	$\displaystyle=\mathbb{E}_{z_{t}}\left[\mathbb{E}_{(w_{t},w)\sim\tilde{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\|z_{t}\right]\right]$
	$\displaystyle\geq\inf_{\hat{\gamma}\in\Gamma(\mu_{t},\mu)}\mathbb{E}_{(w_{t},w)\sim\hat{\gamma}}\left[\\|w_{t}-w\\|_{2}^{2}\right]=W_{2}^{2}(p_{t},p).$

	$\displaystyle{\mathbb{E}}_{w\sim p}\\|\nabla f(w)\\|_{2}^{2}$	$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\\|\nabla f(w)-\nabla f(w_{})\\|_{2}^{2}+2\\|\nabla f(w_{})\\|_{2}^{2}$
		$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\left\\|\left(\int_{0}^{1}\nabla^{2}f\left(\tau w+(1-\tau)w_{}\right)d\tau\right)\left(w-w_{}\right)\right\\|^{2}+2m^{2}\\|w_{*}\\|_{2}^{2}$
		$\displaystyle\leq 2{\mathbb{E}}_{w\sim p}\left\\|H(w-w_{})\right\\|^{2}+2m^{2}\\|w_{}\\|_{2}^{2}.$

	$\displaystyle{\mathbb{E}}_{w_{t}\|\tilde{w}_{t}}\left\\|w_{t}-w\right\\|_{2}^{2}$	$\displaystyle=\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left\langle\nabla f(\tilde{w}_{t}),\tilde{w}_{t}-w\right\rangle+\eta_{t}^{2}{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\\|_{2}^{2}$
		$\displaystyle\leq\left\\|\tilde{w}_{t}-w\right\\|_{2}^{2}-2\eta_{t}\left(f(\tilde{w}_{t})-f(w)\right)+\eta_{t}^{2}{\mathbb{E}}_{\mathcal{S}\|\tilde{w}_{t}}\left\\|\nabla\widetilde{f}(\tilde{w}_{t},\mathcal{S})\right\\|_{2}^{2}.$		(26)

When is the Convergence Time of Langevin Algorithms Dimension Independent? A Composite Optimization Viewpoint

Abstract

1 Introduction

2 Preliminaries

2.1 Two Problem Classes

Assumption on function gg

Assumption on function ff

2.2 Objective Functional and Convergence Criteria

Proposition 1.

Definition 1.

Proposition 2.

3 Related Works

In the Lipschitz continuous case,

In the Lipschitz smooth case,

4 Langevin Algorithms

5 Langevin Algorithms in Lipschitz Convex Case

Assumptions for the Lipschitz Convex Case:

5.1 Full Gradient Langevin Algorithm Convergence in Lipschitz Convex Case

Theorem 1.

Proof of Theorem 1.

Lemma 1 (For Regularized Entropy).

Lemma 2.

Proposition 3.

Proof of Proposition 3.

5.2 Streaming SGLD Convergence in Lipschitz Convex Case

Assumptions on individual loss ℓ\ell

Theorem 2.

Proof of Theorem 2.

Lemma 3.

Lemma 4.

Proposition 4.

6 Langevin Algorithms in Smooth Convex Case

Assumptions for the smooth convex case:

6.1 Full Gradient Langevin Algorithm Convergence in Smooth Convex Case

Theorem 3.

Ridge Separable Case

Assumptions in ridge separable case

Corollary 1.

Proof.

Proof of Theorem 3.

Lemma 5.

Lemma 6.

6.2 SGLD Convergence in Smooth Convex Case

Assumptions on individual loss ℓ\ell

Theorem 4.

Ridge Separable Case

Assumption in ridge separable case on bounded variance

Corollary 2.

Proof of Corollary 2.

Proof of Theorem 4.

Lemma 7.

Lemma 8 (Stochastic Gradient Counterpart of Lemma 6).

7 Proofs of the Supporting Lemmas

7.1 Proofs of Lemmas in the Lipschitz Continuous Case

7.1.1 Proof of Lemma 1

Lemma 9.

Lemma 10.

Proof.

Proof of Lemma 1.

7.1.2 Proof of Lemma 2

Proof of Lemma 2.

7.1.3 Proofs of Lemma 12 and 4 for the streaming SGLD algorithm 2

Proof of Lemma 12.

Proof of Lemma 4.

7.2 Proofs of Lemmas in the Lipschitz Smooth Case

7.2.1 Proofs of Lemmas 5 and 6 for the full gradient Langevin algorithm 1

Proof of Lemma 5.

Proof of Lemma 6.

7.2.2 Proofs of Lemma 7 and 8 for the SGLD algorithm 2

Proof of Lemma 7.

Proof of Lemma 8.

8 Conclusion

References

Assumption on function $g$

Assumption on function $f$

Assumptions on individual loss $\ell$

Assumptions on individual loss $\ell$