On the Theory of Variance Reduction for Stochastic Gradient Monte Carlo

Introduction

One of the major themes in machine learning is the use of stochasticity to obtain procedures that are computationally efficient and statistically calibrated. There are two very different ways in which this theme has played out—one frequentist and one Bayesian. On the frequentist side, gradient-based optimization procedures are widely used to obtain point estimates and point predictions, and stochasticity is used to bring down the computational cost by replacing expensive full-gradient computations with unbiased stochastic-gradient computations. On the Bayesian side, posterior distributions provide information about uncertainty in estimates and predictions, and stochasticity is used to represent those distributions in the form of Monte Carlo (MC) samples. Despite the different conceptual frameworks, there are overlapping methodological issues. In particular, Monte Carlo sampling must move from an out-of-equilibrium configuration towards the posterior distribution and must do so quickly, and thus optimization ideas are relevant. Frequentist inference often involves sampling and resampling, so that efficient approaches to Monte Carlo sampling are relevant.

Variance control has been a particularly interesting point of contact between the two frameworks. In particular, there is a subtlety in the use of stochastic gradients for optimization: Although the per-iteration cost is significantly lower by using stochastic gradients; extra variance is introduced into the sampling procedure at every step so that the total number of iterations is required to be larger. A natural question is whether there is a theoretically-sound way to manage this tradeoff. This question has been answered affirmatively in a seminal line of research [Schmidt et al., 2017, Shalev-Shwartz and Zhang, 2013, Johnson and Zhang, 2013] on variance-controlled stochastic optimization. Theoretically these methods enjoy the best of the gradient and stochastic gradient worlds—they converge at the fast rate of full gradient methods while making use of cheaply-computed stochastic gradients.

A parallel line of research has ensued on the Bayesian side in a Monte Carlo sampling framework. In particular, stochastic-gradient Markov chain Monte Carlo (SG-MCMC) algorithms have been proposed in which approximations to Langevin diffusions make use of stochastic gradients instead of full gradients Welling and Teh . There have been a number of theoretical results that establish mixing time bounds for such Langevin-based sampling methods when the posterior distribution is well behaved [Dalalyan, 2017a, Durmus and Moulines, 2017, Cheng and Bartlett, 2017, Dalalyan and Karagulyan, 2017]. Such results have set the stage for the investigation of variance control within the SG-MCMC framework Dubey et al. , Durmus et al. , Bierkens et al. , Baker et al. , Nagapetyan et al. , Chen et al. . Currently, however, the results of these investigations are inconclusive. Dubey et al. obtain mixing time guarantees for SAGA Langevin diffusion and SVRG Langevin diffusion (two particular variance-reduced sampling methods) under the strong assumption that the log-posterior has the norm of its gradients uniformly bounded by a constant. Another approach that has been explored involves calculating the mode of the log posterior to construct a control variate for the gradient estimate [Baker et al., 2017, Nagapetyan et al., 2017], an approach that makes rather different assumptions. Indeed, the experimental results from these two lines of work are contradictory, reflecting the differences in assumptions.

In this work we aim to provide a unified perspective on variance control for SG-MCMC. Critically, we identify two regimes: we show that when the target accuracy is small, variance-reduction methods are effective, but when the target accuracy is not small (a low-fidelity estimate of the posterior suffices), stochastic gradient Langevin diffusion (SGLD) performs better. These results are obtained via new theoretical techniques for studying stochastic gradient MC algorithms with variance reduction. We improve upon the techniques used to analyze Langevin Diffusion (LD) and SGLD Dalalyan [2017a], Dalalyan and Karagulyan , Durmus and Moulines to establish non-asymptotic rates of convergence (in Wasserstein distance) for variance-reduced methods. We also apply control-variate techniques to underdamped Langevin MCMC [Cheng et al., 2017], a second-order diffusion process (CV-ULD). Inspired by proof techniques for variance-reduction methods for stochastic optimization, we design a Lyapunov function to track the progress of convergence and we thereby obtain better bounds on the convergence rate. We make the relatively weak assumption that the log posteriors are Lipschitz smooth, strongly convex and Hessian Lipschitz—a relaxation of the strong assumption that the gradient of the log posteriors are globally bounded.

We provide sharp convergence guarantees for a variety of variance-reduction methods—SAGA-LD, SVRG-LD, and CV-ULD under the same set of realistic assumptions (see Sec. 4). This is achieved by a new proof technique that yiels bounds on Wasserstein distance. Our bounds allow us to identify windows of interest where each method performs better than the others (see Fig. 1). The theory is verified with experiments on real-world datasets. We also test the effects of breaking the central limit theorem using synthetic data, and find that in this regime variance-reduced methods fare far better than SGLD (see Sec. 5).

Preliminaries

Sum-decomposable: The function $f$ is decomposable, $f(x)=\sum_{i=1}^{N}f_{i}(x)$.

It is worth noting that $M,m$ and $L$ can all scale with $N$.

Wasserstein Distance: We define the Wasserstein distance between a pair of probability measures ($\mu,\nu$) as follows:

where $\Gamma(\mu,\nu)$ denotes the set of joint distributions such that the first set of coordinates has marginal $\mu$ and the second set has marginal $\nu$. (See Appendix A for a more formal definition of $W_{2}$).

Langevin Diffusion: The classical overdamped Langevin diffusion is based on the following Itô Stochastic Differential Equation (SDE):

