Year: 2021

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önwall’s inequality and yields a much tighter error in the 2-Wasserstein (W2) distance. Numerically, we conduct extensive experiments and obtain state-of-the-art results in optimization and uncertainty estimates for synthetic experiments and image data.

heading: 1 INTRODUCTION
text: Stochastic gradient Monte Carlo methods (Welling & Teh, 2011; Chen et al., 2014; Li et al., 2016) are the golden standard for Bayesian inference in deep learning due to their theoretical guarantees in uncertainty quantification (Vollmer et al., 2016; Chen et 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 et 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 et al., 2017) and cyclic learning rates (Zhang et al., 2020), to name a 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 traditional Monte Carlo methods (Swendsen & Wang, 1986; Earl & Deem, 2005), reSGLD (Deng et al., ∗Equal contribution
2020) uses multiple processes based on stochastic gradient Langevin dynamics (SGLD) where interactions between different SGLD chains are conducted in 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-batch settings lead to a 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 reducing 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 et al. (2019) studied the construction of control variates using local modes. Despite the advantages in near-convex problems, a natural discrepancy between theory (Chatterji et al., 2018; Xu et al., 2018; Zou et al., 2019b) and practice (He et al., 2016; Devlin et al., 2019) is whether we should avoid the gradient noise in non-convex 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 gradients 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 dynamics (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., 2018; Zou et al., 2019a) may only work well when the dynamics is sufficiently mixed and the discretization error becomes a major component. Moreover, the existing discretization error of the Langevin-based Markov jump processes (Chen et al., 2019; Deng et al., 2020; Futami et al., 2020) is exponentially dependent on time due to the limitation of Grö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 state-of-the-art performance in both optimization and uncertainty estimates.

heading: 2 PRELIMINARIES
text: A common problem, in Bayesian inference, is the simulation from a posterior P(β|X) ∝ P(β) ∏N i=1 P(xi|β), where P(β) is a proper prior, ∏N i=1 P(xi|β) is the likelihood function and N is the number of data points. When N is large, the standard Langevin dynamics is too costly in evaluating the gradients. To tackle this issue, stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011) was proposed to make the algorithm scalable by approximating the gradient through a mini-batch data B of size n such that
βk = βk−1 − ηk N
n ∑ i∈Bk ∇L(xi|βk−1) + √ 2ηkτξk, (1)
where βk ∈ Rd, τ denotes the temperature, ηk is the learning rate at iteration k, ξk is a standard Gaussian vector, and L(·) := − log P(β|X) is the energy function. SGLD is known to converge weakly to a stationary Gibbs measure πτ (β) ∝ exp (−L(β)/τ) as ηk decays to 0 (Teh et al., 2016). The temperature τ is the key to accelerating the computations in multi-modal distributions. On the one hand, a high temperature flattens the Gibbs distribution exp (−L(β)/τ) (see the red curve in Fig.1(a)) 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 a local region for an exponentially long time, as shown in the blue curve in Fig.1(a,b). To bridge the gap between global exploration and local exploitation, Deng et al. (2020) proposed the replica exchange SGLD algorithm (reSGLD), which consists of a low-temperature SGLD to encourage exploitation and a high-temperature SGLD to support exploration
β (1) k = β (1) k−1 − ηk
N
n ∑ i∈Bk ∇L(xi|β(1)k−1) + √ 2ηkτ (1)ξ (1) k
β (2) k = β (2) k−1 − ηk
N
n ∑ i∈Bk ∇L(xi|β(2)k−1) + √ 2ηkτ (2)ξ (2) k ,
where the invariant measure is known to be π(β(1),β(2)) ∝ exp ( −L(β
(1)) τ(1) − L(β (2)) τ(2) ) as ηk → 0
and τ (1) < τ (2). Moreover, the two processes may swap the positions to allow tunneling between different modes. To avoid inducing a large bias in mini-batch settings, a corrected swapping rate Ŝ is developed such that
Ŝ = exp {( 1
τ (1) − 1 τ (2) ) (N n ∑ i∈Bk L(xi|β(1)k )− N n ∑ i∈Bk L(xi|β(2)k )− ( 1 τ(1) − 1 τ(2) ) σ̂2 F )} ,
where σ̂2 is an estimator of the variance of Nn ∑ i∈Bk L(xi|β (1) k )− N n ∑ i∈Bk L(xi|β (2) k ) and F is the correction factor to balance between acceleration and bias. In other words, the parameters switch the positions from (β(1)k ,β (2) k ) to (β (2) k ,β (1) k ) with a probability r(1∧ Ŝ)ηk, where the constant r is the swapping intensity and can set to 1ηk for simplicity.
From a probabilistic point of view, reSGLD is a discretization scheme of replica exchange Langevin diffusion (reLD) in mini-batch settings. Given a smooth test function f and a swapping-rate function S, the infinitesimal generator LS associated with the continuous-time reLD follows
LSf(β(1),β(2)) = −〈∇β(1)f(β (1),β(2)),∇L(β(1))〉 − 〈∇β(2)f(β (1),β(2)),∇L(β(2))〉
+ τ (1)∆β(1)f(β (1),β(2)) + τ (2)∆β(2)f(β (1),β(2)) + rS(β(1),β(2)) · (f(β(2),β(1))− f(β(1),β(2))),
where the last term arises from swaps and ∆β(·) is the the Laplace operator with respect to β(·). Note that the infinitesimal generator is closely related to Dirichlet forms in characterizing the evolution of a stochastic process. By standard calculations in Markov semigroups (Chen et al., 2019), the Dirichlet form ES associated with the infinitesimal generator LS follows
ES(f) = ∫ ( τ (1)‖∇β(1)f(β (1),β(2))‖2 + τ (2)‖∇β(2)f(β (1),β(2))‖2 ) dπ(β(1),β(2))︸ ︷︷ ︸
vanilla term E(f)
+ r
2 ∫ S(β(1),β(2)) · (f(β(2),β(1))− f(β(1),β(2)))2dπ(β(1),β(2))︸ ︷︷ ︸
acceleration term
,
(2)
which leads to a strictly positive acceleration under mild conditions and is crucial for the exponentially accelerated convergence in the W2 distance (see Fig.1(c)). However, the acceleration depends on the swapping-rate function S and becomes much smaller given a noisy estimate of N n ∑ i∈B L(xi|β) due to the demand of large corrections to reduce the bias.

