N ON - REVERSIBLE P ARALLEL T EMPERING FOR U NCERTAINTY A PPROXIMATION IN D EEP L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Parallel tempering (PT), also known as replica exchange, is the go-to workhorse for simulations of multi-modal distributions.
The key to the success of PT is to adopt efﬁcient swap schemes.
The popular deterministic even-odd (DEO) scheme exploits the non-reversibility property and has successfully reduced the communication cost from O ( P 2 ) to O ( P ) given sufﬁcient many P chains.
However, such an innovation largely disappears in big data problems due to the limited chains and extremely few bias-corrected swaps.
To handle this issue, we generalize the DEO scheme to promote the non-reversibility and obtain an appealing communication cost O ( P log P ) based on the optimal window size.
In addition, we also analyze the bias when we adopt stochastic gradient descent (SGD) with large and constant learning rates as exploration kernels.
Such a user-friendly nature enables us to conduct large-scale uncertainty approximation tasks without much tuning costs.
I NTRODUCTION
Langevin diffusion is a standard sampling algorithm that follows a stochastic differential equation
[Equation]
where β t ∈ R d , ∇ U ( · ) is the gradient of the energy function U ( · ) ,
W t ∈ R
d is a Brownian motion, and τ is the temperature.
The diffusion process converges to a stationary distribution π ( β ) ∝ e − U ( β ) τ and setting τ = 1 yields a Bayesian posterior.
When U ( · ) is convex, the rapid convergence has been widely studied in Durmus & Moulines (2016); Dalalyan (2017); however, when U ( · ) is non-convex, a slow mixing rate is inevitable (Raginsky et al., 2017).
To accelerate the simulation, replica exchange
Langevin diffusion (reLD) proposes to include a high-temperature particle β ( P ) t , where P ∈ N + \{ 1 } , for exploration .
Meanwhile, a low-temperature particle β (1) t is presented for exploitation :
[Equation]
where τ ( P ) > τ (1) and W ( P ) t is independent of W (1) t .
To promote more explorations for the low-temperature particle, the particles at the position ( β (1) , β ( P ) ) ∈ R 2 d swap with a probability aS
[Equation]
where a ∈ (0 , ∞ ) is the swap intensity.
In speciﬁc, the conditional swap rate at time t follows that
[Equation]
In the longtime limit, the Markov jump process converges to the joint distribution π ( β (1) , β ( P ) )
∝ e − U ( β (1))
τ (1) − U ( β ( P ))
τ ( P ) .
For convenience, we refer to the marginal distribution π (1) ( β ) ∝ e − U ( β ) τ (1) and π ( P ) ( β )
∝ e − U ( β ) τ ( P ) as the target distribution and reference distribution , respectively.
P RELIMINARIES
Achieving sufﬁcient explorations requires a large τ ( P ) , which leads to limited accelerations due to a small overlap between π (1) and π ( P ) .
To tackle this issue, one can bring in multiple particles with temperatures ( τ (2) , · · · , τ ( P − 1) ) , where τ (1) < τ (2) < · · · < τ ( P ) , to hollow out “tunnels”.
To maintain feasibility, numerous schemes are presented to select candidate pairs to attempt the swaps.
APE The all-pairs exchange (APE) attempts to swap arbitrary pair of chains (Brenner et al., 2007;
Lingenheil et al., 2009), however, such a method requires a swap time (see deﬁnition in section A.5) of O
( P 3 ) and may not be user-friendly in practice.
ADJ
In addition to swap arbitrary pairs, one can also swap adjacent (ADJ) pairs iteratively from (1 , 2) , (2 , 3) , to ( P − 1 , P ) under the Metropolis rule.
Despite the convenience, the sequential nature requires to wait for exchange information from previous exchanges, which only works well with a small number of chains and has greatly limited its extension to a multi-core or distributed context.
SEO The stochastic even-odd (SEO) scheme ﬁrst divides the adjacent pairs { ( p − 1 , p )
| p = 2 , · · · , P } into E and O , where E and O denote even and odd pairs of forms (2 p − 1 , 2 p ) and (2 p, 2 p + 1) , respectively.
Then, SEO randomly picks E or O pairs with an equal chance in each iteration to attempt the swaps.
Notably, it can be conducted simultaneously without waiting from other chains.
The scheme yields a reversible process (see Figure 1(a)), however, the gains in overcoming the sequential obstacle don’t offset the O ( P 2 ) round trip time and SEO is still not effective enough.
[Figure]
DEO
The deterministic even-odd (DEO) scheme instead attempts to swap even ( E ) pairs at even ( E ) iterations and odd ( O ) pairs at odd ( O ) iterations alternatingly † (Okabe et al., 2001).
The asymmetric manner was later interpreted as a non-reversible PT (Syed et al., 2021) and an ideal index process follows a periodic orbit, as shown in Figure 1(b).
With a large swap rate, Figure 1(c) shows how the scheme yields an almost straight path and a linear round trip time can be expected.
Equi-acceptance
The power of PT hinges on maximizing the number of round trips, which is equivalent to minimizing (cid:80) P − 1 p = 1 1 − r p (Nadler & Hansmann, 2007b), where r p denotes the rejection rate for the chain pair ( p, p + 1) .
Moreover, (cid:80) P − 1 p =1 r p converges to a ﬁxed barrier Λ as P → ∞ (Predescu et al., 2004; Syed et al., 2021).
Applying Lagrange multiplies to the constrained optimization problem leads to r 1 = r 2 = · · · = r P − 1 :
= r , where r is the equi-rejection rate .
In general, a quadratic round trip time is required for ADJ and SEO due to the reversible indexes.
By contrast, DEO only yields a linear round trip time in terms of P as P → ∞ Syed et al.
O PTIMAL N ON - REVERSIBLE SCHEME FOR P ARALLEL TEMPERING
The linear round trip time is appealing for maximizing the algorithmic potential, however, such an advance only occurs given sufﬁciently many chains.
In non-asymptotic settings with limited chains, a pearl of wisdom is to avoid frequent swaps (Dupuis et al., 2012) and to keep the average acceptance rate from 20% to 40% (Kone & Kofke, 2005;
Lingenheil et al., 2009; Atchadé et al., 2011).
Most importantly, the acceptance rates are severely reduced in big data due to the bias-corrected swaps associated with stochastic energies (Deng et al., 2020), see details in section A.1.
As such, maintaining low rejection rates becomes quite challenging and the issue of quadratic costs still exists.
3. G ENERALIZED DEO SCHEME
Continuing the equi-acceptance settings, we see in Figure.2(a)
that the probability for the blue particle to move upward 2 steps to maintain the same momentum after a pair of even and odd iterations is (1 − r ) 2 .
As such, with a large equi-rejection rate r , the blue particle often makes little progress (Figure.2(b-d)).
To handle this issue, the key is to propose small enough rejection rates to track the periodic orbit in Figure.1(b).
Instead of pursuing excessive amount of chains, we resort to a different solution by introducing the generalized even and odd iterations E W and O W , where W ∈ N + ,
E W = {(cid:98) kW (cid:99) mod 2 =
| k
= 1 , 2 , · · · , ∞} and O W = {(cid:98) kW (cid:99)
mod 2 = 1 | k
= 1 , 2 , · · · , ∞} .
Now, we present the generalized DEO scheme with a window size W as follows and refer to it as DEO W : § ◦ Attempt to swap E (or O ) pairs at E W (or O W ) iterations.
◦ Allow at most one swap during each cycle of E W (or O W ) iterations.
[Figure]
As illustrated in Figure.2(e), the blue particle has a larger chance of (1 − r 2 ) 2 to move upwardsteps given W = 2 instead of (1 − r ) 2 when W = 1 , although the window number is also halved.
Such a trade-off inspires us to analyze the expected round trip time based on the window of size W .
3. A NALYSIS OF R OUND T RIP T IME
To bring sufﬁcient interactions between the reference distribution π ( P ) and the target distribution π (1) , we expect to minimize the expected round trip time T (deﬁned in section A.5) to ensure both efﬁcient exploitation and explorations.
Combining the Markov property and the idea of the master equation (Nadler & Hansmann, 2007a), we estimate the expected round trip time E
[ T ] as follows
Lemma 1.
Under the stationary and weak dependence assumptions B1 and B2 in section B, for P ( P ≥ 2 ) chains with window size W ( W ≥ 1 ) and rejection rates { r p } P − 1 p =1 , we have
[Equation]
The proof in section B.1 shows that E [ T ] increases as we adopt larger number of chains P and rejection rates { r p } P − 1 p =1 .
In such a case, the round trip rate P E
[ T ] is also maximized by the key renewal theorem.
In particular, applying W = 1 recovers the vanilla DEO scheme.
3. A NALYSIS OF O PTIMAL W INDOW S IZE AND R OUND T RIP T IME
By Lemma 1, we observe a potential to remove the second quadratic term given an appropriate W .
Such a fact motivates us to study the optimal window size W to achieve the best efﬁciency. Under
(cid:63) and will be studied in section 3.3.
the equi-acceptance settings, by treating the window size W as a continuous variable and taking the derivative of E [ T ] with respect to W , we have
[Equation]
where r is the equi-rejection rate for adjacent chains.
Deﬁne x := r W ∈ (0 , 1) , where W = log r ( x ) = log x log r .
The following analysis hinges on the study of the solution g ( x ) =
(1 − x ) 2 + ( P − 1) x (1 − x + log( x ))
= 0 .
By analyzing the growth of derivatives and boundary values, we can easily identify the uniqueness of the solution.
Then, we proceed to verify that 1 P log P yields an asymptotic approximation such that g ( [Equation] = − [Equation]
P
[Equation]
→ 0 as P → ∞ .
In the end, we have
Theorem 1.
Under Assumptions B1 and B2 based on equi-acceptance settings, if P = 2 , 3 , the maximal round trip time is achieved when W = 1 .
If P ≥ 4 , with the optimal window size
W (cid:63) ≈
[Equation]
(cid:109) , where (cid:100)·(cid:101) is the ceiling function.
The round trip time follows O
( P log P − log r ) .
The above result yields a remarkable round trip time of O
( P log P ) by setting the optimal window size W (cid:63) .
By contrast, the vanilla DEO scheme only leads to a longer time of O
( P 2 ) § .
Denoting by
DEO (cid:63) the generalized DEO scheme with the optimal window size W (cid:63) , we summarize the popular swap schemes in Table.1, where the DEO (cid:63) scheme performs the best among all the three criteria.
T ABLE
[Table]
3. D ISCUSSIONS ON THE O PTIMAL N UMBER OF C HAINS
Note that in practice given P parallel chains, a large P leads to a smaller equi-rejection rate r .
As such [Table] we can further obtain a crude estimate of the optimal P to minimize the round trip time.
Corollary 1.
Under Assumptions B1-B4 and C1 under equi-acceptance settings with the optimal window size [Table] the optimal number of chains follows that P (cid:63) > min p σ p 3 τ ( p ) log( τ ( P ) τ (1) ) [Table]
where σ p is deﬁned in Eq.(15).
The assumptions and proof are postponed in section B.3.
In mini-batch settings, insufﬁcient chains may lead to few effective swaps for accelerations; by contrast, introducing too many chains may be too costly in terms of the round trip time.
This is different from the conclusion in full-batch settings, where Syed et al.
(2021) suggested running the vanilla DEO scheme with as many chains as possible to yield a small enough equi-rejection rate r to maintain the non-reversibility.
Cutoff phenomenon On the one hand, when we only afford at most P chains, where P < P (cid:63) , a large equi-rejection rate r is inevitable and DEO (cid:63) is preferred over DEO; on the other hand, the rejection rate r goes to 0 when P (cid:29) P (cid:63) and DEO (cid:63) recovers the DEO scheme.
In section B.4, we show that P (cid:63) is in the order of thousands for the CIFAR100 example, which is hard to achieve due to the limited computational budget and further motivates us to adopt ﬁnite chains with a target swap rate S to balance between acceleration and accuracy.
U SER - FRIENDLY APPROXIMATE EXPLORATIONS IN B IG D ATA
Despite the asymptotic correctness, SGLD only works well given small enough learning rates and fails in explorative purposes (Ahn et al., 2012).
A large learning rate, however, leads to excessive stochastic gradient noise and ends up with a crude approximation.
As such, similar to Izmailov et al.
; Zhang et al.
, we only adopt SGLD for exploitations.
1 − log( r ) =
O ( r 1 − r ) .
Efﬁcient explorations not only require a high temperature but also prefer a large learning rate.
Such a demand inspires us to consider SGD with a constant learning rate η as the exploration component
[Equation]
where ε ( β k ) ∈
R
d is the stochastic gradient noise.
Under mild normality assumptions on ε (Mandt et al., 2017; Chen et al., 2020), β k converges approximately to an invariant distribution, where the underlying temperature linearly depends on the learning rate η .
Motivated by this fact, we propose an approximate transition kernel T η with P parallel SGD runs based on different learning rates
Exploration:
[Equation]
Exploitation:
[Equation]
where η (1) <
η (2) < · · · < η ( P ) , Ξ k ∼ N (0 , 2 η (1) τ (1) ) , and τ (1) is the target temperature.
Since there exists an optimal learning rate for SGD to estimate the desired distribution through
Laplace approximation (Mandt et al., 2017), the exploitation kernel can be also replaced with SGD based on constant learning rates if the accuracy demand is not high.
Regarding the validity of adopting different learning rates for parallel tempering, we leave discussions to section A.2.
4. A PPROXIMATION ANALYSIS
Moreover, the stochastic gradient noise exploits the Fisher information (Ahn et al., 2012; Zhu et al., 2019; Chaudhari et al., 2017) and yields convergence potential to wide optima with good generalizations (Berthier et al., 2020; Zou et al., 2021).
Despite the implementation convenience, the inclusion of SGDs has made the temperature variable inaccessible, rendering a difﬁculty in implementing the Metropolis rule Eq.(2).
To tackle this issue, we utilize the randomness in stochastic energies and propose a deterministic swap condition for the approximate kernel T η in Eq.(7) such that
Deterministic swap condition: ( β ( p ) , β ( p +1) )
→ ( β ( p +1) , β ( p ) ) if (cid:101) U ( β ( p +1) )
+ C < (cid:101) U ( β ( p ) ) , (8) where p ∈ { 1 , 2 , · · · , P − 1 } , C > 0 is a correction buffer to approximate the Metropolis rule Eq.(2).
Lemma 2.
Assume the energy normality assumption (C1), then for any ﬁxed ∂U p :=
U ( β ( p ) ) −
U ( β ( p +1) )
, there exists an optimal C (cid:63) ∈ (0 , ( 1 τ ( p ) − 1 τ ( p +1) ) σ 2 p ] that perfectly approximates the random event (cid:101) S ( β ( p ) , β ( p +1) ) > u , where σ p deﬁned in Eq.(15) and u ∼ Unif [0 , 1] .
The proof is postponed in section C.1, which paves the way for the guarantee that a deterministic swap condition may replace the Metropolis rule Eq.(2) for approximation tasks.
In addition, the normality assumption can be naturally extended to the asymptotic normality assumption (Quiroz et al., 2019; Deng et al., 2021) given large enough batch sizes.
Admittedly, the approximation error still exists for different ∂U p .
By the mean-value theorem, there exists a tunable C to optimize the overall approximation.
Further invoking the central limit theorem such that ε ( · ) in Eq.(6) approximates a Gaussian distribution with a ﬁxed covariance, we can expect a bounded approximation error for the SGD-based exploration kernels (Mandt et al., 2017).
Theorem 2.
Consider the exact transition kernel T and the proposed approximate kernel T η , which yield stationary distributions π and π η , respectively.
Under smoothness (C2) and dissipativity assumptions (C3) (Mattingly et al., 2002; Raginsky et al., 2017; Xu et al., 2018), T satisﬁes the geometric ergodicity such that there is a contraction constant ρ ∈ [0 , 1) for any distribution µ :
[Equation]
where (cid:107) · (cid:107) TV is the total variation (TV) distance.
Moreover, assume that ε ( · ) ∼ N (0 , M ) for some positive deﬁnite matrix M (C4) (Mandt et al., 2017), then there is a uniform upper bound of the one step error between T and T η such that
[Equation]
where ∆ max ≥ 0 is a constant.
Eventually, the TV distance between π and π η is bounded by
[Equation]
The proof is postponed to section C.2.
The SGD-based exploration kernels no longer require to ﬁne-tune the temperatures directly and naturally inherits the empirical successes of SGD in large-scale deep learning tasks.
The inaccessible Metropolis rule Eq.(2) is approximated via the deterministic swap condition Eq.(8) and leads to a well-controlled approximations by solely tuning η = ( η (1) , · · · , η ( P ) ) and C .
In addition, our proposed algorithm for uncertainty approximation is highly related to non-convex optimization.
For the detailed discussions, we refer interested readers to section A.4.
4. E QUI - ACCEPTANCE PARALLEL TEMPERING ON OPTIMIZED PATHS
Stochastic approximation (SA) is a standard method to achieve equi-acceptance (Atchadé et al., 2011; Miasojedow et al., 2013), however, implementing this idea with ﬁxed η (1) and η ( P ) is rather non-trivial.
Motivated by the linear relation between learning rate and temperature, we propose to adaptively optimize the learning rates to achieve equi-acceptance in a user-friendly manner.
Further by the geometric temperature spacing commonly adopted by practitioners (Kofke, 2002; Earl &
Deem, 2005; Syed et al., 2021), we adopt the following scheme on a logarithmic scale such that
[Equation]
) where p ∈ { 1 , 2 , · · · , P − 1 } , υ ( p ) t = [Equation] [Equation] [Equation]
[Equation]
[Equation] [Equation] [Equation] = (cid:82) H ( p ) ( υ ( p )
[Equation] , β ) π ( p,p +1) ( d β ) is the mean-ﬁeld function, π ( p,p +1) is the joint invariant distribution for the p -th and p + 1 -th processes.
In particular, H ( p ) ( υ ( p ) k , β ) = 1 (cid:101) U ( β ( p +1) )
+ C < (cid:101) U ( β ( p ) )
− S is the random-ﬁeld function to approximate h ( p ) ( υ ( p ) k ) with limited perturbations, υ ( p ) k † implicitly affects the distribution of the indicator function, and S is the target swap rate.
Now consider stochastic approximation of Eq.(9), we have
[Equation]
where γ k is the step size.
Reformulating Eq.(10), we have
[Equation]
where the max operator is conducted explicitly to ensure the sequence of learning rates is nondecreasing.
This means that given ﬁxed boundary learning rates (temperatures) η ( p − 1) k
and η ( p +1) k , applying η ( p ) =
η ( p − 1) + υ ( p ) and η ( p ) =
η ( p +1) − υ
( p +1) for p ∈ { 2 , 3 , · · · , P − 1 } lead to η ( p − 1) k + max [Equation] [Equation]
[Equation]
Adaptive learning rates (temperatures) Now given a ﬁxed η (1) , the sequence η (2) , η (3) , · · · , η ( P ) can be approximated iteratively via the forward sequence of (11); conversely, given a ﬁxed η ( P ) , the backward sequence η
( P − 1) , η ( P − 2) , · · · , η (1) can be decided reversely as well.
Combining the forward and backward sequences, η ( p ) k
+1 can be approximated via
[Equation]
which resembles the binary search in the SA framework.
In particular, the ﬁrst term is the middle point given boundary learning rates and the second term continues to penalize learning rates that violates the equi-acceptance between pairs ( p − 1 , p ) and ( p, p + 1) until an equilibrium is achieved.
This is the ﬁrst attempt to achieve equi-acceptance given two ﬁxed boundary values to our best knowledge.
By contrast, Syed et al.
(2021) proposed to estimate the barrier Λ to determine the temperatures and it easily fails in big data given a ﬁnite number of chains and bias-corrected swaps.
Adaptive correction buffers
In addition, equi-acceptance does not guarantee a convergence to the desired acceptance rate S .
To avoid this issue we propose to adaptively optimize C as follows
[Equation]
As k → ∞ , the threshold and the adaptive learning rates converge to the desired ﬁxed points.
Note that setting a uniform C greatly simpliﬁes the algorithm; in more delicate cases, problem-speciﬁc rules are also recommended.
Now we refer to the approximate non-reversible parallel tempering algorithm with the DEO (cid:63) scheme and SGD-based exploration kernels as DEO (cid:63) -SGD and formally formulate our algorithm in Algorithm 1.
Extensions of SGD with a preconditioner (Li et al., 2016) or momentum (Chen et al., 2014) to further improve the approximation and efﬁciency are both straightforward (Mandt et al., 2017) and are denoted as DEO (cid:63) -pSGD and DEO (cid:63) -mSGD, respectively.
E XPERIMENTS 5. S IMULATIONS OF MULTI - MODAL DISTRIBUTIONS
We ﬁrst simulate the proposed algorithm on a distribution π ( β ) ∝ exp( − U ( β )) , where β = ( β 1 , β 2 ) ,
U ( β ) = 0 .
2( β 21 + β 22 ) − 2(cos(2 πβ 1 ) + cos(2 πβ 2 )) .
The heat map is shown in Figure 3(a) withmodes of different volumes.
To mimic big data scenarios, we can only access stochastic gradient ∇ (cid:101) U ( β ) =
∇ U ( β ) + 2 N (0 , I 2 × 2 ) and stochastic energy (cid:101) U ( β ) =
U ( β ) + 2 N (0 , I ) .
) S
S
S
[Figure]
We ﬁrst run DEO (cid:63) -SGD × P16 based on 16 chains and 20,000 iterations.
We ﬁx the lowest learning rate 0.003 and the highest learning 0.6 and propose to tune the target swap rate S for the accelerationaccuracy trade-off.
Fig.3 shows that ﬁxing S = 0 .
2 or 0.3 is too conservative and underestimates the uncertainty on the corners; S = 0 .
6 results in too many radical swaps and eventually leads to crude estimations; by contrast, S = 0 .
4 yields the best uncertainty approximation among the ﬁve choices.
Next, we select S = 0 .
4 and study the round trips.
We observe in Fig.4(a) that the vanilla DEO only yields 18 round trips every 1,000 iterations; by contrast, slightly increasing W tends to improve the efﬁciency signiﬁcantly and the optimal 45 round trips are achieved at W = 8 , which matches our theory .
In Fig.4(b-c), the geometrically initialized learning rates lead to unbalanced acceptance rates in the early phase and some adjacent chains have few swaps and others swap too much, but as the optimization proceeds, the learning rates gradually converge.
We also observe in Fig.4(d) that the correction is adaptively estimated to ensure the average acceptance rates converge to S
= 0 .
[Figure]
We compare the proposed algorithm with parallel SGLD based on 20,000 iterations and 16 chains (SGLD × P16); we ﬁx the learning rate 0.003 and a temperature 1.
We also run cycSGLD × T16, which is short for a single long chain based on 16 times of budget and cosine learning rates (Zhang et al., 2020) of 100 cycles.
We see in Figure 5(b) that SGLD ×
P16 has good explorations but fails to quantify the uncertainty.
Figure 5(c) shows that cycSGLD × T16 explores most of the modes but overestimates some areas occasionally.
Figure 5(d) demonstrates the DEO-SGD with 16 chains (DEO-SGD × P16) estimates the uncertainty of the centering 9 modes well but fails to deal with the rest of the modes.
As to DEO (cid:63) -SGD × P16, the approximation is rather accurate, as shown in Fig.5(e).
[Figure]
We also present the index process for both schemes in Fig.6.
We see that the vanilla DEO scheme results in volatile paths and a particle takes quite a long time to complete a round trip; by contrast,
DEO (cid:63) only conducts at most one cheap swap in a window and yields much more deterministic paths.
[Figure]
5. U NCERTAINTY APPROXIMATION AND OPTIMIZATION FOR IMAGE DATA
Next, we conduct experiments on computer vision tasks.
We choose ResNet20, ResNet32, and
ResNet56 (He et al., 2016) and train the models on CIFAR100.
We not only report test negative log likelihood (NLL) but also present the test accuracy (ACC).
For each ResNet model, we ﬁrst pre-train 10 ﬁxed models via 300 epochs and then run our algorithm based on momentum SGD (mSGD) for 500 epochs with 10 parallel chains and denote it by DEO (cid:63) -mSGD × P10.
We ﬁx the lowest and highest learning rates as 0.005 and 0.02, respectively.
For a fair comparison, we also include the baseline DEO-mSGD × P10 with the same setup except that the window size is 1; the standard ensemble mSGD × P10 is also included with a learning rate of 0.005.
In addition, we include two baselines based on a single long chain, i.e. we run stochastic gradient Hamiltonian Monte Carlo
(Chen et al., 2014) 5000 epochs with cyclical learning rates and 50 cycles (Zhang et al., 2020) and refer to it as cycSGHMC × T10; we run SWAG ×
T10 (Maddox et al., 2019) under a similar setup.
In particular for DEO (cid:63) -mSGD × P10, we tune the target swap rate S and ﬁnd an optimum at S = 0 .
We compare our proposed algorithm with the four baselines and observe in Table.2 that mSGD
× P can easily obtain competitive results simply through model ensemble (Lakshminarayanan et al., 2017), which outperforms cycSGHMC × T10 and cycSWAG × T10 on ResNet20 and ResNet32 models and perform the worst among the ﬁve methods on ResNet56;
DEO-mSGD × P10 itself is already a pretty powerful algorithm, however, DEO (cid:63) -mSGD
×
P10 consistently outperforms the vanilla alternative.
[Table]
[Figure]
[Table]
To analyze why the proposed scheme performs well, we study the round trips in Figure.7(a) and ﬁnd that the theoretical optimal window obtains around 11 round trips every 100 epochs, which is almost 2 times as much as the vanilla DEO scheme.
In Figure 7(b), we observe that the smallest learning rate obtains the highest accuracy (blue) for exploitations, while the largest learning rate yields decent explorations (red); we see in Figure 7(c-d) that geometrically initialized learning rate fails in producing equi-acceptance, but as the training proceeds, the learning rates converge to ﬁxed points and the acceptance rates for different pairs converge to the target swap rate.
[Figure]
For the visualization of the index process, we observe in Figure.8(a) that the vanilla DEO scheme leads to volatile trajectories wandering back and forth and is rather inefﬁcient; by contrast, the DEO (cid:63) scheme yields well-motivated paths with more deterministic round trips.
Interestingly, this path resembles cyclic learning rates (Zhang et al., 2020), which provides a novel viewpoint to interpret why cyclic learning rates work well empirically .
Nevertheless, the stochastic and parallel manner further improves the margin and eventually leads to the most efﬁcient approximations in this task.
[Figure]
C ONCLUSION
In this paper, we show how to conduct efﬁcient PT in big data problems.
To tackle the inefﬁciency issue of the popular DEO scheme given limited chains, we present the DEO (cid:63) scheme by applying an optimal window size to encourage deterministic paths and obtain in a signiﬁcant acceleration of
O ( P log P ) times .
For a user-friendly purpose, we propose a deterministic swap condition to interact with SGD-based exploration kernels and provide a theoretical guarantee to control the bias solely depending on the learning rate η and the correction buffer C ; we also provide a practical algorithm to adaptively approximate η and C for achieving the optimal efﬁciency in constrained settings .
A B ACKGROUND A. R EPLICA EXCHANGE STOCHASTIC GRADIENT L ANGEVIN DYNAMICS
To approximate the replica exchange Langevin diffusion Eq.(1) based on multiple particles ( β (1)
k +1 , β (2) k +1 , · · · , β ( P ) k +1 ) in big data scenarios, replica exchange stochastic gradient Langevin dynamics (reSGLD) proposes the following numerical scheme:
Exploration:
[Equation]
14) where η ( · ) is the learning rate, ξ ( · ) k is a standard d -dimensional Gaussian noise and each subprocess follows a stochastic gradient Langevin dynamics (SGLD)
(Welling & Teh, 2011).
Further assuming the energy normality assumption C1 in section C.1, there exists a σ p such that
[Equation]
any
[Equation]
15) where p ∈ { 1 , 2 , · · · , P − 1 } , in what follows, Deng et al.
(2020) proposed the bias-corrected swap function as follows
[Equation]
where p ∈ { 1 , 2 , · · · , P − 1 } , (cid:0) 1 τ ( p ) − 1 τ ( p +1) (cid:1) σ 2 p is a correction term to avoid the bias and the swap intensity a can be then set to 1 min { η ( p ) ,η ( p +1) } for convenience.
Namely, given a particle pair at location ( β ( p ) , β ( p +1) ) in the k -th iteration, the conditional probability of the swap follows that
[Equation]
A. I NHOMOGENOUS SWAP INTENSITY VIA DIFFERENT LEARNING RATES
Assume the high-temperature process also applies η ( p +1) ≥ η ( p ) > 0 , where p ∈ { 1 , 2 · · · , P − 1 } .
At time t , the swap intensity can be interpreted as being 0 in a time interval [ t, t + η ( p +1) − η ( p ) ) and being 1 η ( p ) in [ t + η ( p +1) − η ( p ) , t + η ( p +1) ) .
Since the coupled Langevin diffusion process converges to the same joint distribution regardless of the swaps and the swap intensity a varies in a bounded domain, the convergence of numerical schemes is not much affected except the numerical error.
A. A CCURACY - ACCELERATION TRADE - OFF
Despite the exponential acceleration potential, the average bias-corrected swap rate is signiﬁcantly reduced such that E [ (cid:101) S ]
= O (cid:0)
Se − (cid:0) 1 τ ( p ) −
[Equation]
) (cid:1)
[Equation]
(cid:1) (Deng et al., 2021), where S is the swap function following Eq.(2) in full-batch settings.
Including more parallel chains are promising to alleviate this issue, however, it is often inevitable to sacriﬁce some accuracy to obtain more swaps and accelerations, as discussed in section 3.4.
As such, it inspires us to devise the user-friendly
SGD-based approximate exploration kernels.
A. C ONNECTION TO NON - CONVEX OPTIMIZATION
Parallel to our work, a similar SGLD × SGD framework was proposed by Dong & Tong (2021) for non-convex optimization, where SGLD and SGD work as exploration and exploitation kernels, respectively.
By contrast, our algorithm performs exactly in the opposite for uncertainty approximation because SGLD is theoretically more appealing for the exploitations based on small learning rates instead of explorations, while the widely-adopted SGDs are quite attractive in exploration due to its user-friendly nature and ability in exploring wide optima.
If we manipulate the scheme to propose an exact swap in each window instead of at most one in the current version, the algorithm shows a better potential in non-convex optimization.
In particular, a larger window size corresponds to a slower decay of temperatures in simulated annealing (SAA) (Mangoubi & Vishnoi, 2018).
Such a mechanism yields a larger hitting probability to move into a sub-level set with lower energies (losses) and a better chance to hit the global optima.
Nevertheless, the manipulated algorithm possess the natural of parallelism in cyclical fashions.
A. O THERS
Swap time refers to the communication time to conduct a swap in each attempt.
For example, chain pair ( p, p + 1) of ADJ requires to wait for the completion of chain pairs (1 , 2) , (2 , 3) , · · · , ( p − 1 , p ) to attempt the swap and leads to swap time of O ( P ) ; however, SEO, DEO, and DEO (cid:63) don’t have this issue because in each iteration, only even or odd chain pairs are attempted to swap, hence the swap time is O (1) .
Round trip time refers to the time (stochastic variable) used in a round trip.
A round trip is completed when a particle in the p -th chain, where p ∈ [ P ] := { 1 , 2 , · · · , P } , hits the index boundary index at both 1 and P and returns back to its original index p .
B A NALYSIS OF R OUND T RIPS
To facilitate the theoretical analysis, we follow Syed et al.
and make the following assumptions
• (B1)
Stationarity: Each sub-process has achieved the stationary distribution β ( p ) ∼ π ( p ) for any p ∈
[ P ] ; • (B2)
Weak independence: For any ¯ β ( p ) simulated from the p -th chain conditional on β ( j ) , U ( ¯
β j
) ) and U ( β ( j ) ) are independent
B. A NALYSIS OF R OUND T RIP T IME
Proof of Lemma 1.
For t ∈ N , deﬁne Z t ∈
[ P ] = { 1 , 2 , · · · , P } as the index of the chain a particle arrives after t windows.
Deﬁne δ t ∈ { 1 , − 1 } to indicate the direction of the swap a particle intends to make during the t -th window; i.e., the swap is between Z t and Z t
+ 1 if δ t = 1 and is between Z t and Z t − 1 if δ t = − 1 .
Deﬁne U :=
min { t ≥ 0 : Z t = P, δ t = − 1 } and V :=
min { t ≥ 0 : Z t = 1 , δ t = 1 } .
Deﬁne r p :=
P [ reject the swap between Chain p and Chain p +1 for one time ] .
Deﬁne u p,δ :=
E [ U | Z 0 = p, δ
= δ ] for δ ∈ { 1 , − 1 } and v p,δ :=
E [ V | Z 0 = p, δ
= δ ] .
Then the expectation of round trip time T is
[Equation]
By the Markov property, for u p,δ , we have
[Equation] [Equation]
19) where r Wp denotes the rejection probability of a particle in a window of size W at the p -th chain.
According to Eq.(18) and Eq.(19), we have
[Equation] [Equation]
W
[Equation] [Equation]
+ 1 .
Deﬁne α p = u p, 1 − u p − 1 , − 1 .
Then by deﬁnition, Eq.(20), and Eq.(21), we have
[Equation]
which implies that
[Equation]
) for a p :
=
(1 − r Wp − 1 ) α p .
Thus, by Eq.(22), we have
[Equation]
By deﬁnition, u 1 , − 1 = u 1 , 1 + 1 .
According to Eq.(20), for p = 1 , we have
[Equation]
Since r p ∈ (0 , 1) for 1 ≤ p ≤ P , Eq.(24) implies a 2 = − 2 which together with Eq.(23) implies
[Equation]
25) and therefore
[Equation]
According to Eq.(21) and Eq.(26), we have
[Equation]
[Equation]
Since u P, 1 = 1 , we have
[Equation]
Similarly, for v p,δ , we also have
[Equation]
With the same analysis, we have
[Equation] [Equation]
32) where b p :=
(1 − r Wp − 1 ) β p and β p := v p, 1 − v p − 1 , − 1 .
According to Eq.(32), we have
[Equation]
By deﬁnition, v P, 1 = v P, − 1 + 1 .
According to Eq.(32), for p =
P , we have
[Equation]
Since r p ∈ (0 , 1) for 1 ≤ p ≤ P , Eq.(34) implies b P = 2 .
Then according to Eq.(31), we have
[Equation]
, and therefore
[Equation]
By Eq.(33) and Eq.(35), we have
[Equation]
Since v 1 , − 1 = 1 , we have
[Equation]
P
[Equation]
According to Eq.(17), Eq.(30) and Eq.(36), we have
[Equation]
The proof is a generalization of Theorem 1 (Syed et al., 2021); when W = 1 , the generalized DEO scheme recovers the DEO scheme.
For the self-consistency of our analysis, we present it here anyway.
B. A NALYSIS OF O PTIMAL W INDOW
S IZE
Proof of Theorem 1.
By treating W ≥ 1 as a continuous variable and taking the derivative of E [ T ] with respect to W and we can get
[Equation]
Assume that r p = r ∈ (0 , 1) for 1 ≤ p ≤ P .
Then we have
[Equation]
Deﬁne x := r W ∈ (0 , 1) .
Hence W = log r ( x ) =
log x log r and ∂ ∂W E
[ T ] = f
( x ) :
[Equation] [Equation]
Thus it sufﬁces to analyze the sign of the function g ( x ) :
= (1 − x ) 2 + ( P − 1) x (1 − x + log( x )) for x ∈ (0 , 1) .
For g ( x ) , we have
[Equation]
and
[Equation]
Thus, lim x
[Equation]
+ g (cid:48) ( x ) =
−∞ , lim x → 1 g (cid:48)
( x ) = g (cid:48) (1) =
and g (cid:48)(cid:48)
( x ) is monotonically decreasing for x > 0 .
For P > 2 , we know that g (cid:48)(cid:48) ( x ) >
when 0 < x <
P − 1 2( P − 2) and g (cid:48)(cid:48)
( x ) < 0 when x < P
−
1 2( P − 2) .
Therefore, g (cid:48) ( x ) is maximized at P − 1 2( P
−
(a) If P = 3 , by log(1+ y )
<
y for y > 0 , we have g (cid:48)
( x ) = 2 (cid:0) log(1+( x − 1)) − ( x − 1) (cid:1) < 0 for any x ∈ (0 , 1) .
Thus g ( x ) > g (1) = 0 for x ∈ (0 , 1) and therefore ∂∂W E [ T ] >for W ∈ N + .
E [ T ] is globally minimized at W = 1 .
(b) If P > 3 , we have the following lemma and the proof is postponed in section B.2.1.
Lemma 3 (Uniqueness of the solution) .
For P > 3 , there exists a unique solution x ∗ ∈ (0 , 1) such that g ( x ) > 0 for ∀ x ∈ (0 , x ∗ ) and g ( x
) < 0 for ∀ x ∈ ( x ∗ , 1) .
Moreover, W ∗ = log r ( x ∗ ) is the globally minimizer for the round trip time.
For P = 2 , we know that g (cid:48)(cid:48) ( x ) > 0 for x ∈ (0 , 1) .
Thus g (cid:48)
( x ) < g (1) = 0 for x ∈ (0 , 1) and therefore g ( x ) > g (cid:48) (1) = 0 for x ∈ (0 , 1) .
Then according to Eq.(37), ∂∂W E [ T ] >for W ∈ N + .
E [ T ] is globally minimized at W = 1 .
In what follows, we proceed to prove that 1 P log P is a good approximation to x ∗ .
In fact, we have
[Equation]
P
[Equation] [Equation]
P
[Equation]
log(
[Equation]
P ) log P + O (cid:18)
P
[Equation]
= − log(log P ) log P + O (cid:18)log P (cid:19)
Thus, lim P →∞ g ( 1 P log P ) = 0 .
For x = 1 P log P , we have W = log r
[Equation]
Then according to Eq.(38), for
P ≥ 4 ,
[Equation]
In conclusion, for P = 2 , 3 , the maximum round trip rate is achieved when the window size W = 1 .
For P ≥ 4 , with the window size W = log P +loglog P − log r , the round trip rate is at least Ω (cid:16) − log r log P (cid:17) .
Remark:
Given ﬁnite chains with a large rejection rate, the round trip time is only of order
O ( P log P ) by setting the optimal window size W ≈ (cid:108) log P +loglog P − log r (cid:109) .
By contrast, the vanilla DEO scheme with a window of size 1 yields a much longer time of O
( P 2 ) , where
[Equation]
based on Taylor expansion and a large r (cid:29) 0 .
B.2. T ECHNICAL L EMMA
Proof of Lemma 3.
To help illustrate the analysis below, we plot the graphs of g (cid:48)
( x ) and g ( x ) for x ∈ (0 , 1) and P = 5 in Figure 9.
For P > 3 , since g (cid:48) ( x ) is maximized at x = P − 1 2( P − 2) ∈
(0 , 1) with g (cid:48)(cid:48)
( x ) >
when 0 < x < P − 1 2( P − 2) and g (cid:48)(cid:48) ( x ) < 0 when P − 1 2( P − 2) < x < 1 , lim x → 0 + g (cid:48)
( x ) =
−∞ , and g (cid:48) (1) = 0 , we know that g (cid:48)
( P − 1 2( P − 2) ) > 0 and there exists x 0 ∈ (0 , P − 1 2( P − 2) )
such that g (cid:48) ( x ) <
(i.e., g (cid:48)
( x ) is monotonically decreasing) for any x ∈ (0 , x 0 ) and g (cid:48)
( x ) > 0 (i.e., g (cid:48)
( x ) is monotonically increasing) for any x ∈ ( x 0 , 1) .
Then g ( x ) on (0 , 1) is globally minimized at x = x 0 .
Moreover, lim x → 0 + g ( x ) = 1 and g (1) = 0 .
Thus, g ( x 0 ) <
g (1) =
and there exists x ∗ ∈ (0 , x 0 ) (cid:40) (0 , 1) such that g ( x ) > 0 if x ∈ (0 , x ∗ ) and g ( x ) < 0 if x ∈ ( x ∗ , 1) .
Meanwhile, by Eq.(37) and the deﬁnition of x , we know that the sign of ∂∂W E
[ T ] is the same with that of g ( x ) for W = log r ( x ) .
Thus, ∂∂W E [ T ] < 0 when W < W ∗ := log r ( x ∗ ) and ∂∂W E
[ T ] >
when
W > W ∗ , which implies that E [ T ] is globally minimized at W ∗ = log r ( x ∗ ) with some x ∗ ∈ (0 , 1) .
[Figure]
B. D ISCUSSIONS ON THE OPTIMAL N UMBER OF C HAINS IN B IG D ATA
To shed light on the suggestions of the optimal number of chains, we lay out two additional assumptions:
• (B3) Integrability: U 3 is integrable with respect to π (1) and π ( P ) .
• (B4) Approximate geometric spacing: Assume that τ ( p +1) τ ( p ) ≡ (cid:16) τ ( P ) τ (1) (cid:17) 1 P − 1 + o ( 1 P ) for p ∈
[ P − 1] .
Assumption B3 is a standard assumption in Syed et al.
(2021) to establish the convergence of the summation of the rejection rates.
Assumption B4 allows extra perturbations to the geometric spacing assumption (Kofke, 2002) and is empirically veriﬁed in the CIFAR100 example in Figure 7(c).
Proof of corollary 1.
Before we start the proof, we denote by ˜ r and ˜ s the average rejection and acceptance rates in full-batch settings, respectively.
According to section 5.1 in Syed et al.
Assumptions B1, B2 and B3 lead to
[Equation]
where Λ is the barrier, which becomes larger if π (1) and π ( P ) have a smaller probability overlap (Predescu et al., 2004; Syed et al., 2021).
Then acceptance rate in full-batch settings is
[Equation]
By Lemma D4 (Deng et al., 2021) and Eq.(15), a noisy energy estimator with variance ( τ ( p ) − τ ( p +1) )
τ ( p
[Equation] [Equation]
( p
[Equation]
1) 2 σ 2 p yields an average swap rate of O (cid:16)
[Equation]
p
[Equation]
In what follows, the average rejection rate in mini-batch settings becomes
[Equation]
where p ∈ [ P ] and more accurate estimates is studied in proposition 2.4 (Gelman et al., 1997).
Deﬁne ∆ τ = τ ( P ) τ (1) .
By assumption B4 and Taylor’s theorem, we have
[Equation]
Denote γ p = σ p τ ( p +1) .
It follows that
[Equation] [Equation]
Plugging it to Eq.(44), we have
[Equation]
[Equation]
Recall in Eq.(41), the round trip time follows that
[Equation] [Equation]
By log(1 − t ) ≤ − t for t ∈ [0 , 1) and 1 1 − t = 1 + t + O ( t 2 ) for a small t , we have −
[Equation]
[Equation]
(cid:1)(cid:17)(cid:17)(cid:111) ≤
O (cid:16)(cid:0) 1 − Λ P
[Equation]
O (cid:0)
[Equation]
Combining Eq.(49) and Eq.(48), we have
[Equation]
where µ p :=
[Equation]
The optimal P to minimize E [ T ] =
O (cid:16) P log P exp (cid:16)
[Equation]
2 p P 2 (cid:17)(cid:17) is equivalent to the minimizer of h p ( x ) = log (cid:110) x log x exp (cid:16) µ 2 p x 2 (cid:17)(cid:111) =
log x + log log x + µ 2 p x 2 for x ≥ 4 .
Then
[Equation]
Deﬁne s ( x ) = x 2 (1 + 1 log x ) .
Then we have
[Equation]
) for x ≥ 4 .
Thus s ( x ) increases with x for x ≥ 4 .
Obviously, lim x →∞ s ( x ) = ∞ .
Since µ p (cid:29) 1 , we know s (4) < 0 .
Thus, there is a single zero point P (cid:63) ∈ (4 , ∞ ) of s ( x ) such that h (cid:48) p ( x ) < 0 if x ∈ [4 , P (cid:63) ) and h (cid:48) p ( x ) > 0 if x > P (cid:63) .
Therefore, h p ( x ) is minimized at P (cid:63) with
P
[Equation]
(cid:18) 1 +
[Equation]
P
[Equation]
(cid:19)
[Equation]
For any P (cid:63) ≥ 2 , we have (cid:112) 2 / (1 + 1 / log P (cid:63) ) ∈ (1 . 1 , √ 2) .
Thus, we have the optimal number of chains in mini-batch settings following that
[Equation]
B. E MPIRICAL JUSTIFICATION OF THE OPTIMAL NUMBER OF CHAINS
To obtain an estimate of the optimal number of chains for PT in big data problems, we study the standard deviation of the noisy energy estimators on CIFAR100 via ResNet20 models.
We ﬁrst pre-train a model and then try different learning rates to check the corresponding standard deviation of energy estimators.
The reason we abandon the temperature variable is that the widely used data augmentation has drastically affected the estimation of the temperature (Wenzel et al., 2020).
Motivated by the linear relation between the learning rate and the temperature in Eq.(6) and the cold posterior affect with the temperature much smaller than 1 (Aitchison, 2021), applying Figure 10 and
Eq.(53) concludes that P (cid:63) is achieved in an order of thousands in the CIFAR100 example .
The conclusion is different from Syed et al.
(2021) since big data problems require a much smaller swap rate to maintain the unbiasedness of the swaps.
On the one hand, insufﬁcient chains may lead to insigniﬁcant swap rates to generate effective accelerations, but on the other hand, introducing too many chains may be too cost in terms of limited memory and computational budget.
C A PPROXIMATION E RROR FOR SGD- BASED E XPLORATION K ERNELS
This section proposes to analyze the approximation error when we adopt SGDs as the approximate exploration kernels.
Since it is independent of section B, different assumptions may be required.
The ﬁrst subsection shows that for any Gaussian-distributed energy difference estimator (cid:101) U ( β ( p ) ) − (cid:101) U ( β ( p +1) ) , there exists an optimal C (cid:63) such that the deterministic swap condition (cid:101) U ( β ( p +1) )
+ C (cid:63)
<
(cid:101) U ( β ( p ) ) perfectly approximates the random event (cid:101) S ( β ( p ) , β ( p +1) )
> u , where u ∼ Unif [0 , 1] .
C. G AP BETWEEN ACCEPTANCE RATES
To approximate the error when we adopt the deterministic swap condition Eq.(8), we ﬁrst assume the following condition
• (C1) Energy normality: The stochastic energy estimator for each chain follows a normal distribution with a ﬁxed variance.
The above assumption has been widely assumed by Ceperley & Dewing (1999); Bardenet et al.
(2017); Seita et al.
(2017), which can be easily extended to the asymptotic normality assumption depending on large enough batch sizes (Quiroz et al., 2019;
Deng et al., 2021).
Proof of Lemma 2.
Denote ∂T p = 1 τ ( p ) − 1 τ ( p +1) > 0 for any p ∈
[ P − 1] , the energy difference ∂U p =
U ( β ( p ) )
− U ( β ( p +1) ) .
By the energy normality assumption C1 and Eq.(15), the energy difference estimator follows that ∂ (cid:101) U p = ∂U p ( · ) + √ 2 σ p ξ , where ξ ∼ N (0 , 1) .
Recall from Eq.(16) that the bias-corrected acceptance rate follows that
[Equation]
(i) For a uniformly distributed variable µ ∼ Unif [0 , 1] , the base swap function satisﬁes that
[Equation]
(ii)
For the deterministic swaps based on Eq.(8), the approximate swap function follows that
[Equation]
Denote by
[Equation]
u,
[Equation]
55) and
(56), we can easily derive that
[Equation]
Clearly, ∆( C ) is continuous in C and it is straightforward to verify that
[Equation]
Similarly, ∆( −∞ ) = − (cid:82) 1 0 D ( u, C ) du ≤ 0 .
Moreover, the physical construction suggests that the threshold C should be strictly positive to avoid radical swap attempts, which implies that there is an optimal solution C (cid:63) ∈ (0 , ∂T p σ 2 p ] that solves ∆( C (cid:63) )
= 0 .
Remark:
Note that the optimal C (cid:63) may lead to few swaps given a ﬁnite number of iterations and sometimes it is suggested to trade in some accuracy to obtain more accelerations (Deng et al., 2020).
As such, by setting a desired swap rate via the iterate (13), we can estimate the unknown threshold C by stochastic approximation (Robbins & Monro, 1951).
Furthermore, as discussed in Lemma
B2 (Deng et al., 2021), the average swap rate follows that E [ (cid:101) S ( β ( p ) t , β ( p +1) t )]
= O ( e − ∂T 2 pσ 2 p) .
This suggests that for a well-approximated threshold C , the error is at most E [∆( C )]
= O ( e − ∂T 2 pσ 2 p) .
C. A PPROXIMATE U PPER B OUND
The second subsection completes the proof regarding the numerical approximation of SGD-based exploration kernels.
Nevertheless, we still require some standard assumptions.
• (C2)
Smoothness: The function U ( · ) is C -smooth if there exists a positive Lipschitz constant C such that (cid:107)∇ U ( x ) − ∇ U ( y ) (cid:107) 2 ≤ C (cid:107) x − y (cid:107) 2 for every x, y ∈ R d . •
(C3) Dissipativity: The function U ( · ) is ( a, b ) -dissipative if there exist positive constants a and b such that (cid:104) x, ∇ U ( x ) (cid:105) ≥ a (cid:107) x (cid:107) 2 − b .
• (C4)
Gradient normality: The stochastic gradient noise ε ( · ) in Eq.(6) follows a Normal distribution N (0 , M ) with a ﬁxed positive deﬁnite matrix M .
The assumptions C2 and C3 are standard to show the geometric ergodicity of Langevin diffusion and the diffusion approximations for non-convex functions (Mattingly et al., 2002;
Raginsky et al., 2017;
Xu et al., 2018).
The assumption C4 is directly motivated by Mandt et al.
(2017) to track the Ornstein
Uhlenbeck (OU) process with a Gaussian invariant measure.
Proof of Theorem 2.
The transition kernel T of the continuous-time replica exchange Langevin diffusion (reLD) follows that T =
(cid:81) p T ( p ) , where each sub-kernel T ( p ) draws a candidate θ ∈ R d with the following probability
[Equation]
where p, q ∈ { 1 , 2 , · · · , P } , | p − q | = 1 and q is selected based on the DEO (cid:63) scheme (3), Q p ( θ | β ) is the proposal distribution in the p -th chain that models the probability to draw the parameter θ via
Langevin diffusion conditioned on β , and S ( β ( p ) t , β ( q ) t ) follows the swap function in Eq.(2).
Now we consider the following approximations:
[Equation]
I: Numerical approximation via SGLD By Lemma 5, we show that approximating replica exchange Langevin diffusion via reSGLD yields a weak error only depending on the learning rate; similar results depending on the learning rate and the noise in gradient and energy noise have been studied in Lemma 1 of (Deng et al., 2020).
II: Deterministic swap condition By Lemma 2 and the mean-value theorem, there exists an optimal correction buffer C (cid:63) to approximate the random event (cid:101) S ( β ( p ) t , β
( p +1) t
> µ through (cid:101) U ( β ( p ) k ) + C (cid:63) < (cid:101) U ( β ( p +1) k ) in the average sense for any β ( p ) t and β ( p +1) t
III: Laplace approximation via SGD Assumption C4 guarantees that there exists
an optimal learning rate η to conduct Laplace approximation for the target posterior (Mandt et al., 2017).
Further, by the Bernstein-von Mises Theorem, the approximation error goes to 0 as the number of total data points goes to inﬁnity.
Combining the above approximations, there exists a ﬁnite constant ∆ max ≥ 0 depending on the learning rates η and the choice of correction terms such that the total approximation error of the SGD exploration kernels with deterministic swap condition is upper bounded by ∆ max .
For any joint probability density µ , the distance between the distributions generated by one step of the exact transition kernel T and the approximate transition kernel T η in Eq.(7) is upper bounded by
[Equation]
where Q ( ·|· ) = (cid:81) p Q p ( ·|· ) .
Thus, the one-step total variation distance between the exact kernel and the proposed approximate kernel is upper bounded by
[Equation]
The following is a restatement of Lemma 3 in the supplementary ﬁle of Korattikara et al.
Lemma 4.
Consider two transition kernels T and T η , which yield stationary distributions π and π η , respectively.
If T follows a contraction such that there is a constant ρ ∈ [0 , 1) for all probability distributions µ :
[Equation]
and the uniform upper bound of the one step error between T and T η follows that
[Equation]
where ∆ ≥ 0 is a constant.
Then the total variation distance between π and π η is bounded by
[Equation]
Under the smoothness assumption C2 and the dissipative assumption C3, the geometric ergodicity of the continuous-time replica exchange Langevin diffusion to the invariant measure π has been established (Chen et al., 2019).
Combining the exponential convergence of the KL divergence (Deng et al., 2020) and the Pinsker’s inequality (Csiszár & Körner, 2011), we have that for any probability density µ , there exists a contraction parameter ρ ∈
(0 , 1) such that
[Equation]
where ρ depends on the spectral gap established in Raginsky et al.
(2017) and the swap rate.
Applying Lemma 4, the total variation distance can be upper bounded by
[Equation]
which concludes the proof of Theorem 2.
C. W EAK ERROR
In this section, we show the weak error (Lemma 5 below) between the Replica exchange Langevin diffusion ( β (1) t , β (2) t , · · · , β ( P ) t ) and the Replica exchange stochastic gradient Langevin diffusion ( (cid:101) β (1) t , (cid:101) β (2) t , · · · , (cid:101) β ( P ) t ) .
We denote P := P W × N , where P W is the inﬁnite dimensional Wiener measure and N is the Poisson measure independent of P W and has some constant jump intensity.
In our general framework below, the jump process α is introduced by swapping the diffusion matrix of the P Langevin dynamics and the jump intensity is deﬁned through the swap probability in the following sense, which ensures the independence of P W and N in each time interval [ iη, ( i +1) η ] , for i ∈ N + .
Without loss of generality, the replica exchange Langevin diffusion (reLD) is reformulated as below.
For any ﬁxed learning rate η > 0 , we deﬁne
[Equation]
α t ) d W t ,
P ( α ( t ) =
j | α ( t − dt )
= l, β ( (cid:98) t/η (cid:99) η ) = β ) =
aS ( β ) η 1 { t = (cid:98) t/η (cid:99) η } + o ( dt ) , for l (cid:54) =
j, (58) where ∇ U ( β ) :=
(cid:16) ∇ U ( β (1) ) , ∇ U ( β (2) ) , · · · , ∇ U ( β ( P ) ) (cid:17) T , and 1 t = (cid:98) t/η (cid:99) η 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 deﬁned as min { 1 , S ( β (1) , β (2) ) } and S ( β (1) , β (2) ) is deﬁned in (2).
In this case, the Markov Chain α ( t ) is a constant on the time interval [ (cid:98) t/η (cid:99) η, (cid:98) t/η (cid:99) η + η ) with some state in the ﬁnite-state space M = { 0 , 1 , · · · , 2 P − 1 − 1 } and the generator matrix Q follows,
[Equation]
− (cid:98) t/η (cid:99) η )
Q ii = − aS ( β ( i − 1) , ( i ) )
ηδ ( t − (cid:98) t/η (cid:99) η ) − aS ( β ( i ) , ( i +1) )
ηδ
( t − (cid:98) t/η (cid:99) η )
Q P ( P − 1) =
aS ( β ( P − 1) , ( P ) ) ηδ
( t − (cid:98) t/η (cid:99) η ) ,
Q PP = − aS ( β ( P − 1) , ( P ) ) ηδ ( t − (cid:98) t/η (cid:99) η ) , where 1 < i < P , and all the other terms in Q remain zero.
We denote δ ( · ) as a Dirac delta function.
Furthermore, for general adjacent swap with odd iterations (denoted as Q O ), and even iterations (denoted as Q E ), the generator matrices Q O and Q E ∈ R
( Pd ) ×
( Pd ) have the following form,
Q O =  
Q
Q  · · ·
Q− Q     − Q
Q
Q− Q     · · · · · · · · ·   ,
Q E =       0 − Q
Q  · · ·
Q− Q      − Q
Q· · ·
Q− Q · · · · · ·   ,
For the odd iteration case, the diffusion matrix for the ﬁrst two chains β (1) and β (2) has the following form according to the state of Σ( α t ) , i.e.
(Σ(0) , Σ(1)) :
=
    √ 2 τ (1) I d · · · √