where the gradient is evaluated at a fixed point $x_{k}$ (the previous iterate in the chain) and the SDE (2) is integrated up to time $\delta$ (the step size) to obtain

with $\xi_{k}\sim N(0,I_{d\times d})$. Welling and Teh proposed an alternative algorithm—Stochastic Gradient Langevin Diffusion (SGLD)—for sampling from sum-decomposable function where the chain is updated by integrating the SDE:

and where $g_{k}=\frac{N}{n}\sum_{i\in S}\nabla f_{i}(x_{k})$ is an unbiased estimate of the gradient at $x_{k}$. The attractive property of this algorithm is that it is computationally tractable for large datasets (when $N$ is large). At a high level the variance reduction schemes that we study in this paper replace the simple gradient estimate in Eq. (3) (and other variants of Langevin MCMC) with more sophisticated unbiased estimates that have lower variance.

Variance Reduction Techniques

In the seminal work of Schmidt et al. and Johnson and Zhang , it was observed that the variance of Stochastic Gradient Descent (SGD) when applied to optimizing sum-decomposable strongly convex functions decreases to zero only if the step-size also decays at a suitable rate. This prevents the algorithm from converging at a linear rate, as opposed to methods like batch gradient descent that use the entire gradient at each step. They introduced and analyzed different gradient estimates with lower variance. Subsequently these methods were also adapted to Monte Carlo sampling by Dubey et al. , Nagapetyan et al. , Baker et al. . These methods use information from previous iterates and are no longer Markovian. In this section we describe several variants of these methods.

We present a sampling algorithm based on SAGA of Defazio et al. which was developed as a modification of SAG by Schmidt et al. . In SAGA, which is presented as Algorithm 1, an approximation of the gradient of each function $f_{i}$ is stored as $\{g_{k}^{i}\}_{i=1}^{N}$ and is iteratively updated in order to build an estimate with reduced variance. At each step of the algorithm, if the function $f_{i}$ is selected in the mini-batch $S$, then the value of the gradient approximation is updated by setting $g^{i}_{k+1}=\nabla f_{i}(x_{k})$. Otherwise the gradient of $f_{i}$ is approximated by the previous value $g_{k}^{i}$. Overall we obtain the following unbiased estimate of the gradient:

In Algorithm 1 we form this gradient estimate and plug it into the classic Langevin MCMC method driven by the SDE (3). Computationally this algorithm is efficient; essentially it enjoys the oracle query complexity (number of calls to the gradient oracle per iteration) of methods like SGLD but due to the reduced variance of the gradient estimator it converges almost as quickly (in terms of number of iterations) to the posterior distribution as methods such as Langevin MCMC that use the complete gradient at every step. We prove a novel non-asymptotic convergence result in Wasserstein distance for Algorithm 1 in the next section that formalizes this intuition.

The principal downside of this method is its memory requirement. It is necessary to store the gradient estimator for each individual $f_{i}$, which essentially means that in the worst case the memory complexity scales as $\mathcal{O}(Nd)$. However in many interesting applications, including some of those considered in the experiments in Sec. 5, the memory costs scale only as $\mathcal{O}(N)$ since each function $f_{i}$ depends on a linear function in $x$ and therefore the gradient $\nabla f_{i}$ is just a re-weighting of the single data point $z_{i}$.

2 SVRG Langevin MC

3 Control Variates with Underdamped Langevin MC

Another approach is to use control variates [Ripley, 2009] to reduce the variance of stochastic gradients. This technique has also been previously explored both theoretically and experimentally by Baker et al. and Nagapetyan et al. . Similarly to SAGA and SVRG the idea is to build an unbiased estimate of the gradient $g(x)$ at a point $x$:

where the set $S$ is the mini-batch and $\hat{x}$ is a fixed point that is called the centering value. Observe that taking an expectation over the choice of the set $S$ yields $\nabla f(x)$. A good centering value $\hat{x}$ would ensure that this estimate also has low variance; a natural choice in this regard is the global minima of $f$, $x^{*}$. A motivating example is the case of a Gaussian random variable where the mean of the distribution and $x^{*}$ coincide.

A conclusion of previous work that applies control variate techniques to stochastic gradient Langevin MCMC is the following—the variance of the gradient estimates can be lowered to be of the order of the discretization error. Motivated by this, we apply these techniques to underdamped Langevin MCMC where the underlying continuous time diffusion process is given by the following second-order SDE:

The discretization of SDE (6) (which we can simulate efficiently) is

In Algorithm 3 the updates of the gradients are dictated by,

Convergence results

In this section we provide convergence results of the algorithms presented above, which improve upon the convergence guarantees for SGLD. Dalalyan and Karagulyan show that for SGLD run for $T$ iterations:

under assumptions (A2)-(A4) with access to stochastic gradients with bounded variance – $\sigma^{2}d$. The term involving the variance – $\sigma\sqrt{\delta d}/{2\sqrt{m}}$ dominates the others in many interesting regimes. For sum-decomposable functions that we are studying in this paper this is also the case as the variance of the gradient estimate usually scales linearly with $N^{2}$. Therefore the performance of SGLD sees a deterioration when compared to the convergence guarantees of Langevin Diffusion where $\sigma=0$. To prove our convergence results we follow the general framework established by Dalalyan and Karagulyan , with the noteworthy difference of working with more sophisticated Lyapunov functions (for Theorems 4.1 and 4.2) inspired by proof techniques in optimization theory. This contributes to strengthening the connection between optimization and sampling methods raised in previous work and may potentially be applied to other sampling algorithms (we elaborate on these connections in more detail in Appendix B). This comprehensive proof technique also allows us to sharpen the convergence guarantees obtained by Dubey et al. on variance reduction methods like SAGA and SVRG by allowing us to present bounds in $W_{2}$ and to drop the assumption on requiring uniformly bounded gradients. We now present convergence guarantees for Algorithm 1.