heading: 3 VARIANCE REDUCTION IN REPLICA EXCHANGE STOCHASTIC GRADIENT LANGEVIN DYNAMICS
text: The desire to obtain more effective swaps and larger accelerations drives us to design more efficient energy estimators. A naı̈ve idea would be to apply a 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.
A natural idea to propose more effective swaps is to reduce the variance of the noisy energy estimator L(B|β(h)) = Nn ∑ i∈B L(xi|β(h)) for h ∈ {1, 2}. Considering an unbiased estimator L(B|β̂(h))
for ∑N i=1 L(xi|β̂(h)) and a constant c, we see that a new estimator L̃(B|β(h)), which follows
L̃(B|β(h)) = L(B|β(h)) + c ( L(B|β̂(h))−
N∑ i=1 L(xi|β̂(h))
) , (3)
is still the unbiased estimator for ∑N i=1 L(xi|β(h)). By decomposing the variance, we have
Var(L̃(B|β(h))) = Var ( L(B|β(h)) ) + c2Var ( L(B|β̂(h)) ) + 2cCov ( L(B|β(h)), L(B|β̂(h)) ) .
In such a case, Var(L̃(B|β(h))) achieves the minimum variance (1 − ρ2)Var(L(B|β(h))) given c? := −Cov(L(B|β
(h)),L(B|β̂(h))) Var(L(B|β̂(h))) , where Cov(·, ·) denotes the covariance and ρ is the correlation
coefficient of L(B|β(h)) and L(B|β̂(h)). To propose a correlated control variate, we follow Johnson & Zhang (2013) and update β̂(h) = β(h)
mb km c every m iterations. Moreover, the optimal c? is often
unknown in practice. To handle this issue, a well-known solution (Johnson & Zhang, 2013) is to fix c = −1 given a high correlation |ρ| of the estimators and then we can present the VR-reSGLD algorithm in Algorithm 1. Since the exact variance for correcting 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 β̂(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; a largem, however, may yield a less correlated control variate and setting c = −1 fails to reduce the variance. In spirit of the adaptive variance in Deng et al. (2020) to estimate the unknown variance, we explore the idea of the adaptive coefficient c̃k = (1− γk)c̃k−m + γkck such that the unknown optimal c? is well approximated. We present the adaptive VR-reSGLD in Algorithm 2 in Appendix E.2 and show empirically later that the adaptive VR-reSGLD leads to a significant improvement over VR-reSGLD 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(Nd), 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. (2019) 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 reduction 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 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-time process in the early stage, in addition to the improvement on the discretization error. In other words, our algorithm is not only theoretically sound but also more empirically appealing for a wide variety of problems in non-convex learning.
Algorithm 1 Variance-reduced replica exchange stochastic gradient Langevin dynamics (VRreSGLD). 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.
Input The initial parameters β(1)0 and β (2) 0 , learning rate η, temperatures τ (1) and τ (2), correction factor F and smoothing factor γ. repeat
Parallel sampling Randomly pick a mini-batch set Bk of size n.
β (h) k = β (h) k−1 − η
N
n ∑ i∈Bk ∇L(xi|β(h)k−1) + √ 2ητ (h)ξ (h) k , for h ∈ {1, 2}. (4)
Variance-reduced energy estimators Update L̂(h) = ∑N i=1 L ( xi ∣∣∣β(h) mb k
m c
) every m iterations.
L̃(Bk|β(h)k ) = N
n ∑ i∈Bk [ L(xi|β(h)k )− L ( xi ∣∣∣β(h) mb km c )] + L̂(h), for h ∈ {1, 2}. (5)
if k mod m = 0 then Update σ̃2k = (1− γ)σ̃2k−m + γσ2k, where σ2k is an estimate for Var ( L̃(Bk|β(1)k )− L̃(Bk|β (2) k ) ) . end if Bias-reduced swaps Swap β(1)k+1 and β (2) k+1 if u < S̃η,m,n, where u ∼ Unif [0, 1], and S̃η,m,n follows
S̃η,m,n = exp {( 1
τ(1) − 1 τ(2) )( L̃(Bk+1|β(1)k+1)− L̃(Bk+1|β (2) k+1)− 1 F ( 1 τ(1) − 1 τ(2) ) σ̃2 mb k
m c
)} .
(6) until k = kmax. Output: The low-temperature process {β(1)iT } bkmax/Tc i=1 , where T is the thinning factor.

heading: 4 THEORETICAL PROPERTIES
text: 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 1 (Variance-reduced energy estimator) Under the smoothness and dissipativity assumptions 1 and 2 in Appendix A, the variance of the variance-reduced energy estimator L̃(B|β(h)), where h ∈ {1, 2}, is upper bounded by
Var ( L̃(B|β(h)) ) ≤ min { O ( m2η
n
) ,Var (N n ∑ i∈B L(xi|β(h)) ) + Var (N n ∑ i∈B L(xi|β̂(h)) )} ,
where the detailed O(·) constants is shown in Lemma B1 in the appendix.
The analysis shows the variance-reduced estimator L̃(B|β(h)) yields a 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 Sη,m,n = min{1, S̃η,m,n} still satisfies the “stochastic” detailed balance given an unbiased swapping-rate estimator S̃η,m,n (Deng et al., 2020) †, it doesn’t mean the efficiency of the swaps is not affected. By contrast, we can show that the number of swaps may become exponentially smaller on average.
Lemma 2 (Variance reduction for larger swapping rates) Given a large enough batch size n, the variance-reduced energy estimator L̃(Bk|β(h)k ) yields a truncated swapping rate that satisfies
E[Sη,m,n] ≈ min { 1, S(β(1),β(2)) ( O ( 1 n2 ) + e −O ( m2η n + 1 n2 ))} , (7)
†Andrieu & Roberts (2009); Quiroz et al. (2019) achieve a similar result based on the unbiased likelihood estimator for the Metropolis-hasting algorithm. See section 3.1 (Quiroz et al., 2019) for details.
where S(β(1),β(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 require the normality assumption. As n goes to infinity, where the asymptotic normality holds, the RHS of (7) changes to
min { 1, S(β(1),β(2))e −O ( m2η n )} , which becomes exponentially larger as we use a smaller up-
date frequency m and learning rate η. Since the continuous-time reLD induces a jump operator in the infinitesimal generator, the resulting Dirichlet form potentially leads to a much larger acceleration term which linearly depends on the swapping rate Sη,m,n and yields a faster exponential convergence. Now we are ready to present the first main result.
Theorem 1 (Exponential convergence) Under the smoothness and dissipativity assumptions 1 and 2, the probability measure associated with reLD at time t, denoted as νt, converges exponentially fast to the invariant measure π:
W2(νt, π) ≤ D0 exp { −t ( 1 + δSη,m,n ) /cLS } , (8)
where D0 is a constant depending on the initialization, δSη,m,n := inft>0 ESη,m,n (
√ dνt dπ )
E( √ dνt dπ )
− 1 ≥ 0
depends on Sη,m,n, ESη,m,n and E are the Dirichlet forms based on the swapping rate Sη,m,n and are defined in (2), cLS is the constant of the log-Sobolev inequality for reLD without swaps.
We detail the proof in Theorem.1 in Appendix B. Note that Sη,m,n = 0 leads to the same performance as the standard Langevin diffusion and δSη,m,n is strictly positive when dνt dπ is asymmetric (Chen et al., 2019); given a smaller η and m or a large n, the variance becomes much reduced according to Lemma 1, yielding a much larger truncated swapping rate by Lemma 2 and a faster exponential convergence to the invariant measure π compared to reSGLD.
Next, we estimate the upper bound of the 2-Wasserstein distance W(µk, νkη), where µk denotes the probability measure associated with VR-reSGLD at iteration k. We first bypass the Grönwall inequality and conduct the change of measure to upper bound the relative entropy DKL(µk|νkη) following (Raginsky et al., 2017). In addition to the approximation in the standard Langevin diffusion Raginsky et al. (2017), we 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 DKL(µk|νkη) to the Wasserstein distanceW2(µk, νkη) via a weighted transportation-cost inequality of Bolley & Villani (2005).
Theorem 2 (Diffusion approximation) Assume the smoothness, the dissipativity and the gradient assumptions 1, 2 and 3 hold. Given a large enough batch size n, a small enough m and η, we have
W2(µk, νkη) ≤ O ( dk3/2η ( η1/4 + δ1/4 + (m2 n η )1/8)) , (9)
where δ is a constant that characterizes the scale of noise caused in mini-batch settings and the detail is given in Theorem 2 in Appendix C . Here the last term O (( m2 n η )1/8)
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 η. Note that our upper bound is linearly dependent on time approximately, which is much tighter than the exponential dependence using the Grö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 forW2(µk, νkη) andW2(νkη, π) leads to the final result
Theorem 3 Assume the smoothness, the dissipativity and the gradient assumptions 1, 2 and 3 hold. Given a small enough learning rate η, update frequencym and a large enough batch size n, we have
W2(µk, π) ≤ O ( dk3/2η ( η1/4 + δ1/4 + (m2 n η )1/8)) +O ( e −kη(1+δSη,m,n ) cLS ) .
This theorem implies that increasing the batch size n or decreasing the update frequency m not only reduces the numerical error but also potentially leads to a faster exponential convergence of the continuous-time dynamics via a much larger swapping rate Sη,m,n.

heading: 5 EXPERIMENTS

heading: 5.1 SIMULATIONS OF GAUSSIAN MIXTURE DISTRIBUTIONS
text: We first study the proposed variance-reduced replica exchange stochastic gradient Langevin dynamics algorithm (VR-reSGLD) on a Gaussian mixture distribution (Dubey et al., 2016). The distribution follows from xi|β ∼ 0.5N(β, σ2) + 0.5N(φ − β, σ2), where φ = 20, σ = 5 and β = −5. We use a training dataset of size N = 105 and propose to estimate the posterior distribution over β. 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), we present trace plots and kernel density estimates (KDE) of samples generated from VR-reSGLD with m = 40, τ (1) = 10 †, τ (2) = 1000, η = 1e− 7, and F = 1; reSGLD adopt the same hyper-parameters except for F = 100 because a smaller F may fail to propose any swaps; SGLD uses η = 1e − 7 and τ = 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 β0 = 30, is attracted to the nearest mode and fails to escape throughout the run; reSGLD manages to jump between the two modes, however, F is chosen as large as 100, which induces a 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 it 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 σ̃2 as a function of the ratio n/N and the learning rate η; for this estimate we average out 10 realizations of VR-reSGLD, and our results agree with the theoretical analysis in Lemma 1.

heading: 5.2 NON-CONVEX OPTIMIZATION FOR IMAGE DATA
text: We further test the proposed algorithm on CIFAR10 and CIFAR100. We choose the 20, 32, 56-layer residual networks as the training models and denote them by ResNet-20, ResNet-32, and ResNet56, respectively. Considering the wide adoption of M-SGD, stochastic gradient Hamiltonian Monte Carlo (SGHMC) is selected as the baseline. We refer to the standard replica exchange SGHMC algorithm as reSGHMC and the variance-reduced reSGHMC algorithm as VR-reSGHMC. We also include another baseline called cyclical stochastic gradient MCMC (cycSGHMC), which proposes a cyclical learning rate schedule. To make a fair comparison, we test the variance-reduced replica exchange SGHMC algorithm with cyclic learning rates and refer to it as cVR-reSGHMC.
We run M-SGD, SGHMC and (VR-)reSGHMC for 500 epochs. For these algorithms, we follow a setup from Deng et al. (2020). We fix the learning rate η(1)k = 2e-6 in the first 200 epochs and decay it by 0.984 afterwards. For SGHMC and the low-temperature processes of (VR-)reSGHMC, we anneal the temperature following τ (1)k = 0.01/1.02
k in the beginning and keep it fixed after the burn-in steps; regarding the high-temperature process, we set η(2)k = 1.5η (1) k and τ (2) k = 5τ (1) k . The initial correction factor F0 is fixed at 1.5e5. The thinning factor T is set to 256. In particular for †We choose τ (1) = 10 instead of 1 to avoid peaky modes for ease of illustration.
cycSGHMC, we run the algorithm for 1000 epochs and choose the cosine learning rate schedule with 5 cycles; η0 is set to 1e-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-reSGHMC 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 2 standard deviations.
We evaluate the performance of variance reduction using VR-reSGHMC and compare it with reSGHMC. We first increase the batch size n from 256 to 512 for reSGHMC and notice that the reduction of variance is around 2 times (see the red curves in Fig.3(c,d)). Next, we try m = 50 and n = 256 for the VR-reSGHMC algorithm, which updates the control variates every 50 iterations. As shown in Fig.3(a,b), during the first 200 epochs, where the largest learning rate is used, the variance of VR-reSGHMC is slightly reduced by 37% on CIFAR100 and doesn’t make a difference on CIFAR10. However, as the learning rate and the temperature decrease, the reduction of the variance gets more significant. We see from Fig.3(c,d) that the reduction of variance can be up to 10 times on CIFAR10 and 20 times on CIFAR100. This is consistent with our theory proposed in Lemma 1. The reduction of variance based on VR-reSGHMC starts to outperform the baseline with n = 512 when the epoch is higher than 370 on CIFAR10 and 250 on CIFAR100. We also try m = 392, which updates the control variates every 2 epochs, and find a similar pattern.
For computational reasons, we choose m = 392 and n = 256 for (c)VR-reSGHMC and compare 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 CIFAR10 and 1% improvement on CIFAR100 (Table.1) and the number of swaps is increased to around a hundred under the same setting. We also try cyclic learning rates and compare cVR-reSGHMC with cycSGHMC, we see cVR-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.
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, we release the code at https://github.com/WayneDW/ Variance_Reduced_Replica_Exchange_Stochastic_Gradient_MCMC.

heading: 5.3 UNCERTAINTY QUANTIFICATION FOR UNKNOWN SAMPLES
text: A reliable model not only makes the right decision among potential candidates but also casts doubts on irrelevant choices. For the latter, we follow Lakshminarayanan et al. (2017) and evaluate the uncertainty on out-of-distribution samples from unseen classes. To avoid over-confident predictions on unknown classes, the ideal predictions should yield a higher uncertainty on the out-of-distribution samples, while maintaining the accurate uncertainty for the in-distribution samples.
Continuing the setup in Sec.5.2, we collect the ResNet20 models trained on CIFAR10 and quantify
the entropy on the Street View House Numbers (SVHN) dataset, which contains 26,032 RGB testing images of digits instead of objects. We compare cVR-reSGHMC with M-SGD, SGHMC, reSGHMC, and cSGHMC. Ideally, the predictive distribution should be the uniform distribution and leads to the highest entropy. We present the empirical cumulative distribution function (CDF) of the entropy of the predictions on SVHN and report it in Fig.4. As shown in the left figure,
M-SGD shows the smallest probability for high-entropy predictions, implying the weakness of stochastic optimization methods in uncertainty estimates. By contrast, the proposed cVR-reSGHMC yields the highest probability for predictions of high entropy. Admittedly, the standard ResNet models are poorly calibrated in the predictive probabilities and lead to inaccurate confidence. To alleviate this issue, we adopt the temperature-scaling method with a scale of 2 to calibrate the predictive distribution (Guo et al., 2017) and present the entropy in Fig.4 (right). In particular, we see that 77% of the predictions from cVR-reSGHMC yields the entropy higher than 1.5, which is 7% higher than reSGHMC and 10% higher than cSGHMC and much better than the others.
For more discussions of uncertainty estimates on both datasets, we leave the results in Appendix F.

heading: 6 CONCLUSION
text: 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 Grönwall inequality to avoid the crude numerical error and consider a change of Poisson measure in the generalized Girsanov theorem to 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 hyper-parameter setting can be also naturally imported, which greatly facilitates the training of deep neural works.

heading: ACKNOWLEDGMENT
text: We would like to thank Maxim Raginsky and the anonymous reviewers for their insightful suggestions. Liang’s research was supported in part by the grants DMS-2015498, R01-GM117597 and R01-GM126089. Lin acknowledges the support from NSF (DMS-1555072, DMS-1736364), BNL Subcontract 382247, W911NF-15-1-0562, and DE-SC0021142.

heading: A PRELIMINARIES
text: Notation We denote the deterministic energy based on the parameter β by L(β) = ∑N i=1 L(xi|β) using the full dataset of size N . We denote the unbiased stochastic energy estimator by N n ∑ i∈B L(xi|β) using the mini-batch of data B of size n. The same style of notations is also applicable to the gradient for consistency. We denote the Euclidean L2 norm by ‖ · ‖. To prove the desired results, we need the following assumptions:
Assumption 1 (Smoothness) The energy function L(xi|·) is CN -smoothness if there exists a constant CN > 0 such that ∀β1,β2 ∈ Rd, i ∈ {1, 2, · · · , N}, we have
‖∇L(xi|β1)−∇L(xi|β2)‖ ≤ CN‖β1 − β2‖. (10) Note that the above condition further implies for a constant C = NCN and ∀β1,β2 ∈ Rd, we have
‖∇L(β1)−∇L(β2)‖ ≤ C‖β1 − β2‖. (11)
The smoothness conditions (10) and (11) are standard tools in studying the convergence of SGLD in (Xu et al., 2018) and Raginsky et al. (2017), respectively.
Assumption 2 (Dissipativity) The energy function L(·) is (a, b)-dissipative if there exist constants a > 0 and b ≥ 0 such that ∀β ∈ Rd, 〈β,∇L(β)〉 ≥ a‖β‖2 − b.
The dissipativity condition implies that the Markov process is able to move inward on average regardless of the starting position. It 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 3 (Gradient oracle) There exists a constant δ ∈ [0, 1) such that for any β, we have
E[‖∇L̃(β)−∇L(β)‖2] ≤ 2δ(C2‖β‖2 + Φ2), (12) where Φ is a positive constant. The same assumption has been used in Raginsky et al. (2017) to control the stochastic noise from the gradient.

heading: B EXPONENTIAL ACCELERATIONS VIA VARIANCE REDUCTION
text: We aim to build an efficient estimator to approximate the deterministic swapping rate S(β(1),β(2))
S(β(1),β(2)) = e
( 1
τ(1) − 1 τ(2)
) ( ∑N i=1 L(xi|β (1))− ∑N i=1 L(xi|β
(2))). (13)
In big data problems and deep learning, it is too expensive to evaluate the energy ∑N i=1 L(xi|β) for each β for a large N . To handle the computational issues, a popular solution is to use the unbiased stochastic energy Nn ∑ i∈B L(xi|β) for a random mini-batch data B of size n. However, a näive
replacement of ∑N i=1 L(xi|β) by N n ∑ i∈B L(xi|β) leads to a large bias to the swapping rate. To remove such a bias, we follow Deng et al. (2020) and consider the corrected swapping rate
Ŝ(β(1),β(2)) = e
( 1
τ(1) − 1 τ(2) )( N n ∑ i∈B L(xi|β (1))−Nn ∑ i∈B L(xi|β (2))− ( 1 τ(1) − 1 τ(2) ) σ̂2 2 ) , (14)
where σ̂2 denotes the variance of Nn ∑ i∈B L(xi|β(1))− N n ∑ i∈B L(xi|β(2)). ∗ Empirically, σ̂2 is quite large, resulting in almost no swaps and insignificant accelerations. To propose more effective swaps, we consider the variance-reduced estimator
L̃(Bk|βk) = N
n ∑ i∈Bk ( L(xi|βk)− L ( xi ∣∣∣βmb km c))+ N∑ i=1 L ( xi ∣∣∣βmb km c) , (15) where the control variate βmb km c is updated everym iterations. Denote the variance of L̃(B|β
(1))− L̃(B|β(2)) by σ̃2. The variance-reduced stochastic swapping rate follows
S̃η,m,n(β (1),β(2)) = e
( 1
τ(1) − 1 τ(2)
)( L̃(B|β(1))−L̃(B|β(2))− ( 1
τ(1) − 1 τ(2)
) σ̃2
2 ) . (16)
∗We only consider the case of F = 1 in the stochastic swapping rate for ease of analysis.
Using the strategy of variance reduction, we can lay down the first result, which differs from the existing variance reduction methods in that we only conduct variance reduction in the energy estimator for the class of SGLD algorithms.
Lemma B1 (Variance-reduced energy estimator) Under the smoothness and dissipativity assumptions 1 and 2, the variance of the variance-reduced energy estimator L̃(Bk|β(h)k ), where h ∈ {1, 2}, is upper bounded by
Var ( L̃(Bk|β(h)k ) ) ≤ m 2η
n D2R
( 2η
n (2C2Ψd,τ(2),C,a,b + 2Q 2) + 4τ (2)d
) . (17)
whereDR = CR+maxi∈{1,2,··· ,N}N‖∇L(xi|β?)‖+ Cba andR is the radius of a sufficiently large ball that contains β(h)k for h ∈ {1, 2}.
Proof
Var ( L̃(Bk|β(h)k ) ) =E N n ∑ i∈Bk [ L(xi|β(h)k )− L ( xi ∣∣∣β(h) mb k m c )] + N∑ j=1 L ( xj ∣∣∣β(h) mb k m c ) − N∑ j=1 L(xj |β(h)k )
2 =E N n ∑ i∈Bk [ L(xi|β(h)k )− L ( xi ∣∣∣β(h) mb k m c ) + 1 N ( N∑ j=1 L ( xj ∣∣∣β(h) mb k m c ) − N∑ j=1 L(xj |β(h)k )
)]2 = N2
n2 E ∑ i∈Bk [ L(xi|β(h)k )− L ( xi ∣∣∣β(h) mb k m c ) + 1 N ( N∑ j=1 L ( xj ∣∣∣β(h) mb k m c ) − N∑ j=1 L(xj |β(h)k ) )]2 = N2
n2 ∑ i∈Bk E (L(xi|β(h)k )− L(xi∣∣∣β(h)mb k m c ) − 1 N [ N∑ j=1 L(xj |β(h)k )− N∑ j=1 L ( xj ∣∣∣β(h) mb k m c )])2 ≤N 2
n2 ∑ i∈Bk E [( L(xi|β(h)k )− L ( xi ∣∣∣β(h) mb k m c ))2]
≤D 2 R
n E [∥∥∥∥β(h)k − β(h)mb k m c ∥∥∥∥2 ] ,
(18)
where the last equality follows from the fact that E[( ∑n i=1 xi) 2] = ∑n i=1 E[x2i ] for independent variables {xi}ni=1 with mean 0. The first inequality follows from E[(x − E[x])2] ≤ E[x2] and the last inequality follows from Lemma D1, whereDR = CR+maxi∈{1,2,··· ,N}N‖∇L(xi|β?)‖+ Cba and R is the radius of a sufficiently large ball that contains β(h)k for h ∈ {1, 2}.
Next, we bound E [∥∥∥β(h)k − β(h)mb km c∥∥∥2 ] as follows
E [∥∥∥∥β(h)k − β(h)mb k m c ∥∥∥∥2 ] ≤ E  ∥∥∥∥∥∥∥ k−1∑ j=mb k
m c
(β (h) j+1 − β (h) j ) ∥∥∥∥∥∥∥ 2 ≤ m k−1∑
j=mb k m c
E [∥∥∥(β(h)j+1 − β(h)j )∥∥∥2] . (19)
For each term, we have the following bound
E [∥∥∥β(h)j+1 − β(h)j ∥∥∥2] =E ∥∥∥∥∥ηNn ∑ i∈Bk ∇L(xi|β(h)k ) + √ 2ητ (h)ξk ∥∥∥∥∥ 2 
≤2η 2N2
n2 ∑ i∈Bk E [∥∥∥∇L(xi|β(h)k )∥∥∥2]+ 4ητ (2)d
≤2η 2
n (2C2E[‖β(h)k ‖ 2] + 2Q2) + 4ητ (2)d
≤2η 2
n (2C2Ψd,τ(2),C,a,b + 2Q 2) + 4ητ (2)d,
(20)
where the first inequality follows by E[‖a+ b‖2] ≤ 2E[‖a‖2] + 2E[‖b‖2], the i.i.d of the data points and τ (1) ≤ τ (2) for h ∈ {1, 2}; the second inequality follows by Lemma D2; the last inequality follows from Lemma D3.
Combining (18), (19) and (20), we have Var ( L̃(Bk|β(h)k ) ) ≤ m 2η
n D2R
( 2η
n (2C2Ψd,τ(2),C,a,b + 2Q 2) + 4τ (2)d
) . (21)
Since Var ( L̃(Bk|β(h)k ) ) ≤ Var ( N n ∑ i∈B L(xi|βk) ) + Var ( N n ∑ i∈B L ( xi ∣∣∣βmb km c)) by definition, Var ( L̃(Bk|β(h)k ) ) is upper bounded by O ( min{σ̂2, m 2η n } ) , which becomes much smaller using a small learning rate η, a shorter period m and a large batch size n.
Note that S̃η,m,n(β(1),β(2)) is defined on the unbounded support [0,∞] and E[S̃η,m,n(β(1),β(2))] = S(β(1),β(2)) regardless of the scale of σ̃2. To satisfy the (stochastic) reversibility condition, we consider the truncated swapping rate min{1, S̃η,m,n(β(1),β(2))}, 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 variance increases.
Lemma B2 (Variance reduction for larger swapping rates) Given a large enough batch size n, the variance-reduced energy estimator L̃(Bk|β(h)k ) yields a truncated swapping rate that satisfies
E[min{1, S̃η,m,n(β(1),β(2))}] ≈ min { 1, S(β(1),β(2)) ( O ( 1
n2
) + e −O ( m2η n + 1 n2 ))} . (22)
Proof By central limit theorem, the energy estimator Nn ∑ i∈B L(xi|βk) converges in distribution to a normal distributions as the batch size n goes to infinity. In what follows, the variance-reduced estimator L̃(Bk|βk) also converges to a normal distribution, where the corresponding estimator is denoted by L̃(Bk|βk). Now the swapping rate Sη,m,n(·, ·) based on normal estimators follows
Sη,m,n(β(1),β(2)) = e ( 1 τ(1) − 1 τ(2) )( L̃(B|β(1))−L̃(B|β(2))− ( 1 τ(1) − 1 τ(2) ) σ̄2 2 ) , (23)
where σ̄2 denotes the variance of L̃(B|β(1)) − L̃(B|β(2)). Note that Sη,m,n(β(1),β(2)) follows a log-normal distribution with mean logS(β(1),β(2)) − ( 1 τ(1) − 1 τ(2) )2 σ̄2 2 and variance(
1 τ(1) − 1 τ(2)
)2 σ̄2 on the log-scale, and S(β(1),β(2)) is the deterministic swapping rate defined
in (13). Applying Lemma D4, we have
E[min{1,Sη,m,n(β(1),β(2))}] = O ( S(β(1),β(2)) exp { − ( 1 τ(1) − 1 τ(2) )2 σ̄2
8
}) . (24)
Moreover, σ̄2 differs from σ̃2, the variance of L̃(B|β(1)) − L̃(B|β(2)), by at most a bias of O( 1n2 ) according to the estimate of the third term of (S2) in Quiroz et al. (2019) and σ̃ 2 ≤
Var ( L̃(Bk|β(1)k ) ) + Var ( L̃(Bk|β(2)k ) ) , where both Var ( L̃(Bk|β(1)k ) ) and Var ( L̃(Bk|β(2)k ) ) are upper bounded by m
2η n D 2 R ( 2η n (2C 2Ψd,τ(2),C,a,b + 2Q 2) + 4τd ) by Lemma B1, it follows that
E[min{1,Sη,m,n(β(1),β(2))}] ≤ S(β(1),β(2))e −O
( m2η n + 1 n2 ) . (25)
Applying min{1,A + B} ≤ min{1,A}+ |B|, we have E[min{1, S̃η,m,n(β(1),β(2))}]
=E [ min {
1, S̃η,m,n(β (1),β(2))− Sη,m,n(β(1),β(2))︸ ︷︷ ︸
B
+Sη,m,n(β(1),β(2))︸ ︷︷ ︸ A }] ≤E
[∣∣∣S̃η,m,n(β(1),β(2))− Sη,m,n(β(1),β(2))∣∣∣]︸ ︷︷ ︸ I +E[min{1,Sη,m,n(β(1),β(2))}]︸ ︷︷ ︸ see formula (25)
(26)
By the triangle inequality, we can further upper bound the first term I E [∣∣∣S̃η,m,n(β(1),β(2))− Sη,m,n(β(1),β(2)∣∣∣]
≤ ∣∣∣E[S̃η,m,n(β(1),β(2))]− S(β(1),β(2))∣∣∣︸ ︷︷ ︸
I1
+ ∣∣∣S(β(1),β(2))− E[Sη,m,n(β(1),β(2))]∣∣∣︸ ︷︷ ︸
I2 = S(β(1),β(2))O ( 1
n2
) + S(β(1),β(2))O ( 1
n2
) ,
(27)
where I1 and I2 follow from the proof of S1 without and with normality assumptions, respectively (Quiroz et al., 2019).
Combining (26) and (27), we have E[min{1, S̃η,m,n(β(1),β(2))}] ≈ min { 1, S(β(1),β(2)) ( O ( 1
n2
) + e −O ( m2η n + 1 n2 ))} . (28)
This means that reducing the update period m (more frequent update the of control variable), the learning rate η and the batch size n significantly increases min{1, S̃η,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 1 (Exponential convergence) Under the smoothness and dissipativity assumptions 1 and 2, the replica exchange Langevin diffusion associated with the variance-reduced stochastic swapping rates Sη,m,n(·, ·) = min{1, S̃η,m,n(·, ·)} converges exponential fast to the invariant distribution π given a smaller learning rate η, a smaller m or a larger batch size n:
W2(νt, π) ≤ D0 exp { −t ( 1 + δSη,m,n ) /cLS } , (29)
where D0 = √ 2cLSD(ν0||π), δSη,m,n := inft>0 ESη,m,n ( √ dνt dπ )
E( √ dνt dπ )
− 1 is a non-negative constant de-
pending on the truncated stochastic swapping rate Sη,m,n(·, ·) and increases with a smaller learning rate η, 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 ×Rd → R, the infinitesimal generator LSη,m,n associated with the replica exchange Langevin diffusion with the swapping rate Sη,m,n = min{1, S̃η,m,n} follows LSη,m,nf(β(1),β(2)) =− 〈∇β(1)f(β(1),β(2)),∇L(β(1))〉 − 〈∇β(2)f(β(1),β(2)),∇L(β(2))〉
+ τ (1)∆β(1)f(β (1),β(2)) + τ (2)∆β(2)f(β (1),β(2))
+ rSη,m,n(β (1),β(2)) · (f(β(2),β(1))− f(β(1),β(2))),
(30)
where∇β(h) and ∆β(h) are the gradient and the Laplace operators with respect to β(h), respectively. Next, we model the exponential decay ofW2(νt, π) using the Dirichlet form
ESη,m,n(f) = ∫ ΓSη,m,n(f)dπ, (31)
where ΓSη,m,n(f) = 1 2 · LSη,m,n(f 2)− fLSη,m,n(f) is the Carré du Champ operator. In particular for the first term 12LSη,m,n(f 2), we have
1 2 LSη,m,n(f(β(1),β(2))2)
=− 〈f(β(1),β(2))∇β(1)f(β(1),β(2)),∇β(1)L(β(1))〉+ τ (1)‖∇β(1)f(β(1),β(2))‖2
+ τ (1)f(β(1),β(2))∆β(1)f(β (1),β(2))
− 〈f(β(1),β(2))∇β(2)f(β(1),β(2)),∇β(2)L(β(2))〉+ τ (2)‖∇β(2)f(β(1),β(2))‖2
+ τ (2)f(β(1),β(2))∆β(2)f(β (1),β(2))
+ r
2 Sη,m,n(β
(1),β(2))(f2(β(2),β(1))− f2(β(1),β(2))).
Combining the definition of the Carré du Champ operator, (30) and (B), we have
ΓSη,m,n(f(β (1),β(2)))
= 1
2 LSη,m,n(f 2(β(1),β(2)))− f(β(1),β(2))LSη,m,n(f(β (1),β(2)))
=τ (1)‖∇β(1)f(β (1),β(2))‖2 + τ (2)‖∇β(2)f(β (1),β(2))‖2
+ r
2 Sη,m,n(β
(1),β(2))(f(β(2),β(1))− f(β(1),β(2)))2.
(32)
Plugging (32) into (31), the Dirichlet form associated with operator LSη,m,n follows ESη,m,n(f) = ∫ ( τ (1)‖∇β(1)f(β (1),β(2))‖2 + τ (2)‖∇β(2)f(β (1),β(2))‖2 ) dπ(β(1),β(2))︸ ︷︷ ︸
vanilla term E(f)
+ r
2
∫ Sη,m,n(β
(1),β(2)) · (f(β(2),β(1))− f(β(1),β(2)))2dπ(β(1),β(2))︸ ︷︷ ︸ acceleration term ,
(33)
where f corresponds to dνt dπ(β(1),β(2)) . Under the asymmetry conditions of νt π(β1,β(2)) and Sη,m,n > 0, the acceleration term of the Dirichlet form is strictly positive and linearly dependent on the swapping rate Sη,m,n. Therefore, ESη,m,n(f) becomes significantly larger as the swapping rate Sη,m,n increases significantly. According to Lemma 5 (Deng et al., 2020), there exists a constant
δSη,m,n = inft>0 ESη,m,n (
√ dνt dπ )
E( √ dνt dπ )
− 1 depending on Sη,m,n that satisfies the following log-Sobolev
inequality for the unique invariant measure π associated with variance-reduced replica exchange Langevin diffusion {βt}t≥0
D(νt||π) ≤ 2 cLS
1 + δSη,m,n ESη,m,n( √ dνt dπ ),
where δSη,m,n increases rapidly with the swapping rate Sη,m,n. By virtue of the exponential decay of entropy (Bakry et al., 2014), we have
D(νt||π) ≤ D(ν0||π)e−2t(1+δSη,m,n )/cLS ,
where cLS is the standard constant of the log-Sobolev inequality asscoiated with the Dirichlet form for replica exchange Langevin diffusion without swaps (Lemma 4 as in Deng et al. (2020)). Next, we upper boundW2(νt, π) by the Otto-Villani theorem (Bakry et al., 2014)
W2(νt, π) ≤ √ 2cLSD(νt||π) ≤ √ 2cLSD(µ0||π)e−t(1+δSη,m,n)/cLS ,
where δSη,m,n > 0 depends on the learning rate η, the period m and the batch size n.
In the above analysis, we have established the relation that δSη,m,n = inft>0 ESη,m,n (
√ dνt dπ )
E( √ dνt dπ )
− 1
depending on Sη,m,n may increase significantly with a smaller learning rate η, a shorter period m and a large batch size n. For more quantitative study on how large δSη,m,n is on related problems, we refer interested readers to the study of spectral gaps in Lee et al. (2018); Dong & Tong (2020); Futami et al. (2020).

heading: C DISCRETIZATION ERROR
text: Consider a complete filtered probability space (Ω,F ,F = (Ft)t∈[0,T ],P) which supports all the random subjects considered in the sequel. With a little abuse usage of notation, the probability measure P (component wise if P is joint probability measure with mutually independent components) would always denote the Wiener measure under which the process (Wt)0≤t≤T is a P-Brownian motion. To be precise, in what follows, we shall denote P := PW × N, where PW is the infinite dimensional 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 α is introduced by swapping the diffusion matrix of the two Langevin dynamics and the jump intensity is defined through the swapping probability in the following sense, which ensures the independence of PW and NS in each time interval [iη, (i + 1)η], for i ∈ N+. The precise definition of the Replica exchange Langevin diffusion (reLD) is given as below. For any fixed learning rate η > 0, we define dβt = −∇G(βt)dt+ Σ(αt)dWt,P (α(t) = j|α(t− dt) = l,β(bt/ηcη) = β) = rS(β)η1{t=bt/ηcη} + o(dt), for l 6= j, (34) where ∇G(β) := ( ∇L(β(1)) ∇L(β(2)) ) , and 1t=bt/ηcη is the indicator function, i.e. for every t = iη with i ∈ N+, given β(iη) = β, we have P (α(t) = j|α(t− dt) = l) = rS(β)η, where S(β) is defined as min{1, S(β(1),β(2))} and S(β(1),β(2)) is defined in (13). In this case, the Markov Chain α(t) is a constant on the time interval [bt/ηcη, bt/ηcη+η) with some state in the finite-state space {0, 1} and the generator matrix Q follows
Q = ( −rS(β)ηδ(t− bt/ηcη) rS(β)ηδ(t− bt/ηcη) rS(β)ηδ(t− bt/ηcη) −rS(β)ηδ(t− bt/ηcη) ) ,
where δ(·) is a Dirac delta function. The diffusion matrix Σ(αt) is thus defined as (Σ(0),Σ(1)) :={(√ 2τ (1)Id 0
0 √ 2τ (2)Id
) , (√ 2τ (2)Id 0
0 √ 2τ (1)Id
)} . 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 [iη, (i+ 1)η). The distribution of process (βt)0≤t≤T is denoted as νT := PG ×NS which is absolutely continuous with respect to the reference measure P := PW × N, under which W is Brownian motion and α(·) is a Poisson process with some constant jump intensity. This fact follows from the result in Gikhman & Skorokhod (1980)[VII, Section 6, Theorem 2] and Yin & Zhu (2010)[Section 2.5, formula (2.40)]. The motivation of only considering the positive swapping rate in iη, for i ∈ N+, and zero elsewhere is due to our construction of the discretized process β̃ as shown below (see equation 35). A simple illustration of the idea can be seen from the auxiliary process construction in Yin & Zhu (2010)[Section 2.5], following which we want to make sure the stopping time of β and β̃ happening at the same time. Otherwise, it is unlikely (and also unreasonable) to derive the Radon-Nikodym derivative of the two process β and β̃. Thus, we should think of the process is concatenated on the time interval [iη, (i + 1)η) up to time horizon T . Similarly, we consider the following Replica exchange stochastic gradient Langevin diffusion, for the same learning rate η > 0 as above, we have dβ̃ηt = −∇G̃(β̃ η bt/ηcη)dt+ Σ(α̃bt/ηcη)dWt, P ( α̃(t) = j|α̃(t− dt) = l, β̃(bt/ηcη) = β̃ ) = rS̃(β̃)η1{t=bt/ηcη} + o(dt), for l 6= j, (35)
where∇G̃(β) := ( ∇L̃(β(1)) ∇L̃(β(2)) ) and S̃(β̃) = min{1, S̃η,m,n(β̃(1), β̃(2))} and S̃η,m,n(β̃(1), β̃(2)) is
shown in (16). The distribution of process (β̃t)0≤t≤T is denoted as µT := PG̃ ×NS̃ , where α̃ is a Poisson process with jump intensity rS̃(β̃)ηδ(t−bt/ηcη) on the time interval [bt/ηcη, bt/ηcη+η). Note that β and β̃ are defined by using the same P-Brownian motion W , but with two different jump intensity on the time interval [bt/ηcη, bt/ηcη + η). Notice that, if there is no jump, the construction of β̃ based on β follows from the fact that they share the same marginal distributions as shown in Gyöngy (1986), where one can find the details in Raginsky et al. (2017). Given the jump process α and α̃ introduced into the dynamics of β and β̃, the construction is more complicated. Thanks to Bentata & Cont (2009), we can carry on the similar construction in our current setting. We then introduce the following Radon-Nikodym density for dνT /dµT . 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 A ∈ Rn, we denote ‖A‖2 := A∗A. Furthermore, we introduce a sequence of stopping time based on our definition of process β and β̃. For j ∈ N+, we denote ζ ′js as a stopping times defined by ζj+1 := inf{t > ζj : α(t) 6= α(ζj)} and N(T ) = max{n ∈ N : ζn ≤ T}. It is easy to see that for any stopping time ζi, there exists l ∈ N+ such that ζj = lη. Similarly, we have the stopping time for the process β̃ by ζ̃j+1 := inf{t > ζ̃j : α̃(t) 6= α̃(ζj)} and α̃(t) follows the same trajectory of α(t). To serve the purpose of our analysis, one should think of the process β as the auxiliary process to the process β̃, see similar constructions in Yin & Zhu (2010)[Section 2.5, formula (2.39)]. The difference is that both of our process β and β̃ are associated with jump process jumping at time iη, for some integer i ∈ N+, instead of jumping at any continuous time. We combine approximation method from Yin & Zhu (2010)[Section 2.7] for non-constant generator matrix Q and the density representation for Markov process in Gikhman & Skorokhod (1980)[VII, Section 6, Teorem 2] to get the following
Lemma C1 Let {ζj |j ∈ {0, 1, · · · , N(T )}} be a sequence of stopping time defined by α. Let k ∈ N+ be an fixed integer such that kη ≤ T ≤ (k + 1)η. For each fixed learning rate η > 0 and for any ε > 0, the Radon-Nikodym derivative of dµT /dνT is given as below,
dµT dνT = exp (N(T )∑ j=0 ∫ ζj+1∧T ζj [ Σ−1(α̃(ζj))∇G̃(βt)− Σ−1(α(ζj))∇G(βt) ] dWGt
− 1 2 N(T )∑ j=0 ∫ ζj+1∧T ζj ∥∥∥Σ−1(α̃(ζj))∇G̃(βt)− Σ−1(α(ζj))∇G(βt)∥∥∥2dt)
× exp − N(T ))∑ j=0 ∫ ζj+1∧T−ε ζj rδ(t− bt/ηcη)[S̃(β̃bt/ηcη)− S(βbt/ηcη)]ηdt ×ΠN(T )j=0 S̃(β̃ζj )S(βζj ) . Proof Recall that ζj is stopping time defined by α (same as defined by α̃), i.e. ζj+1 := inf{t > ζj : α(t) 6= α(ζj)}, for j = 0, 1, · · · , N(T ), and for each ζj , there exists l ∈ {0, 1, · · · , k} such that ζj = lη. We now follow Gikhman & Skorokhod (1980)[VII, Section 6, Theorem 2] to derive the Radon-Nikodym density for dµT /dνT . In this case, if the generator matrixQ is constant, i.e. the jump intensity is constant, we can follow the similar construction from Yin & Zhu (2010)[Formula (2.40)], see also Eizenberg & Freidlin (1990)[Formula(3.13)]. Next, we adjust our setting to the case that we can treat our generator matrix as constant on each time interval [ζj , ζj+1), then the existing results apply to our case for the density with respect to the Poisson measure (jump process α and α̃), i.e. dNS/dNS̃ . Furthermore, once the generator matrix Q is constant, then the measure PG ( or PG̃) is independent to NS (or NS̃). We show the following steps to give a clear outline of our proof.
Step 1: For each stopping time interval [ζj , ζj+1), no jump would occur after the initial point at time ζj and the diffusion matrix Σ and Σ̃ 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 dNS/dNS̃ and dPG/dPG̃, we need the independent property of the two measures on the same time interval, then we directly get the density
following from Gikhman & Skorokhod (1980)[VII, Section 6, Theorem 2]. Different from the work mentioned above, we will first write all the density on each time interval [iη, (i + 1)η) to incorporate the independent requirement mentioned above. Notice that the relative change of density for dNS/dNS̃ would only depends on the left end point, since the jump intensity would change its values at the initial value of interval [iη, (i + 1)η), which is a standard idea to deal with generator matrix Q depending on the initial value instead of a constant matrix case. (See Yin & Zhu (2010)[Section2.7] for similar treatments).
Step 3: In general, the stopping time interval could contain several time interval with length η, however the jump intensity should only depend on the left end point for each time interval [iη, (i+ 1)η). Based on the above set up, we now derive the Radon-Nikodym derivative. First notice that, on each period [ζj , ζj+1), the matrix Σ is fixed and is evaluated at Σ(α(ζj)), which is the same for Σ(α̃(ζj)). In particular, Σ(α(ζj)) = Σ(α̃(ζj)) is a constant diagonal matrix. According to our definition dνT = dPG × dNS and dµT = dPG̃ × dNS̃ , we write the Radon-Nikodym derivative on each of the time interval [iη, (i+ 1)η) and concatenate them together. We consider the swapping of the diffusion matrix first where a similar construction can be found in Yin & Zhu (2010)[Formula (2.40)], we get the following Radon-Nikodym derivative, for any ε > 0,
dNS̃ dNS = exp { − N(T )∑ j=0 ∫ (j+1)η∧T−ε jη rδ(t− bt/ηcη)(S̃(β̃βbt/ηcη )− S(ββbt/ηcη ))ηdt }
(36)
×ΠN(T )j=0 S̃(β̃ζj )
S(βζj ) .
Next, we show the density for dPG/dPG̃ as below. On each interval [ζj , ζj+1), given initial value (βj , β̃j), the matrix Σ(α(ζj)) and Σ(α̃(ζj)) are always the same, since no jump would happen. In particular, in this continuous case the integral on [ζj , ζj+1) and [ζj , ζj+1] are the same. Thus we have the following Radon-Nikodym derivative
dPG̃ dPG = exp (N(T )∑ j=0 ∫ ζj+1∧T ζj [ Σ−1(α̃(ζj))∇G̃(βt)− Σ−1(α(ζj))∇G(βt) ] dWGt
− 1 2 N(T )∑ j=0 ∫ ζj+1∧T ζj ∥∥∥Σ−1(α̃(ζj))∇G̃(βt)− Σ−1(α(ζj))∇G(βt)∥∥∥2dt). (37) Notice that matrix Σ is diagonal square matrix, thus we have Σ = Σ∗. Recall that W is a PBrownian motion, assuming there is no jump in the dynamic for β, then according to the Girsanov theorem (see an example in Theorem 8.6.6 and Example 8.6.9 (Øksendal, 2003)) with RadonNikodym derivative dPG/dP, we have the PG-Brownian motion, denoted asWG, which follows
WGt := Wt + ∫ t 0 Σ−1(αs)(∇G(βs))ds. (38)
This fact holds true on each of the time interval [ζj , ζj+1]. Multiplying the two density dPG/PG̃ and dNS/dNS̃ , we complete the proof.
Remark 1 Notice that, if we keep the constant diffusion matrix without jump, then the RandonNikodym derivative dµT /dνT has been used in the stochastic gradient descent setting, for example Raginsky et al. (2017). 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 Σ, we have
dPG̃ dPG = exp (∫ T 0 [ Σ−1∇G̃(β̃s)− Σ−1∇G(βs) ] dWGs
− 1 2 ∫ T 0 ∥∥∥Σ−1∇G̃(β̃s)− Σ−1∇G(βs)∥∥∥2ds), (39) whereWG is a PG-Brownian motion as shown in equation 38, not a P-Brownian motionW .
Remark 2 The density dµTdνT that we derived above is so far the best we can do. If one would like to use the continuous time control α(t) with continuous jump intensity S(β(t)) instead of jumping at the initial point with a fixed rate, then we can not even write the Randon-Nikodym derivative anymore, since α(t) and α̃(t) will define different stopping time, i.e. jump at different time and µT is not absolutely continuous with respect to νT .
Based on the above lemma, we further get the following estimates.
Lemma C2 Given a large enough batch size n or a small enough m and η, we have the bound of the KL divergence of DKL(µT |νT ) as below,
DKL(µT |νT ) ≤ (Φ0 + Φ1η)kη +N(T )Φ2, with
Φ0 = O ( m√ n √ ηd ) + rδΦ2 4τ (1) ,
Φ1 = ( C2d τ (2)
τ (1) + C2δkd 2τ (1) [τ (1) + τ (2)]
) ,
Φ2 = O ( m√ n √ ηd ) .
Proof By the very definition of the KL-divergence, we have DKL(µT |νT ) =− ∫ dνT log
dµT dνT
=− EνT [ log(dµT /dνT ) ∣∣∣(β, β̃) = (β, β̃)].
We shall keep the convention below and denote EνT ,β = EνT [·|(β, β̃) = (β, β̃)], where β = (β(1), β(2)) ∈ R2d and β̃ = (β̃(1), β̃(2)) ∈ R2d denotes the values at each time iη, i = 0, 1, · · · , k. Plugging Lemma C1 in the above equation and we unify the notation by using time intervals of the type [iη, (i+ 1)η]. To be precise, we get
dPG̃ dPG = exp ( k−1∑ i=0 ∫ (i+1)η iη [ Σ−1(α̃(iη))∇G̃(βt)− Σ−1(α(iη))∇G(βt) ] dWGt
+ ∫ T kη [ Σ−1(α̃(kη))∇G̃(βt)− Σ−1(α(kη))∇G(βt) ] dWGt
− 1 2 k−1∑ i=0 ∫ (k+1)η iη ∥∥∥Σ−1(α̃(iη))∇G̃(βt)− Σ−1(α(iη))∇G(βt)∥∥∥2dt − 1
2 ∫ T kη ∥∥∥Σ−1(α̃(kη))∇G̃(βt)− Σ−1(α(kη))∇G(βt)∥∥∥2dt). (40) The above equality follows from the fact that each time interval [ζj , ζj+1] always contain exactly some sub-interval [iη, (i+ 1)η]. Namely, we have [ζj , ζj+1] = [iη, (i+ 1)η]∪ [(j+ 1)η, (j+ 2)η]∪ · · · ∪ [lη, (l + 1)η], for some i, l ∈ {0, 1, · · · , k}. In particular, the matrix Σ keep the same on each interval [iη, (i + 1)η], for some i ∈ {0, 1, · · · , k}. Similarly, we expand the Radon-Nokodym derivative for dN S̃
dNS on the time interval of length η. Based on our definition of jump intensity, we get
dNS̃ dNS = exp { − N(T )∑ j=0 ∫ (j+1)η∧T−ε jη rδ(t− bt/ηcη)(S̃(β̃bt/ηcη)− S(βbt/ηcη))ηdt
− ∫ T kη rδ(s− bs/ηcη)(S̃(β̃kη)− S(βkη))ηds } ×ΠN(T )j=0 ( S̃(β̃ζj ) S(βζj ) )
= exp { −
k∑ i=0 r(S̃(β̃iη)− S(βiη))η } ×ΠN(T )j=0 ( S̃(β̃ζj ) S(βζj ) ) . (41)
Without loss of generality, we shall only consider the sum ∑k−1 i=0 and skip the interval [kη, T ]. Notice that on each time interval [iη, (i + 1)η), the control α(iη) and α̃(iη) are fixed, thus the two component of the measure dνT,β are independent. Taking into account the fact that WG is PG-Brownian motion, thus we apply the martingale property and arrive at
DKL(µT |νT )
=EνT ,β [1
2 k−1∑ i=0 ∫ (i+1)η iη ∥∥∥Σ−1(α̃(iη))∇G̃(βt)− Σ−1(α(iη))∇G(βt)∥∥∥2dt]
+ EνT ,β [ k−1∑ i=0 [S̃(β̃iη)− S(βiη)]η − N(T )∑ j=0 ( log S̃(β̃ζj )− logS(βζj ) )]
≤ 1 2 k−1∑ i=0 EνT ,β [ ∫ (i+1)η iη ∥∥∥Σ−1(α̃(iη))∇G̃(βt)− Σ−1(α(iη))∇G(βt)∥∥∥2dt]︸ ︷︷ ︸ I
+ k−1∑ i=0 EνT ,β [ r|S̃(β̃iη)− S(βiη)|η ] ︸ ︷︷ ︸
J
+ N(T )∑ j=0 EνT ,β [ | log S̃(β̃ζj )− logS(βζj )| ] ︸ ︷︷ ︸
K
.
We then estimates the three terms I,J ,K in order as below. Estimate of I: Due to the fact that every interval [iη, (i + 1)η) ⊂ [ζj , ζj+1) for some j ∈ {0, 1, · · · , N(T )}, we know that the control α and α̃ are the same in the interval [iη, (i + 1)η] and the diffusion matrix Σ is just constant matrix. Thus, we know that matrix Σ−1(α̃(iη)) = Σ−1(α(iη)), which takes one of the form from (Σ−1(0),Σ−1(1)) :={(
1√ 2τ(1) Id 0
0 1√ 2τ(2) Id
) , ( 1√ 2τ(2) Id 0
0 1√ 2τ(1) Id
)} . If Σ−1(α(iη)) = Σ−1(0), we get
∥∥∥Σ−1(α̃(iη))∇G̃(βt)− Σ−1(α(iη))∇G(βt)∥∥∥2 =
d∑ j=1 1 2τ (1) |∇jG̃(βt)−∇jG(βt)|2 + 2d∑ j=d+1 1 2τ (2) |∇jG̃(βt)−∇jG(βt)|2
≤ 1 2τ (1) 2d∑ j=1 |∇jG̃(βt)−∇jG(βt)|2 ≤ 1 2τ (1) ‖∇G̃(βt)−∇jG(βt)‖2.
Here ∇G(β) := ( ∇L(β(1)) ∇L(β(2)) ) and ∇G̃(β) := ( ∇L̃(β(1)) ∇L̃(β(2)) ) . The other matrix form of Σ−1(1) will result in the same estimates. We thus get
I ≤ 1 4τ (1) k−1∑ i=0 EνT ,β [ ∫ (i+1)η iη ∥∥∥∇G̃(βt)−∇G(βt)∥∥∥2dt] On each fixed interval, for t ∈ [kη, (k + 1)η) , we have PG-Brownian motion and PG̃-Brownian motion (see examples in Theorem 8.6.6 and Example 8.6.9 (Øksendal, 2003)),
dWGt =dWt + Σ −1(αt)(∇G(βt))dt.
dW G̃t =dWt + Σ −1(αt)(∇G̃(β̃t))dt.
Plugging the PG (and PG̃)-Brownian motions to the original dynamics (34) and (35), we have
dβt = Σ(αt)dW G t , and dβ̃t = Σ(αt)dW G̃ t .
On each interval [iη, (i+1)η), Σ(αt) is a constant matrix, thus we know that the probability distribution of {βt}t∈[kη,(k+1)η) and {β̃t}t∈[kη,(k+1)η) are the same and we denote as L(βt) = L(β̃t). The difference is that βt is driven by PG-Brownian motion and β̃t is driven by PG̃-Brownian motion, which implies that, for t ∈ [iη, (i+ 1)η), we have
EνT ,β [∥∥∥∇G̃(βt)−∇G(βt)∥∥∥2] = EµT ,β̃[∥∥∥∇G̃(β̃t)−∇G(β̃t)∥∥∥2]. (42)
Thus, we have the following estimates,
I ≤ 1 4τ (1) k−1∑ i=0 EµT ,β̃ [ ∫ (i+1)η iη ∥∥∥∇G(β̃t)−∇G(β̃bt/ηcη)∥∥∥2dt] + 1
4τ (1) k−1∑ i=0 EµT ,β̃ [ ∫ (i+1)η iη ∥∥∥∇G(β̃bt/ηcη)−∇G̃(β̃bt/ηcη)∥∥∥2dt] ≤ C 2
4τ (1) k−1∑ i=0 EµT ,β̃ [ ∫ (i+1)η iη ∥∥∥β̃t − β̃iη)∥∥∥2dt] · · · I1 + 1
4τ (1) k−1∑ i=0 EµT ,β̃ [ ∫ (i+1)η iη ∥∥∥∇G(β̃bt/ηcη)−∇G̃(β̃bt/ηcη)∥∥∥2dt] · · · I2. We now estimate the two terms I1 and I2 separately. Notice that, following our notation of PG̃Brownian motion, for t ∈ [iη, (i+ 1)η), we have
β̃t − β̃iη = Σ(αt)(W G̃t −W G̃iη ) = Σ(αt)(W G̃t −W G̃iη ),
which implies that (recall that dµT = dPG̃ ×NS̃ and Σ ∈ R2d×2d),
Eµ T,β̃ [‖β̃t − β̃iη‖2] ≤ 2τ (1)dη + 2τ (2)dη ≤ 4τ (2)dη.
We thus conclude that,
I1 ≤ C2 τ (2)
τ (1) kdη2.
As for the term I2, according to Assumption 3, we obtain that
I2 ≤ ηδ
4τ (1) k−1∑ i=0 EµT ,β̃ [ C2‖β̃iη‖2 + Φ2 ] .
Now, we just need to estimate EµT ,β̃[‖β̃kη‖ 2] †. On each interval [iη, (i+ 1)η], under the measure dµT,β̃, we have
β̃(i+1)η = β̃iη + Σ(α(iη))(W G̃ (i+1)η −W G̃ iη ),
which implies that
EµT ,β̃[‖β̃(i+1)η‖ 2]
=EµT ,β̃[‖β̃iη‖ 2] + EµT ,β̃[〈β̃iη,W G̃ (i+1)η −W G̃ iη 〉] + EµT ,β̃[‖W G̃ (i+1)η −W G̃ iη‖2]
=EµT ,β̃[‖β̃iη‖ 2] + [2τ (1) + 2τ (2)]dη
The last equality follows from the independence of β̃kη and W G̃(k+1)η −W G̃ kη , and W G̃ is a PG̃Brownian motion. By induction, we get
EµT ,β̃[‖β̃iη‖ 2] ≤ 2id[τ (1) + τ (2)]η ≤ 2kd[τ (1) + τ (2)].
†In principle, the Wiener measureW under PG̃ is not a Brownian motion, thus the uniform L2 bound used in Lemma.3 may not be appropriate. Instead, we estimate the upper bound using a slightly weaker result.
We conclude that,
I2 ≤ kηδ
4τ (1)
( 2C2[τ (1) + τ (2)]kdη + Φ2 ) ,
which implies that
I ≤ kη 4τ (1)
( 2δC2[τ (1) + τ (2)]kdη + δΦ2 ) + C2 τ (2)
τ (1) kdη2.
Estimate J : According to our definition of the swapping probability, we have, for each i,
S̃(β̃iη) = min{1, S̃η,m,n(β̃(1)iη , β̃ (2) iη )}, S(βiη) = min{1, S(β (1) iη ,β (2) iη )},
which means |S̃(β̃iη)− S(βiη)| ≤ 1. Denote Cτ = | 1τ(1) − 1 τ(2) |, we have
S̃η,m,n(β̃ (1) iη , β̃ (2) iη ) = exp ( Cτ (L̃(Biη|β(1)iη )− L̃(Biη|β (2) iη ))− C 2 τ σ̃2
2 ) S(β
(1) iη ,β (2) iη ) = exp ( Cτ (L(β (1) iη )− L(β (2) iη )) ) .
Applying Taylor expansion for the exponential function at Cτ (L(β (1) kη )− L(β (2) kη )), we have
EνT ,β [ |S̃η,m,n(β̃(1)iη , β̃ (2) iη )− S(β (1) iη ,β (2) iη )| ] .EνT ,β [ S(β (1) iη ,β (2) iη ) ∣∣∣Cτ (L̃(Bkη|β(1)iη )− L̃(Bkη|β(2)iη ))− C2τ σ̃22 − Cτ (L(β(1)iη )− L(β(2)iη ))∣∣∣+ higher order term]
≤EνT ,β [∣∣∣Cτ (L̃(Biη|β(1)iη )− L̃(Biη|β(2)iη ))− C2τ σ̃22 − Cτ (L(β(1)iη )− L(β(2)iη ))∣∣∣+O(σ̃2)]
where the last inequality follows from S(β(1)iη ,β (2) iη ) ≤ 1. Combining Lemma B1, we thus get the following estimates, J = k−1∑ i=0 EνT ,β [ r|S̃η,m,n(β̃iη)− S(βiη)|η
] ≤ rη
k−1∑ i=0 EνT ,β [∣∣∣Cτ (L̃(Biη|β(1)iη )− L̃(Biη|β(2)iη ))− C2τ σ̃22 − Cτ (L(β(1)iη )− L(β(2)iη ))∣∣∣+O(σ̃2)]
≤ rkηO(Cτ σ̃ + σ̃2) = rkηO
(( m2
n η
)1/2 d ) where the last inequality follows from the Jensen’s inequality and the last order holds given a large enough batch size n or a small enough m and η.
Estimate K: We now estimate the last term K, we have
K = N(T )∑ j=0 EνT ,β [ | log S̃η,m,n(β̃ζj )− logS(βζj )| ]
≤Cτ N(T )∑ j=0 EνT ,β [∣∣∣[L̃(Bζj |β(1)ζj )− L̃(Bζj |β(2)ζj )− Cτ σ̃22 ]− [L(β(1)ζj )− L(β(2)ζj )]∣∣∣] ≤N(T )C2τEνT ,β[σ̃2/2] + Cτ N(T )∑ j=1 Var[L̃(Bζj |β (1) ζj )− L̃(Bζj |β (2) ζj )]1/2
≤N(T )C 2 τ σ̃ 2
2 +N(T )Cτ σ̃
Combining Lemma B1 again, we conclude with
K ≤ C2τ N(T )σ̃2
2 +N(T )Cτ σ̃ = N(T )O
(( m2
n η
)1/2 d ) .
Combining the estimates of I, J , and K, we complete the proof.
Remark 3 After the change of measure, the expectation is under the new measure PG (or PG̃) instead of the Wiener measure P. In the estimate of term I, similar L2 estimates of the term EµT ,β̃[‖β̃(i+1)η‖ 2] has been obtained in Raginsky et al. (2017)[Proof of Lemma 7] when there is no swap. The difference is we write the dynamic of β̃(i+1)η with respect to the PG̃-Brownian motion W G̃ instead of the P-Brownian motion W . In principle, W under PG̃ is not a Brownian motion.
We then extend the distance of relative entropy DKL(µT |νT ) to the Wasserstein distance W2(µT , νT ) via a weighted transportation-cost inequality of Bolley & Villani (2005).
Theorem 2 Given a large enough batch size n or a small enough m and η, we have
W2(µT , νT ) ≤ O ( dk3/2η ( η1/4 + δ1/4 + ( m2
n η
)1/8)) . (43)
Proof Before we proceed, we first show in Lemma.D5 that νT has a bounded second moment; the L2 upper bound of µT is majorly proved in Lemma.C2 (Chen et al., 2019) except that the slight difference is that the constant in the RHS of (C.38) Chen et al. (2019) is changed to account for the stochastic noise. Then applying Corollary 2.3 in Bolley & Villani (2005), we can upper bound the two Borel probability measures µT and νT with finite second moments as follows
W2(µT , νT ) ≤ Cν [√ DKL(µT |νT ) + ( DKL(µT |νT )
2
)1/4] , (44)
where Cν = 2 infλ>0 ( 1 λ ( 3 2 + log ∫ Rd e λ‖w‖2ν(dw) ))1/2 . Applying Lemma D6, we have
W22 (µT , νT ) ≤ ( 12 + 8 ( κ0 + 2b+ 4dτ (2) ) kη )( DKL(µT |νT ) + √ DKL(µT |νT ) ) .
Combining Lemma.C2 and √ N(T ) ≤ N(T ) and taking η ≤ 1, kη > 1, and λ = 1, we have
W22 (µT,β̃, νT,β) ≤ ( 12 + 8 ( κ0 + 2b+ 4dτ (2) )) kη ( (Φ̃0 + Φ̃1 √ η)kη +N(T )Φ̃2 ) ,
where Φ̃i = Φi + √ Φi for i ∈ {0, 1, 2}. In what follows, we have
W22 (µT,β̃, νT,β) ≤ (Ψ0 + Ψ1 √ η) (kη)2 + Ψ2kηN(T ),
where Ψi = ( 12 + 8 ( κ0 + 2b+ 4dτ (2) ))
Φ̃i for i ∈ {0, 1, 2}. By the orders of Φ0, Φ1 and Φ2 defined in Lemma.C2, we have
W22 (µT,β̃, νT,β) ≤ O
( d2k3η2 ( η1/2 + δ1/2 + ( m2
n η
)1/4 + N(T )
kη
( m2
n η
)1/4))
≤ O ( d2k3η2 ( η1/2 + δ1/2 + ( m2
n η
)1/4)) ,
where N(T )kη can be interpreted as the average swapping rate from time 0 to T and is of order O(1). Taking square root to both sides of the above inequality lead to the desired result (43).

heading: D PROOF OF TECHNICAL LEMMAS
text: Lemma D1 (Local Lipschitz continuity) Given a d-dimensional centered ball U of radiusR, L(·) is DR-Lipschitz continuous in that |L(xi|β1) − L(xi|β2)| ≤ DRN ‖β1 − β2‖ for ∀β1,β2 ∈ U and any i ∈ {1, 2, · · · , N}, where DR = CR+ maxi∈{1,2,··· ,N}N‖∇L(xi|β?)‖+ Cba .
Proof
For any β1,β2 ∈ U , there exists β3 ∈ U that satisfies the mean-value theorem such that
|L(xi|β1)− L(xi|β2)| = 〈∇L(xi|β3),β1 − β2〉 ≤ ‖∇L(xi|β3)‖ · ‖β1 − β2‖,
Moreover, by Lemma D2, we have
|L(xi|β1)− L(xi|β2)| ≤ ‖∇L(xi|β3)‖ · ‖β1 − β2‖ ≤ CR+Q
N ‖β1 − β2‖.
Lemma D2 Under the smoothness and dissipativity assumptions 1, 2, for any β ∈ Rd, it follows that
‖∇L(xi|β)‖ ≤ C N ‖β‖+ Q N . (45)
where Q = maxi∈{1,2,··· ,N}N‖∇L(xi|β?)‖+ bCa .
Proof According to the dissipativity assumption, we have
〈β?,∇L(β?)〉 ≥ a‖β?‖2 − b, (46)
where β? is a minimizer of∇L(·) such that∇L(β?) = 0. In what follows, we have ‖β?‖ ≤ ba . Combining the triangle inequality and the smoothness assumption 1, we have
‖∇L(xi|β)‖ ≤CN‖β − β?‖+ ‖∇L(xi|β?)‖ ≤ CN‖β‖+ CNb
a + ‖∇L(xi|β?)‖. (47)
Setting CN = CN as in (11) and Q = maxi∈{1,2,··· ,N} ‖∇L(xi|β?)‖+ bC a completes the proof.
The following lemma is majorly adapted from Lemma C.2 of Chen et al. (2019), except that the corresponding constant in the RHS of (C.38) is slightly changed to account for the stochastic noise. A similar technique has been established in Lemma 3 of Raginsky et al. (2017).
Lemma D3 (Uniform L2 bounds on replica exchange SGLD) Under the smoothness and dissipativity assumptions 1, 2. Given a small enough learning rate η ∈ (0, 1 ∨ aC2 ), there exists a positive constant Ψd,τ(2),C,a,b <∞ such that supk≥1 E[‖βk‖2] < Ψd,τ(2),C,a,b.
Lemma D4 (Exponential dependence on the variance) Assume S is a log-normal distribution with mean u − 12σ 2 and variance σ2 on the log scale. Then E[min(1, S)] = O(eu−σ 2
8 ), which is exponentially smaller given a large variance σ2.
Proof For a log-normal distribution S with mean u − 12σ 2 and variance σ2 on the log scale, the probability density fS(S) follows that 1S √ 2πσ2 exp { − (logS−u+ 1 2σ 2)2 2σ2 } . In what follows, we have
E[min(1, S)] = ∫ ∞ 0 min(1, S)fS(S)dS = ∫ ∞ 0 min(1, S) 1 S √ 2πσ2 exp { − (logS − u+ 1 2 σ2)2 2σ2 } dS
By change of variable y = logS−u+ 1 2σ 2
σ where S = e σy+u− 12σ 2 and y = −uσ + σ 2 given S = 1, it
follows that
E[min(1, S)]
= ∫ 1 0 S 1 S √ 2πσ2 exp { − (logS − u+ 12σ 2)2 2σ2 } dS + ∫ ∞ 1
1
S √ 2πσ2 exp
{ − (logS − u+ 12σ 2)2
2σ2
} dS
= ∫ −uσ+σ2 −∞ 1√ 2πσ2 e− y2 2 σeu− 1 2σ 2+σydy + ∫ ∞ −uσ+ σ 2 1√ 2πσ2 e−σy−u+ 1 2σ 2 e− y2 2 σeu− 1 2σ 2+σydy
=eu ∫ −uσ+σ2 −∞ 1√ 2π e− (y−σ)2 2 dy + 1 σ ∫ ∞ −uσ+ σ 2 1√ 2π e− y2 2 dy
=eu ∫ ∞ u σ+ σ 2 1√ 2π e− z2 2 dz + 1 σ ∫ ∞ −uσ+ σ 2 1√ 2π e− y2 2 dy
≤eu ∫ ∞ −uσ+ σ 2 1√ 2π e− z2 2 dz + 1 σ ∫ ∞ −uσ+ σ 2 1√ 2π e− y2 2 dy
≤ ( eu + 1
σ
) e− (−u σ +σ 2 )2 2 . eu− σ2 8 ,
where the last equality follows from the change of variable z = σ− y and the second last inequality follows from the exponential tail bound of the standard Gaussian variable P(y > ) ≤ e− 2 2 .
Lemma D5 (Uniform L2 bound on replica exchange Langevin diffusion) For all η ∈ (0, 1 ∧ a 4C2 ), we have that
E[‖(β(1)t ,β (2) t )‖2] ≤ E[e‖β (1) 0 ,β (2) 0 ‖ 2 ] + b+ 2dτ (2)
a .
Proof Consider Lt(βt) = ‖βt‖2, where βt = (β(1)t ,β (2) t ) ∈ R2d. The proof is marjorly adapted from Lemma 3 in Raginsky et al. (2017), except that the generalized Itô formula (formula 2.7 in page 29 of Yin & Zhu (2010)) is used to handle the jump operator, which follows that
dLt =− 2〈βt,∇G(βt)〉+ 2d(τ (1) + τ (2))dt+ 2βTt Σ(αt)dW (t)
+ rSη,m,n(β (1) t ,β (2) t ) · (Lt(β(2),β(1))− Lt(β (1) t ,β (2) t ))︸ ︷︷ ︸
Jump-inducing drift
+M1(t) +M2(t),
where ∇G(β) := ( ∇L(β(1)) ∇L(β(2)) ) and M1(t) and M2(t) are two martingales defined in formula 2.7
in Yin & Zhu (2010)). Due to the definition of Lt(βt), we have Lt(β (1) t ,β (2) t ) = Lt(β (2) t ,β (1) t ), which implies that the Jump-inducing drift actually disappears. Taking expectations and applying the margingale property of the Itô integral, we have the almost the same upper bound as Lemma 3 in Raginsky et al. (2017). Combining E[‖β0‖2] ≤ logE[e‖β0‖ 2 ] completes the proof.
Lemma D6 (Exponential integrability of replica exchange Langevin diffusion) For all τ ≤ 2a , it follows that
logE[e‖(β (1) t ,β (2) t )‖ 2 ] ≤ logE[e‖(β (1) 0 ,β (2) 0 )‖ 2 ]︸ ︷︷ ︸ κ0 +2(b+ 2dτ (2))t.
Proof The proof is marjorly adapted from Lemma 4 in Raginsky et al. (2017). The only difference is that the generalized Itô formula (formula 2.7 in Yin & Zhu (2010)) is used again as in Lemma D5. Consider L(t,βt) = e‖βt‖ 2 , where β = (β(1)t ,β (2) t ) ∈ R2d. Due to the special structure that L(t,βt) is invariant under the swaps of (β (1) t ,β (2) t ), the generator of L(t,βt) 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. (2017).
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 γ captures the trend better but becomes less robust.
Input Initial parameters β(1)0 and β (2) 0 , learning rate η and temperatures τ (1) and τ (2), correction factor F .
repeat Parallel sampling Randomly pick a mini-batch set Bk of size n.
β (h) k = β (h) k−1 − η
N
n ∑ i∈Bk ∇L(xi|β(h)k−1) + √ 2ητ (h)ξ (h) k , for h ∈ {1, 2}.
Variance-reduced energy estimators Update L̂(h) = ∑N i=1 L ( xi ∣∣∣β(h) mb k
m c
) every m iterations.
L̃(Bk|β(h)k ) = N
n ∑ i∈Bk L(xi|β(h)k ) + c̃k · N n ∑ i∈Bk L ( xi ∣∣∣β(h) mb k m c ) − L̂(h)  , for h ∈ {1, 2}. if k mod m = 0 then
Update σ̃2k = (1− γ)σ̃2k−m + γσ2k, where σ2k is an estimate for Var ( L̃(Bk|β(1)k )− L̃(Bk|β (2) k ) ) .
Update c̃k = (1− γ)c̃k−m + γck, where ck is an estimate for − Cov ( L(B|β(h) k ),L ( B|β(h) mb k m c )) Var ( L ( B|β(h)
mb k m c )) . end if Bias-reduced swaps Swap β(1)k+1 and β (2) k+1 if u < S̃η,m,n, where u ∼ Unif [0, 1], and S̃η,m,n follows
S̃η,m,n = exp {( 1
τ(1) − 1 τ(2) )( L̃(Bk+1|β(1)k+1)− L̃(Bk+1|β (2) k+1)− 1 F ( 1 τ(1) − 1 τ(2) ) σ̃2 mb k
m c
)} .
until k = kmax. Output: {β(1)iT } bkmax/Tc i=1 , where T is the thinning factor.

heading: E MORE EMPIRICAL STUDY ON IMAGE CLASSIFICATION
text: E.1 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 m = 50 is even more time-consuming 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 m = 392, which still obtains 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 is only increased by 8% while the variance reduction can be still at most 6 times on CIFAR10 and 10 times on CIFAR100.
E.2 ADAPTIVE COEFFICIENT
We study the correlation coefficient of the noise from the current parameter β(h)k , where h ∈ {1, 2}, and the control variate β(h)
mb km c . As shown in Fig.5, the correlation coefficients are only around -0.5
due to the large learning rate in the early period. This implies that VR-reSGHMC may overuse the noise from the control variates and thus fails to 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 time-varying correlation coefficients and present it in Algorithm 2.
As a result, we can further improve the performance of variance reduction by as much as 40% on CIFAR10 and 30% on CIFAR100 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 is still 10% improvement compared to the standard VR-reSGHMC.
In a 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.

heading: F MORE EMPIRICAL STUDY ON UNCERTAINTY QUANTIFICATION
text: 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.2, we see that the scores before calibration in the seen CIFAR10 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 MSGD still performs the worst for frequently making over-confident predictions; SGHMC performs better but is 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 highly-optimized BS score based on reSGHMC from 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 seen dataset is still maintained. Together with the lowest BS score in the unseen SVHN dataset, cVR-reSGHMC shows its strength in uncertainty quantification.

heading: G MODIFIED EXAMPLE 5.1
text: We revisit Example 5.1, and re-run the procedures with temperature τ (1) = 1.0. In Fig. 6, we present trace plots and kernel density estimates (KDE) of samples generated from VR-reSGLD, reSGLD, and SGLD. In particular, we run VR-reSGLD with m = 40, τ (1) = 1, τ (2) = 500, η = 1e− 5, and F = 1; reSGLD with the same hyper-parameters as VR-reSGLD except for F = 500; and SGLD with η = 1e − 5 and τ = 1. Note that here, we run reSGLD with a greater F than in Example 5.1 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
†BS = 1 N ∑N i=1 ∑R j=1(fij − oij)
2, where fi is the predictive probability and oi is actual output of the event which is 1 if it happens and 0 otherwise; N is the number of instances and R is the number of classes.
with long enough iterations. In Figs 6(a) and 6(b), we observe that, even though the distribution of interest has a pickier density, our 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.

Review summaries:
The 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. 

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.