[Equation]
I d
[Equation] [Equation]
[Equation]
· 

[Equation]
d
[Equation] [Equation]
[Equation] [Equation]
[Equation]
· · · · 
 ( Pd
[Equation]
( Pd )   , where the rest of the diffusion matrix remain unchanged.
Based on the observation of the transition matrix of Q O , and Q E , there is no interaction between every pairs (each pair has two chains) in each iteration, we will prove the weak error for only two chains.
We can then add different square blocks together as shown in Q O and Q E .
From now on, we restrict to the two chain case with β =
( β (1) , β (2) ) .
Next, we denote (cid:101) β as the Replica exchange stochastic gradient Langevin diffusion, for the same learning rate η > 0 as above, we have
[Equation]
) where ∇ (cid:101) U ( β ) := (cid:18) ∇ (cid:101) U (
[Equation]
[Equation]
) ∇ (cid:101) U ( β (2) )
(cid:19) and (cid:101) S ( (cid:101) β ) =
min { 1 , (cid:101) S η,m,n ( (cid:101) β (1) , (cid:101) β (2) ) } and (cid:101) S η,m,n ( (cid:101) β (1) , (cid:101) β (2) ) is shown in Eq.(16).
The distribution of process ( (cid:101) β t ) 0 ≤ t ≤ T is denoted as µ T :
= P (cid:101) G × N (cid:101) S , where (cid:101) α is a
Poisson process with jump intensity a (cid:101) S ( (cid:101) β ) ηδ ( t − (cid:98) t/η (cid:99) η ) on the time interval [ (cid:98) t/η (cid:99) η, (cid:98) t/η (cid:99) η + η ) .
For any open set U ⊂ R 2 d , which is contained in a large enough open ball with radius R centered at β (0) , we deﬁne stopping time
[Equation]
Lemma 5.
Assume the smoothness assumption C2 holds.
For any large enough open set U , we have
[Equation]
for any continuous differentiable and polynomial growth function h .
Proof.
We denote β t =
( β (1) t , β (2) t ) , which follows the dynamics below,
[Equation]
Denote u ( β t = x, t, i ) =
E x,i [ h ( β T ∧ τ U , α T
∧ τ U ) | ( β t , α t ) =
( x, i )] , i ∈ M .
For t <
T
∧ τ U , applying Feynman-Kac formula (see e.g. Baran et al.
, we have
[Equation]
Applying generalized Itô formula, for each i ∈ M we have
[Equation]
where M 1 + M 2 is a local martingale.
See details in (Yin & Zhu, 2010; Skorokhod, 2009).
We thus get, for kη ≤ t < ( k + 1) η < T
∧ τ U ,
[Equation]
For 0 < t 1 < t 2 < · · · , t N = T ∧ τ U , with t k = kη , we next derive the Itô formula for h ( (cid:101) β t , α t ) .
Note that β and (cid:101) β are deﬁned by using the same P -Brownian motion W , but with two different jump intensity on the time interval [ (cid:98) t/η (cid:99) η, (cid:98) t/η (cid:99) η + η ) .
We have the following approximation,
[Equation]
For stochastic noise ξ t k − 1 , with E [ ξ t k − 1 ] = 0 , we rewrite the above approximation as below,
[Equation]
For notation convenience, for t k − 1 ≤ t < t k , we keep the following convention,
[Equation]
For any t k − 1 ≤ t < t k , we apply Itô formula for u ( (cid:101) β t , t, i ) ,
[Equation]
where the last equality follows from (61)
[Equation]
and for t k − 1 ≤ t < t k ,
[Equation]
Evaluate the integral on time interval [0 , T ∧ τ U ]
[Equation]
and take expectation E x,i , we have
[Equation]
Where we have
[Equation]
In the next step, we introduce a sequence of stopping time based on our deﬁnition of process β and (cid:101) β .
For j ∈ N + , we denote ζ (cid:48)
j s as a stopping times deﬁned by ζ j +1 :
= inf { t > ζ j : α ( t ) (cid:54) =
α ( ζ j ) } and N ( T )
=
max { n ∈ N : ζ n ≤ T } .
According to the deﬁnition of process β , we have ζ j = lη , for some l ∈ N + .
Similarly, we have stopping time (cid:101) ζ j +1 :
=
inf { t > (cid:101) ζ j : (cid:101) α
( t ) (cid:54)
=
(cid:101) α ( ζ j ) } associated with process (cid:101) β , and ˜ α ( t ) follows the same trajectory of α ( t ) .
We thus have the following estimates for the terms I 1 , I 2 , and I 3 .
Estimate of I 1 : as for I 1 , −∇ U ( · ) does not depend on α , we get
I 1 = −
[Equation]
where we use the assumption E x,i [ ξ t k ] = 0 , and independent property of the noise ξ t k .
For t k ≤ t < t k +1 , let
[Equation]
applying Itô formula (same as (63) above), taking expectation, and using the fact (cid:101) Q ρ 1 ( (cid:101) β t , t )
= a ˜ S ( (cid:101) β (1) t , (cid:101) β (2) t )
[ ρ 1 (( β (1) t , β (2) t ) , t ) − ρ 1 (( β (2) t , β (1) t ) , t )]
= 0 , we get d E x,i [ ρ 1 ( (cid:101) β t , t )]
[Equation]
where the bound follows from the Weierstras theorem, and the same idea is also adopted in Sato &
Nakagawa (2014) for diffusion process without jumps.
We end up with,
[Equation]
which implies
[Equation]
Estimate of I 2 : for each time interval t k ≤ t < t k +1 , falling between two stopping time, the matrix Σ( α ( t ))
=
(cid:101) Σ( (cid:101) α ( t )) remains constant, and tr (ΣΣ T ∇ 2 u ) =
tr ( (cid:101) Σ (cid:101) Σ T ∇ 2 u ) =
(cid:80) Pj =
τ j ∆ (cid:101) β
j u , we have
I 2 = 0 .
Estimate of I 3 : we rewrite the time interval [0 , T ∧ τ U ] as the summation of the stopping times, i.e. (cid:82) T ∧ τ U =
[Equation]
N
[Equation]
1 l =
[Equation]
l , we thus have
[Equation]
Recall that,
[Equation] [Equation]
[Equation]
)) , and we denote (cid:101) Q ij as the components of the generator matrix (cid:101) Q of (cid:101) β .
For ζ l ≤ t < ζ l +1 , we denote ρ l 3 ( (cid:101) β t ,
[Equation]
In particular, for the two chain case, M = { 0 , 1 } .
For any t ≤ T ∧ τ U
[Equation]
there exists large enough open ball with radius R , denoted as U ⊂ B ( β (0) ,R )
[Equation]
such that U ⊂ B ( β (0) ,R )
[Equation]
which implies the uniform bound of u ( β t , t ) ≤
C
[Equation]
for t ≤ T ∧ τ U .
We thus get the estimates
[Equation]
where the last inequality follows from the fact that S, (cid:101) S ≤ 1 , and N ( T ∧ τ U ) ≤ K .
Combing all the estimates in I 1 , I 2 , I 3 , the proof is thus completed.