Let assumptions (A1)-(A4) hold. Let $p^{(T)}$ be the distribution of the iterate of Algorithm 1 after $T$ steps. If we set the step size to be $\delta<\frac{n}{8MN}$ and the batch size $n\geq 9$ then we have the guarantee:

For the sake of clarity, only results for small step-size $\delta$ are presented however, it is worth noting that convergence guarantees hold for any $\delta\leq\frac{1}{8M}$ (see details in Appendix B.2). If we consider the regime where $\sigma,M,L$ and $m$ all scale linearly with the number of samples $N$, then for SGLD the dominating term is $\mathcal{O}(\sigma\sqrt{\delta d/m})$. If the target accuracy is $\epsilon$, SGLD would require the step size to scale as $\mathcal{O}(\epsilon^{2}/d)$ while for SAGA a step size of $\delta=\mathcal{O}(\epsilon/d)$ is sufficient. The mixing time $T$ for both methods is roughly proportional to the inverse step-size; thus SAGA provably takes fewer iterations while having almost the same computational complexity per step as SGLD. Similar to the optimization setting, theoretically SAGA Langevin diffusion recovers the fast rate of Langevin diffusion while just using cheap gradient updates. Next we present our guarantees for Algorithm 2.

Let assumptions (A1)-(A4) hold. Let $p^{(T)}$ be the distribution of the iterate of Algorithm 2 after $T$ steps.

If we set $\delta<\frac{1}{8M}$, $n\geq 2$, $\tau\geq\frac{8}{m\delta}$ and run $\mathsf{Option\ I}$ then for all $T$ mod $\tau=0$ we have

If we set $\delta\!<\!\frac{\sqrt{n}}{4\tau M}$ and run $\mathsf{Option\ II}$ for $T$ iterations then,

For Option I, if we study the same regime as before where $M,m$ and $L$ are scaling linearly with $N$ we find that the discretization error is dominated by the term which is of order $\mathcal{O}(\sqrt{\delta Nd/n})$. To achieve target accuracy of $\epsilon$ we would need $\delta=\mathcal{O}(\epsilon^{2}n/Nd)$. This is less impressive than the guarantees of SAGA and essentially we only gain a constant factor as compared to the guarantees for SGLD. This behavior may be explained as follows: at each epoch, a constant decrease of the objective is needed in the classical proof of SVRG when applied to optimization. When the step-size is small, the epoch length is required to be large that washes away the advantages of variance reduction.

For Option II, similar convergence guarantees as SAGA are obtained, but worse by a factor of $\sqrt{n}$. In contrast to SAGA, this result holds only for small step-size, with the constants in Eq. (12) blowing up exponentially quickly for larger step sizes (for more details see proof in Appendix B.1). We also find that experimentally SAGA routinely outperforms SVRG both in terms of run-time and iteration complexity to achieve a desired target accuracy. However, it is not clear whether it is an artifact of our proof techniques that we could not recover matching bounds as SAGA or if SVRG is less suited to work with sampling methods. We now state our results for the convergence guarantees of Algorithm 3.

Let assumptions (A1)-(A3) hold. Let $p^{(T)}$ be the distribution of the iterate of Algorithm 3 after $T$ steps starting with the initial distribution $p^{(0)}(x,v)=1_{x=x^{*}}\cdot 1_{v=0}$. If we set the step size to be $\delta<1/M$ and run Algorithm 3 then we have the guarantee that

We initialize the chain in Algorithm 3 with $x^{*}$, the global minimizer of $f$ as we already need to calculate it to build the gradient estimate. Observe that Theorem 4.3 does not guarantee the error drops to when $\delta\to 0$ but is proportional to the standard deviation of our gradient estimate. This is in contrast to SAGA and SVRG based algorithms where a more involved gradient estimate is used. The advantage however of using this second order method is that we get to a desired error level $\epsilon$ at a faster rate as the step size can be chosen proportional to $\epsilon/\sqrt{d}$, which is $\sqrt{d}$ better than the corresponding results of Theorem 4.1 and 4.2 and without Assumption (A4) (Hessian Lipschitzness).

Here we compare the theoretical guarantees of Langevin MCMC [LD, Durmus and Moulines, 2016], Underdamped Langevin MCMC [ULD, Cheng et al., 2017], SGLD [Dalalyan and Karagulyan, 2017], stochastic gradient underdamped Langevin diffusion [SGULD, Cheng et al., 2017], SAGA-LD (Algorithm 1), SVRG-LD (Algorithm 2 with Option I and II), Control Variate Langevin diffusion [CV-LD, Baker et al., 2017] and Control Variate underdamped Langevin diffusion (CV-ULD, Algorithm 3). We always consider the scenario where $M,m$ and $L$ are scaling linearly with $N$ and where $N\gg d$ (tall-data regime). We note that the memory cost of all these algorithms except SAGA-LD is $\mathcal{O}(nd)$; for SAGA-LD the worst-case memory cost scales as $\mathcal{O}(Nd)$. Next we compare the mixing time ($T$), that is, the number of steps needed to provably have error less than $\epsilon$ measured in $W_{2}$ and the computational complexity, which is the mixing time $T$ times the query complexity per iteration. In the comparison below we focus on the dependence of the mixing time and computational complexity on the dimension $d$, number of samples $N$, condition number $\kappa$, and the target accuracy $\epsilon$. The mini-batch size has no effect on the computational complexity of SGLD, SGULD and SAGA-LD; while for SVRG-LD, CV-LD and CV-ULD the mini-batch size is chosen to optimize the upper bound.

