Accelerating Convergence of Replica Exchange Stochastic Gradient MCMC via Variance Reduction
Authors: Wei Deng
Qi Feng
Georgios P. Karagiannis
Guang Lin
Faming Liang
Venue: ICLR 2021 Poster (2021)
Keywords: variance reduction
replica exchange
parallel tempering
stochastic gradient Langevin dynamics
uncertainty quantification
change of measure
generalized Girsanov theorem
Dirichlet form
Markov jump process

Abstract: Replica exchange stochastic gradient Langevin dynamics (reSGLD) has shown promise in accelerating the convergence in non-convex learning; however, an excessively large correction for avoiding biases from noisy energy estimators has limited the potential of the acceleration. To address this issue, we study the variance reduction for noisy energy estimators, which promotes much more effective swaps. Theoretically, we provide a non-asymptotic analysis on the exponential convergence for the underlying continuous-time Markov jump process; moreover, we consider a generalized Girsanov theorem which includes the change of Poisson measure to overcome the crude discretization based on the Gr\"{o}wall's inequality and yields a much tighter error in the 2-Wasserstein ($\mathcal{W}_2$) distance. Numerically, we conduct extensive experiments and obtain state-of-the-art results in optimization and uncertainty estimates for synthetic experiments and image data.

Published as a conference paper at ICLR 2021
ACCELERATING CONVERGENCE OF REPLICA EX- CHANGE STOCHASTIC GRADIENT MCMC VIA VARI- ANCE REDUCTION
Wei Deng
Qi Feng
Department of Mathematics Purdue University
Department of Mathematics
University of Southern Californa Los Angeles, CA, USA qi F@usc.edu
West Lafayette, IN, USA
we weideng05 gma com
Georgios Karagiannis Department of Mathematical Sciences
Guang Lin
Departments of Mathematics & School of Mechanical Engineering Purdue University
Durham University Durham, UK
os kas.kaaiaamnisebrrhhen
West Lafayette, IN, USA guanglin@purdue.
Faming Liang
Departments of Statistics Purdue University West Lafayette, IN, USA fml imllang@purdue.een edu
ABSTRACT
Replica exchange stochastic gradient Langevin dynamics (reSGLD) has shown promise in accelerating the convergence in non-convex learning: however, cessively large correction for avoiding biases from noisy energy estimators has limited the potential of the acceleration. To address this issue, study the vari- ance reduction for noisy energy estimators, which promotes much more effective swaps. Theoretically, convergence for the underlying continuous-time Markov jump process; moreover, we consider generalized Girsanov theorem which includes the change of Poisson measure to overcome the crude discretization based on the Grönwall's inequality and yields a much tighter eтrоr in the 2- Wasserstein (W2) distance. Numerically, we conduct extensive experiments and obtain state-of-the-arr results in optimiza- tion and uncertainty estimates for synthetic experiments and image data.
provide non-asymptotic analysis on the exponential
INTRODUCTION
Stochastic gradient Monte Carlo methods (Welling Teh 2011 Chen et al. 2014 et 2016
are the golden standard for Bayesian inference in deep learning due to their theoretical guarantees uncertainty quantification (Vollmer et al.. 2016 Chen al. 2015) and non-convex optimization (Zhang et al.) 2017). However, despite their scalability with respect to the data size, their mixing rates are often extremely slow for complex deep neural networks with rugged energy landscapes (Li ct al. 2018). To speed up the convergence, several techniques have been proposed in the literature in order to accelerate their exploration of multiple modes on the energy landscape, for example,
dynamic temperatures Ye al. 2017) and cyclic learning rates (Zhang al. 2020). name few. However, such strategies only explore contiguously a limited region around a few informative modes. Inspired by the successes of replica exchange, also known as parallel tempering, in tradi- tional Monte Carlo methods (Swendsen & Wang 1986 Earl & Deem 2005). reSGLD (Deng et al.
"Equal contribution

---

Published as a conference paper at ICLR 2021
2020) uses multiple processes based on stochastic gradient Langevin dynamics (SGLD) where inter- actions between different SGLD chains are conducted a manner that encourages large jumps. In addition to the ideal utilization of parallel computation, the resulting process is able to jump to more informative modes for more robust uncertainty quantification. However, the noisy energy estimators in mini-hatch settings lead large bias in the na've swaps, and a large correction is required to reduce the bias, which yields few effective swaps and insignificant accelerations. Therefore, how to reduce the variance of noisy energy estimators becomes essential in speeding up the convergence.
A long standing technique for variance reduction is the control variates method. The key to reduc- ing the variance is to properly design correlated control variates so as to counteract some noise. Towards this direction, Dubey et al. (2016); Xu et al. (2018) proposed to update the control variate periodically for the stochastic gradient estimators and Baker al. (219)) studied the construction of control variates using local modes. Despite the advantages in near-convex problems, natural discrepancy between theory (Chatterji et al. 2018 Xu et al., 2018 Zou al., 2019b) and practice should avoid the gradient noise in non-convex
(He et al. 2016 Devlin al., 2019) is whether problems. To fill in the gap, we only focus on the variance reduction of noisy energy estimators to exploit the theoretical accelerations but no longer consider the variance reduction of the noisy gra- dients so that the empirical experience from stochastic gradient descents with momentum (M-SGD) can be naturally imported.
In this paper we propose the variance-reduced replica exchange stochastic gradient Langevin dynam- ics (VR-reSGLD) algorithm to accelerate convergence by reducing the variance of the noisy energy estimators. This algorithm not only shows the potential of exponential acceleration via much more effective swaps in the non-asymptotic analysis but also demonstrates remarkable performance in practical tasks where a limited time is required; while others (Xu et al.] 20188) ZZou et al.) 2019a) may only work well when the dynamics sufficiently mixed and the discretization error becomes a major component. Moreover, the existing discretization error of the Langevin- sangevin-based Markov jump processes (Chen al. 2019 Deng et al. 2020 Futami et al. 2020) exponentially dependent on time due to the limitation of Grönwall's inequality. To avoid such a crude estimate, we consider the generalized Girsanov theorem and a change of Poisson measure. As a result, we obtain a much tighter discretization error only polynomially dependent on time. Empirically, we test the algorithm through extensive experiments and achieve trat-off-the-a performance in both optimization and uncertainty estimates.
0.2
VR -LGGGL
Swap
Explotation
(a) Gibbs measures at three temperatures T.
(b) Sample trajectories on energy landscape.
(c) Faster exponential convergence in Wu
Figure 1: An illustration of replica exchange Monte Carlo algorithms for non-convex learning.
2 PRELIMINARIES
A common problem, in Bayesian inference, is the simulation from a posterior P(3]X) oc P(3) IT P(x; ( where P(θ) is proper prior, IT, P(x.(33 is the likelihood function and N the number of data points. When is large, the standard Langevin dynamics is too costly in evaluating the gradients. To tackle this issue, stochastic gradient Langevin dynamics (SGLD)
(Welling & Teh 201T) was proposed to make the algorithm scalable by approximating the gradient through a mini-batch data B of size n such that
N
Ba = Bk-1 1 T7K
Σ
VL(x((Bu-1) + 22TTTTER
(1)
n
2

---

Published as a conference paper at ICLR 2021
where 3y I denotes the temperature, 7] is the learning rate at iteration k, Ex is standard Gaussian vector, and L(-) == log P(3 is the energy function. SGLD known to converge weakly to stationary Gibbs measure T ( exp (xp(-((3)/T) 7h decays to 0 (Teh al. 2016).
The temperature is the key to accelerating the computations multi-modal distributions. On the one hand, a high temperature flattens the Gibbs distribution exp (-I(3)/T) (see the red curve in aa) and accelerates mixing by facilitating exploration of the whole domain, but the resulting distribution becomes much less concentrated around the global optima. On the other hand, a low temperature exploits the local region rapidly; however, it may cause the particles to stick in local region for an exponentially long time, as shown in the blue curve in Fig[T..)). To bridge the gup between global exploration and local exploitation, Deng et al. (2020) proposed the replica exchange SGLD algorithm (reSGID), which consists of low-temperature SGLD to encourage exploitation and hig--temperature SGLD to support exploration
( k = (1))
N
Σ
vL(x;1( + ( (1)
k-1
n
E
/ = k-1
(2)
N
VL(x;13 k-1 + v 22) (2)
where the invariant measure is known to be T x exp
[(a
L(a T(1)
र(अ)
as T/ +
and (2) < 72)), Moreover, the two processes may swap the positions to allow tunneling between different modes. To avoid inducing large bias in mini-batch settings, a corrected swapping rate $ is developed such that
3 = exp ( (HD ((2)
1 7(2)
N
Σ B
1(x133
N T
(xx)
1 F()
1 7(2)
0²
F
B
where G² is an estimator of the variance of N ΣiEB, ((x,13))) NN EiEB, (xx and F is the correction factor to balance between acceleration and bias. In other words, the parameters switch the positions from (.) ( to (3( ((( with a probability r(1 A 3) where the constant v is the swapping intensity and can set to v/a
for simplicity.
From a probabilistic point of view, reSGLD is discretization scheme of replica exchange Langevin diffusion (reLD) in mini-batch settings. Given a smooth test function f and swapping- rate function S, the infinitesimal generator Ls associated with the continuous-ting reLD follows
(sf(8 (0))) oV f( (())), (L(B"))) (V f ,(())) VL(3 )+ (2)) aren ((3"), ()) + ( (1)(()_)),101) (f(3()) ( - "(
A f ,
where the last term arises from swaps and Ag is the the Laplace operator with respect (0)). Note that the infinitesimal generator closely related to Dirichlet forms in characterizing the evolution of stochastic process. By standard caleulations in Markov semigroups (Chen al. 2019). the Dirichlet form &g associated with the infinitesimal generator Lg follows
"1)00,10))) Es(f) / ( (100) |V gan f( 11 + + ( gan oin (() ,g()))) dr
vanilla term E(f)
(2)
s(1 ,g22) -1133 (0(38)80)) - ( dn (3())) 9(
acceleration lerm
which leads to a strictly positive acceleration under mild conditions and is crucial for the expo- nentially accelerated convergence in the W2 distance (see [i).. However, the acceleration depends on the swapping-rate function and becomes much smaller given noisy estimate of N Eien [(x:33) due to the demand of large corrections to reduce the bias
3

---

Published as a conference paper at ICLR 2021
3 VARIANCE REDUCTION IN REPLICA EXCHANGE STOCHASTIC GRADIENT LANGEVIN DYNAMICS
The desire to obtain more effective swaps and larger accelerations drives us to design more efficient energy estimators. A naive idea would be to apply large batch size n, which reduces the variance of the noisy energy estimator proportionally. However, this comes with a significantly increased memory overhead and computations and therefore is inappropriate for big data problems.
natural idea to propose more effective swaps is to reduce the variance of the noisy energy estimator L(B] = N' Eien ((x,)3()) for h E (1,2)- Considering an unbiased estimator ((B)
for ΣN ((x,) and a constant c, we see that new estimator (B1,3())) which follows
i(B|g()) L(B) +o L(BBBU)) -
(3)
is still the unbiased estimator for Ei= Lxx]g(W) By decomposing the variance, have Var(T((B(a()) = Var (I( + Var (I(B)B3))) + 2cCov (T((B)B(h), (((13)))))
In such a case, Var ( achieves the minimum variance (1 (1-))))r(L(BBB(())) given CovLL(B ((1)))())..)))))))) L(BB)_))))))) where Cov(;,) denotes the covariance and p is the correlation (ac(L(B)@(hr)
coefficient of L(B) ((h)) and I(B] To propose correlated control variate, we follow Johnson & Zhang (013) and update 3(h) = (/hh) m every m iterations. Moreover, the optimal c* is often unknown practice. To handle this issue, well-known solution (Johnson & Zhang 2013) is to fix c -1 given a high correlation |o| of the estimators and then we can present the VR-reSGLD algorithm in Algorithm [ Since the exact variance for conrecting the stochastic swapping rate is unknown and even time-varying, we follow Deng et al. (2020) and propose to use stochastic approximation (Robbins & Monro 1951) to adaptively update the unknown variance.
Variants of VR-reSGLD The number of iterations m to update the control variate g(h) gives rise to a trade-off in computations and variance reduction. A small m introduces a highly correlated control variate at the cost of expensive computations; large m, however, may yield less correlated control variate and setting e = fails to reduce the variance. spirit the adaptive variance in Deng et al. 2020) to estimate the unknown variance, we explore the idea of the adaptive coefficient
C = (1 Th c Ck-m + Ty&C such that the unknown optimal c* is well approximated. We present the adaptive VR-reSGLD in Algorithm 2 in Appendix E.22 and show empirically later that the adaptive VR-reSGLD leads to a significant improvement over VR-reSGL.D for the less correlated estimators.
A parallel line of research is to exploit the SAGA algorithm (Defazio et al. 2014) in the study of variance reduction. Despite the most effective performance in variance reduction (Chatterji et al.] 2018), the SAGA type of sampling algorithms require an excessively memory storage of O(Na), which is too costly for big data problems. Therefore, we leave the study of the lightweight SAGA algorithm inspired by Harikandeh et al. (2015): Zhou et al. (209) for future works.
Related work Although our VR-reSGLD is, in spirit, similar to VR-SGLD (Dubey et al. 2016 Xu et al., 2018). it differs from VR-SGLD in two aspects: First, VR-SGLD conducts variance re- duction on the gradient and only shows promises in the nearly log-concave distributions or when the Markov process is sufficiently converged; however, our VR-reSGLD solely focuses on the variance reduction of the energy estimator to propose more effective swaps, and therefore we can import the empirical experience in hyper-parameter tuning from M-SGD to our proposed algorithm. Second, VR-SGLD doesn't I accelerate the continuous-time Markov process but only focuses on reducing the discretization error; VR-reSGLD possesses a larger acceleration term in the Dirichlet form (2) and shows a potential in exponentially speeding up the convergence of the continuous-ting process in the early stage, in addition to the improvement the discretization error. In other words, our al- gorithm is not only theoretically sound but also more empirically appealing for a wide variety of problems in non-convex learning.
4

---

Published as a conference paper at ICLR 2021
Algorithm 1 Variance-reduced replica exchange stochastic gradient Langevin dynamics (VR- reSGLD). The learning rate and temperature can be set to dynamic to speed up the computations. A larger smoothing factor captures the trend better but becomes less robust. T is the thinning factor to avoid a cumbersome system.
The initial parameters (6 and ( learning rate 17, temperatures 7(1) and 7(2), correction factor F and smoothing factor 7.. repeat Parallel sampling Randomly pick a mini-batch set By of size n.
Input
N 17 n ieB
E (L(xx
+ ( for h (1,2).
(4)
Variance-reduced energy estimators Update T,(N) Em Zimml
(x) mL
every iterations.
L(B
Σ [(xx - (
+ Ih) for {1,2}. (5)
BB
if' mod m then Update σं (1 - - rat. where of is an estimate for Var (L(BalB))) - L(Bulss)))
end if
Bias-reduced swaps Swap ( and 3 if < Smm.n. where Unif [0, 1]. and Smm,, follows
Spma = exp (rdss ( - I(Bu
(I(Bu1118
13 (+11) F (+d p(a) 0
(6)
until = kmaxx
Output: The low-temperature process (B Jim1
[Mmax/]]
where T is the thinning factor.
4 THEORETICAL PROPERTIES
The large variance of noisy energy estimators directly limits the potential of the acceleration and significantly slows down the convergence compared to the replica exchange Langevin dynamics. As a result, VR-reSGLD may lead to a more efficient energy estimator with a much smaller variance.
Lemma (Variance-reduced energy estimator) Under the smoothness and dissipativity assump- tions I and in Appendix A the variance of the variance- reduucdd energy estimator L(B) ((h)), where h E {I, 2}. is upper bounded by
Var (T(BIB(A)) min {o
m²
Var
N
MI (xxxB(())
N + Var
Exxx| ((M)))(1)))
VEB
iEB
where the detailed O-) constants is shown in Lemma
Bl in the appendix.
The analysis shows the variance-reduced estimator (h)) yields much-reduced variance given a smaller learning rate and a smaller m for updating control variates based on the batch size n. Although the truncated swapping rate Sy,m, = min(1, (1,mmmm..)) still satisfies the "stochastie" detailed balance given unbiased swapping-rate estimator S,,,, (Deng et al. 2020 it doesn't 't mean the efficiency of the swaps is not affected. By contrast, we can show that the number swaps may become exponentially smaller average.
Lemma (Variance reduction for larger swapping rates) Given a large enough batch size n, the (1) yields truncated swapping rate that satisfies
variance-reduced energy estimator ((Br)B()))
m. R min (1, ((B ((()) (0 (
: Quiroz et al. 2019 achieve similar result hased on the unbiased likelihood estimator for the Metropolis - algorithm. See section 3.1 QQiroz et al. 2019) for details.
Andrieu & Roberts 2009
5

---

Published as a conference paper at ICLR 2021
where S(B(), ((2 is the deterministic swapping rate defined in Appendix B The proof is shown in Lemma B2 in Appendix B Note that the above lemma doesn't 't require the normality assump- tion. As goes to infinity, where the asymptotic normality holds, the RHS of 7 changes to
min 1, 1,5() g())e of P(=")] }. which becomes exponentially larger as we use a smaller up-
date frequeney m and learning rate 17. Since the continuous- time reLD induces a jump operator in the infinitesimal generator, the resulting Dirichlet form potentially leads to much larger accel- eration term which linearly depends on the swapping rate Sp,m,, and yields faster exponential convergence. Now we are ready to present the first main result.
Theorem (Exponential convergence) Under the smoothness and dissipativity assumptions Lnd 2 the probability measure associated with reLD at time t, denoted as Vt. converges exponentially Jast to the invariant measure
Wa(ve, S Do exp {-t (1 +8 8s,. /css),
where Do is a constant depending the initialization, 85n,m.. inf Egg
1 0
E(v un
depends on m.m.n Esmmm.. and & are the Dirichlet forms based on the swapping rale Spmmn and are defined in (2) Cts is the constant of the log- tog-Sobolev inequality for rel.D without swaps.
We detail the proof in Theorem T in Appendix B Note that Sp.m.n = 0 leads to the same perfor-
mance as the standard Langevin diffusion and 8s..m.. is strictly positive when du is asymmetric (Chen et al. 2019) given a smaller 7 and m or a large n, the variance becomes much reduced according to Lemma yielding much larger truncated swapping rate by Lemma 2 and faster exponential convergence to the invariant measure и compared to reSGI.D
dr
Next, we estimate the upper bound of the 2- Wasserstein distance W(Js, Pkn), where 14 denotes the probability measure associated with VR-reSGLD at iteration k. We first bypass the Grönwall inequality and conduct the change of measure upper bound the relative entropy D (px (vtn) following (Raginsky et al. 2017). In addition to the approximation the standard Langevin diffu- sion Raginsky et al. (2017), 7), also consider the change of Poisson measure following Yin & Zhu (2010); Gikhman & Skorokhod (1980) to handle the error from the stochastic swapping rate. We then extend the distance of relative entropy D (Mg. Pkty) to the Wasserstein distance W2 (pt. Pkk)) via a weighted transportation-com inequality of Bolley & Villani (2005).
Theorem (Diffusion approximation) Assume the smoothness, the dissipativity and the gradient assumptions 7 and hold. Given large enough batch size n, a small enough m and n, we have
Wa Pkn) s 0 (lkkk (n22212(11 (1 + 5//4 +
m
1/8
where 5 is a constant that characterizes the scale of noise caused in mini-batch settings and the detail is given in Theorem in Appendix C. Here the last term O(( (m=)) n) 2/$)) comes from the error induced by the stochastic swapping rate, which disappears given a large enough batch size n or a small enough update frequency m and learning rate 7). Note that our upper bound linearly dependent on time approximately, which is much tighter than the exponential dependence using the Grönwall inequality. Admittedly, the result without swaps is slightly weaker than the diffusion approximation (3.1) in Raginsky et al. (2017) and we refer readers to Remark 3 in Appendix C
Applying the triangle inequality for Wa(us, Mkn,) and Wz (Vin F) leads to the final result
Theorem 3 Assume the smoothness, the dissipativity and the gradient assumptions 2 and 3 hold. Given a small enough learning rate update frequency m and a large enough batch size n, we have
m² (=))) ))))/ +
-kn(1+6s
Wa(ue, m) S O (1133/2 (1 + 82/4 +
This theorem implies that increasing the batch size or decreasing the update frequency not only reduces the numerical error but also potentially leads to a faster exponential convergence of the continuous-time dynamics via much larger swapping rate Sq,man

---

Published as a conference paper at ICLR 2021
Sreand
SOLD
0.06
AVV
9-
10
0.02
-00
0.14 0.66"
400
0.3
460
200
600
1900
epozh
epoch
spach
loguol
Trace plot for (3))
(b) Trace plot for (
(c) Paths of 1og
(d) Contour of log10 7²
Figure 2: Trace plots, KDEs of ( and sensitivity study of 82 with respect to m, and n.
5 EXPERIMENTS
5.1
SIMULATIONS OF GAUSSIAN MIXTURE DISTRIBUTIONS
We first study the proposed vericcee-reduced replica exchange stochastic gradient Langevin dynam- ies algorithm reR-D)) on Gaussian mixture distribution (Dubey et al., 2016). The distribu- tion follows from 2: 2 (,, (3, o²) + 0.5N( 3, o²), where Φ 20, σ and 8 = 55 We use training dataset of size N 10 and propose to estimate the posterior distribution over 3. We compare the performance of VR-reSGLD against that of the standard stochastic gradient Langevin dynamics (SGLD), and replica exchange SGLD (reSGLD).
In Figs 2(a) and 2(b 2(b) we present trace plots and kernel density estimates (KDE) of samples generated from VR-reSGLD with m 40, r(1) 10
1 7(2) 1000, 7) 1e 7, and F = reSGLD adopt the same hyper-parameters except for F = 100 because smaller F may fail to propose any swaps; SGLD uses // = le - 7 and r 10. As the posterior density is intractable, we consider a ground truth by running replica exchange Langevin dynamics with long enough iterations. We observe that VR-reSGLD is able to fully recover the posterior density, and successfully jump between the two modes passing the energy barrier frequently enough. By contrast, SGLD, initialized at 8g 30, is attracted to the nearest mode and fails to escape throughout the run; reSGL. manages to jump between the two modes, however, F is chosen as large 100, which induces large bias and only yields three to five swaps and exhibits the metastability issue. In Figure 2(c) we present the evolution of the variance for VR-reSGLD over a range of different m and compare with reSGLD.
We see that the variance reduction mechanism has successfully reduced the variance by hundreds of times. In Fig 2(d) we present the sensitivity study of ō² as a function of the ratio n/N and the learning rate for this estimate we average out 10 realizations of VR-re and our results agree with the theoretical analysis Lemma]] 1
5.2 NON-CONVEX OPTIMIZATION FOR IMAGE DATA
We further test the proposed algorithm on CIFAR10 10 and CIFAR100. We choose the 20, 32, 56-layer residual networks as the training models and denote them by ResNet-20, ResNet-32, and ResNet- 56, respectively. Considering the wide adoption of M-SGD, stochastic gradient Hamiltonian Monte Carlo (SGHMC) is selected as the bascline. We refer to the standard replica exchange SGHMC algorithm as reSGHMC and the variance-reducedd reSGHMC algorithm as VR-reSGHMC. We also include another baseline called cyclical stochastic gradient MCMC (cycSGHMC), which proposes cyclical learning rate schedule. To make fair comparison, we test the variance-reduced replica exchange SGHMC algorithm with cyclic learning rates and refer to it as eVR-reSGHMC.
We run M-SGD, SGHMC and (ERG)M for 500 epochs. For these algorithms, we follow a setup from Deng et al. (2220). We fix the learning rate v/(2) 2c-6 in the first 200 epochs and decay by 0.984 afterwards. For SGHMC and the low-temperature processes of (VR-) (RG-)MMGG
we anneal the temperature following = (1) = 0.01/1.000000
in the beginning and keep it fixed after the
burn-in steps; regarding the high-temperature process, we set (2) 1.5mg (1) and = (2) 57 (1) The
initial correction factor Fg is fixed at 1.5e5. The thinning factor T is set to 256. In particular for
We choose (1)) 10 instead of to avoid peaky modes for ease of illustration.

---

Published as a conference paper at ICLR 2021
m0 n56
40
m= & n2256 mwitf nw256
Varterce retiationss (imee
10.0
m0
n52
20
m5
& n225
Varbacc (nndloor
mwinf 56
Vertence (nimoooos)
30
mm302
n=256
m=302
nw256
7.5
m=inf 5 522
Varfance not coouudtoonn 0X
mmmin
6 n1152
10.
-0
3.0
5.
100
400
500
100
300
100
200 Epochs
400
500
200 300 Epochs
500
Epochs
Epocrs
(a) CIFARIO: Original
(h) CIFARIO0 Original (c)
Variance reduction (d)
Variance reduction
v.s. proposed (m=50)
v.s. proposed (m=50)
setups CIFARIO
setups
CIFAR 100
Figure 3: Variance reduction on the noisy energy estimators on CIFARIO & CIFAR100 datasets.
cycSGHMC, we run the algorithm for 1000 epochs and choose the cosine learning rate schedule with 5 cycles; is set to le-5; we fix the temperature 0,001 and the threshold 0.7 for collecting the samples. Similarly, we propose the cosine learning rate for cVR-reGGHMC with 2 cycles and run it for 500 epochs using the same temperature 0.001. We only study the low-temperature process for the replica exchange algorithms. Each experiment is repeated five times to obtain the mean and standard deviations.
We evaluate the performance of variance reduction using VR-reSGHMC and compare with reS- GHMC. We first increase the batch size from 256 to 512 for reSGHMC and notice that the duction variance around 2 times (see the red curves in fig]),,)), Next, we try m = 50 and 256 for the VR-reSGHMC algorithm, which updates the control variates every 50 iterations. As shown in Fig 3(a,)), during the first 200 epochs, where the largest learning rate is used, the vari- ance of VR-reSGHMC is slightly reduced by 37% on CIFAR100 and doesn't 't make difference CIFARIO. However, as the learning rate and the temperature decrease, the reduction of the variance gets more significant. We see from Fig (,,d) that the reduction of variance can be up 10 times on CIFARIO and 20 times on CIFARIOO. This is consistent with our theory proposed in Lemma] I The reduction of variance based on VR-reSGHMC starts outperform the baseline with 512 when the epoch is higher than 370 on CIFAR10 and 250 on CIFAR 100. We also try m = 392, which updates the control variates every 2 epochs, and find similar pattern.
For computational reasons, we choose them with the baseline algorithms. With the help of swaps between two SGHMC chains, reSGHMC already obtains remarkable performance (Deng et al.) 2020) and five swaps often lead to an optimal performance. However, VR-reSGHMC still outperforms reSGHMC by around 0.2% on CIFARIO CIFAR100 [ and the number of swaps is increased to around
392 and
256 for reSGMCC and compare
and 1% improvement hundred under the same setting. We also try cyclic learning rates and compare cVR-reSGHMC with cycSGHMC, we see cVH-reSGHMC outperforms cycSGHMC significantly even if cycSGHMC is running 1000 epochs, which may be more costly than cVR-reSGHMC due to the lack of mechanism in parallelism. Note that cVR-reSGHMC keeps the temperature the same instead of annealing it as in VR-reSGHMC, which is more suitable for uncertainty quantification.
TABLE 1: PREDICTION ACCURACIES (%) BASED ON BAYESIAN MODEL AVERAGING. IN PAR- TICULAR, M-SGD AND SGHMC RUN 500 EPOCHS USING SINGLE CHAIN; CYCSGHMC RUN 1000 EPOCHS USING A SINGLE CHAIN; REPLICA EXCHANGE ALGORITHMS RUN 500 EPOCHS
USING TWO CHAINS WITH DIFFERENT TEMPERATURES.
CIFARIO
CIFAR100
METHOD
RESNET20
RESNET32
RESNET56
RESNET20
RESNET32
RESNET56
M-SGD
94.07=0.11
95.11+0.07
96.05+0.21
71.93+0.13
74.65+0.20
78.7640.22
SGHMC
0..16.0.13
95.1710.08
96.04.10.18
12.0910.14
74.80.10.19
10.25
TeSGHMC
94.56=0.23
95.4440.16
96.1540.17
73.9440.34
76.38:0.23
79.86+0.26
VR-ISGGHMC
94.84=0.11
95.62:10.09
96.32-1..15
0.
77.4010.27
80.62-10.22
сусвGHMC
94.61=0.15
92.
96. 16.19
74.2110.22
76.60:10.25
80.3940.22
eVR ReGIMI
94.91=0.10
95.64+0.113
96.36+0.116
75.02+40.19
77.58+0.21
80.50+20225
8

---

Published as a conference paper at ICLR 2021
Regarding the training cost and the treatment for improving the performance of variance reduction using adaptive coefficients in the early period, we refer interested readers to Appendix E
For the detailed implementations, release the code at /
Variance Bariance_Reduced Repleca__pppiaaaeaaeaaa R Exchaage_Stochhheeeeeee SawnnngeeStchhhttiicccoooo ett.chtttti_rarrrenettt.taeeeeaa MCMC
5.3 UNCERTAINTY QUANTIFICATION FOR UNKNOWN SAMPLES
A reliable model not only makes the right decision among potential candidates but also casts doubts on irrelevant choices. For the latter, we follow Lakshminrryynaan et al. (2017) and evaluate the uncertainty on outto--distribution samples from unseen classes. To avoid ever-confident predictions on unknown classes, the ideal predictions should yield a higher uncertainty on the cuttof-distribution samples, while maintaining the accurate uncertainty for the in-distribution samples.
Continuing the setup in S.2 we collect the ResNet20 models trained on CIFAR 10 and quantify the entropy on the Street View House Numbers (SVHN) dataset,
1.00
1.00
which contains 26,032 RGB test- ing images of digits instead of ob-
Empingal CDD
0.75
Emplinal CDE
0.75
jects. We compare cVR-reGGHMC with M-SGD, SGHMC, reSGHMC, and cSGHMC. Ideally, the predic- tive distribution should be the uni- form distribution and leads to the highest entropy. We present the em- pirical cumulative distribution func- tion (CDF) of the entropy of the pre- dictions SVHN and report it in Fig 4 As shown the left figure, M-SGD shows the smallest probability for high-entropy predictions, implying the weakness of stochastic optimization methods uncertainty estimates. By contrast, the proposed cVR-reGGHMC yields the highest probability for predictions of high entropy. Admittedly, the standard ResNet mod- els are poorly calibrated in the predictive probabilities and lead to inaccurate confidence. To alleviate this issue, we adopt the temperature-scae scaalin method with scale of calibrate the predictive dis- tribution (Guo et al.) 2017) and present the entropy in Fig 4 (right). In particular, we see that 77% of the predictions from C-G yields the entropy higher than 1.5, which is 7% higher than reSGHMC and 10% higher than cSGHMC and much better than the others.
0.50
3.500
M-SGD
M-SGD
8GHMC
SGHMC
0.25
reSGHMC
0.25
cycSCHMC
cycSGHMC
0.00
cVR-reSGHMC
0.00
eVH-reSGHMC
2.0
1.5
0.5
0.0
2.0
1.5
1.0
0.0
Entropy (betore calibration)
Entropy (after calibration)
Figure 4: CDF of entropy for predictions on SVHN via CI- FARIO models. A temperature scaling is used in calibrations.
For more discussions of uncertainty estimates on both datasets, we leave the results in Appendix
6 CONCLUSION
We propose the variance-reduced replica exchange stochastic gradient Langevin dynamics algorithm to accelerate the convergence by reducing the variance of the noisy energy estimators. Theoretically, this is the first variance reduction method that yields the potential of exponential accelerations instead of solely reducing the discretization error. In addition, we bypass the Grönwall inequality to avoid the crude numerical error and consider a change of Poisson measure in the generalized Girsanov theorem obtain a much tighter upper bound. Since our variance reduction only conducts on the noisy energy estimators and is not applied to the noisy gradients, the standard setting can be also naturally imported, which greatly facilitates the training of deep neural works.
ACKNOWLEDGMENT
We would like to thank Maxim Raginsky and the anonymous reviewers for their insightful sugges- tions. Liang's research supported in part by the grants DMS-2015498, R01-GM117597 and R01-GM126089, Lin acknowledges the support from NSF Subcontract 382247, W91 and DE-SC0021144 142.
DMS- 173-61333644, BNL
9

---

Published as a conference paper at ICLR 2021
REFERENCES
Christophe Andrieu and Gareth O. Roberts. The Pseudo-Marginal Approach for Efficient Monte Carlo Computations. Annals of Statistics, 37:697-725, 2009.
Jack Baker, Paul Fearnhead, Emily B. Fox, and Christopher Nemeth. Control Variates for Stochastic Gradient MCMC. Statistics and Computing, 29:599-615, 2019.
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and Geometry of Markov Diffusion Operators. Springer, 2014.
Amel Bentata and Rama Cont. Mimicking the Marginal Distributions of a Semimartingale. arXiv preprint arXiv:0910.3992, 2009.
François Bolley and Cédric Villani. Weighted Kierrar-Kulbacc-PPP Inequalities and Ap- plications to Transportation Inequalities. Annales de la Faculté des sciences de Toulouse Mathématiques, Serie. 6, 14(3):331-35,2200 200..
Niladri Chatterji, Nicolas Flammarion, Yi-An Ma, Peter Bartlett, and Michael Jordan. On the The- ory of Variance Reduction for Stochastic Gradient Monte Carlo. In Proc. of the International Conference Machine Learning (ICML), 2018.
Changyou Chen, Nan Ding, and Lawrence Carin. On the Convergence of Stochastic Gradient MCMC Algorithms with High-order Integrators. In Advances in Neural Information Process- ing Systems (NeurIPS). pp. 2278-2286, 2015.
Tianqi Chen, Emily B. Fox. and Carlos Guestrin. Stochastic Gradient Hamiltonian Monte Carlo, In Proc. of the International Conference on Machine Learning (ICML), 2014.
Yi Chen, Jinglin Chen, Jing Dong, Jian Peng, and Zhaoran Wang. Accelerating Nonconvex Learning via Replica Exchange Langevin Diffusion. In Proc. of the International Conference on Learning Representation (ICLR), 2019.
Aaron Defazio, Francis Bach, and Simon Lacoste-Julien. SAGA: Fast Incremental Gradient Method with Support for Non-Strongly Convex Composite Objectives. In Advances in Neural Information Processing Systems (NeuriPS). 2014.
Wei Deng, Qi Feng, Liyao Gao, Faming Liang, and Guang Lin. Non-Convex Learning via Replica Exchange Stochastic Gradient MCMC. In Proc. of the International Conference Machine Learning (ICML), 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proc. of the Annual Meeting of the Association Computational Linguistics (ACL), 2019.
Jing Dong and Xin T. Tong. Spectral Gap of Replica Exchange Langevin Diffusion on Mixture Distributions. ArXiv 16193v2, July 2020.
Avinava Dubey, Sashank Reddi, Barnahás Póezos, Alexander Smola, Eric P. Xing, and Sinead A. Williamson. Variance Reduction in Stochastic Gradient Langevin Dynamies. In Ad- vances Neural Information Processing Systems (NeurIPS), 2016.
David J. Earl and Michael W. Deem. Parallel Tempering: Theory, Applications, and New Perspec-
tives. Phys. Chem. Chem. Phys., 7:3910-3916, 2005.
A. Eizenberg and M. Freidlin. On the Dirichlet Problem for a Class of Second Order PDE Systems with Small Parameter. Stochastics and Stochastic Reports, 33:111-1453 1990.
Futoshi Futami, Issei Sato, and Masashi Sugiyama. Accelerating the Diffusion-based Ensemble Sampling by Non-reversible Dynamics. In Proc. of the International Conference on Machine Learning (ICML), 2020.
losif I. Gikhman and Anatoli V. Skorokhod. The Theory of Stochastic Processes Springer, 1980.
10

---

Published as a conference paper at ICLR 2021
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On Calibration of Modern Neural Networks. Proc. of the Interaational Conference on Machine Learning (ICML), 2017.
István Gyöngy. Mimicking the One-dimensional Marginal Distributions of Processes Having an Itō differential. Probability theory and related fields, 7 7144)5001-11, 1986.
Reza Harikandeh, Mohamed Osama Ahmed, Alim Virani, Mark Schmidt, Jakub Konečný, and Scott Sallinen. Stop Wasting My Gradients: Practical SVRG. In Advances in Neural Information Processing Systems (NeurIPS), 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Rie Johnson and Tong Zhang. Accelerating Stochastic Gradient Descent using Predictive Variance Reduction. In Advances Neural Information Processing Systems (NeurIPS). 2013.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensemble. In Advances in Neural Information Processing Systems (NeurlPS), 2017.
Holden Lee, Andrej Risteski, and Rong Ge. Beyond Log-concavity: Provable Guarantees for Sam- pling Multi-modal Distributions using Simulated Tempering Langevin Monte Carlo. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks. In Proc. of the National Conference on Artificial Intelligence (AAAI), pp. 1788-1794, 2016.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the Loss Land- scape of Neural Nets. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
J.C. Mattingly, A.M. Stuartb, and D.J. Highame. Ergodicity for SDEs and Approximations: Locally Lipschitz Vector Fields and Degenerate Noise. Stochastic Processes and their Applications, 101: 185-232, 2002.
B. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Springer, 2003.
Matias Quiroz, Robert Kohn, Mattias Villani, and Minh-Ngoc Tran. Speeding Up MCMC by Eff- cient Data Subsampling. Journal of the American Statistical Association, 114:831-843 2019.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex Learning via Stochastic Gradient Langevin Dynamics: a Nonasymptotic Analysis. In Proc. of Conference on Learning Theory (COLT), June 2017.
Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Mathe- matical Statistics, 2203):400-407 1951.
Robert H. Swendsen and Jian-Sheng Wang. Replica Monte Carlo Simulation of Spin-Glasses Phys- ical Review Letters, 57:2607-2600 1986.
Yee Whye Teh, Alexandre Thiery, and Sebastian Vollmer. Consistency and Fluctuations for Stochas- tic Gradient Langevin Dynamics. Journal of Machine Learning Research, 17:1-33, 2016.
Sebastian J. Vollmer, Konstantinos C. 7ygalakis, and Yee Whye Teh. Exploration of the (Non-) Asymptotic Bias and Variance of Stochastic Gradient Langevin Dynamics. Journal of Machine Learning Research, 17(159):1-48, 2016.
Max Welling and Yee Whye Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In Proc. of the International Conference on Machine Learning (ICML), pp. 681-688, 2011.
Pan Xu, Jinghui Chen, Difan Zou, and Quanquan Gu. Global Convergence of Langevin Dynamics Based Algorithms for Nonconvex Optimization. In Advances in Neural Information Processing Systems (NeurlPS), 2018.
11

---

Published as a conference paper at ICLR 2021
Nanyang Ye, Zhanxing Zhu, and Rafal K.Mantiuk. Langevin Dynamics with Continuous Tempering for Training Deep Neural Networks. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
George Yin and Chao Zhu, Hybrid Switching Diffusions: Properties and Applications. Springer, 2010.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning. In Proc. of International Conference on Learning Representation (ICLR), 2020.
Yuchen Zhang, Percy Liang, and Moses Charikar. Hitting Time Analysis of Stochastic Gradient Langevin Dynamics. In Proc. of Conference on Learning l'heory (COLT), pp. 1980-2022, 2017.
Dongruo Zhou, Pan Xu, and Quanquan Gu. Stochastic Nested Variance Reduction for Nonconvex Optimization. Journal of Machine Learning Research, 20: 1-47, 2019.
Difan Zou, Pan Xu, and Quanquan Gu. Sampling from Non-Log Con-Log-Concave Distributions via Reduuec Gradient Langevin Dynamics. In Proc. of the International Conference Artificial Intelligence and Statistics (AISTATS), 2019a.
Difan Zou, Pan Xu, and Quanquan Gu. Stochastic Gradient Hamiltonian Monte Carlo Methods with Recursive Variance Reduction. In Advances in Neural Information Processing Systems (NeurIPS), 2019b.
12

---

Published as a conference paper at ICLR 2021
A PRELIMINARIES
Notation We denote the deterministic energy based on the parameter 3 by using the full dataset of size N. We denote the unbiased stochastic energy estimator by N LiEB ((x,(3) using the mini-batch data B of size n. The same style notations is also applicable to the gradient for consistency. We denote the Euclidean I norm by || Il. To prove the desired results, we need the following assumptions:
(xx;13)
Assumption (Smoothness) The energy function L(x,)-) (-) is Sm-smootssess if there exists a con- stant CN such that vB1. .2 E Rd (1.2. NN. we have
[VV[(x))))) /=(x & Cx 1| - B211-
(10)
Note that the above condition furthee implies constant C' = NCv and 1/1, E R we have [VL(()) - VI(3)) & Ci - -211-
(11)
The smoothness conditions (10) and [I] are standard tools in studying the convergence of SGLD in (Xu et al., 2018) and Raginsky et (2017). respectively.
Assumption (Dissipativity) The energy function L(-) is (a, didisiiie if there exist constants and b 2 0 such that VB E R2 (B. V(()) 2 a||B||2 b.
The dissipativity condition implies that the Markov process is able to move inward on average regardless of the starting position. has been widely used in proving the geometric ergodicity of dynamic systems (Mattingly et al. 2002 Raginsky et al. 2017 Xu et al. 2018).
Assumption (Gradient oracle) There exists constant ó E [0, 1) such that for any B, we have
|| E|| [ VL(B) = 25(C 2 + 2²),
(12)
where is a positive constant. The same assumption has been used in Raginsky et al. (2077)
to
control the stochastic noise from the gradient.
B EXPONENTIAL ACCELERATIONS VIA VARIANCE REDUCTION
We aim to build an efficient estimator to approximate the deterministic swapping rate S(3
(( -e (
Sabr-=by -b))) (EA, L(x, - ((x:19(9)))
(13)
In big data problems and deep learning, is too expensive to evaluate the energy Σ (xx133 for each 33 for large N. To handle the computational issues, a popular solution is to use the unbiased stochastic energy N Eien L(x:13) for random mini-batch data B of size However, näive
13 replacement E/ L(x; (3) by N EiE
JiB (3x,13) leads to a large bias to the swapping rate. To remove such bias, follow Deng et al. (2020) and consider the corrected swapping rate
$100),000) (, | (
Eien L(x, - E.e (xx.1118) ) (.b (
(14)
where 2 denotes the variance of Eier L(x, ((x)) NN- Eie ((2)) Empirically, d² is quite large, resulting in almost no swaps and insignificant accelerations. To propose more effective swaps, we consider the variance-reduced estimator
L(Bo\Ba)
=
N T B
Σ
((xx(Bb)) - L (xx BmL )m)++))
Σ
x1 B (1)
(15)
where the control variate PmL is updated every m iterations. Denote the variance of ((B)3(1))- (())3(2) by σ2. The variance-reduced stochastic swapping rate follows
(g()) ((2)) c( ب =
( (
(16)
We only consider the of F in the stochastic swapping rate for ease of analysis.
13

---

Published as a conference paper at ICLR 2021
Using the strategy of variance reduction, can lay down the first result, which differs from the isting variance reduction methods that we only conduct variance reduction in the energy estimator for the class of SGLD algorithms.
Lemma (Variance-redueed energy estimator) Under the smoothness and dissipativity as- sumptions and 2 the variance of the variance-reduced energy estimator T(BBLBB (),, where (1,2), upper bounded by
s Var (I(BAlo()))
m D3
2n (2C d,,,0,c.. + +2Q²) + 47(0))
(17)
where DR = C--CC++aan max. [ Cb and the radius of a sufficiently large ball that contains (h)) for h E {1,2}.
Proof
Var (I(B(I{())
((xz18((0) '
N
((x))0()))
(xxx(0 3
ml an
+ +
x5
(1)
#n
Σ
72
BEB
E
N
((xa1p())
- L x: (())
+ N
>
Σ
(x
m
j=1
N²
F.
Σ
(b))
Σ T
x5|
|2t ml
Σ ((x)1,{00)
2
ml m +
N
ق=1
N2
(301) )1)
Σ
TR
((x)13()) To
(k)
Σ
Txx,
Σ T.
x5
mi
1EDx
N22
1 )
E
((x
x:
(h)
n
D3
g
3
ml
(18)
where the last equality follows from the fact that E{E 11 Στ E[r;3 for independent
variables {r,)/=1 with mean 0. The first inequality follows from E( - -[=)) E|z²] and the
last inequality follows from Lemma D1 where DR = CR+ m1,., .// [z) Ch
and R is the radius of sufficiently large ball that contains (a4 for (1,2).
Next, we bound E
|
g ml $
follows
E
(())
SE
Σ
(3 - (
S m Σ
E
-
- -3001 (19)
ml
m
14

---

Published as a conference paper at ICLR 2021
For each term, we have the following bound
B p
=R
11
N
(l(x:13())) + 21
B
2nj²
N
E
V (3 (1))
+ 12)
(20)
iEB
2n1
(M) 1 +22002) + 41)
2nF
2202y (,,(9),C,ar + =Q2) 4
n
where the first inequality follows by [IIaa + bll2] v 2F[||a| [b]] the i.i.d of the data points
and (11) < 7(2) for h E {1, the second inequality follows by Lemma D2 the last inequality
follows from Lemma D3
Combining 188. 19 and (20). we have
Var ([(Balo(
I m
DR
2m (20 d + +202) + 4r(3)d)
(21)
n
Since Var (B(B)(( I Var ( Eien ((x113a) + Var ( Eien 3m by defi-
nition, Var [[BBB is upper bounded by O mi?, " m²²)). which becomes much smaller using a small learning rate y,, shorter period m and large batch size n.
n
Note that mp (B(1), ((2)) is defined on the unbounded support [0, oo] and m (3(1), ((22) =
regardless of the scale of σ². To satisfy the (stochas- tic) reversibility condition, we consider the truncated swapping rate min S (. (3(1), ,g(())), which still targets the same invariant distribution (see section 3.1 (Quiroz et al. 2019) for details). We can show that the swapping rate may even decrease exponentially as the variacee increases.
Lemma B2 (Variance reduction for larger swapping rates) Given large enough batch size n, the variance-reduced energy estimator (() yields truncated swapping rate that satisfies
E (, ((2))) so min [1. ((30)) ,(2)
(o
o 20m=a++))) ) (22)
n² +e
Proof
N By central limit theorem, the energy estimator Σ iEB L(x|B3x) converges in distribution normal distributions as the batch size n goes to infinity. In what follows, the variance-reduced estimator L L(Bx (3x) also converges to a normal distribution, where the corresponding estimator is denoted by L(Bx (Bx). Now the swapping rate S.).m,((..) based on normal estimators follows
( - b (1) ( - )
mnm,n (301,
(23)
where 8² denotes the variance of ((BB)() L(B Note that Sy,m (B ( fol- lows a log-normal distribution with mean log (( (2)) ( y 2 and variance
x² on the log-scale, and S(g(1, (1g("),33 ((2)) is the deterministic swapping rate defined Applying Lemma D4 we have
)
E/minç1, [1min(1,S,.m.i. Sp,m (1(1, 3(2)))) = o ( (s(30),312)) exp
( SOo 02 1).
exp
(24)
) by at most a bias of ( according to the estimate of the third term of (S2) in Quiroz et al. (2019) and 5² <
Moreover, σ² differs from ~2, the variance i(BB3())
[((B)(()
15

---

Published as a conference paper at ICLR 2021
Var (i(BalB())) + Var (i(Bal(()) where both Var (B(Bal(()) and Var (i(Bals{))
are upper bounded by m²n DD ( (20 d,r(2),C,a, + 2Q²) + 4rd) by Lemma B1 it follows that
n
0 _
Elmin (1, S n (B(1) 3 H S
(25)
Applying A B < min {1 1, A} + BB, we have
Elmin S m, (3 ,()))))))))))))))))))) (()
-B| min {1, pmm. (3 3) m , S mn ((())
(26)
E mn ,
m ((2))
E/min (1, (p ((2)
see formala 25
By the triangle inequality, we can further upper bound the first term I S m. no (B(1), g(2)) S m,m,n (8 g
E
< [ S n.m. (g(1) ,300 - 5(1()),,)))))00)))))) ,(0)) + S((()) ((())) B BS y. m. (3(1) (B(1)) (3
(27)
I
=(((0) ((2))
+ 1
1
n
n
where I, and Ig follow from the proof of S1 without and with normality assumptions, respectively
(Quiroz et al. 2019).
Combining (26 and (27). we have
E1min Sy (. (( xe min (L3) ,38) (0
1 - 0 m 1 ) }
)
(28)
This means that reducing the update period m (more frequent update the of control variable), the learning rate 7) and the batch size n significantly increases min(1, 3p,m.n on average.
The above lemma shows a potential to exponentially increase the number of effective swaps via variance reduction under the same intensity r. Next, we show the impact of variance reduction in speeding up the exponential convergence of the corresponding continuous-time replica exchange Langevin diffusion.
Theorem (Exponential convergence) Under the smoothness and dissipativity assumptions Aand 2 the replica exchange Langevin diffusion associated with the variance-reduced stochastic swap- ping rates (p min,L, m), converges exponential fast to the invariant distribu- tion a given smaller learning rate smaller m. or larger batch size n:
Wy(ve, π) S Du exp {-t +++(1++ (1+ (1+ 85,m /cLs).
(29)
m.
where Do 2ccs (vol [),, 0s.,m.. == infr>0 Em (V des
1 is a non-negative constant de-
E(V
pending on the truncated stochastic swapping rate Symm, (.))) and increases with smaller learning rate n,, a shorter period m and a large batch size n. CLs is the standard constant of the log-Sobolev inequality asscoiated with the Dirichlet form for replica exchange Langevin diffusion without swaps.
Proof Given a smooth function f Rd x Rd + R, the infinitesimal generator Ls,, associated with the replica exchange Langevin diffusion with the swapping rate S.m,, min(1, 3 m,, follows
f( ((3(2) = (V gan (()(), 3(2),, VL( - (V mm) aw ( VL(
+ (1) (1)) A gm f(g ((2) + ((2) A f((1) 3(2)
+ r tr (1 ,(2) ( f( 3(2))
(30)
16

---

Published as a conference paper at ICLR 2021
where V yga and Agih) are the gradient and the Laplace operators with respect to (th,, respectively. Next, we model the exponential decay of (2 x) using the Dirichlet form
Es (f) = f I. (f)dT
(31)
where I m.mm. - -Ssmm.n (f2) f.mmmmn (f) is the Carré du Champ operator. In particular for the first term &.Ssy.... (f²), we have
1/2
2 S. m. (1130),000113
1
= - -(((())) ( gas (
io (eee v L + - gu, s(3 ,(301,,
- (( ( gu (B3 3(3)
-U(3(1), g s 3(2)) v sm L(3(2)) + - gm S , 3(
(2) (( ) S (B ,(0))
T S + 2 m. m. (3 )(0)) (3 g0)) (B ((2))
Combining the definition of the Carré du Champ operator, 30 and B we have
I's m. (f ((())
2 mm mn (3 ,(0))) ((8 ( Ls m. ((B ((()))
(32)
= "|V am ,(())) + 11
| ((8 !
2 Snm.n ( ((4)
f(3(1) ,)
Plugging 32
into 31.. the Dirichlet form associated with operator f. follows
Es ( P f(3 3 I1 (2) I A ,3) do d3 ,3()))
vanilla term E(f)
(33)
dr 2 / Sy (g()) ((a)) ( - s(B() )())) d (3 ,()
,
necellration temm
dyc where f corresponds to (((B")_ga)) Under the asymmetry conditions of ा| , and S.,m,n >> O, the accelcration term of the Dirichlet form is strictly positive and linearly dependent on the swapping rate S.m,.. Therefore, Es (f) becomes significantly larger the swapping rate Sp,m, increases significantly. According to Lemma 5 (Deng et al. 2020). there exists a constant
8...m. =inf E (V nt
depending on S.,m, that satisfies the following log-Sobolev
E(V
des 3
inequality for the unique invariant measure associated with variance-reduced replica exchange Langevin diffusion (B.
D(vol(=)= S
CLS Es.,
dv
1
+ dsm,m.n
dm
where ōs. increases rapidly with the swapping rate Sy By virtue of the exponential decay of entropy (Bakry et al. 2014), we have
D(vIlTT) S D(Pol(T)e - 20
where CLS is the standard constant of the log- Sog-oooolev inequality asscoiated with the Dirichlet form for replica exchange Langevin diffusion without swaps (Lemma 4 as in Deng et al. (2020)) Next, we upper bound Wa(ue, T) by the Villaiilani theorem (Bakry et al. 2014)
Wa(u: 7) / 2 D (T) V 201s D (1ol| ( (1165m.mm.))) /
17

---

Published as a conference paper at ICLR 2021
where ös Sy.m.. 0 depends on the learning rate 77, the period m and the batch size
& m E(V Tt
In the above analysis, we have established the relation that ds inf
dus
depending on mmm, may increase significantly with a smaller learning rate 17, a shorter period m and large batch size n. For more quantitative study on how large óg 0pmma. is related problems, we refer interested readers to the study of spectral gaps in Lee et al. (2018): Dong & TTng 20200;; Futami al. 2020).
C DISCRETIZATION ERROR
Consider complete filtered probability space (R, (Fe) which supports all the ran- dom subjects considered in the sequel. With a little abuse usage of notation, the probability measure P (component wise if IP is joint probability measure with mutually independent components) would always denote the Wiener measure under which the process (We)osisT is a IP-Brownian motion. To be precise, in what follows, we shall denote IP == pw x N. where pw is the infinite dimen- sional Wiener measure and N is the Poisson measure independent of pW and has some constant jump intensity. In our general framework below, the jump process a is introduced by swapping the diffusion matrix of the two Langevin dynamics and the jump intensity is defined through the swap- ping probability in the following sense, which ensures the independence of F.W and N in each time interval [in,. (i 1)m] for N+ The precise definition of the Replica exchange Langevin diffusion (reLD) is given as below. For any fixed learning rate 17 > 0, we define
d3, = dt + )d,,
(34)
P ((a(() jla(t di) = L, (,,((L/mm))) = 3) = ( + oldt), for 1 + j.
where VG( (L((82)) and [=(//n1 is the indicator function, i.e. for every 1 in with
(L(
E N+, given B(in) = B, we have IP (a(() = j|a((t di) = 1) rS(3),, where S(B) is defined and (1((1),g(9)) is defined in (33 In this case, the Markov Chain (t))
as (1, ( is a constant on the time interval [Lt/n] [t/n/nn+/ n) with some state in the finite-state space {0, 1)
and the generator matrix Q follows
- (t rS(B)) - [t/nlm) -rS(B)no(t - (ttltt(((/nn))) -
- LL/1m))(
Q =
where 5(.) is a Dirac delta function. The diffusion matrix E(ae) is thus defined as ((00)E(1) {
(
(2
V2r("II,)) From our definition and following Yin & Zhu 2010 (Section 2.7]. the generator matrix Q will depend on the initial value at each time interval en,, + 1)7). The distribution of process (B:)ssiST denoted as VT == ==P x N which is absolutely continuous with respect the reference measure PP PW N, under which W is Brownian motion and a(-) is Poisson process with some constant jump intensity. This fact follows from the result Gikhman & Skorokhod (VII, Section Theorem 2] and Yin & Zhu [Sectin 2.5, formula (2.40)). The motivation of only considering the positive swapping rate in in, for i N+, and zero elsewhere is due to construction of the discretized process 3 as shown below (see equation (3). A simple illustration of the idea can be seen from the auxiliary process construction in Yin & Zhu (210)]Section 2.5], following which we want make sure the stopping time of / and |3 happening the same time. Otherwise, is unlikely (and also unreasonable) to derive the Nikoody derivative of the two process B3 and B. Thus, we should think of the process is concatenated on the time interval [in, (i + 1)m) up to time horizon T". Similarly, we consider the following Replica exchange stochastic gradient Langevin diffusion, for the same learning rate 17 > 0 as above, have
d3 = (
S(G
(35)
P (a(t) jlā(t dt) = 1, 1,8((t/lm)) = =ō) = ( + +odd)), for 1 j.
18

---

Published as a conference paper at ICLR 2021
V and 5( = min S m. )(3)) } and n,m, mn (3 is (I((())
where Vā(3) =
shown in 16). The distribution of process (Be)o is denoted as pr = ph x N where ã is Poisson process with jump intensity r ( [t/n|n) on the time interval [[t/n|n. lt/n 17). Note that 33 and [ are defined by using the same P-Brownian motion W. but with two different jump intensity
the time interval [Lt/n/n, [t/n]n + 11). Notice that, if there is no jump, the construction of /3 based on (3 follows from the fact that they share the same marginal distributions as shown in Gyöngy ( (1986), where one can find the details in Raaiinkk et al. (2017). Given the jump process a and ā introduced into the dynamics of ( and B3, the construction is more complicated. Thanks to |Bentata & Cont (209)), we can carry on the similar construction in our current setting. We then introduce the following Radon-Nikodym density for dur/dur. In the current setting, the change of measure can be seen as the combination of two drift-diffusion process and two jump process simultaneously. We first introduce some notation. For each vector E R", we denote ||l1I := A A. Furthermore, we introduce a sequence of stopping time based on our definition of process (3 and B. For j E N+ we denote Ç/s as a stopping times defined by (j+1 (t)) * aG))) and ((T) max E N Sn S T). It is easy to see that for any stopping time S. there exists 1 E N+ such that E = In. Similarly, we have the stopping time for the process / by 2J+1 == inf(t > [j ā(t) + x(G))) and ā(t) follows the same trajectory of att). To serve the purpose of our analysis, one should think of the process ,3 as the auxiliary process to the process B3, see similar constructions in Yin & Zhu [Seetin 2.5, formula (2.39)]. The difference is that both of our process & and B are associated with jump process jumping at time in, for some integer i E N instead of jumping at any continuous time. We combine approximation method from Yin & Zhu 20100|Section 2.7] for non-constant generator matrix Q and the density representation for Markov process in Gikhman & Skorokhod (V80) Section 6, Teorem 2] to get the following
inf G :
Lemma C1 Let {Gilj E 10,1,... N(T))} be sequence of stopping time defined by Let E N be an fixed integer such that kn S S (k + 1)m, For each fixed learning rate r >> 0 and for any >> 0, the Kadon-Nikodym derivative of / is given below,
duT
N(T)
STTIAT
Σ
Σ (((G))) - x (a(()))) awo
exp
dut
1=
N(T)
SHIAT
1
[=(())))) - ((((())) || dtt)
2 1=
N(T)
SH-INT-E
N(T) S((c,)
x exp
Σ
rō(t [ttm S(B(
x xII S(BB))
Recall that Cj is stopping time defined by a (same as defined by ã), i.e. Cj+1 := infst G a{t) * a(())), for j = 1. NT),,, and for each Çj. there exists 1 E 1. kbb such that Cj bn. We now follow Gikhman & Skorokhod [VII, Section 6, Theorem 2] to derive the Radon-Nikodym density for durr duy. In this case, if the generator matrix Q is constant, i.e. the jump intensity is constant, we can follow the similar construction from Yin & Zhu (FomO) (2.40)], see also Eizenberg & Freidlin that we can treat our generator matrix as constant on each time interval [G, {1+1), then the existing results apply to our case for the density with respect to the Poisson measure (jump process and σ), i.e. dN dN3/dN's / Furthermore, once the generator matrix Q is constant, then the measure PG (or pC) is independent to S (or ) We show the following steps to give a clear outline of our proof.
Proof
[Formula()
Next, we adjust our setting to the case
Step I: For each stopping time interval Ki (i+), no jump would occur after the initial point at time &j and the diffusion matrix Σ and E keep the same, thus we can apply the generalized Girsanov theorem to get the Randon-Nikodym derivative for dPG /dPG.
Step 2: In order to combine the the two density of dN /dN and d / we need the inde- pendent property of the two measures on the same time interval, then we directly get the density
19

---

Published as a conference paper at ICLR 2021
following from Gikhman & Skorokhod (III Section 6, Theorem 2]. Different from the work mentioned above, we will first write all the density on each time interval [in, (i + 1)7) to incor- porate the independent requirement mentioned above. Notice that the relative change of density for dN /dNN would only depends on the left end point, since the jump intensity would change its values at the initial value of interval [ånz, (i + 1)m), which is standard idea to deal with gen- erator matrix Q depending on the initial value instead constant matrix case. (See Yin & Zhu [Scctint for similar treatments).
general, the stopping time interval could contain several time interval with length y, however the jump intensity should only depend on the left end point for each time interval [in, (i + 17)) Based on the above set up, we now derive the Radon-Nikodym derivative. First notice that, on each period [Si, (j+1), the matrix Σ is fixed and is evaluated at E(a((;),, which the same for (.)( In particular, M(a({;) = [()())) is constant diagonal matrix. According to our definition dvy dPG x dN and dur = dP x áN we write the Radon-Nikode each of the time interval [án, (17)) and concatenate them together. We consider the swapping of the diffusion matrix first where a similar construction can be found Yin Zhu (2010) (2.40)), we get the following Radon-Nikodys derivative, for any 0,
Step 3:
derivative on
N(T) Σ 1=0
(+1)AA-=
dNs = exp {
rō(t - (tnnLtt/nn)(((pmmaal) S(3p (nm
(36)
X 11 N(T) ((
S(Bc,)
Next, we show the density for dPG /dI@ as below. On each interval [S,+ Sj+1), given initial value (3,, By), the matrix M(a((j) and E(ā(())) are always the same, since no jump would happen. In particular, in this continuous case the integral on [Gi (i+1) and [er, Gi+1] are the same. Thus have the following Radon-Nikodyn derivative
dPG
N(T)
SH41AT
dPG = cx
N(T)
CHHAA
E (a(())) - E "a(((((()) dt (37)
2
Notice that matrix Σ is diagonal square matrix, thus we have Σ - Σ' Recall that W is P- Brownian motion, assuming there is no jump in the dynamic for 3, then according to the Gir- sanov theorem (see an example in Theorem 8.6.6 and Example 8.6.9 (Øksendal, 2003)) with Radon- Nikodym derivative dP /dP, we have the pG Brownian motion, denoted as WG which follows
WC == W. + S E (a.)))))
(38)
This fact holds true each of the time interval [C,, S4+1). Multiplying the two density dpG /pC and 5 / we complete the proof.
Remark Notice that, we keep the constant diffusion matrix without jump, then the Randon- Nikodym derivative dur / dvr has been used in the stochastic gradient descent setting, for example Raginsky et al.) (0017). However, the notation of the Brownian motion has been used freely, we try to make it consistent in the current setting. Namely, for constant diffusion matrix L, we have
dpē dPG
= -xxp (//"
X - awo
1
E - ds)
-
(39)
2 2
where WG is PG -Boowiinan motion as shown in equation
38 not P- Broowiian motion W.
20

---

Published as a conference paper at ICLR 2021
Remark 2 The density der that we derived above is far the best do. one would like dur use the continuous time control at)) with continuous jump intensity (((()) instead of jumping at the initial point with a fixed rate, then we can not even write the Gandon-Nikodyyl derivative anymore, since att) and a(t) will define different stopping time, i.e. jump different time and pr is not absolutely continuous with respect to Vr.
Based on the above lemma, we further get the following estimates.
Lemma C2 Given large enough batch size n small enough m and T3, we have the bound of the KL divergence of DK (pTlur) as below,
(KI(UTUTT) (Do +
with
ro42
m
20 0
vid
+
Vn
47(1)²
(2)
c2
Φ (caa
(T(1) [T(2] + +111))
+
-(1)
2 O
Vidd
Proof By the very definition of the KI.-divergence, we have
/ dur log duT dut
DKI(PT|UT)
= -ur log(dpr/der/)) (1.
We shall keep the convention below and denote Euy = E (1(,,3) = (3, where 3 = (8(1) ,g(2) 2dd and B (3(1), (()) ((2)) Rdd denotes the values at each time in, 0, kk.
Plugging Lemma CI in the above equation and we unify the notation by using time intervals of the type (i1,, (i 1)m] To be precise, we get
dpp
(1+1)m
Σ (a(in) ((m)))((GGG)-=="(m))))))))))) VC((
dIpG
exp
6n
Σ ((
(k+1)n
Σ (a(()))))
- -aat))))
in
1
Σ (a(k))
-(1 (-ack))))))))))))))))))))))))))) (a(x)) di).
(40)
kn
The above equality follows from the fact that each time interval [(,, G++1] always contain exactly some sub-interval [én,, (i 1)m]. Namely, we have (G, (i+1] = [ir,, (i+ 1)m] [(n+ (j ...U [In. (1 + 1)m). for some i E 1, kk.. In particular, the matrix Σ keep the same on
each interval [in7, (i + 1)m], for some E 10 .k.. Similarly, we expand the Radon-Nokodym
f the time interval of length n, Based on our definition of jump intensity, we get
derivative for aN
N(T)
dN
Σ
rō(t (t/ll ((B(e/n)z)))dd
exp
S
in
((3(,)
r8(s (3)m) S(3kg) x 1
N(T)
S(Br,)
km
(S((B.)) -((w))) x 11 N(T) 5((ct)
exp
S(B,,)
(41)
D=0
21

---

Published as a conference paper at ICLR 2021
Without loss of generality, we shall only consider the sum Σ and skip the interval [kn, J]- Notice that on each time interval [in, (i + 1)7), the control a (i7) and a(ir) are fixed, thus the two component of the measure dv independent. Taking into account the fact that W is Brownian motion, thus we apply the martingale property and arrive at
DKI(UT(UT)
((+1)7
(a(in)) - Vc( [t]
B .3
==
in
N(T)
EuT, 9
(S(nn) 5((mm)
(l3,) - log S(((,)
-log
(1+1)
E Bur
Σ (aii)) - (a((m)) d1
2
1=0
I
N(T)
Evr.s @ ((S((Bnn) S((:)))
Evr
1og 3( - log S(G
We then estimates the three terms I, J, K. in order as below.
Estimate of Σ: Due to the fact that every interval [ön, (i + 1)m) C 1s (j+1) for some j E {0. 1,...... ,N.,,,,NN .N))),, we know that the control a and are the same the interval lin, (i 1)n] and the diffusion matrix E is just constant matrix. Thus, we know that ma- (ar()()))) which takes one of the form from (O-)(0), E1))
trix E-1 (a(in))
La
O
La
V2T
IFE E-)(0), get
0
V2
V2
(a(m)) - Σー ->-)))-~~))((()) (a(i))
2d
V,G(B,) V,G(B,)
Σ
2r(2) V,ū(B,) (1)8))-V,((B)) G(3))))
[T
1=1
j=d+1
1
V VG.(B)) v,((3,)
VV V,G(3)]
(1)
2-(1)
Here VC(
(L(((2 and Vō(3)
VL(BO))
( The other matrix form of E-1(11
will result in the same estimates. We thus get
Vā(B.) V
I
Eur 3
On each fixed interval, for t [Kn,, (k + 1)m) we have PC Brownian motion and IP -Brownian motion (see examples in Theorem 8.6.6 and Example 8.6.9 Øksendal 2003)).
dw. dW, + (at)
dw -.: +=>+++
Plugging the PG (and B -Brownian motions to the original dynamics 34 and (35). we have
dB. E(ax)aWG and d3 E(or)aW
22

---

Published as a conference paper at ICLR 2021
On each interval (én), (i+ 1)m), E(ar) is constant matrix, thus know that the probability distribu- tion of (e)))e((n,( and {Se (k)))) are the same and we denote as ((}.) [(3.). The difference is that Be is driven by pG Pr-Brownian motion and Be is driven by PPo-Brownian motion, which implies that, for & E [in,, (i + 1)7), we have
Vā())) - VC( -B Byr -
E Euy 3
(42)
Thus, we have the following estimates,
R-1 I I
VG(3.) - VG ( de
47
PT
3
1=0
in
(i+1)n
VG((Le/n))) V0(B (000)|/0)0)
E
4-(1)
Mr
iin
1=0
C2
(++1)
33 Bunll dt) .I.
47
PT
h-1
(1+1)
VG( Vā (
dt
Iz.
4-
1=0
We now estimate the two terms I and In separately. Notice that, following our notation of PG_ Brownian motion, for t E [in, (i + 1)m), we have
3. - Bin = E(a+)
W = E(az)(W, w
which implies that (recall that dur dPp N and R2dx2d),
E |1| - < 2r 27 dn + 2 2- dn S 47 dn.
We thus conclude that,
I < C C (()) 711) kdn
As for the term Iz, according to Assumption
we obtain that
no
k-1
C B,, 11 1 Φ2
I2
< (1))
E
3
i=
Now, we just need to estimate
E pr
On each interval (in,, (i + 1)m], under the measure
dp ... we have
Bi1t2 = - + Ela((n)) (1+1) - w&),
which implies that
E "T.
@ 111 1,,
= kn. 3 e1,1 | + E, Bl. w
(1+1)m w + E - @ (++1) w&
= Fr.B
|1| || + (2) + (3)dn
The last equality follows from the independence of Bkn and W (k+1)) W and W is P6_ Brownian motion. By induction, we get
μr. a 11, < (id(7(1) + ((2) S 2kd(7 + +2)].
In principle, the Wiener measure W under pa is not Brownian motion, thus the uniform T,² bound used Lemma.3 may not be appropriate. Instead, we estimate the upper hound using slightly weaker result.
23

---

Published as a conference paper at ICLR 2021
We conclude that, which implies that
knp
(20 + kld + $2).
I2 S 4r(1]
kn (25)2 + 712 + sT02) + C² 7(2) 4r(1)
I IS
7(1) kddp
Estimate J: According to our definition of the swapping probability, we have, for each i,
5(Bm,) min Sy m, (B%)) ;, S(Bir) min S(0.)) ( )}
which means (3()).) s(Br,)) 1. Denote C, ||) - we have
C 02
m (3!) (((2 = exp (I (8! - [(B.nll?
S(B!) / exp (L((( - 11) 1))
Applying Taylor expansion for the exponential function at C+( )) )), we have
B mm (3 ( S( _
C 52 ---6- - C,(L - -
<Eur (( ( C. (B ( [( ((
+ higher order term
a| [I i(, { - c² /²
EEvr
) - | +0
where the last inequality follows from S( ) < 1. Combining Lemma B1 we thus get the following estimates,
)
k-1 J
Σ ==
Evr. (l (B,n) 5((3,))11)
C 0² -C-(T((!))
<
rn
Eur ,
CT ((
-((3)) +
1=0
mm² -1
< rkkno(CC, = -knn
d
where the last inequality follows from the Jensen's inequality and the last order holds given large enough batch size n or small enough m and n7.
Estimate Kc: We now estimate the last term KC, we have
N(T)
K
E Bpr l log S (m log S(3;))
3=0
N(T)
Cr 02 ) - -,00
<C,
Eur
-
(88)
N(T)
<N(T)C? / +2]+ C
Σ
Var Var[ -0,-(,,,,,)))11
N(T)CTO2
CN 2
24

---

Published as a conference paper at ICLR 2021
Combining Lemma B1 again, conclude with
c N(T)Ø² + ((T)CT=-N(T))
(mm²))
(m² 1/2
2
Combining the estimates of I. J, and XC, we complete the proof.
Remark After the change of measure, the expectation is under the new measure PG (or pC) instead of the Wiener measure P. In the estimate of term I, similar I estimates of the term uur.a (.r.alll/(++>)]l has been obtained in Raginsky et al. (017)) of Lemma 7] when there is no swap. The difference is we write the dynamic of B(t+T with respect the pd -Brownian motion @ instead of the P-Brownian motion W. In principle, W under Tp not Brownian motion.
We then extend the distance of relative entropy (Kz(UTleT) to the Wasserstein distance W2(Pr, ur) via a weighted transportation-cost inequality of Bolley & Villani (0055).
Theorem 2 Given a large enough batch size nu small enough and we huve
m )
1/8
W2(UT, PT) S O d1312 1/14 + 88/14
(43)
72
Proof Before we proceed, we first show in Lemma D5 that ur has a bounded second moment; the L22 upper bound of frr is majorly proved in Lemma.C2 (Chen et al. 2019) except that the slight et al. (2019) is changed to account for the
difference is that the constant the RHS of (C.38) stochastic noise. Then applying Corollary 2.3 in Bolley & Villani (2005), we can upper bound the two Borel probability measures fer and Py with finite second moments as follows
(KI(UTIUT) (TT 1/4
We(HT,PT) s Cv
VDKLHrIPR)) +
(44)
2
(x (3 + log Sae Mlllll v(dw)) 1/2
where C, = =iinf>>
Applying Lemma D6 we have
W 12 (No + 2b + 42) kn) (DII(WTTTT) (
Combining Lemma C2 and /N(T) N(T) and taking 17 1, kn 1, and À have
(3 (1,3) < (12+8 (oo + 26 + 4db(2)) km +
where + Φ; + / Φ. for i E (0, 1, In what follows, we have
(2 VT.3) (vo + VIvT) (ko))² N2)
where v, (12 8 (No 2b + (2) . for i {0, 2.).
By the orders of 2o. Φ and D2 defined Lemma C2 we have
m=
1/4
N(T)
m²
1/4
(T,)) 0
2
1/2 40/22 +
+
-17
77
Kn
mm
1/4
< 0
1///2 + 81/2+
+0/22
where N(T) can be interpreted the average swapping rate from time 0 to T and is of order (())
Taking square root both sides of the above inequality lead to the desired result 43
25

---

Published as a conference paper at ICLR 2021
D PROOF OF TECHNICAL LEMMAS
Lemma D1 (Local Lipschitz continuity) Given a d-dimensional centered ball U of radius R, L()) DR- -ipiscchit contimuous in that (x) - (xx & Da [| - 32 for v 32 E U and any i E (1,2, ,N,, where DR CR + m1.. N [x; + c.
Proof
For any B1. U. there exists 33 U that satisfies the mean-value theorem such that
[T(x,13)) (x) (VL 32) [V ||3 -321,
Moreover, by Lemma D2 we have
CR N
[xx](a)) [(x [VLx -13, - B2ll S
113 -32l1-
Lemma D2 Under the smoothness and dissipativity assumptions
for any B E Rd it follows
that
C
Q N
VL(x;13)| - N
|
(45)
where Q = mmx.ge1..,..... (xx) DC a
Proof
According to the dissipativity assumption, we have
(B.. VI(3.) a|| b,
(46)
where B, a minimizer of VL(-) such that VL(3.) = 0. what follows, we have ||3. S 8.
Combining the triangle inequality and the smoothness assumption |
we have
Cub + [VL((
[V[xx1B) SCv 3 3+1l [Vxxx S Cn B|
(47)
Setting ON
N in and Q m (1.2,...)) [VL(x(BB)) + LC completes the proof.
The following lemma is majorly adapted from Lemma C.2 of Chen et al. (2019), except that the corresponding constant the RHS of (C.38) is slightly changed to account for the stochastic noise. A similar technique has been established in Lemma of Raginsky et al. (2017).
Lemma D3 (Uniform 2 bounds on replica exchange SGLD) Under the smoothness and dissi-
putivity assumptions
7 2 Given small enough learning rate 1) (0,1 V (x), there exists
positive constant V (,,, < such that sup E B1] Warre)
Lemma D4 (Exponential dependence the variance) Assume S is log-normal distribution with mean u 3 and variance o² on the log scale. Then E[min(1, 0(e"~** )), which is exponentially smaller given large variance o².
Proof For log-normal distribution S with mean J02 and variance 22 on the log scale, the (log s $0222 In what follows, have
probability density fs(S) follows that SVZngz exp
20a
B(min(1,S)) (1,),))]// / (ff(1, ( min(1,
(log S - 2000=
exp
ds
S 22
20
26

---

Published as a conference paper at ICLR 2021
By change of variable y = log ---+-uth to² follows that
where S = @ and y = y/4 + given = 1,
E
[1m
1
(log S + 2
/
(log S - u + 02)
S
exp
dS +
exp
ds
S
2mo2
20²
SV2md=
20²
++
1
u²
o0
tay by +
1
1²
06 2 to dy
e
c
e
2TT02
2
(u-o12
1
=
dy +
2m
7
2T
1
22
d2
+
e
dy
2m
V
0
1
2
1
1
2
Zea
dz
dy
2
2
e"
6
where the last equality follows from the change of variable a - y and the second last inequality
follows from the exponential tail bound of the standard Gaussian variable Pyy e) Sez _,a
Lemma DS (Uniform L bound on replica exchange Langevin diffusion) For all n E (0,1A 40 ), we have that
E ( (. < Blello! 1 sP b+ (2)
a
Proof Consider (B,)
B3 2 where B, (8 ( E Rdd The proof is marjorly adapted
from Lemma in Raginsky et al. (2017), except that the generalized Itô formula (formula 2.7 in page 29 of Yin & Zhu (2010) is used to handle the jump operator, which follows that
dh. 2B3,, VG((.)) + 2d(7(") + + + 233 Elar)dW
+ S. (B. ((())) (T.(3 ( Lx( . ((())) +M) + M2(1),
Jump-inducing drift
VL(B()) V ( and Mi()) and M2(() are two martingales defined in formula 2.7
where VG(
in Yin & Zhu 2. Due to the definition of (33). we have (((() (=) ,3(T)).
which implics that the Jump-inducing drift actually disappears. Taking expectations and applying the margingale property of the Itō integral, we have the almost the same upper bound as Lemma in Raginsky et al. 2017). Combining E[||@o [1a] < log Efel (all") completes the proof.
Lemma D6 (Exponential integrability of replica exchange Langevin diffusion) For all < a' it follows that
(a( log El (a s(Di) & log Elee P "p +2(b + 2d (2)tt
Proof
The proof is marjorly adapted from Lemma 4 in Raginsky et al. (017)). The only difference is that the generalized Itō formula (formula 2.7 in Yin & Zhu (2010) used again in Lemma
DS Consider L(t, 3,) = ell where / = (33 (.(3)) E Rd Due to the special structure that L(t, B:) is invariant under the swaps of (33(1) 11 _.()). the generator of L(t, Bi) with swaps is the same as the one without swaps. Therefore, the desired result follows directly by repeating the steps
from Lemma 4 in Raginsky et al. (201)).
27

---

Published as a conference paper at ICLR 2021
Algorithm 2 Adaptive variance-reduced replica exchange SGLD. The learning rate and temperature can be set to dynamic to speed up the computations. A larger smoothing factor y captures the trend better but becomes less robust.
Input Initial parameters ( and g learning rate y and temperatures =()) and r(?), correction factor F.
repeat
Parallel sampling Randomly pick mini-batch set Be of size n.
3 hb)
1 (h)
N
Σ
vL(x;
for h {1,2]. ( + 22
BE
Variance-reduced energy estimators Update P,U)) ='
x1| (0/mmia mL every m iterations.
N Σ
+
N Σ
x:
3/h)) ml
d
F(m)
for h E {1,2}-
π VED
if & mad m 0 then
Update o = (1 - --)0k- + yo where of is an estimate for Var (T(Bal3.2) [(Balo()))
Update č = (1 - 7)kk- + 7ck. where o is an estimate for ( (11(
Var(I
(B1S mL
end if
Bias-reduced swaps Swap (()) and [ if Sn.m.n. where u Unif [0, 1]. and S.,m. follows
Spma = exp { (=dr r (I(Be+1)81h) I(Bu41 (3 (zds
82 m 1
until A = kimax.
Output: (3(11) ) )i
where is thinning factor.
E MORE EMPIRICAL STUDY ON IMAGE CLASSIFICATION
E.
TRAINING COST
The batch size of n = 512 almost doubles the training time and memory, which becomes too costly in larger experiments. A frequent update of control variates using cinecconu and is not acceptable in practice. The choice of m gives rise to a tradeoff between computational cost and variance reduction. As such, we choose significant reductions of the variance at the cost of 40% increase on the training time. Note that when we set m 2000, the training cost only increased by 8% while the variance reduction can be still at most 6 times on CIFARIO and 10 times on CIFARI00.
50 is even more
392, which still obtains
E.2 ADAPTIVE COEFFICIENT
We study the correlation coefficient of the noise from the current parameter (/a) where h {1, 2). and the control variate (h) - As shown in Fig the correlation coefficients are only around -0.5 3mLa
due to the large learning rate the early period. This implies that VR-reSGHMC may overuse the noise from the control variates and thus fails fully exploit the potential in variance reduction. In spirit to the adaptive variance, we try the adaptive correlation coefficients to capture the pattern of the varying correlation coefficients and present it in Algorithm 2
As a result, we can further improve the performance variance reduction by as much as 40% on CIFAR10 and 30% CIFAR 100 in the first 200 epochs. As the training continues and the learning rate decreases, the correlation coefficient is becoming closer to -1. In the late period, there still 10% improvement compared to the standard VR-reGGHMC.
28

---

Published as a conference paper at ICLR 2021
In nut shell, we can try adaptive coefficients in the early period when the absolute value of the correlation is lower than 0.5 or just use the vanilla replica exchange stochastic gradient Monte Carlo to avoid the computations of variance reduction.
-0.5
3.0.
-44
Tary sdudaa tnreel
-0.
-0.05
150
200 300
450
500
205 100
302
400
50b
100
200
100
400
500
1.6g
100
2dc
300
ato
500
(a) CIFARIO & m=50
(b) CIFARI00 & m=50
(c) CIFARIO & m=392
(d) CIFARIOO & m=392
Figure 5: A study of variance reduction techniques using adaptive coefficient and non-adaptive coefficient on CIFARIO & CIFAR 100 datasets.
F MORE EMPIRICAL STUDY ON UNCERTAINTY QUANTIFICATION
To avoid sacrificing the prediction power for the known classes, we also include the uncertainty estimate on CIFAR10 using the Brier score (BS) [ and compare it with the estimates on SVHN. The optimal BS scores on the seen CIFAR10 dataset and the unseen SVHN dataset are 0 and 0.1, respectively. As shown in Table we see that the scores before calibration in the seen CIFARIO I0 is much lower than the ones in the unseen SVHN. This implies that all the models perform quite well in terms of what it knows, although cSGHMC are slightly better than the alternatives. To alleviate this issue, we propose to calibrate the predictive probability through the temperature scaling (Guo et al., 2017) and obtain much better results. Regarding the BS score on the unseen dataset, we see that M- SGD still performs the worst for frequently making over-confident predictions; SGHMC performs better but far away from satisfying. reSGHMC obtains much better performance by allowing interactions between different chains. However, the large correction term affects the efficiency of the swaps significantly. In the end, our proposed algorithm increases the efficiency of the swaps via variance reduction and further improves the righly-optimize 0.29 to 0.27, which is much closer to the ideal 0.1. Note that the accurate uncertainty estimates of cVR-reSGHMC on the dataset is still maintained. Together with the lowest BS score in the unseen SVHN dataset, eVR-reSGMMC shows its strength in uncertainty quantification.
BS score based reSGHMC from
TABLE 2: UNCERTAINTY ESTIMATES ON SVHN USING CIFARIO MODELS.
BRIER SCORE (before calibration)
BRIER SCORE (after calibration)
METHOD
CI (seen)
SVHN (unseen)
CIFARIO (seen)
SVHN (unscen)
M-SGD 0.09010.001
0.48
0.09840.001
0.3310.02
SGHMC
0.089+0.001
0.47+0.02
0.099
0.31+0.02
reSGHMC
0.086
0.41110.03
0.097100001
0.29
cSGHMC
0.084+0.001
0.43+0.02
0.092
0.30+0.02
cVR-reSGHMC
0.085+0.001
0.38+0.02
0.094#0.001
0.27+0.02
G MODIFIED EXAMPLE 5.1
We revisit Example 5.1 and re-run the procedures with temperature 7(1) 1.0. Fig. we present trace plots and kernel density estimates (KDE) of samples generated from VR-reSGI.D, reSGLD, and SGLD. In particular, we run VR-reSGLD with m = 40, -(1) = 1, (2) = 500, le 5, and F 1; reSGLD with the same hyper-parameters as VR-reSGLD except for F 500; and SGLD with 7 1e 5 and - = 1. Note that here, we run reSGLD with greater F than in Example 5.I in order to prohibit the drastic reduction of the swapping rate which is caused by the pickier target density. As in Example 5.1 for the ground truth, we run replica exchange Langevin dynamics
IBS = DN (Jog o, where fs is the predictive probability and o, is actual output of the which is if happens and 0 otherwise; N is the number of instances and is the number of classes.
29

---

Published as a conference paper at ICLR 2021
40
40
30
Ground bruth
30
Ground truth
SGLD
20
20
10
10
-00
400 800 epoch
400
800
epoch
(a) Trace plot VR-reSGLD and
(b) plot SGLD, reSGLD ground truth
ground truth
Figure 6: Trace plots and KDEs of (8()
with long enough iterations. Figs 6(a) and 6(b) we observe that, even though the distribution of interest has pickier density, proposed algorithm VR-reSGLD was able to detect both modes and acceptably jump between them. On the other hand, the competitor algorithm SGLD was trapped in the first mode visited and never escaped. reSGLD was able to jump some times between modes only after considering a substantial factor F 500 which, according to the theory, introduces bias.
30


Reviews:
[{"review_id": "xx78WKR2Rzl", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a variant of the Replica Exchange Stochastic Gradient Langevin Dynamics (reSGLD) for non log-concave sampling by using a variant reduction technique on the estimation of the swapping rate. Assuming that the log-density is a finite sum. the authors apply classical variance reduction techniques to the energy estimator necessay to compute the swapping rate. They show that applying such technique yields a higher swapping frequency and faster convergent rate of both the continuous time SDE and its dicretization scheme. Finally, the authors perform numerical experiments on both synthetic and real world data, and show that VR indeed reduces the variance of energy estimator by several orders of magnitude, hence inducing faster convergence.\n\nIn the algorithm, what is the role of the thining factor T? Is it involved in the convergence guarantee given by Theorem 3?\n\nFor the synthetic experiments, can you explain this choice for different values of F (1 for VR-reSGLD and 100 for reSGLD)? I can believe that variance reduction lkeads to a different bias/variance trade-off, but it would be good to explain whythese values of F were used.\n\nThm 1: Operator E is not defined in the main text (only in the appendix).", "extra_scores": {}, "extra_text": {}}, {"review_id": "xmyKe1d0JQW", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##  Summary of the paper\nThis paper extends the replica stochastic gradient MCMC by incorporating the new variance reduction technique. The author analyzed the non-asymptotic theoretical behavior of the proposed method, which shows the proposed method provides a better energy landscape.\n\n## Strong and weak points of the paper\n### Strong points\n- Although the applied variance reduction technique is not new, the analysis of the swapping rate in Lemma 2 seems novel compared to the past replica-exchange MCMC work.\n- Provided an asymptotic analysis for the variance reduction in Theorem 4.\n\n### Weak points \n- Proposed variance reduction is not new, I think. It is very similar to the standard control variates methods, which had been extensively applied into MCMCs and other machine learning tasks.\n\n## Rating\n- Clarity: Well written, easy to read, although I did not check the proofs in detail.\n- Correctness: I did not check all the proof in detail.\n- Novelty: The idea seems not novel but the swapping rate analysis seems new and interesting.\n\n## Comments and Questions\n- Q) All the experimental results and theoretical analysis suggest that using a smaller $m$ is better. But how can we tune this $m$? Does this variance reduction work even I use $m=1$?  But I also think that sufficiently large $m$ is required to estimate the control covariate appropriately as the work of  Rie Johnson and Tong Zhang suggest in their paper. Does this intuition wrong?\n\n- Q) In Sec5.2, as for Fig.3 c and d, the effect of variance reduction becomes significant after the sufficient epochs and the author conjectured that this is because that the learning rate is decreased. I think that this suggests that the proposed variance reduction is significantly affected by the step size which changes during the exploration stage in SG-MCMCs. And thus I thought that using the constant $m$ during the exploration stage might not be a good idea.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "c_35NgVa005", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary\n\nThe paper presents VR-reSGLD, a method to accelerate replica exchange stochastic gradient Langevin diffusion (reSGLD), which has been proposed recently to tackle non-convex learning problems. reSGLD suffers from two major sources of error resulting in low swapping rates: minibatch noise and the discretization error of Langevin diffusion. The idea of the paper is to use control variates to reduce the variance of energy estimators and thereby improve the swapping rate (which should lead to an accelerated convergence). Unlike previous modifications of SGLD, the variance reduction proposed in the paper aims at improving the energy estimators rather than the gradient estimators. The paper presents non-asymptotic results backing the intended acceleration of the Markov jump process. Numerical experiments illustrate the performance gain achieved by VR-reSGLD. These tests include a one-dimensional example (learning the component mean of a bimodal mixture of Gaussians) and Bayesian training of DNNs based on CIFAR imaging data. \n\n# Assessment\n\nOverall, I think the paper could be accepted in principle, but requires a thorough revision. \n\nreSGLD has been proposed quite recently (ICML 2020); and the use of variance reduction (VR) techniques is common to account for minibatch noise. So the combination of reSGLD and VR seems like an obvious idea. However, the current paper adds more than just a minor modification of reSGLD in that it provides interesting theoretical results and illustrates the validity of these results by numerical experiments. \n\n## Pros\n\n1. Convergence in 2-Wasserstein distance, no asymptotic normality assumed.\n\n2. Significant acceleration, tighter discretization error.\n\n3. Improved uncertainty quantification.\n\n## Cons\n\n1. The theoretical part of the paper is hard to follow for a non-expert.\n\n2. There are quite a number of algorithmic parameters (e.g. $m$, $n$, $\\eta$, $\\gamma$, $F$, etc.) that seem to require some tweaking. \n\n3. The improvement presented on imaging data (Table 1) seems to be rather marginal. \n\n# Comments and Questions:\n\n- How much tweaking do parameters such as $F$ etc. require? Can you provide some intuition on how to set these parameters? \n\n- Have you considered learning rates $\\eta_k$ that depend on the temperature?\n\n- Have you considered using more than two temperatures? \n\n- Please define \"acceleration\" (relative to what?)\n\n- Page 2: \"To avoid the crude estimate, ...\" is unclear. What is \"the crude estimate\"? The upper bound provided by Gronwall's inequality?\n\n- Page 3: The sentence starting with \"The underlying Dirichlet form ...\" is a bit obscure. Please try to clarify. \n\n- Page 3: Regarding the constant $c$ achieving minimum variance: Is the sign of $c^\\star$ correct? It seems that the denominator should be $\\text{Var}(B|\\hat\\beta^{(h)})$ rather than $\\text{Var}(B|\\beta^{(h)})$. \n\n- Page 4, Algorithm 1: How to set $\\gamma$? The update rule for $\\tilde \\sigma_k^2$ is unclear to me. What is $\\sigma_k^2$? Should this be $\\sigma^2$?\n\n- Page 5, Theorem 1: What do you mean by \"exponentially faster\", compared to reSGLD? What are the $\\mathcal{E}$s? \n\n- Page 6: Your comment after Theorem 3 (\"This theorem implies ...\"). Why does Theorem 3 imply a much larger swapping rate?\n\n- Page 6: You set $\\tau^{(1)}=10$ \"to avoid peaky modes\", but it would be interesting to see the performance of VR-reSGLD for a sharply peaked posterior (which is more realistic when thinking about high-dimensional parameter spaces and a large number of data...)\n\n- Page 7: Why do you anneal $\\tau^{(1)}$ in the CIFAR experiment?\n\n# Minor:\n\n## Typos\n\n- Replace \"Gröwall\" with \"Grönwall\" or \"Gronwall\"\n\n## Grammar and wording\n\n* In \"obtain the state-of-the-art results\" delete \"the\"\n\n* Try to improve: \"the noisy energy estimators in mini-batch settings render the naive implementation a large bias...\"\n\n* Page 2: \"it may cause the process to stuck in ...\"\n\n* There a probably more problematic phrases. Please improve the quality of writing. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "5R9M62gbVow", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Accelerating convergence of replica exchange stochastic gradient mcmc via variance reduction\n\nSummary:\n\nThe paper presents a variance reduction technique to achieve more efficient swaps in replica exchange stochastic gradient Langevin dynamics MCMC. The paper provides detailed analysis of the method as well as empirical evaluation on some standard deep learning tasks.\n\nPositive:\n\n1. Overall I would say that the paper is well written and and it is fairly easy to follow the presentation and details in the derivations.\n2. The topic is very timely and the method appears to be very useful. As an attractive method for minibatched Bayesian inference, stochastic gradient Langevin Dynamics samplers are of high interest, but tuning the algorithm can be somewhat finicky in my experience. Replica exchange is sometimes extremely useful, and finding good defaults for these types of methods is important.\n3. Experimental validation is reasonable (although a bit limited) and the methods chosen for comparison are resonable.\n4. A comprehensive set of appendices are included to provide further details. Although I did not go through the appendices in detail, I find it appealing that further information is provided for readers wishing to apply these methods in practice.\n\nNegative:\n\n1. The authors do not provide a reference software implementation. This makes it more difficult for readers to verify the results and might limit the impact of the paper. I would highly appreciate that the authors would create and share a minimal implementation.\n2. The novelty / originality is limited: A well known type of variance reduction applied in a new way/context where it makes perfect sense though.\n\nRecommendation:\n\nGood paper. Accept.\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work aims at doing Bayesian inference via Langevin dynamics with data subsampling. This builds on previous work with "replica exchange" where parallel chains are run at different temperatures and can be swapped to encourage moving between modes. The main technical novelty here is a scheme to reduce variance. This is done in the style of SGRD by periodically computing the gradient on all data and then using those values as control variates. This is shown to reduce variance.

Reviewers generally felt that this represented a sensible combination of known ideas aimed at an important and timely problem with sufficient empirical evaluation. There was consensus the paper was clearly written. I concur that even if the combination is "expected" to work, the presence of guarantees for performance represent sufficient technical novelty. I particularly applaud the fact that the paper does not over-claim and generously gives credit to related work. This is helpful to the reader and encourages the flow of ideas. For these reasons I recommend acceptance of the paper.

In reading the paper, I had a couple questions about the experiments:

1. It's not obvious to me from the experiments how specific the method is to the replica exchange setting. The main control variate idea appears to be applicable without replica exchange. I would very much like to see a "VR-SGHMC" row in Table 1 unless there is a good reason that this cannot be done. It would be very beneficial to understand the contributions of these different algorithmic components.

2. The CIFAR experiments directly test variance. That's fine, the paper is aimed at reducing variance, after all. However, I would like to see more tests of the follow-on improvements in optimization speed. It has been my experience that improvements in variance sometimes produce surprisingly small improvements in optimization speed. My intuition for this is that reduced variance mostly helps by making it possible to use a larger step-size without the same penalty in the stationary dist. In practice, the step-size typically ends up being imperfect, meaning that changes in variance have small changes.

Decision: Accept (Poster)