As illustrated in Fig. 1 we see a qualitative difference in behavior of variance reduced algorithms compared to methods like SGLD. In applications like calculating higher order statistics or computing confidence intervals to quantify uncertainty it is imperative to calculate the posterior with very high accuracy. In this regime when the target accuracy $\epsilon<\mathcal{O}(\sqrt{d/N})$, the computational complexity of SGLD starts to grow larger than $\mathcal{O}(N)$ at rate $\mathcal{O}(\sqrt{d}/\epsilon^{2})$ whereas the computational cost of variance reduced methods is lower. For SAGA-LD the computational cost is $\mathcal{O}(N)$ up until when $\epsilon=\mathcal{O}(\sqrt{d}/N)$ after which it grows at a rate $\mathcal{O}(\sqrt{d}/\epsilon)$. CV-ULD also has a computational cost of $\mathcal{O}(N)$ up until the point where $\epsilon=\mathcal{O}(\sqrt{d}/N^{5/6})$ after which it starts to grow as $\mathcal{O}(d^{3/2}/(N^{3/2}\epsilon^{3}))$. When $\mathcal{O}(\sqrt{d}/N^{5/6})\leq\epsilon<\mathcal{O}(\sqrt{d}/\sqrt{N})$ our bounds predict both SAGA-LD and CV-ULD to have comparative performance ($\mathcal{O}(N)$) and in some scenarios one might outperform the other. For higher accuracy our results predict SAGA-LD performs better than CV-ULD. Note that Option II of SVRG performs also well in this regime of small $\epsilon$ but not as well as SAGA-LD or CV-ULD.

At the other end of the spectrum for most classical statistical problems accuracy of $\epsilon=\mathcal{O}(\sqrt{d/N})$ is sufficient and less than a single pass over the data is enough. In this regime when $\epsilon>\mathcal{O}(\sqrt{d/N})$ and we are looking to find a crude solution quickly, our bounds predict that SGLD is the fastest method. Other variance reduction methods need at least a single pass over the data to initialize.

Our sharp theoretical bounds allow us to classify and accurately identify regimes where the different variance reduction algorithms are efficient; bridging the gap between experimentally observed phenomenon and theoretical guarantees of previous works. Also noteworthy is that here we compare the algorithms only in the tall-data regime which grossly simplifies our results in Sec. 4, many other interesting regimes could be considered, for example the fat-data regime where $d\approx N$, but we omit this discussion here.

Experiments

In this section we explore the performance of SG-MCMC with variance reduction via experiments. We compare SAGA-LD, SVRG-LD (with option II), CV-LD, CV-ULD and use SGLD as the baseline method.

We make use of three datasets available at the UCI machine learning repository. The first two datasets describe the connections between heart disease and diabetes with various patient-specific covariates. The third dataset captures the generation of supersymmetric particles and its relationship with the kinematic properties of the underlying process. We use part of the datasets to obtain a mean estimate of the parameters and hold out the rest to test their likelihood under the estimated models. Sizes of the datasets being used in Bayesian estimation are $100$, $600$, and $1e5$, respectively.

Performance is measured by the log probability of the held-out dataset under the trained model. We first find the optimal log held-out probability attainable by all the currently methods being tested. We then target to obtain levels of log held-out probability increasingly closer to the optimal one with each methods. We record number of passes through data that are required for each method to achieve the desired log held-out probability (averaged over $30$ trials) for comparison in Fig. 2. We fix the batch size $n=10$ as constant, to explore whether the overall computational cost for SG-MCMC methods can grow sub-linearly (or even be constant) with the overall size of the dataset $N$. A grid search is performed for the optimal hyperparameters in each algorithm, including an optimal scheduling plan of decreasing stepsizes. For CV-LD, we first use a stochastic gradient descent with SAGA variance reduction method to find the approximate mode $x^{*}$. We then calculate the full data gradient at $x^{*}$ and initialize the sampling algorithm at $x^{*}$.

From the experiments, we recover the three regimes displayed in Fig. 1 with different data size $N$ and accuracy level with error $\epsilon$. When $N$ is large, SGLD performs best for big $\epsilon$. When $N$ is small, CV-LD/ULD is the fastest for relatively big $\epsilon$. When $N$ and $\epsilon$ are both small so that many passes through data are required, SAGA-LD is the most efficient method. It is also clear from Fig. 2 that although CV-LD/ULD methods initially converges fast, there is a non-decreasing error (with the constant mini-batch size) even after the algorithm converges (corresponding to the last term in Eq. (13)). Because CV-LD and CV-ULD both converge fast and have the same non-decreasing error, their performance overlap with each other. Convergence of SVRG-LD is slower than SAGA-LD, because the control variable for the stochastic gradient is only updated every epoch. This attribute combined with the need to compute the full gradient periodically makes it less efficient and costlier than SAGA-LD. We also see that number of passes through the dataset required for SG-MCMC methods (with and without variance reduction) is decreasing with the dataset size $N$. Close observation shows that although the overall computational cost is not constant with growing $N$, it is sublinear.

2 Breaking CLT: Synthetic Log Normal Data

Many works using SG-MCMC assume that the data in the mini-batches follow the central limit theorem (CLT) such that the stochastic gradient noise is Gaussian. But as explained by Bardenet et al. , if the dataset follows a long-tailed distribution, size of the mini-batch needed for CLT to take effect may exceed that of the entire dataset. We study the effects of breaking this CLT assumption on the behaviors of SGLD and its variance reduction variants.

We use synthetic data generated from a log normal distribution: $f_{X}(x)=\dfrac{1}{x}\cdot\dfrac{1}{\sigma\sqrt{2\pi}}\exp\left(-\dfrac{(\ln x-\mu)^{2}}{2\sigma^{2}}\right)$ and sample the parameters $\mu$ and $\sigma$ according to the likelihood $p({\bf x}|\mu,\sigma)=\prod_{i=1}^{N}f_{X}(x_{i})$. It is worth noting that this target distribution not only breaks the CLT for a wide range of mini-batch sizes, but also violates assumptions (A2)-(A4).

To see whether each method can perform well when CLT assumption is greatly violated, we still let mini-batch size to be $10$ and grid search for the optimal hyperparameters for each method. We use mean squared error (MSE) as the convergence criteria and take LD as the baseline method to compare and verify convergence.

From the experimental results, we see that SGLD does not converge to the target distribution. This is because most of the mini-batches only contain data close to the mode of the log normal distribution. Information about the tail is hard to capture with stochastic gradient. It can be seen that SAGA-LD and SVRG-LD are performing well because history information is recorded in the gradient so that data in the tail distribution is accounted for. As in the previous experiments, CV-LD converges fastest at first, but retains a finite error. For LD, it converges to the same accuracy as SAGA-LD and SVRG-LD after $10^{4}$ number of passes through data. The variance reduction methods which uses long term memory may be especially suited to this scenario, where data in the mini-batches violates the CLT assumption.

It is also worth noting that the computation complexity for this problem is higher than our previous experiments. Number of passes through the entire dataset is on the order of $10^{2}\sim 10^{3}$ to reach convergence even for SAGA-LD and SVRG-LD. It would be interesting to see whether non-uniform subsampling of the dataset Schmidt et al. can accelerate the convergence of SG-MCMC even more.

Conclusions

In this paper, we derived new theoretical results for variance-reduced stochastic gradient MC. Our theory allows us to accurately classify two major regimes. When a low-accuracy solution is desired and less than one pass on the data is sufficient, SGLD should be preferred. When high accuracy is needed, variance-reduced methods are much more powerful. There are a number of further directions worth pursuing. It would be of interest to connect sampling with advances in finite-sum optimization. specifically advances in accelerated gradient [Lin et al., 2015] or single-pass methods [Lei and Jordan, 2017]. Finally the development of a theory of lower bounds for sampling will be an essential counterpart to this work.

References

Organization of the Appendix

In Appendix A we formally define the Wasserstein distance. In Appendix B we introduce the notations required to prove Theorems 4.1 and 4.2. In Appendix B.1 we prove Theorem 4.2 and then in Appendix B.2 we prove Theorem 4.1. Finally in Appendix C we prove Theorem 4.3.

Appendix A Wasserstein Distance

We define the Wasserstein distance of order two between a pair of probability measures as follows:

Finally we denote by $\Gamma_{opt}(\mu,\nu)$ the set of transference plans that achieve the infimum in the definition of the Wasserstein distance between $\mu$ and $\nu$ [see, e.g., Villani, 2008, for more properties of $W_{2}(\cdot,\cdot)$].

In this section we will prove Theorem 4.1 and Theorem 4.2 and include details about Algorithms 1 and 2 that were omitted in our discussion in the main paper. Throughout this section we assume that assumptions (A1)-(A4) holds. First we define the continuous time (overdamped) Langevin diffusion process defined by the Itô SDE:

Throughout this section we will denote by $\{x_{k}\}_{k=1}^{T}$ the iterates of Algorithm 1 or Algorithm 2. Also we will define the distribution of the $k^{th}$ iterate of Algorithm 1 or Algorithm 2 by $p^{(k)}$. With this notation in place we are now ready to present the proofs of Theorem 4.2 and Theorem 4.1.

In both the proof of Theorem 4.1 and Theorem 4.2 we draw from and sharpen techniques established in the literature of analyzing Langevin MCMC methods and variance reduction techniques in optimization. In both the proofs we use Lyapunov functions that are standard in the optimization literature for analyzing these methods; we use them to define Wasserstein distances and adapt methods from analysis of sampling algorithms to proceed.

B.1 Stochastic Variance Reduced Gradient Langevin Monte Carlo

In the proof of SVRG for Langevin diffusion, it is common to consider the Lyapunov function to the standard 2-norm. We define a Wasserstein distance with respect to distance – $\lVert x_{k}-y_{k}\rVert_{2}^{2}$.

where the Brownian motion is independent of $(x_{k},y_{k},S)$. Thus integrating the above SDE we get upto time $\delta$ (the step-size),

where $\xi_{k}\sim N(0,I_{d\times d})$. Note that since $y_{k}\sim p^{*}$, we have that $y_{k+1}\sim p^{*}$. Similarly we also have that the next iterate $x_{k+1}$ is given by

where $\xi_{k}$ is the same normally distributed random variable as in (17). Let us define $\Delta_{k}:=y_{k}-x_{k}$ and $\Delta_{k+1}:=y_{k+1}-x_{k+1}$. Also define

Now that we have the notation setup, we will prove the first part of this Theorem. We procede in 5 steps. In Step 1 we will express $\lVert\Delta_{k+1}\rVert_{2}^{2}$ in terms of $\lVert\Delta_{k+1}+V_{k}\rVert_{2}^{2}$ and $\lVert V_{k}\rVert_{2}^{2}$, in Step 2 we will control the expected value of $\lVert V_{k}\rVert_{2}^{2}$. In Step 3 we will express $\Delta_{k+1}$ in terms of $\Delta_{k}$ and other terms, while in Step 4 we will use the characterization of $\Delta_{k+1}$ in terms of $\Delta_{k}$ combined with the techniques established by Dubey et al. to bound the expected value of $\lVert\Delta_{k+1}+V_{k}\rVert_{2}^{2}$. Finally in Step 5 we will put this all together and establish our result. First we prove the result for Algorithm 2 run with Option I.

Step 1: By Young’s inequality we have that $\forall a>0$,

We will choose $a$ at a later stage in the proof to minimize the bound on the right hand side.

Step 2: By Lemma 6 of Dalalyan and Karagulyan we have the bound,

Step 3: Next we will bound the other term in (19), $\lVert\Delta_{k+1}+V_{k}\rVert_{2}^{2}$. First we express $\Delta_{k+1}$ in terms of $\Delta_{k}$,

Step 4: Using the above characterization of $\Delta_{k+1}$ in terms of $\Delta_{k}$ established above, we get

where in the second equality we used the definition of $U_{k}$ and $\zeta_{k}$, while in the third equality we used the definition of $g_{k}$. By Young’s inequality we now have,

Now we bound each of the 4 terms. First for $\omega_{1}$ by $M$-smoothness of $f$ we get,

By the bounds we have established on $\omega_{1},\omega_{2},\omega_{3}$ and $\omega_{4}$ we get that $\Omega_{2}$ is upper bounded by,

Next we control $\Omega_{1}$ by using the convexity of $f$,

Define the contraction rate to be $\alpha:=\delta m\left(1-2\delta M\left(1+\frac{2}{n}\right)\right)$, then we have

and with $a=\alpha/(1-\alpha)$ then we have

We now sum this inequality from $k=u\tau$ to $(u+1)\tau-1$ we get,

Using the strong convexity and smoothness of $f$ we have,

Using this in (29) and rearranging terms we get,

for any $u\in\{0,\ldots,\left\lfloor T/\tau\right\rfloor\}$. Unrolling the equation above for $\left\lfloor T/\tau\right\rfloor$ steps we get,

where we denote by $\rho=1-\frac{1-\alpha}{\alpha\tau}-\frac{4mM\delta^{2}}{\alpha n}$. Finally we use again the strong convexity and smoothness of $f$ to obtain

Now for $n>2$, and $\delta<\frac{1}{8M}$ we have $\alpha\geq\delta m/2$, and $\frac{4mM\delta^{2}}{\alpha n}\leq\frac{8M\delta}{n}\leq 1/n\leq 1/4$. Then consider $\tau=4/\alpha$, in order to have $\rho=1/2$.

Finally our choice of $n>2$, and $\delta<\frac{1}{8M}$ and $\tau=4/\alpha$ ensures that

for all $T>0$ such that $T$ mod $\tau=0$, which completes the proof of part 1.

Proof for Option 2: To prove part 2 of the theorem Steps 1-3 are same as above. The technique to control $\Omega_{2}$ is going to differ which leads to a different bound.

where the last step is by Young’s inequality. First we claim a bound on $\omega_{3}$,

by the same argument as used in (B.1). Next we control $\omega_{1}$ by

by the $M$-smoothness of $f$. Finally we control $\omega_{2}$,

where the first inequality follows by Jensen’s inequality. Further we have,

by using discrete Grönwall lemma [see, e.g., Clark, 1987] we get,

Combined with the bound on $\omega_{1}$ and $\omega_{3}$ this yields a bound on $\Omega_{2}$ which is

As before, using strong convexity of $f$, we get that $\Omega_{1}$ is bounded by

Having established these bounds on $\Omega_{1}$ and $\Omega_{2}$ we get,

Define the contraction rate $\alpha:=\delta m\left(1-\frac{3M\square\delta}{2}\right)$. Further we choose $a=\alpha/(1-2\alpha)$, then we get,

Taking the global expectation, we obtain by a direct expansion:

Let us use the fact that $\tau<\sqrt{n}/(2\tau M)$ then we have $\exp(4\tau^{2}M^{2}\delta^{2}/n)<3$, $\square\leq 4/3$ and $\alpha\geq\delta m(1-2M\delta)$. We assume also that $\delta\leq 1/(4M)$ then $\alpha\geq\delta m/2$ and

Using that $\triangle\leq M^{2}d\left(1+\frac{9\tau^{2}}{n}\right)$ and $\circ\leq L^{2}d^{2}+M^{3}d$ we finally obtain

B.2 SAGA Proof

We will proceed as in the proof of Theorem 4.2 and borrow notation established in the proof of Theorem 4.2. For $k\in\{1,\ldots,T\}$, we denote by $\{h_{k}^{i}\}_{i=0}^{N}$ analogously to $\{g_{k}^{i}\}_{i=0}^{N}$ (updated at the same point in the sequence $\{y_{k}\}$). We consider the Lyapunov function $T_{k}=c\sum_{i}^{N}\|g_{k}^{i}-h_{k}^{i}\|_{2}^{2}+\|x_{k}-y_{k}\|_{2}^{2}$ for some constant $c>0$ to that will be chosen later. The first 4 steps of the proof are exactly the same as the proof above, in Step 5 we shall control the norm of the other part of the Lyapunov function involving $g_{k}^{i}$ and $h_{k}^{i}$. In the rest of the steps we gather these bounds on the different parts of the Lyapunov function and establish convergence.

Step 1: By Young’s inequality we have that $\forall a>0$,

We will choose $a$ at a later stage in the proof to minimize the bound on the right hand side.

Step 2: By Lemma 6 in [Dalalyan and Karagulyan, 2017] we have the bound,

Step 3: Next we will bound the other term in (31). First we express $\Delta_{k+1}$ in terms of $\Delta_{k}$,

Step 4: Using the above characterization of $\Delta_{k+1}$ in terms of $\Delta_{k}$, we now get

Now we take expectation with respect to all sources of randomness conditioned (Brownian motion and the randomness in the choice of $S$) on $x_{k}$ and $y_{0}$ (thus $\Delta_{k}$ is fixed). Recall that conditioned on $\Delta_{k}$, $\Psi_{k}$ and $\zeta_{k}$ are zero mean, thus we get

where in the second equality we used the definition of $U_{k}$ and $\zeta_{k}$, while in the third equality we used the definition of $g(x_{k})$. By Young’s inequality we now have,

Let us define the random variable $\beta^{(i)}=\nabla f_{i}(x_{k})-g_{k}^{i}-\frac{n}{N}\left(\nabla f(x_{k})+\sum_{i=1}^{n}g_{k}^{i}\right)$. Observing that $\{\beta^{(i)}\}$ are zero mean (taking expectation over $i$) and independent, we have

Using the following decomposition $\lVert\nabla f_{i}(x_{k})-g_{k}^{i}\rVert_{2}^{2}\leq 3\lVert\nabla f_{i}(x_{k})-\nabla f_{i}(y_{k})\rVert_{2}^{2}+3\lVert\nabla f_{i}(y_{k})-h_{k}^{i}\rVert_{2}^{2}+3\lVert h_{k}^{i}-g_{k}^{i}\rVert_{2}^{2}$, we have shown $\Omega_{2}$ may be bounded as follow

Step 5: We will now bound the different term in Equation 34. First, using a similar technique as Dubey et al. , we bound the term $\|h_{k}^{i}-\nabla f_{i}(y_{k})\|^{2}_{2}$. Let $p=1-(1-1/N)^{n}$ be the probability to chose an index, then

where we have use lemma and the bound on $p\geq\frac{n}{2N}$ from Eq. (19) of Dubey et al. . Therefore

Step 6: We can combine now the previous bound to first obtain

where we have denoted by $\square=2M^{2}d+\frac{72dNM^{2}}{n^{2}}\left[\frac{\delta MN}{n}+1\right]$ and $\triangle=\frac{1}{2}\left(L^{2}d^{2}+M^{3}d\right)$.

Step 7: We expand now the first part of $T_{k}$ We directly obtain:

Step 8: We are now able to upper-bound $T_{k}$:

With the strong-convexity of $f$ we obtain

Step 9: We will now fix the different values of $c$ and $a$ in order to obtain the final recursive bound on $T_{k}$. With $c=\frac{12(1+a)\delta^{2}N^{2}}{n^{2}}$ we obtain

Assume now that $a=\frac{\alpha}{2(1-\alpha)}$ where $\alpha=2m\delta\left(1-{3}(1+\frac{9}{n})\delta M\right)$

Let us assume that $n>9$ and $\delta<1/12M$, then

where we denote by $\rho=\min\{\frac{n}{3N},m\delta/2\}$.

Step 10: We are able now to solve this recursion to obtain an upper bound on $T_{k}$. We obtain a recursive argument,

We have that $\frac{1}{\rho}=\max\{\frac{3N}{n},\frac{2}{m\delta}\}\leq\frac{3N}{n}+\frac{2}{m\delta}$:

Using the fact that $x_{k}$ and $y_{k}$ are optimally coupled, the results follows. ∎

Appendix C Control Variates with Underdamped Langevin MCMC

In this section we will prove Theorem 4.3 and also include details regarding Algorithm 3 that was omitted in Section 3. Throughout this section we will assume that assumptions (A1)-(A3) holds. Crucially in this section we will not assume the Hessian of $f$ to be Lipschitz ((A4)).

Underdamped Langevin Markov Chain Monte Carlo [see, e.g., Cheng et al., 2017] is a sampling algorithm which can be viewed as discretized dynamics of the following Itô stochastic differential equation (SDE):

With these definitions and notation in place we can first show that $p_{t}$ contracts exponentially quickly to $p^{*}$ measured in $W_{2}$.

Let $p_{0}$ be a distribution with $(x_{0},v_{0})\sim p_{0}$. Let $q_{0}$ and $q_{t}$ be the distributions of $(x_{0},x_{0}+v_{0})$ and $(x_{t},x_{t}+v_{t})$, respectively (i.e., the images of $p_{0}$ and $p_{t}$ under the map $g(x,v)=(x,x+v)$). Then

The next lemma establishes a relation between the Wasserstein distance between $q$ to $q^{*}$ and between $p$ to $p^{*}$.

The triangle inequality for the Euclidean norm implies that

Thus we also get convergence of $p_{t}$ to $p^{*}$:

We will now present a discretization of the dynamics in 35. A natural discretization to consider is defined by the SDE,

Corollary C.1 (contraction of the continuous time-process) coupled with the result presented as Theorem C.4 (discretization error bound; stated and proved in Appendix C.1) will now help us prove Theorem 4.3.

For any random variable $(x,v)$ with distribution $p$, let $q$ denote the distribution of the random variables $(x,x+v)$. From Corollary C.1, we have that for any $i\in\{1,\ldots,T\}$

By the discretization error bound in Theorem C.4 and Lemma C.2, we get

Let us define $\eta=e^{-m\delta/2}$. Then by applying (44) $T$ times we have:

where the second step follows by summing the geometric series and by applying the upper bound Lemma C.2. By another application of 37 we get:

In Lemma C.7 we establish a bound on $W_{2}^{2}(p^{(0)},p^{*})\leq 2d/m$. This motivates our choice of $T>\frac{1}{m\delta}\log\left(\frac{12d}{\epsilon m}\right)$, $\delta=\frac{\epsilon}{M^{2}}\sqrt{5m^{3}/(425984\cdot d)}$ and $n=(249M)^{2}d/(m^{3}\epsilon^{2})$ which establishes our claim. ∎

C.1 Discretization Error Analysis

In this section, we will repeatedly use the following inequality:

which follows from Jensen’s inequality using the convexity of $\|\cdot\|_{2}^{2}$.

This completes the bound for the velocity variable. Next we bound the discretization error in the position variable:

where the first line is by coupling through the initial distribution $p_{0}$, the second line is by Jensen’s inequality and the third inequality uses the preceding bound. Setting $s=M\delta$ and by our choice of $u=1/M$ we have that the squared Wasserstein distance is bounded as

Given our assumption that $\delta$ is chosen to be smaller than $1/M$, this gives the upper bound:

Taking square roots establishes the desired result. ∎

C.2 Auxiliary Results

In this section, first we establish an explicit bound on the kinetic energy $\mathcal{E}_{K}$ in (46) which is used to control the discretization error at each step.

Let $p^{(0)}(x,v)=1_{x=x^{*}}\cdot 1_{v=0}$— the Dirac delta distribution at $(x^{*},0)$. Further let $p^{(i)}$ be defined as in (42) for $i=1,\ldots T$, with step size $\delta$ and number of iterations $T$ as specified in Theorem 4.3. Then for all $i=1,\ldots T$ and for all $t\in[0,\delta]$, we have the bound

We first establish an inequality that provides an upper bound on the kinetic energy for any distribution $p$. Step 1: Let $p$ be any distribution over $(x,v)$, and let $q$ be the corresponding distribution over $(x,x+v)$. Let $(x^{\prime},v^{\prime})$ be random variables with distribution $p^{*}$. Further let $\zeta\in\Gamma_{opt}(p,p^{*})$ such that,

where for the second and the third inequality we have used Young’s inequality, while the final line follows by optimality of $\zeta$.

Step 3: For our initial distribution $p_{0}(q_{0})$ we have the bound

Putting all this together along with (48) we have

Step 4: By Corollary C.1, we know that $\forall t>0$,

for all $t>0$ and $i\in\{0,1,\ldots,T\}$. ∎

Let $p^{(0)}(x,v)=1_{x=x^{*}}\cdot 1_{v=0}$— the Dirac delta distribution at $(x^{*},0)$. Further let $p^{(i)}$ be defined as in (42) for $i=1,\ldots T$, with step size $\delta$ and number of iterations $T$ as specified in Theorem 4.3. Then for all $i=1,\ldots T$ and for all $t\in[0,\delta]$, we have the bound

We first establish an inequality that provides an upper bound on the kinetic energy for any distribution $p$. Step 1: Let $p$ be any distribution over $(x,v)$, and let $q$ be the corresponding distribution over $(x,x+v)$. Let $(x^{\prime},v^{\prime})$ be random variables with distribution $p^{*}$. Further let $\zeta\in\Gamma_{opt}(p,p^{*})$ such that,

where for the second and the third inequality we have used Young’s inequality, while the final line follows by optimality of $\zeta$.

Step 3: For our initial distribution $p_{0}(q_{0})$ we have the bound

where the first inequality is an application of Young’s inequality, the equality in the second line follows as $p^{*}(v)\propto\exp(-M\lVert v\rVert_{2}^{2}/2)$ and the second inequality follows by again applying the bound from Theorem D.1. Combining these we have the bound, Putting all this together along with (49) we have

Step 4: By Corollary C.1, we know that $\forall t>0$,

Next we prove that the distance of the initial distribution $p^{(0)}$ to the optimum distribution $p^{*}$ is bounded.

Let $p^{(0)}(x,v)=1_{x=x^{*}}\cdot 1_{v=0}$— the Dirac delta distribution at $(x^{*},0)$. Then

As $p^{(0)}(x,v)$ is a delta distribution, there is only one valid coupling between $p^{(0)}$ and $p^{*}$. Thus we have

Next we calculate integral representations of the solutions to the continuous-time process (35) and the discrete-time process (38).

The solution $(x_{t},v_{t})$ to the underdamped Langevin diffusion (35) is

Appendix D Technical Results

We state this Theorem from Durmus and Moulines used in the proof of Lemma C.5.