Non-reversible Parallel Tempering for Uncertainty Approximation in Deep Learning
Authors: Wei Deng
Qian Zhang
Qi Feng
Faming Liang
Guang Lin
Venue: ICLR 2022 Submitted (2022)
Keywords: replica exchange
parallel tempering
non-reversibility
stochastic approximation
round trip rate
deep learning

Abstract: 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 efficient 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 sufficient many $P$ chains. However, such an innovation largely disappears given limited chains in big data problems due to the extremely few bias-corrected swaps. To handle this issue, we generalize the DEO scheme to promote the non-reversibility and obtain an optimal communication cost $O(P\log P)$. 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.

Under review as a conference paper at ICLR 2022
PARALLEL TEMPERING FOR UNCERTAINTY APPROXIMATION IN DEEP LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
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 adopt efficient swap schemes. The popular deterministic even-odd (DEO) scheme exploits the non-reversibilityy property and has successfully reduced the communication cost from O(P2) to O(P) given sufficient many P chains. However, such an innovation largely disappears 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 exploration kernels. Such user-friendly nature enables conduct large-scale uncertainty approximation tasks without much tuning costs.
INTRODUCTION
Langevin diffusion is standard sampling algorithm that follows stochastic differential equation
d3: = (U,) + 22dW
where Bg E Rd, (U(-)) is the gradient the energy function U(),, We E Rd a Brownian motion, and is the temperature. The diffusion process converges stationary distribution x (Bf) xx and setting т = 1 yields a Bayesian posterior. When (-)) is convex, the rapid convergence has been widely studied in Durmus & Moulines (2016); Dalalyan (2017); however, when U/-)) is non-convex, slow mixing rate is inevitable (Raginsky et al., 2017). To accelerate the simulation, replica exchange Langevin diffusion (reLD) proposes include a high-temperature particle (! where P N+ ((1),
for exploration. Meanwhile, a low-temperature particle (( is presented for exploitation:
d - vU( dt + (PW)) (P)
db(1 - -U(( Jtt + d2 (1)
where 7PP) > 7(1) and (( is independent of w( To promote more explorations for the low-temperature particle, the particles at the position (g(1), g(P)) E 2dd swap with a probability
aS(3(),g(P))) - (Inelto - (1( -((6)))
(2)
where a E (0, is the swap intensity. In specific, the conditional swap rate at time t follows that P( = (B(P), ( (8(1,,(P)) ((P))) _ (())) P( = (80))g(())) = (3())g(P)) = - aa(()) g(P))at.
In the longtime limit, the Markov jump process converges to the joint distribution x ((8()),((P))
(18())
((S(P))
(18)
200
P(P) For convenience, we refer to the marginal distribution (1() x =00 and
uLg -(P)
(B)
as the target distribution and reference distribution, respectively.

---

Under review as a conference paper at ICLR 2022
2 PRELIMINARIES
Achieving sufficient explorations requires large r(P), which leads to limited accelerations due to a small overlap between x+() and x(P), To tackle this issue, one can bring in multiple particles with temperatures (r(2),-. ... (((-1)), where 7(1) 7(2)
< -(P), hollow out "tunnels". To maintain feasibility, numerous schemes are presented to select candidate pairs to attempt the swaps.
The all-pairs exchange (APE) attempts to swap arbitrary pair of chains (Brenner et al., 2007; Lingenheil et al., 2009), however, such method requires swap time (see definition in section A.5) of O(PS) and may not be user-friendly in practice.
APE
ADJ In addition to swap arbitrary pairs, one can also swap adjacen! (ADJ) pairs iteratively from (1,2), 3), to (P- 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 multi-core distributed context.
SEO The stochastic even-odd (SEO) scheme first divides the adjacent pairs {(p - 1, 1)p ,.....PP into and O, where and O denote even and odd pairs of forms (2p 1, 2p) and (2p, 2p + 1), respectively. Then, SEO randomly picks E or O pairs with 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 (()),, however, the gains in overcoming the sequential obstacle don't offset the ((P2) round trip time and SEO is still not effective enough.
Raforance ditteaoon
Tru
(a) Reversible indexes
(b) Nom-revensible indexes
(c) Non-reversible chains
Figure 1: Reversibility v.s. ronrrsversibll,ty In (a), reversible index takes O(P²) time to communicate; in (b), ideal non-reversible index moves along a periodic orbit, where dark and light arrows denote even and odd iterations, respectively; (c) shows
non-reversible chain conducts round trip with the DEO scheme.
DEO The deterministic even-odd (DEO) scheme instead attempts to swap even (E) pairs at even (E) iterations and odd (O) pairs 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 periodic orbit, as shown in Figure (b)). With large swap rate, Figure ((c) shows how the scheme yields an almost straight path and linear round trip time can be expected.
The power of hinges on maximizing the number of round trips, which equivalent to minimizing Ep=1 1-r, (Nadler & Hansmann, 2007b), where rp denotes the rejection
Equi-acceptance
1
rate for the chain pair (p. p + 1). Moreover, Ep=1 Tp converges to fixed barrier A P - oo (Predescu et al., 2004; Syed et al., 2021). Applying Lagrange multiplies to the constrained optimization problem leads T1 r2
P-1
rP-1 := r, where is the equi-rejection rate. In general, quadratic round trip time is required for ADJ and SEO due to the reversible indexes. By contrast, DEO only yields linear round trip time in terms of P as oc Syed al. (2021).
OPTIMAL NON- NON-REVERSIBLE SCHEME FOR PARALLEL TEMPERING
The linear round trip time is appealing for maximizing the algorithmic potential, however, such an advance only occurs given sufficiently many chains. In non asymptotic settings with limited chains, a pearl wisdom is to avoid frequent swaps (Dupuis al., 2012) and keep the average acceptance rate from 20% to 40% (Kone & Kofke, 2005; Lingenheil et al., 2009; Atchadé et al., 2011).
E shown iterations means even iterations; atherwise, even pairs
Ingic applies to
2

---

Under review as a conference paper at ICLR 2022
Most importantly, the acceptance rates severely reduced in big data due to the bias- 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.1 GENERALIZED DEO SCHEME
Continuing the equi-acceptane settings, we see in Figure.20 (()) 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- As such, with large equi-rejection rate r, the blue particle often makes little progress ((bd)) To handle this issue, the key is to propose small enough rejection rates to track the periodic orbit in Figure. (b)). Instead of pursuing excessive amount of chains, we resort to different solution by introducing the generalized even and odd iterations Ew and Ow, where W E N
Ew = [ mod 2 Ojk- 1, oo and Ow [[W] mod 2 = 1|k 1,2, 20). Now, we present the generalized DEO scheme with a window size W follows and refer to it DEOw:
o Attempt to swap E (or O) pairs at Fw (or Ow) iterations.
(3)
Allow most one swap during each cycle of Ew (or Ow) iterations.
(a) DEO
(b) Bad case 1
(c) Bad case
(d) Bad case
(e) DEO2
Figure 2: Illustration of DEO and DEO, In (a), we show an ideal case of DEO: (b-d) show bad cases of DEO based large equi-rejection rate r; (e) show how the generalized DEO scheme of window size tackles the issue with large r. The x-axis and y-axis denote (generalized) E (or O) iterations and E (or O) pairs, respectively. The dashed line denotes the failed swap attempts; the gray shaded areas are frozen to refuse swapping odd pairs at even iterations (or vice versa); the blue shaded area ensures at most one swap in a window.
As illustrated in Figure.2(e), the blue particle has a larger chance of (1 - +2)² to move upward 2 steps given W instead of (1 r)² when W 1, although the window number is also halved. Such trade-off inspires us to analyze the expected round trip time based on the window of size W. Although allowing at most one swap introduces the stopping time and may affect the distribution, the bias is rather mild due to the noisy energy estimators in big data. Check section C.2 for the details.
3.2 ANALYSIS OF ROUND TRIP TIME
To bring sufficient interactions between the reference distribution (P) and the target distribution (1),, we expect to minimize the expected round trip time (defined in section A.5) to ensure both efficient 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 BI and B2 in section B, for P (P 2 2) chains with window size W (W 1) and rejection rates P=t we have
P-1
E [ = 2WP + PWP
(4)
W
p=1
The proof in section B. shows that [[[]] increases as we adopt larger number of chains P and P rejection rates (rp)p=] In such a case, the round trip rate उ is also maximized by the key renewal theorem. In particular, applying W- recovers the vanilla DEO scheme.
3.3 ANALYSIS OF OPTIMAL WINDOW SIZE AND ROUND TRIP TIME
By Lemma 1, we observe a potential to remove the second quadratic term given an appropriate W. Such fact motivates to study the optimal window size W to achieve the best efficiency. Under
'he generalized DEO with the optimal window size is denoted by DEO, and will be studied in section 3.3.
3

---

Under review as a conference paper at ICLR 2022
the equi-acceptance settings, by treating the window size W as continuous variable and taking the derivative of [[] with respect to W, we have
a E
2P
(1 rh- {(1 - (P- r. W Woggr);
(5)
where r is the equi-rejection rate for adjacent chains. Define x := rWW E (0, 1), where W = log, (xx log
logr' The following analysis hinges the study of the solution g(x) (1 x)² (P- 1)x(1 - + log(z)) = O. By analyzing the growth of derivatives and boundary values, we can easily identify the uniqueness of the solution. Then, we proceed to verify that Plog P yields an asymptotic
1
approximation such that g( Ping p) log(log P)
- as P co. In the end, we have (1ng,) P) +0
Ing P
Theorem 1. Under Assumptions Bl and B2 based
aeci-acceptance settings, P 2.3. the
maximal round trip time is achieved when W = 1. If P 2 4, with the optimal window size where [-] the ceiling function. The round trip time follows O Plog
W. log P+log
log
log
The above result yields a remarkable round trip time of O(P log P) by setting the optimal window size W. By contrast, the vanilla DEO scheme only leads to a longer time of ((P2) $ Denoting by DEO, the generalized DEO scheme with the optimal window size W., we summarize the popular swap schemes in Table. 1, where the DEO, scheme performs the best among all the three criteria.
TABLE 1: ROUND TRIP TIME AND SWAP TIME FOR DIFFERENT SCHEMES. THE APE SCHEME REQUIRES
EXPENSIVE SWAP TIME O(P²) AND IS NOT COMPARED,
ROUND TRIP TIME (NON-ASYMPTOTIC)
ROUND TRIP (ASYMPTOTIC)
SWAP TIME
ADJ
O(P2) (NADLER & HANSMANN, 2007A)
O(P)) (NADLER & HANSMANK, 2007A)
O(P)
SEO
O(P2) (SYED ET AL.. 2021)
O(P))
(SYED ET AL., 2021)
O(1)
DEO
O(P2) (SYED ET Al., 2021)
O(P)
(SYED ET AL.,, 2021)
0(1)
DEO.
O P)
O(P)
O(1)
3.4
DISCUSSIONS ON THE OPTIMAL NUMBER OF CHAINS
Note that in practice given P parallel chains, a large P' leads to smaller equi-rejection rate r. As such, 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, the optimal number of chains follows that Ps > minp 3T(5 defined in (g.(1
og
(P)) log( (())) where op is
The assumptions and proof are postponed in section B.3. In mini-batch settings, insufficient 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 many chains as possible to yield a small enough equi-rejection rate to maintain the non-reversibility.
Cutoff phenomenon On the one hand, when we only afford at most P chains, where < P.. a large equi-rejection rate r is inevitable and DEO, is preferred over DEO; on the other hand, the rejection rate r goes when P > P. and DEO, recovers the DEO scheme.
In section B.4, we show that P. is in the order of thousands for the CIFAR example, which is hard to achieve due to the limited computational budget and further motivates us adopt finite chains with target swap rate S to balance between acceleration and accuracy.
4 USER-FRIENDE APPROXIMATE EXPLORATIONS IN BIG DATA
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. (2018); Zhang et al. (2020), we only adopt SGLD for exploitations.
By Taylor expansion, given large rejection rate r, log(r) r, which means -og(r] -((Fe))
4

---

Under review as a conference paper at ICLR 2022
Efficient explorations not only require a high temperature but also prefer large learning rate. Such demand inspires us to consider SGD with constant learning rate 17 as the exploration component
B =B 77 (B)) + <(B.)) 3x (/VU(B))
)2n(")
(6)
where s(3x) E Rd is the stochastic gradient noise. Under mild normality assumptions on e (Mandt al., 2017; Chen al., 2020), 3y converges approximately
invariant distribution, where the underlying temperature linearly depends on the learning rate Motivated by this fact, we propose an approximate transition kernel Tn with P parallel SGD runs based on different learning rates
( k+1 = (3 -
Exploration:
(7)
optional
Exploitation:
3 ( - (
E
where 1/()) < 1)(2) < (P) N(0, 2ny (1)) and 7 (1) 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.1 APPROXIMATION 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 difficulty in implementing the Metropolis rule (2). To tackle this issue, we utilize the randomness in stochastic energies and propose deterministic swap condition for the approximate kernel Ty in Eq.(7) such that
Deterministic swap condition: (gi1), (g("),
g((+)) - (ggp+t),g((>) if (((+++))) + < U(((P), (8)
where p (1, ... . P 1}, > 0 is a correction buffer to approximate the Metropolis rule Eq. Lemma 2. Assume the energy normality assumption (C1), then for any fixed OU, == U(B(P))
U(G(++1)), there exists an optimal C. E (0, (-b. - Joū] that perfectly approximates the random event ((g(r), 3(++1)) ((3t)),/(++111 > u, where op defined (1)) and Unif [0. 1].
The proof is postponed in section C.I. which paves the way for the guarantee that deterministic swap condition may replace the Metropolis rule (2. for approximation tasks. In addition, the normality assumption can be naturally extended to the asymptotic normality assumption (Quiroz al., 2019; Deng al., 2021) given large enough batch sizes.
Admittedly, the approximation error still exists for different aU, By the mean-value theorem, there exists a tunable C to optimize the overall approximation. Further invoking the central limit theorem such that E(-) in (6) approximates Gaussian distribution with fixed covariance, can expect 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 Typ which yield stationary distributions æ and Ty, respectively. Under smoothness (C2) and dissipativity assumptions (C3) (Mattingly al., 2002; Raginsky et al., 2017; Xu al., 2018), T satisfies the geometric ergodicity such that there is contraction constant E [0, 1) for any distribution μ.:
ll - -T|IIV= S pllu- Silu--llre Tllr.
where || 1|rv is the total variation (TV) distance. Moreover, assume that e(-) NO, M) for some positive definite matrix M (C4) (Mandi el al., 2017), then there a uniform upper bound of the one step error between T and Tap such that
1 TTT A Amax, Vp,
5

---

Under review as a conference paper at ICLR 2022
where maxx 2 0 is constant. Eventually, the TV distance between and Ty is bounded by A
|| Tnllor
max
I-P
The proof is postponed to section C.2. The SGD-based exploration kernels no longer require to fine-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 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.2 EQUT-ACCEPTANCE PARALLEL TEMPERING OPTIMIZED PATHS
Stochastic approximation (SA) is standard method to achieve 2011; Miasojedow et al., 2013), however, implementing this idea with fixed (7 and ry(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 logarithmic scale such that
(Atchadé et al.,
a Olgg(u) (())) h (1) (p)
where p E {1,2, , 13. u. = m (m+1) ni (p) ( ())p(,)))))))))))))))) (d3) is the mean-field function, π (p.P+1) is the joint invariant distribution for the p-th and + 1-th processes. In particular, H()(((() ,)) -000 S is the random-fied function to approximate (P(P)) (())) with limited perturbations, ux implicitly affects the distribution of the indicator function, and S is the target swap rate. Now consider stochastic approximation of (9),)), we have
log(u k+1) log(u (p) + + (p) (uk))
(p)
(p) B²),
(10)
where is the step size, Reformulating (q.() we have
K+1 max(0, (p) )en
(p)
where the max operator is conducted explicitly to ensure the sequence of learning rates is non- decreasing. This means that given fixed boundary learning rates (temperatures) 1/ (2-1) and flk (p+1) applying 10(0) = 1f 11- vd0) and 1760) 17(0+1) v((+1) for porpe (2,3, ..,,,.... P lead to
(p-1) + (0, v ee 11441 mB
(p) (p+1) - max(0, up )e
(1I)
forward sequence
backward sequence
Adaptive learning rates (temperatures) Now given fixed 17t), the sequence (2(2), 17(3), r)(P) can be approximated iteratively via the forward sequence of (11); conversely, given fixed (P) the backward sequence (PP-1)), (,(""-), ...,, (7 can be decided reversely as well. Combining the (p)
forward and backward sequences, Th+ can be approximated via
(p-1) + m (p+1) m0 (p) (e max(0, of (p+1) e
1/6
(12)
2
which resembles the binary search in the SA framework. In particular, the first 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- p) and (p.p + 1) until an equilibrium is achieved.
This is the first attempt to achieve aqui-acceptance given two fixed boundary values to our best knowledge. By contrast, Syed al. (2021) proposed to estimate the barrier to determine the temperatures and it easily fails in big data given finite number of chains and bias-corrected swaps.
'For convenience, (p) denotes the continuous-time diffusion at time t and vx represents the discrete approximations at iteration

---

Under review as a conference paper at ICLR 2022
Algorithm I Non-reversible parallel tempering with SGD-based exploration kernels (DEO,-SGD).
Input Number of chains P 2 3, boundary learning rates 17 and r7 ry(P), target swap rate 3, Input Optimal window size W log log log
total iterations K, and step sizes {7 1
log
for l K do _ ~ (T((3) following Eq.(7)
> Exploration / exploitation phase (parallelizable) Generalized even/odd iterations
P {V' 21 2...
P} mod 2 w
mod 2}.
for p = 1,2 to 1 do
A(r) = " +CE
gln) IA mod w-o-
Open the gate allow swaps
if p P and g(v) and A then
Swap: 3(m and g(a41)
b. Communication phase (parallelizable)
Freeze: g(v) = O.
Close the gate to refuse swaps
end if
if p > then
Update learning rate (eeperature) (q.(12)
end if end for Adaptive correction buffer: Cx+1 = Ck 7h (P'z Ex= Dp-140 A(r) S)
end for Output Models collected from the target temperature 11 Ik-1.
Adaptive correction buffers In addition, equi-acceptance does not guarantee convergence to the desired acceptance rate S. To avoid this issue, we propose to adaptively optimize follows
C++1 CC 7k
(13)
P P-
As k -> oc, the threshold and the adaptive learning rates converge to the desired fixed points. Note that setting a uniform C greatly simplifies the algorithm; in more delicate cases, problem-specifi rules are also recommended. Now we refer to the approximate non-reversible parallel tempering algorithm with the DEO, scheme and SGD-based exploration kernels as DEO,-SGD and formally formulate our algorithm in Algorithm Extensions of SGD with a preconditioner (Li al., 2016) or momentum (Chen et al., 2014) to further improve the approximation and efficiency are both straightforward (Mandt et al., 2017) and are denoted as DEOL-PSGD and DEO,-mSGD, respectively.
EXPERIMENTS
5.1 SIMULATIONS OF MULTI-MODAL DISTRIBUTIONS
We first simulate the proposed algorithm on distribution TB3) x exp(-U where 3 = (3 32). U(B) 0.2 + 82) ((c + cos(2 The heat map shown in Figure 3(a) with 25 modes of different volumes. To mimic big data scenarios, we can only access stochastic gradient (i3 VU(G) + (0, T2x2) and stochastic energy U(3) = U(3) + 2N(0,
(a) Ground truth
(b) = 0.2
(c) S = 0.3
(d) 8 = 0.4
(c) S = 0.5
(f) 8 = 0.6
Figure 3: Study of different target swap rate $ via DEO.-SGD SGD, where SGLD is the exploitation kemnel.
We first run DEO, -E0.- based on 16 chains and 20,000 iterations. We fix the lowest learning rate 0.003 and the highest learning 0.6 and propose to tune the target swap rate & for the acceleration- accuracy trade-off. Fig.3 shows that fixing S 0.2 or 0.3 is too conservative and underestimates the
7

---

Under review as a conference paper at ICLR 2022
uncertainty the corners; 8 = 0.6 results in too many radical swaps and eventually leads to crude estimations; by contrast, = 0.4 yields the best uncertainty approximation among the five choices.
Next, we select & 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 efficiency significantly 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 adaptively estimated to ensure the average acceptance rates converge to S 0.4.
Boung Bard
0.6
1.0
1.0
0.1
0.8
0.00
0.4
0.0
0.4
0.003
0.0
0.2
5000 10000 15000 20000
10000
2000
10000
20000
Window size
(a) Round trips
(b) Learning rates
(c) Acceptance rates
(d) Corrections
Figure 4: Study of window sizes, learning rates, acceptance rates, and the corrections.
We compare the proposed algorithm with parallel SGLD based on 20,000 iterations and 16 chains ; we fix the learning rate 0.003 and temperature We also run cycSGLD T16, which is short for single long chain based on 16 times of budget and cosine learning rates (Zhang al., 2020) of 100 cycles. We see Figure 5(b) that SGL.xxI6 has good explorations but fails to quantify the uncertainty. Figure 5(c) shows that xxSGGIDX116 explores most of the modes but overestimates some areas occasionally. Figure 5(d) demonstrates the DEO-SGD with 16 chains (DEO-SGD> x P16) estimates the uncertainty of the centering 9 modes well but fails to deal with the rest of the modes. As to DEO,-SGD P16, the approximation rather accurate, as shown in 5ie.
(a) Ground truth (b) SGLDxP16
(c) eycSGLDxT16
(d) DEO-SGDXPI6
(c)
Figure 5:
Simulations of the multi-modal distribution through different sampling algorithms.
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 long time to complete a round trip: by contrast, DEO, only conducts at most one cheap swap in window and yields much more deterministic paths.
(a) HEO-SGDxPI
(b) DEO,-SGDxPI
Figure 6: Dynamics of the index process. The red path denotes the round trip path for a particle.
5.2 UNCERTAINTY 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 first pre-train 10 fixed models via 300 epochs and then run our algorithm based on momentum SGD (mSGD) for 500 epochs with 10 parallel chains and denote by IEO..mSGDxx10 We fix the lowest and highest learning rates as 0.005 and 0.02, respectively. For a fair comparison, we also include the baseline DEO-mSGDxXn P10 with the same setup except that the window size is I; the standard ensemble mSGD x P00 also included with a learning rate of 0.005. In addition, we include two baselines based on single long chain, i.e. we run stochastic gradient Hamiltonian Monte Carlo
8

---

Under review as a conference paper at ICLR 2022
(Chen al., 2014) 5000 epochs with cyclical learning rates and 50 cycles (Zhang al., 2020) and refer to as cycSGHMCxTI0; we run SWAGx T10 (Maddox et al., 2019) under a similar setup.
In particular for DEO, -MSGDx1000 tune the target swap rate S and find optimum at S = 0.005. We compare our proposed algorithm with the four baselines and observe in Table.2 that mSGD xP10 easily obtain competitive results simply through model ensemble (Lakshminarayanan et al., 2017), which outperforms cycSGHMCX T10 and cycSWAGx on ResNet20 and ResNet32 models and perform the worst among the five methods on ResNet56; DEO-mSGD xP10 itself is already pretty powerful algorithm, however, DEO,-m x PIO consistently outperforms the vanilla alternative.
TABLE UNCERTAINTY APPROXIMATION AND OPTIMIZATION ON CIFARIO0 VIA 10x BUDGET.
ResNet20
ResNet32
ResNe156
MODEL
NLL
(%)
NLL
(%)
NLL
ACC (%)
sySSCMMCx110
8198159
76.2610.18
7401128
78.54_0.15
6460121
81.7820.08
eycSAAGXT10
8164138
76.1310.21
73891:32
78.62=0.13
6486:1229
81.60=0.11
mSGDxPIO
7902:164
76.59:0.11
7204:29
79.02=0.09
6553:15
81.49=0.09
DEO-mSGDxP10
7964122
76.84+0.12
7152#41
79.34=0.15
65341:26
81.72-0.1
DEO, mSGD
7741+67
77.37+0016
7019435
79.54 0.12
6439+-32
02.02-00.5
To analyze why the proposed scheme performs well, we study the round trips in Figure.7(a) and find 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 the training proceeds, the learning rates converge to fixed points and the acceptance rates for different pairs converge to the target swap rate.
0.00
0.000
Aittii My MYMON
d
60
Cran1
C
0.000
Crain 10
0.005
0.000
350
94h 1000
1500 1850
Spoch
Spsch
Esoch
(a) Round trips
(b) Accuracies
(c) Learning rates
(d) Acceptance rates
Figure 7: Study of window sizes, accuracies, learning rates, and acceptance rates ResNet20.
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 inefficient; by contrast, the DEO, scheme yields well-motivated paths with more deterministic round trips. Interestingly, this path resembles cyclic learning rates (Zhang et al., 2020), which provides novel viewpoint to interpret why cyclic learning rates work well empirically. Nevertheless, the stochastic and parallel manner further improves the margin and eventually leads the most efficient approximations in this task.
(a) DEO-DSGDD PP00mS
(b)
Figure 8: Dynamics of the index process of ResNet20 models in the last 200 epochs.
6 CONCLUSION
In this paper, we show how to conduct efficient PT big data problems. tackle the inefficiency issue of the popular DEO scheme given limited chains, we present the DEO, scheme by applying an optimal window size to encourage deterministic paths and obtain a significant acceleration of times. For user-friendly purpose, propose deterministic swap condition interact
iing ( with SGD-based exploration kernels and provide theoretical guarantce control the bias solely depending on the learning rate and the correction buffer C: we also provide a practical algorithm to adaptively approximate 77 and C for achieving the optimal efficiency in constrained settings.
O

---

Under review as a conference paper at ICLR 2022
REFERENCES
Sungjin Ahn, Anoop Korattikara, and Max Welling, Bayesian Posterior Sampling via Stochastic Gradient Fisher Scoring. In Proc. of the International Conference on Machine Learning (ICML), 2012.
Laurence Aitchison. A Statistical Theory of Cold Posteriors in Deep Neural Networks. In Proc. of the International Conference on Learning Representation (ICLR), 2021.
Yves F. Atchadé, Gareth O. Roberts, and Jeffrey S. Rosenthal. Towards Optimal Scaling of Metropolis- coupled Markov Chain Monte Carlo. Statistics and Computing, 21:555-568, 2011.
Nicholas Baran, George Yin, and Chao Zhu. Feynman-kac formula for switching diffusions: connections of systems of partial differential equations and stochastic differential equations. Advances in Difference Equations. 2013(1):1-13, 2013.
Rémi Bardenet, Arnaud Doucet, and Chris Holmes, On Markov Chain Monte Carlo Methods for Tall Data. Journal of Machine Learning Research, 18:1-43, 2017.
Raphaël Berthier, Francis Bach, and Pierre Gaillard. Tight Nonparametrie Convergence Rates for Stochastic Gradient Descent under the Noiseless Linear Model. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Paul Brenner, Christopher R. Sweet, Dustin VonHandorf, and Jesús A. Izaguirre. Accelerating the Replica Exchange Method through Efficient All-pairs Exchange. The Journal of Chemical Physics, 126:074103, 2007.
David Ceperley and Mark Dewing. The Penalty Method for Random Walks with Uncertain Energies. The Journal of Chemical Physics, 110:9812-9820 9820, 1999.
Pratik Chaudhari. Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing Gradient Descent into Wide Valleys. Proc. of the International Conference on Learning Representation (ICLR),
2017.
Tianqi Chen, Emily B. Fox, and Carlos Guestrin. Stochastic Gradient Hamiltonian Monte Carlo. In Proc. of the International Conference on Machine Learning (ICML), 2014.
Xi Chen, Jason D. Lee, Xin T. Tong, and Yichen Zhang. Statistical Inference for Model Parameters in Stochastic Gradient Descent. Annals of Statistics, 48(1)2251- 2020.
Yi Chen, Jinglin Chen, Jing Dong, Jian Peng, and Zhaoran Wang. Accelerating Nonconvex Learning via Replica Exchange Langevin Diffusion. In Proc. of the International Conference on Learning Representation (ICLR), 2019.
Imre Csiszár and János Körner. Information Theory; Coding Theorems for Discrete Memoryless Systems. Cambridge University Press, 2011.
Amak S Dalalyan. Theoretical Guarantees for Approximate Sampling from Smooth and Log-concave Densities, Journal of the Royal Statistical Society: Series B, 79(3):651-676, 2017.
Wei Deng, Qi Feng, Liyao Gao, Faming Liang, and Guang Lin. Non-Convex Learning via Replica Exchange Stochastic Gradient MCMC. In Proc. of the International Conference Machine Learning (ICML), 2020.
Wei Deng, Qi Feng, Georgios Karagiannis, Guang Lin, and Faming Liang. Accelerating Conver- gence Replica Exchange Stochastic Gradient MCMC via Variance Reduction. In Proc. of the International Conference on Learning Representation (ICLR), 2021.
Jing Dong and Xin T. Tong. Replica Exchange for Non-Convex Optimization. arXi::200.0833604 2021.
Paul Dupuis, Yufei Liu, Nuria Plattner, and D. Doll. On the Infinite Swapping Limit for Parallel Tempering. SIAM J. Multiscale Modeling Simulation, 10, 2012.
10

---

Under review as a conference paper at ICLR 2022
Alain Durmus and Éric Moulines. Sampling from Strongly Log-concave Distribution with the Unadjusted Langevin Algorithm. 1rXiv:1605.01559, 2016.
David J. Earl and Michael W. Deem. Parallel Tempering: Theory, Applications, and New Perspectives.
Phys. Chem. Chem. Phys., 7:3910-3916. 2005.
A. Gelman, W. Gilks, and G. O. Roberts. Weak Convergence and Optimal Sealing of Random Walk Metropolis Algorithms. Annals of Applied Probability, 7:1 7:110-120, 1997.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Pavel Izmailov, Dmitry Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging Weights Leads to Wider Optima and Better Generalization. In Proc. of the Conference on Uncertainty in Artificial Intelligence ( (UAI), 2018
David A. Kofke. On the Acceptance Probability of Replica-Fxchange Monte Carlo Trials. The Journal of Chemical Physics, 117, 2002.
Aminata Kone and David A. Kofke. Selection of Temperature Intervals for Parallel-temperig Simulations. The Journal of Chemical Physics, 122:206101, 2005.
Anoop Korattikara, Yutian Chen, and Max Welling, Austerity in MCMCLand Cutting the Metropolis- Hastings Budget. In Proc. of the International Conference on Machine Learning (ICML), 2014.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and Scalable Predictive Uncertainty Estimation using Deep Ensemble. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks. In Proc. of the National Conference on Artificial Intelligence (AAAI), pp. 1788-1794, 2016.
Martin Lingenheil, Robert Denschlag. Gerald Mathias, and Paul Tavan. Efficiency of Exchange Schemes in Replica Exchange. Chemical Physics Letters, 478:80-84, 2009.
Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A Simple Baseline for Bayesian Uncertainty in Deep Learning. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic Gradient Descent as Approxi- mate Bayesian Inference. Journal of Machine Learning Research, 18:1-35, 2017.
Oren Mangoubi and Nisheeth K. Vishnoi. Convex Optimization with Unbounded Nonconvex Oracles using Simulated Annealing. In Proc. of Conference on Learning Theory (COLT), 2018.
J.C. Mattingly, A.M. Stuartb, and D.J. Highamc. Ergodicity for SDEs and Approximations: Locally Lipschitz Vector Fields and Degenerate Noise. Stochastic Processes and their Applications, 101: 185-232 2002.
Btazej Miasojedow, Éric Moulines, and Matti Vihola. An Adaptive Parallel Tempering Algorithm. Journal of Computational and Gruphical Statistics, 22(3):649-664, 2013.
Walter Nadler and Ulrich H. E. Hansmann. Generalized Ensemble and Tempering Simulations: A Unified View. Phys. Rev. E, 75:026109, 2007a.
Walter Nadler and Ulrich H. E. Hansmann. Dynamics and Optimal Number of Replicas in Parallel Tempering Simulations. Phys. Rev. E, 76:065701, 2007b.
Tsuneyasu Okabe, Masaaki Kawata, Yuko Okamoto, and Masuhiro Mikami. Replica Exchange Monte Carlo Method for the 2001.
Ensemble. Chemical Physics Letters, 335:435-44
11

---

Under review as a conference paper at ICLR 2022
Cristian Predescu, Mihaela Predescu, and Cristian V. Ciobanu. The Incomplete Beta Function Law for Parallel Pempering Sampling of Classical Canonical Systems. Chemical Physics Letters, 120: 4119-4128, 2004.
Matias Quiroz, Robert Kohn, Mattias Villani, and Minh-Ngoe Tran. Speeding Up MCMC by Efficient Data Subsampling. Journal of the American Statistical Association, 1 11:8831-43,, 2019.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex Learning via Stochastic Gradient Langevin Dynamies: Nonasymptotic Analysis. In Proc. of Conference on Learning Theory (COLT), June 2017.
Herbert Robbins and Sutton Monro. A Stochastic Approximation Method. The Annals of Mathemati- cal Statistics, 22(3):400-407, 1951.
Issci Sato and Hiroshi Nakagawa. Approximation Analysis of Stochastic Gradient Langevin Dynamics by Using Fokker-Planck Equation and Ito Process. In Proc. of the International Conference on Machine Learning (ICML), 2014.
Daniel Seita, Xinlei Pan, Haoyu Chen, and John Canny. An Efficient Minibatch Acceptance Test for Metsopolis-Hastings. In Proc. of the Conference on Uncertainty Artificial Intelligence (UAI),
2017.
Anatolií Vladimirovich Skorokhod. Asymptotic methods in the theory of stochastic differential equations, volume 78. American Mathematical Soc., 2009.
Saifuddin Syed, Alexandre Bouchard-Côréé George Deligiannidis, and Arnaud Doucet. Non- Reversible Parallel Tempering: Scalable Highly Parallel MCMC scheme. arX.v:/905.020: 00293944 2021.
Max Welling and Yee Whye Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In Proc. of the International Conference on Machine Learning (ICML), pp. 681-688, 2011.
Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Swiatkowski. Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How Good is the Bayes Posterior in Deep Neural Networks Really? In Proc. of the International Conference Machine Learning (ICML), 2020.
Pan Xu, Jinghui Chen, Difan Zou, and Quanquan Gu. Global Convergence of Langevin Dynamics Based Algorithms for Nonconvex Optimization. In Advances in Neural Information Processing Systems (NeurIPS), 2018
George Yin and Chao Zhu. Hybrid Switching Diffussins: Properties and Applications. Springer, 2010.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning. In Proc. of the International Conference on Learning Representation (ICLR), 2020.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects. In Proc. of the International Conference on Machine Learning (ICML), 2019.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M. Kakade. Benign Overfitting of Constant-Stepsize SGD for Linear Regression. In Proc. of Conference on Learning Theory (COLT), 2021.
12

---

Under review as a conference paper at ICLR 2022
A BACKGROUND
A.1 REPLICA EXCHANGE STOCHASTIC GRADIENT LANGEVIN DYNAMICS
To approximate the replica exchange Langevin diffusion Eq.(1) based on multiple particles ( - ((P)) in big data scenarios, replica exchange stochastic gradient Langevin dy- namics (reSGLD) proposes the following numerical scheme:
+ + (2) -,,0
Exploration:
B ( 12 v( + _2
B (( (v + (2n
Exploitation:
(14)
where 1/1)) is the learning rate, c is standard d-dimensional Gaussian noise and each subprocess follows stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011) 1). Further assuming the energy normality assumption CI in section C.I, there exists op such that
Ü(g()) - -((((=)) ~ N(U())))
(U(((()))-(((((0 U(Br+t)) 207), for any (66) ~ π (P),
-
(15)
where p E (1,2, ... P 1]. in what follows, Deng et al. (2020) proposed the bias-corrected swap function as follows
aBt aB ((ptt)) p(ott) a. (Inelii) -
(pो., -( - ( )o%)
(16)
where p (1,2,.... P- 1}, (7b) f0 )oz is correction term avoid the bias and the swap intensity a can be then set to min(r("), ,pp:: dy for convenience. Namely, given a particle pair at location ((0+)))))))))))))))))))))))))))) in the kitth iteration, the conditional probability the swap follows that P(B (,(et)), 1 = (g(2), ((et2)))) ((3(0),g(tt)) P(2x+1 = (8(r),B(0t)) 3(++))) = (g(6) B(o+1))) = 5(3(r), glot
A.2 INHOMOGENOUS SWAP INTENSITY VIA DIFFERENT LEARNING RATES
Assume the high-temperature process also applies v7 (P+11 > r)(P) > O, where p (1,2..., 1}. At time the swap intensity can be interpreted being 0 in time interval t n((+1) 1)(1)) and being 7
1 in [t + r7(0+1) -)(2), .t +++ n((++))). Since the coupled Langevin diffusion process converges to the same joint distribution regardless of the swaps and the swap intensity a varies in bounded domain, the convergence of numerical schemes is not much affected except the numerical error.
A.3 ACCURACY-ACCHLLRRATTII TRADE-OFF
Despite the exponential acceleration potential, the average birs-cordetted swap rate significantly
reduced such that EJ = - m 22
(Deng et al., 2021), where S is the swap function following (2) in full-batch settings. Including more parallel chains are promising to alleviate this issue, however, it is often inevitable to sacrifice some accuracy to obtain more swaps and accelerations, discussed in section 3.4. As such, inspires us to devise the user-friendly SGD-based approximate exploration kernels.
A.4 CONNECTION TO NON-CONVEX OPTIMIZATION
Parallel to our work, a similar SGLDxSGD framework was proposed by Dong & Tong (2021) for non-convex optimization, where SGLD and SGD work as exploration and exploitation kernels,
13

---

Under review as a conference paper at ICLR 2022
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 manipulate the scheme to propose an exact swap in each window instead of at most 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 mechanism yields a larger hitting probability to move into 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.5 OTHERS
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, (p 1, 1, to attempt the swap and leads to swap time of O()); however, SEO, DEO, and DEO, don't 't have this issue because in each iteration, only even or odd chain pairs attempted swap, hence the swap time is (())
Round trip time refers to the time (stochastic variable) used in a round trip. A round trip is completed when a particle in the y-th chain, where p E [P] == (1,2,... P), hits the index boundary index at both 1 and P and returns back its original index p.
B ANALYSIS OF ROUND TRIPS
To facilitate the theoretical analysis, we follow Syed et al. (2021) and make the following assumptions
(B1) Stationarity: Each sub-process has achieved the stationary distribution 3(P) ~ ((P) for any p E [P]:
(B2) Weak independence: For any @ simulated from the p-th chain conditional on (50, U(Bi) and U(B))) are independent.
B.1 ANALYSIS OF ROUND TRIP TIME
Proof of Lemma 1. For t E N. define Z, E [ (1,,.... as the index of the chain a particle arrives after windows. Define d, E {1, 11 indicate the direction of the swap a particle intends to make during the 1-th window; i.e., the swap is between Z₁ and Z, + 1 if 6 = 1 and is between Z, and Z 1 if & = 1.
Define U == mi 2 0: Z. P. = 1 and V == minft > 0 2, = 1,4, 1}. Define
Tp : PP| reject the swap between Chain p and Chain p + for one time). Define up.d == E =
6o = 8) for 5 E (1, 11) and up. = EE(V p, 5n = 8). Then the expectation of round trip
time T is
BTT ((1.1 + uP.-
(17)
By the Markov property, for Up.5. we have
Up r (up.-1 + 1) + (1 +1 W (1 (up+1,1 +
(18)
Up - W. (Upp11 - 1) + (1 - 1- (up-1,-1 + 1),
(19)
where r denotes the rejection probability of a particle in a window of size W at the p-th chain. According to (18) and Eq.(1)), we have
Up+111 - Up = Tp (up+1,1 - Up,-1) -
(20)
Up 1 1, 1.-1 p-1 (Mpp1 1p-11-1) + 1.
W
(21)
14

---

Under review as a conference paper at ICLR 2022
Define Op Upp1 Up-1,-1- Then by definition, (q.() and Eq.( we have
apt1 Op = (Up+1,1 - Up.-1) - (11p.1 Up 1,-1)
(Up+1,1 [ppp1) (Upp-1 1 1p-11,11)
ry ap+1 -p- 1 1p - 2,
which implies that
ap+1 Op 22
(22)
for ap == (1 pW pp-1app. Thus, by (22), have
ap 02 - 2(pp- - 22).
(23)
By definition, 21,-1 U1 + 1. According to (2.). for p 1, we have
a2 (1 r (u2,1 -1,-1)
= (1 - r1 |ur,1 - -11--1 + 11 (u2,1 - -1,-1)-
(24)
= - r1 + W az.
Since rp E 10 for S p P. Eq. implies a2 = 22 which together with (23 implies
(1 p )a = Op 22(p- 1).
(25)
and therefore
p- p--1 ap -2(p- 1) W 1
W
(26)
pp-1
According to (2.( and (26), we have
UP, $1
Σ
p
W
11.-1) - (1,+1,1-Upp.11
(PP-)
(27)
Σ
W Aptl -PP-1) p
(28)
P-1
W
Σ
p
pWy P- - (P 11)
(29)
p=1
Since up,1 = 1, we have
P-1 W
1- w. P + (P 1) $1 = up,1 +
p=1
(30)
P-1 W
P+2 Σ
w D.
Similarly, for Up.6, we also have
Up, r (1pp 1 1) + (1 (1-r,, ) (1pt11, + 1)
W
up 1 + +1) (1 (1-1 11 1).
With the same analysis, we have
bp+1 - bp -2
(31)
W Up. Up-1,-1 = r p (Up.1 1p-1,-1) + 1,
(32)
where by : (1- - - 18 and 3p := Up.1 1p-11-1- According to Eq.(32), we have
P-1 vp U1 =
Σ
p Bpt1 + 11)
(33)
p=1
15

---

Under review as a conference paper at ICLR 2022
By definition, 2p,1 UP + 1. According to Eq.(2)) for p P, we have
bp = (1 r W
(1p11 - -1p-1.-1)
W = (1 - W [ur,1 -18-11 + --1 (UP, -r-1,,1)+
(34)
W -pp-1bp- 2(1 -rp-1)
Since rp E (0, for _ v P. (34) implies bp = 2. Then according to Eq.(31,, we have
bp = 2(P - + 1),
and therefore
W
pp 1pp
W
p-1
(P- p 1).
(35)
W
p-1
By E33( and (35() we have
P-1
pp
er- 11,
2
(P- - p) + (P 1).
W
p=1
Since -1 1, we have
PP =
Σ
p) P.
(36)
p
According to (q7) (3.) and (3,(3 we have
W r
E B W (m1,1 up,-1)
W P
The proof is generalization of Theorem (Syed 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.2 ANALYSIS OF OPTIMAL WINDOW SIZE
Proof of Theorem I. By treating W 2 1 as a continuous variable and taking the derivative of B[T] with respect to W and get
P-1
W
P-1
W log Tp
ET = 2P
1+
+
W
aw
W
W )
p=1
p=1
W
2W
+ W W
log Tp
2P
1
p
(37)
(1
p )2
P-1
W 9
W log Tp
2P
14
Σ
'p
W
1 )2
p=1
Assume that rp =r e 10 for 1 <p - P. Then we have
E 2WP 2WP(P 11) 1) 1 ww
(38)
8
2P
E
wwyz {(1 W (P 1)r (1 W + W Woogr))
(39)
OW
Define T. (0, 1). Hence W log, (x)
log and Tog r
a
2P
T3 T f(x) (1 r) {11 x)2 + (P- 1)x(1 + log
(40)
aw
16

---

Under review as a conference paper at ICLR 2022
Thus it suffices analyze the sign of the function g(æ) (1-2) (1-x)² + (P 1)r(1 + +log(x)) for I E (0, 1). For g(z), we have
g((x) (4-2P) 1) (P- 1) log(2)
and
d'(x)= =4 2P P-1
x
Thus, lim, g'(x) = -00, lim lim.--1 = g' = and gz is monotonically decreasing for 0.
For P> 2, we know that g"(x) when < x P_y and g=(x)
when
P-1 I < 2)-22
P-1 Therefore, g (z) is maximized 2)
(a) IfP = 3, by log(1+y) < for y > 0, we have gf (z) = 2( (xg)(1( < (0. Thus g(x) > g(1) = for E (0, 1) and therefore a 8w E > 0
O for any
for W E 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 unique solution
x² E (0, 1) such that g(æ) > 0 for Vx E (0,z*) and g(x) < 0 for Va E 11",1). Moreover, W log, (xx) the globally minimizer for the round trip time.
2. For P = we know that gx) > for x E (0, 1). Thus g' < g(1) for (0, 1) and therefore g(x) g((1) = 0 for x E (0, 1). Then according to (37)) aw E[T]
for W E N+ [[[] is globally minimized W
In what follows, we proceed to prove that Plog is a good approximation to x". In fact, we have
1
9 (
(1
P
log P - log(log P))
Plog P
Plog P
Plog P
Plog P
1
1
log(log P)
1
1+0
+
+0
P
log P
log P
P
log(log P)
+0
log P
log P
Thus, lim o r) O.
For x Plng p. we have W log,
1-+ log Ing
Then according to Eq.(38), for
P log p
P 4,
E T = 2P
log P + log log P
1+ (P - 1)
Plog P
- llggr r
1
Plog
1
2
(Plog P + P log log P)
log P
logr
Plog P
(41)
1
(Plog P P Poo log P)
log 1
1
log r
4log
4
(Plog P + P Plog log P)
log
conclusion, for = 2, 3, the maximum round trip rate achieved when the window size For P 2 with the window size W
log P log P. the round trip rate is least Ω log r
logr
Tog
Remark: Given finite chains with a large rejection rate, the round trip time is only of order
logg P) by setting the optimal window size W ~
log P+1og log
By contrast, the vanilla DEO
log
17

---

Under review as a conference paper at ICLR 2022
scheme with window of size yields much longer time of O(P2), where -lo((r) (
based on Taylor expansion and large > O.
B.2.1 TECHNICAL LEMMA
Proof of Lemma 3. To help illustrate the analysis below, we plot the graphs of g1x) and g(x) for (0, 1) and P 5 in Figure 9. For P 3, since g' (z) is maximized at z 2(P-2) (0,
P-1
P-1 and gx) < 0 when 2(P-2) < < 1, limm g (x) = with g² > O when 0 < < 2(P-2)
p-1
-oo, and g1) = O, we know that 9 (((P-) g((x) 0 (i.... g (ax) is monotonically decreasing) for any E (0, xo) and g'(x) > 0 (i.e., (xxx) is monotonically increasing) for any x E (ro, 1). Then g(x) on (0.1) 1) is globally minimized at ZT.- Moreover, - g(x) and g(1) 0. Thus, g(mo) < {((1) 0 and there exists
> O and there exists 2o E (0, 2(P-2) such that
<
x E (0, & (0, 1) such that g(x) > 0f E (0, and g(Φ) < 0f (1",
Meanwhile, by (37) and the definition z, we know that the sign of aw E[] is the same with that a
of (x) for W = log, Thus, 0 [[T] < 0 when W < W == log, (x") and aw E[T]] > 0 when By W > W', which implies that [[]] globally minimized W log, (x*) with some x E (0. 1).
1.00
400
025
000-
1.00
025
650
135
075
(a)
(xx
(b) g(x)
Figure 9: Illustration of functions g' g((x) and g(x) for x E (0, 1) and P
B.3 DISCUSSIONS ON THE OPTIMAL NUMBER OF CHAINS IN BIG DATA
To shed light on the suggestions of the optimal number of chains, we lay out two additional assump- tions:
(B3) Integrability: 3 is integrable with respect to (1) and or (P))
(B4) Approximate geometric spacing: Assume that
(P)
+ + +
+ for E PP
T
Assumption B3 is 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 verified in the CIFAR100 example in Figure 7(c).
Proof of corollary I. Before we start the proof, we denote by F and ă the average rejection and acceptance rates in full-batch settings, respectively. According to section 5.1 in Syed et al. (2021), Assumptions B1, B2 and B3 lead to
1 p3
$ +
0
(42)
P
18

---

Under review as a conference paper at ICLR 2022
where A is the barrier, which becomes larger if x(1) and xP)) have smaller probability overlap
(Predescu et al., 2004; Syed et al., 2021). Then acceptance rate full-batch settings
A P
s= -=
(43)
p3
By Lemma D4 (Deng et al., 2021) and (q,(1 a noisy energy estimator with variance (r(0) (-60)~((r1)) 1ya 1(*)" p yields an average swap rate of O "xxp (- (r())_r611))) (((0)")(++1)"
oz) In what follows,
the average rejection rate in mini-bater settings becomes
(scxp (-
(T10)- 22+ 6r(0))+(++1)
1-
o
(44)
where [P] and more accurate estimates is studied in proposition 2.4 (Gelman et al., 1997).
.(P)
Define A, Fay By assumption B4 and Taylor's theorem, we have
(T(tt1) - 7(7) 2
1 }) ))
2 log Ar 2
AF
-(p)
P P-1
+0
p2
(45)
Denote 7/- 1-- It follows that
(7(p) - 1(0
exp
7(=+)) o
(log A2 8(P -122 +o
1
of
exp
p2
-(0)
(46)
(7p log 42 8(P - 1)2 p2
(
exp
+o
Plugging it to Eq.(44), we have
r r=1-0 ((1- A
1 p3
(- (7p log A,
1
0
exp
8(P 1)2
+o
p2
(47)
Recall in (4.) the round trip time follows that
P Plgg P -logr
E[T]- O
(48)
By log(1 t) -t fort E [0, 1) and 1 = 1 + +(t((t for small have
log r
log
1-0 - +0 exp
(Tp A.
+ +0 (
8(P 112
1
(49)
0 0 A +0 exp
(3, A,
+ t (TtT
8(P-18
0
(1+ P
1
(7p log 22)
+
o
exp
+
p3
8(P 1)2
p2
Combining (49) and (4.(48), we have
A
212 (P- 1)2
ET =0 Plog P
1
+
+ 0
P
p8
p2
(50)
-0
Plog P exp
pz
19

---

Under review as a conference paper at ICLR 2022
where Up == Tp log A. The optimal P to minimize B[T] = O P exp
is equivalent to
2/2
p2
the minimizer of (x) log log legogrexp eeggrepp exp
log x + log log x +
x²
M
for x 2 4. Then
2p2 x
2 +
2
(x)))) I log x
log
(51)
x
Define s(x) x2 + log
Then have
2x log x 1 > 0,
s (xx) = 2x
(52)
(log x²2
for x 2 4. Thus s(x) increases with x for x 2 4. Obviously, limg-po s(z) co. Since M4 > 1, we know s(4) < Thus, there is single zero point P. E (4, 0o) of s(x) such that R2(x) [4, P.) and (x) > Oif P.. Therefore, ht(()) is minimized P. with
p2 1 log P. ) E 2 min n 2 max v.))
For any P. 2 2, have (1 log P.) E (1.1, 22) Thus, have the optimal number of chains mini-batch settings following that
P. E (1.1 min Up: /2 pm)xxpp)=> E min 37
dp
log A,, max op 2-(++1) log A.).
(53)
P
B.4 EMPIRICAL. JUSTIFICATION OF THE OPTIMAL NUMBER OF CHAINS
To obtain an estimate of the optimal number chains for PT in big data problems, we study the standard deviation of the noisy energy estimators on CIFAR 100 via ResNet20 models. We first pre-train a model and then try different learning rates to check the corresponding standard deviation of energy estimators. The reason 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 posterior affect with the temperature much smaller than 1 (Aitchison, 2021), applying Figure 10 and (53) concludes that P. is achieved in an order of thousands in the CIFARI00 example.
and the temperature in (6) and the cold
The conclusion is different from Syed et al. (2021) since big data problems require a much smaller swap rate maintain the unbiasedness of the swaps. On the one hand, insufficient chains may lead to insignificant swap rates to generate effective accelerations, but on the other hand, introducing too many chains may be too cost terms of limited memory and computational budget.
Standard deviation energy
6000
2000
8
Learning
1e-7
Figure 10: Analysis of standard deviation of energy estimators with respect to different learning rates through ResNet20 on CIFAR 100 dataset. Note that we have transformed the average likelihood into the sum of likelihood with total number of 50,000 datapoints. Hence, the learning rate is also reduced by 50,000 times.
20

---

Under review as a conference paper at ICLR 2022
C APPROXIMATION ERROR FOR SGD-BASED EXPLORATION KERNELS
This section proposes analyze the approximation error when adopt SGDs the approximate exploration kernels. Since it is independent of section B, different assumptions may be required.
The first subsection shows that for any Gaussian-distnn distributed energy difference estimator (((P)) U(B(rt1), there exists optimal C, such that the deterministic swap condition U (B(r+1) C. (((r)) perfectly approximates the random event S( B(++1)) u, where Unif(0. 1].
C.I GAP BETWEEN ACCEPTANCE RATES
To approximate the error when we adopt the deterministic swap condition (8),8), we first assume the following condition
(C1) Energy normality: The stochastic energy estimator for each chain follows normal distribu- tion with a fixed variance.
The above assumption has been widely assumed by Ceperley & Dewing (1999); Bardenet 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 0T, rin) -(++j O for any p E [P - 1]. the energy difference aU, U(B(P) - U(3(r+1)). By the energy normality assumption Cl and Eq,(15), the energy difference estimator follows that aU,- aU,()) 22pp, where & (0 Recall from (q.( that the bias-corrected acceptance rate follows that
= 1 (1 -8T
(54)
(i) For a uniformly distributed variable μ ~ Unif(0, 1], the base swap function satisfies that
, ((1-1)) u1aU,)
P(S g(ptt)) > vl0U,))du
P 07, p 01300 > (U
(55)
S P &>>
p
-OU, oT,O² 0 +
loR
8f, du
V20p
(ii) For the deterministic swaps based on (8), the approximate swap function follows that
P (((((ttD) + (((B("))0U,) = P(C < OU, + v20, = =P(5> -OU, +
V20,
(56)
21

---

Under review as a conference paper at ICLR 2022
Denote by D(u C) = P p <>
log -UU -
- PP<>>
-OU,-C v20,,
[0,,2) 2). Combining (55)
v20,
and (56), easily derive that
A(C) = (( 3 >
4/OU,)
- (((8(
+ +C (
02338+0000
-UU, + pT
log 0Ip
- <>
-dUp + C 220p
E>
du
v 22p
((u,C)) small enough
0 + p + 8T.
P E
-0Up + C
P p
du
V20p
V20p
D(u,C) large enough u.
[. J--or----once D(u, C)du
D(u, C)du -
Clearly, A(C) is continuous in C and is straightforward to verify that
or3e3 1
or
Duu, oT S(
A(0T,00))
3T3-3 073-3 D(u, No,,d
= l D(u, 8T, o2 20
Similarly, A(-20) threshold C should be strictly positive to avoid radical swap attempts, which implies that there an optimal solution C, E (0, of that solves A(C.) 0.
D(u, C)du s O. Moreover, the physical construction suggests that the
Remark 1. Note that the optimal C. may lead to few swaps given finite 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 (( .33 , O
or
This
wr 2
suggests that for well-approximated threshold C, the error is at most B[A(C) O(e
C.2 APPROXIMATE UPPER BOUND
The second subsection completes the proof regarding the numerical approximation of SGD-based exploration kernels. Nevertheless, we still require standard assumptions.
(C2) Smoothness: The function U(- is C-smooth there exists a positive Lipschitz constant C such that (x))) VU (2) S C'|æ - vll2 for every æ, y E IRd.
(C3) Dissipativity: The function U(-) is (a, diissippte there exist positive constants a and such that (x. VU(x)) 2 a||>=| 2 b.
(C4) Gradient normality: The stochastic gradient noise e()) in Eq.(6) follows a Normal distribution (0, with fixed positive definite 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 (Matiiggyy 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 Gaussian invariant measure.
22

---

Under review as a conference paper at ICLR 2022
Proof of Theorem 2. The transition kernel T of the continuous-time replica exchange Langevin diffusion (rel.D) follows that T IIp T with the following probability
where each sub-kernel T(P) draws a candidate A E IR
T (0) ((0) (1- 1,0 - s(B! 2,,
where p.9 E (1,2. ... P). p- = and is selected based on the DEO, scheme (3), Qp the proposal distribution in the p-th chain that models the probability to draw the parameter 0 via Langevin diffusion conditioned on B, and ((3 ,3(()) follows the swap function (2.
Now we consider the following approximations:
reLDD 4 reSGLD DEO,-SGLD 7 DEO,-SGLD IN DEO,-SGD,
where DEO,-SGLD resembles the DEO, scheme except that there are no restrictions on conducting most one swap.
Numerical approximation via SGLD By Lemma 5, we show that approximating replica ex- change Langevin diffusion via reSGLD yields weak error only depending the learning rate; similar results depending on the learning rate and the noise in gradient and energy noise have been studied in Lemma of (Deng et al., 2020).
I: Restrictions on at most one swap Adopting at most one swap the DEO, scheme includes stopping time, which may affect the underlying distribution. Fortunately, the bias can be controlled and becomes smaller in big data problems, which is detailed as follows.
Note that DEO,-SGLD differs from DEO, in that DEO, -SLD freezes at most W attempts of swaps each neighboring chains in each window, while DEO,-SGLD keeps all of these positions open. Recall that the bias-corrected swap rate follows that 5 = (Se Following similar technique in analyzing the bias through the difference of acceptance rates in Lemma 2, we can show that the bias is at most S P(Acceptance rate given frozen swaps)
by invoking (1.(
((e-))))) In other words, the restrictions at most one swap only lead to mild bias in big data problems and is less of a concern compared to the local trap problems. Empirically, we have verified the relation between the bias and the variance of noisy energy estimator in simulated example in section C.2.1.
III: Deterministic swap condition By Lemma 2 and the mean-value theorem, there exists an optimal correction buffer C, to approximate the random event S(3(2) (.(ptt) / through Ü(((P)) + C. < ((((+))) in the average sense for any ((P) and g(++1)
IV: Laplace approximation via SGD Assumption C4 guarantees that there exists optimal learning rate y) 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 the number of total data points goes to infinity.
Combining the above approximations, there exists a finite constant Amax 0 depending the learning rates 7,, the variance of noisy energy estimators, and the choice correction terms such that the total approximation еитог of the SGD exploration kernels with deterministic swap condition is upper bounded by A Amax- For any joint probability density p. the distance between the distributions generated by step of the exact transition kemel T and the approximate transition kernel T, in
23

---

Under review as a conference paper at ICLR 2022
Eq.(7) is upper bounded by
/ d2 (T(()) (T,((3)
- / d2(3))
/ du( 20) Amx (Q(3(0) -
1..00
s A Amax
dS2(3)
I du(d. vb) 1b.b
(0(31))+
2A 2Auxx;
where Q(---) = ПI, 2p(-).). Thus, the one-step total variation distance between the exact kernel and the proposed approximate kernel is upper bounded by
(BT (TT
dSu(B) (3) (TTT((3)
Amax.
(57)
The following a restatement of Lemma 3 in the supplementary file of Korattikara al. (2014).
Lemma 4, Consider two transition kernels T and Typ which yield stationary distributions т and Tay respectively, If follows contraction such that there constant p E [0, 1) for all probability distributions μ.
lu - - S pllm-x||m T||TV,
and the uniform upper bound of the one step error between 'T and Tr follows that
I| pTyll S A,
where 2 0 a constant. Then the total variation distance between π and Tn is bounded by
A
- TTllIv 1
1-p
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 E (0, 1) such that plle [|lu-=]|ıv.
IP
where P depends on the spectral gap established in Raginsky al. (2017) and the swap rate.
Applying Lemma 4, the total variation distance can be upper bounded by
||t TTylTV S max
1-p"
which concludes the proof of Theorem 2.
C.2.1 BIAS ANALYSIS FOR GENERALIZED WINDOWS IN BIG DATA PROBLEMS.
Following the empirical setup in section 5.1 (Deng al., 2020), we run parallel tempering based on two gradient Langevin dynamics (GLD) with noisy energy estimators of different variances and generalized windows of different sizes. We fix the learning rate 0.003 and set the temperatures and 10, respectively. The swap condition follows from (16). We run the algorithm 3,000,000 iterations and also include baseline by running the low-temperature process with 6,000,000 iterations.
We try different window sizes (2 and 5) and difference standard deviations (1, 2, and 3) for the noisy energy estimators. We observe in Fig. 11 that simply running GLD (red curves) suffers from the local trap issues and consistently lead to the worst performance; the exact energy estimators, as shown in the cyan curves, yield crude estimations as well due to the inclusion of stopping time. By contrast, as the variance increases, we see a pattern that the bias tends to decrease and a good trade-off is obtained sd 3, where the blue curves show remarkable approximations.
Further increasing sd (or the variance) may yield fewer swaps and less accelerations.
24

---

Under review as a conference paper at ICLR 2022
PT with Windowes
hark
Grunddnu (abr))
Figure 11: Bias analysis based on extended windows and nosiy energy estimators.
C.3 WEAK ERROR
In this section, we show the weak error (Lemma below) between the Replica exchange Langevin diffusion (B( .3.2) (B(P)) and the Replica exchange stochastic gradient Langevin diffusion (3 3(2) measure and N is the Poisson measure independent of PW and has some constant jump intensity. In our general framework below, the jump process a is introduced by swapping the diffusion matrix of the P Langevin dynamics and the jump intensity is defined through the swap probability in the following sense, which ensures the independence of pW and N in each time interval [án, (i+1 1)], for i N Without loss of generality, the replica exchange Langevin diffusion (reLD) is reformulated as below. For any fixed learning rate 1) > O, we define
B P((P) We denote
pw N, where pW is the infinite dimensional Wiener
d3: = - (B.) + )d.
(58)
P (a(() jla(t dt) 1,,3(1////m))) 3) = (8(3)) (1-(0)1)))) o(dt), for 1 j.
where VU(B) (V (
and 1==l//n1n the indicator func- tion, i.c. for every in with N+, given B(ir) 3, we have (a(t) = jaa(t dt)
VU(B
VU(B(P))
rS(m), where S(B) defined as min(1, (( ) } and S(3(2), ,g(()) defined in (2). In this case, the Markov Chain att) is constant on the time interval [[t/m] [t/n|n 1 7) with some state in the finite-state space M 1, I................ 2" 1} and the generator matrix follows.
Qn
- -aS(3())))) mnō(1 - [4/n]n), Q12 (2 [1/n]n)
Qiii1)) aS(3(i-1) )no(t - [t/n]n), Q((+11) = (S(3()) (t/n]n) -aS(B(--D.00 )n(( [1/n]n) -8(((()((+))))) -1//n)n) [t/n]n),
Qú
QP(P-1) -aS(3(P-1.(P)))) )n5(t [t/n/n). Qpp -S(3P--)(
where <i< P. and all the other terms in Q remain zero. We denote 5(.) as a Dirac delta function. Furthermore, for general adjacent swap with odd iterations (denoted Q²), and even iterations (denoted as )²), the generator matrices Qº and QF (Pd)x(PP)) have the following form,
Qu
Q12
0
0
0
0
Q21
Q21
0
Q23
Q23
0
0
Qo
0
0
Q4
Q44
QE
0
Q
Q
0
0
0
0
Q43
Q43
0
0
0
Q45
Q45
0
0
0
0
0
0
0
Q54
-Q54
0
For the odd iteration case, the diffusion matrix for the first two chains (3(1)
and
((2) has the following form according to the state of E(ar), i.e. (Σ(0), E(1)
(22
12
V 12
2
where
(Pd)x(Pd)
(Pd)x(P))
the rest of the diffusion matrix remain unchanged. Based on the observation of the transition matrix of Q", and Q there is no interaction between every pairs (each pair has two chains) in
25

---

Under review as a conference paper at ICLR 2022
each iteration, will prove the weak error for only two chains. We can then add different square blocks together as shown Qo and Q", From now on, we restrict to the two chain case with 3 = (3(1) (. Next, we denote 3 as the Replica exchange stochastic gradient Langevin diffusion, for the same learning rate 7j > 0 as above, we have
dB - - - (t/1/// dt + (((t//))))
(59)
(ā(1)) = j|ā(t - di) = L, (1,((4/m)l)) = 3) =((3)n + o(dt), for s
( and S(B) min (B
where Vū(B) (
and m, (3(1) is
shown in (q.( The distribution of process (B.)) is denoted as pr PG x N where ā is Poisson process with jump intensity ( - [t/n]n) on the time interval [[t/n]n. Lt/m]y + n).
For any open set U C IR2dd which is contained in large enough open ball with radius R centered at (3)), we define stopping time
Tu inf 2 0((3()),(()) & x M)).
(60)
Lemma 5. Assume the smoothness assumption C2 holds. For any large enough open set we have
E (M((3 A mu), aT A Au)) (M3(( A ru), aT A mn))]l O()). for any continuous differentiable and polynomial frowth.fucctinn h.
Proof. We denote Br = (1 (3(())) which follows the dynamics below,
dB, VU E(a))dW,
Denote
a (B, = (r,1),, i M. u( = -r,,,i (M
For < T A nu applying Feynman- Kac formula (see e.g. Baran et al. (2013)), we have
{
0,u(x,t, Lu(x,
(61)
Cu(x,t,8) trEEEt Vun( 1)) (u((,,, VU) + ((,
Applying generalized Itô formula, for each i M we have
h(3,, o.)) - h(Bo. aa) = -[[ Cg(B,, a, + (t) + M2(t),
(62)
where M1 + M2 is local martingale. See details in (Yin & Zhu, 2010; Skorokhod, 2009). We thus get, for Kn t (k 1)n T Tud
E [1 ar) - 1((x, = E - ( Ch(B,,
For < 11 l2
< ... IN T Tu. with 4 = knj, we next derive the Itô formula for h( at).
Note that ,3 and / are defined by using the same P--Brownian motion W, but with two different jump intensity the time interval [[t/nJn: Ll/njn 11). We have the following approximation,
3i2 -B__, = - vi (B...)) I(a,,_, n).
For stochastic noise E with E(Erx-11 = O, we rewrite the above approximation as below,
B.. Bon- = - + E )m + E(o
For notation convenience, for tx-1 S t < tk. we keep the following convention,
((G)) )(0()/)))) E(3% vū( (E(())))))))) (V
26

---

Under review as a conference paper at ICLR 2022
For any to-1 St tke apply Itō formula for u ((B,,t, t,),,
Du (B.) + d.i(() + dM2(t);
du t, = (d.u - Lu(B,, ()d d) + [(ā
(63)
dr
du (3)d + dM) (M2(()))
=
t, d) + [(a
0z
where the last equality follows from (61), and for t S l.. Lu(B,,t, 1, = tr V (VÜ (3.): Vu) + 113 (3.
Evaluate the integral on time interval [0, A Tu). and take expectation E we have
TTm
E (ulBABB T TA u(So, 0.i) =
E B=
J
= I + In Iz.
Where we have
TATU
E (VV((V vu.)). (VU( vu(B, d)]
==
TATu
I2 =
E Vz u t)) tr "REE" v (i, dt)
TAAT t)((B,tt) 2(Br, t),
I3 =
In the next step, we introduce a sequence of stopping time based on our definition of process 3 and G. For N we denote Ç/ as stopping times defined by Si+1 := inf(t Cj axl)) & a(G)) and N(7) max E N: &n S T). According to the definition of process B, we have G Iy, for some E N+ Similarly, we have stopping time Çj+1 := inf(t E, ā(t) * ā(1))) associated with process 3, and @(t) follows the same trajectory of att). We thus have the following estimates for the terms In Iz, and Iz.
Estimate of Iy: as for Iı, (-) does not depend on a, we get
K
Z =
Σ
E
((VU + - vU(B)), Vu(3r, [)
K=0
Mk11
-
E.
VU f)), Vu(3., t))
where we use the assumption 41-1. let
( = 0, and independent property of the noise En. For ty S <
pi (B., 1) (VVU (ū.)), Vu(3,
applying Itō formula (same (63) above), taking expectation, and using the fact (Br,t)
(5())))) ,3(2) ) (p((((()) ,3(2)), 1) - pi (()(2) ((()),)) dE (Br t)] dt
()] = 0, we get
E 1 (B., + tr T v² p1 (3,, (vū( (pI(3,, < Ck.
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,
E (3r. S Ckn
which implies
In < TnCmman-
27

---

Under review as a conference paper at ICLR 2022
Estimate of Iz: for each time interval ty t < tk+1, falling between two stopping time, the matrix [(a(t)) [(((t) remains constant, and vr(22>) tr(>>>/'u) have
u,
12 0.
Estimate of Iz: rewrite the time interval [0, T A Tu] the summation of the stopping times, i.c. So [T
Ei=0 (TArl))- f(i41 we thus have
TATU
Ia =
E
12, ()(,,,,, 2(3r, t), J]tt
Jo
N(TAPA--
Σ
- 1),B,, Jdl..
Recall that,
(1)))) = E - u(x,))).
jEMJati
and we denote Qij as the components of the generator matrix of
of B. For Gi S < C1+1, we denote
()(3, = (),,,,) (((T))
Σ (B. (6.)) (0,
JEM.)A(())
In particular, for the two chain case, M {0, 1}. For any 1 < T' nu. there exists large enough open ball with radius R, denoted as C Bia( R2), such that C B)a((). which implies the uniform bound of u uBe, 1) S C, for S nu. We thus get the estimates
N(TAm)) -1
C++1
I
Σ
E (3,)
1=0
N(TAnd)
6441
20
Σ
Σ
(1Q Qa(n))
1=0
JEMJZa(T)
N(TAAu)-
2Can
Σ
Σ
E (((((((())) Sa(m))
1=0
jEM.jta(())
2CKan.
where the last inequality follows from the fact that S, < 1, and N(T A Tu) s K. Combing all the estimates in I, Iz, Iz. the proof is thus completed.
28


Reviews:
[{"review_id": "mVw57pv9m3", "reviewer": "Reviewer_pLfL", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Pros** \n1. The contribution is very concrete\n2. $O(P)$ improvement is significant \n3. The proposed algorithm is based on a slight modification of an existing algorithm \n4. The paper is very well written \n5. The proof is quite straight words\n6. Literature is well-reviewed in section 2.  \n**My concerns**\n1.  Is possible to modify the rejection probability in DEO to get the same rounding time (for example use the rejection $O(r^{\\log(p)})$? \n2. I could not understand why the paper focuses on the gradient Langevin dynamics. Is it possible to use a different sampling algorithm? \n3. The notion of optimality is rather confusing in the abstract. Is it possible to show that it is not possible to get a communication cost lower than $O(P\\log(P))$? I think that the optimality is about the optimal window size for the adjusted DEO. \n4. I wondered whether authors use the real SGD algorithm or they use gradient descent + a Gaussian noise? \n5. The transition from the main result to applications for SGD can be more detailed and more clear. This application can be motivated more and explain why this particular case of parallel tempering is studied in the paper. The difference between SGD and Euler's discretization of Langevin dynamics can be explained in more detail. \n6. I recommend presenting experimental validations for Thm. 1 after this Thm. before talking about SGD. ", "summary_of_the_paper": "The paper studies the communicational complexity of parallel tempering. One example of parallel tempering is when one uses Langevin dynamics with different temperatures to balance exploit and exploration for sampling from a multi-modal distribution. The balance is made by swapping running particles with different temperatures according to a particular schedule. The round trip time (RTT) is a computational complexity metric for measure for the swapping schedule.  It is known that the best-known algorithm suffers from $O(P^2)$ RTT for $P$ parallel particles. By adjusting this algorithm, the authors improve the RTT to $O(P\\log(P))$. Then, they demonstrate the application of nonconvex with stochastic gradient descent. ", "summary_of_the_review": "I liked the contribution and narrative of this paper, and my opinion about the result is rather positive. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "gCLHog3lKJm", "reviewer": "Reviewer_mbau", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. The generalized DEO scheme is interesting and the O(P logP) expected round trip time for finite number of chains and non-diminishing rejection rate is novel. However, DEO* doesn't converge to the target distribution. The following example constitutes a counterexample of soundness of DEO*.\n\nConsider an energy function $f(x) = 8 x^2$ on the domain $[-1,1]$. We construct only two chains with temperatures $T_1=1$, $T_2=\\infty$. Two posterior should be $p_1$ as a truncated Gaussian distribution highly concentrated around origin, $p_2$ be a uniform distribution. Let local exploration kernel to be identity (e.g. SGLD with $0$ step size). Let window size to be large enough, such that swap will happens in one window with high probability. Initialize two chains with target distributions $p_1$ and $p_2$, then run DEO* with one even window and one odd window. The output distribution is no longer $p_1$ and $p_2$.\n\n2. This paper is clear until section 4 where DEO*-SGD method is proposed. Although practical, DEO*-SGD contains multiple approximations that require careful investigation:\n+ a. While using SGD as approximated local exploration kernel in (6), it seems only gradient noise is considered. However, discretization error is another source of noise. More specifically, for Langevin dynamics $d\\beta_t = -\\nabla U(\\beta_t) + \\sqrt{2} dW_t$, the discretization error for Euler discretization is $\\int_0^\\eta \\nabla U(\\beta_t) dt-\\eta\\nabla U(\\beta_k)$. How does this discretization error affect the distribution and temperature?\n+ b. on second line at page 5, the authors state that \"$\\beta_k$ converges approximately to an invariant distribution, where the underlying temperature linearly depends on the learning rate $\\eta$\". What is this invariant distribution? Does different step size induce different invariant distribution?\n+ c. With update rule in (12), can we ensure that $\\eta^{(p)}$ are always ordered?\n\n3. The experiments also need to be amended.\n+ a. Important baseline (E.g. Deep Ensemble) is missing. \n+ b. No uncertainty approximation quality (e.g. expected calibration error) is reported.\n+ c. Figure 7.c has incorrect y-axis value.\n+ d. experiment code for CIFAR100 experiment is not provided.", "summary_of_the_paper": "This paper proposes generalized DEO scheme (DEO*) that extends DEO with a window where at most one exchange could happen within one window. This paper also extend existing analysis on expected round trip time into DEO* and authors show that O(P logP) expected round trip time could be achieved even for finite number of chains and non-diminishing rejection rate. This paper also proposes multiple approximation to make DEO* practical: The exploration chains substitute gaussian noise with gradient noise, where noise magnitude is controlled with step size. Deterministic swap condition is used to avoid specific temperature. Adaptive learning rates and adaptive correction buffers are adopted to tune acceptance rate. The authors evaluate DEO* on ResNet model over CIFAR100.", "summary_of_the_review": "The proposed method does reduce expected round trip time but fails to converge. The practical method is not well understood and experiments need to be amended.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "9gm7t3isrTY", "reviewer": "Reviewer_ofJx", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I thank the authors for their article, which addresses the important topic of choice of swap schemes in parallel tempering. They build upon the current deterministic even-odd (DEO) swap scheme and develop modifications and generalisations which have more favourable theoretical and empirical performance. \n\nOriginality: The authors propose a generalisation of the existing DEO swap scheme. The details of their theoretical analysis and their methodological contributions are original and significant. Their investigation of stochastic approximations for big data, optimal number of chains, analysis of round trip time are detailed and through. \n\nClarity: The article is clearly written, even for readers who are less familiar with the swap schemes for parallel tempering. \n\nI have minor comments about the paper. They are:\n1) The approximation analysis theorem (theorem 2) assumes uniform ergodicity, which is stronger than “geometric ergodicity” as currently stated. Secondly, note that the bound given for theorem 2 is vacuous if (1-\\rho)<\\Delta. Some discussion of this point (that we require small \\delta and \\rho much less than 1 for the given upper bound \\Delta/(1-\\rho) to be useful) would be helpful.\n\n2) There should be some discussion about the choice of W (or target swap rate S) in practice, what algorithms/ heuristics are recommended when choosing the parameter W, and what additional computational cost this incurs. For example, at the moment the optimal window size is given as a function of the number of chains and the target swap rate- but how is the target swap rate chosen? I could be misunderstanding something here.\n\n3) If possible, the authors could present a version of Cor 1 which does not assume that the optimal W is chosen, but a generic W. Alternatively, some discussion about the importance of the choice of W should be included. ", "summary_of_the_paper": "The manuscript considers parallel tempering for sampling from multi-modal distributions. The choice of the swap scheme can great impact the performance of parallel tempering algorithms. The authors propose a modification of the existing deterministic even-odd (DEO) swap scheme. Theoretical results are established which show the proposed scheme improves communication cost from O(P^2) to O(P*logP) for P chains. Simulation studies are performed on a stylised multimodal distribution and on CIFAR100 datasets showing empirical improvements.  ", "summary_of_the_review": "This manuscript proposes new swap schemes for parallel tempering. The methodology is supported by theoretical analysis, detailed discussion of methodological details and simulations. It will be of interest to the ICLR community.\n\n\nEDIT: Having read the review of \"mbau\" and the corresponding discussion, I agree with the reviewer that the DEO* scheme may not converge to the target distribution of interest. I also agree the current argument in Section C.2 needs to be made more rigorous. Based on this point, I am inclined to revise my initial review score for the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9XDjp3-UIqJ", "reviewer": "Reviewer_kvcG", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea considered is very interesting and it demonstrates a lot of promises. The theoretical components seem to be strong too. I’ll be happy to be corrected and increase my rating, but I’d love to see more integration of various theoretical components, as well as how their implications get reflected in the empirical results. For example:\n* Is there any theoretical guide on the choice of $a$? By the way, what is the relation between $a$ and $\\mathbb{S}$?\n* If I understood correctly, Theorem 2 is just for one chain, not all the $P$ chains. When there are $P>1$ chains, don't errors add up? How does that scale with $P$?\n* How can Theorem 1 and Theorem 2 be combined to produce a theoretical guarantee of the final algorithm? It will be great if notions built around `round trip’ get translated into final iteration complexity.\n* Can wall clock counts be added to the experiments? Since improved efficiency seems to be a key contribution, it is a little odd to only show improved accuracy.\n* Can results for various W values be compared in the experiments, so that the theoretically optimal value could be justified?\n* Reading `We choose ResNet20, ResNet32, and ResNet56 (He et al., 2016) and train the models on CIFAR100.’ can I assume Bayesian neural network is used based on these architectures? What is the prior?\n", "summary_of_the_paper": "This paper considers sampling multi-modal distributions by the powerful parallel tempering methodology. An existing exchange scheme, DEO (deterministic even-odd), is improved by generalizing even and odd chain indices to even and odd windows of width W. Some nontrivial theoretical and empirical results are provided.", "summary_of_the_review": "This is a very interesting paper. Since the treatment is rather nontrivial and I feel it deserves appreciation, I’d like to see just a little more about how different theoretical and empirical ingredients of this work mesh together.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes the algorithm which they call DEO*-SGD, which is a combination of the ideas of the generalized DEO scheme, denoted by DEO*, to facilitate exploration (Section 3.1), adoption of stochastic gradient descent (SGD) in the exploration chains (i.e., those chains except the one with the lowest temperature) (Section 4), and use of adaptive tuning of learning rates (Section 4.2). The proposal is applied experimentally in Section 5 to demonstrate superiority of the proposal over existing approaches.

The initial review scores of the four reviewers were one positive and three negatives. Most reviewers positively evaluated the proposal, including the proposal of DEO* and its theoretical analysis, as well as its empirical usefulness in deep learning for a computer-vision task. On the other hand, some reviewers showed concern about soundness of the proposal. Upon reading the reviews and the author responses, as well as the paper itself, I think that this paper lacks a clear statement on its objective.

* **What does "uncertainty approximation" mean?:** The paper title would imply that the objective of the proposal in this paper is for "uncertainty approximation," but I could not find any concrete description on what it exactly is.
* **Sampling versus optimization:** The methods of Langevin dynamics, or more generally Markov-chain Monte-Carlo methods, have been used for two distinct purposes: sampling and optimization. In any case fast relaxation towards equilibrium would be of practical importance. For sampling purposes it is also important to assure that the stationary distribution of the Markov chain corresponds to the target distribution (In Langevin dynamics the target distribution would be the canonical ensemble defined by the energy $U(\cdot)$ and the temperature $\tau$). For optimization purposes, however, the assurance of the stationary distribution to be equal to the target distribution would be less of concern. It seems that the authors' interest would be in optimization rather than in sampling, but it is not clearly stated.
* **Soundness issue:** As Reviewer mbau pointed out, DEO* does not have a guarantee of convergence to the target distribution. I thought that if the objective of this paper would be in optimization rather than in sampling, the existence of approximation already in DEO* might be thought of as a minor problem, as the proposal already has other approximations introduced in Section 4. The authors claim that this problem does not affect the main body of the paper, but I feel that it would affect the overall organization of the paper, as the current organization seems to presume that approximation only resides in the adoption of the SGD-based exploration kernels with deterministic swap. In any case, this problem has been acknowledged by the authors themselves, as well as Reviewer ofJx.

In particular, the detailed discussion between the authors and Reviewer mbau has been very fruitful in clarifying technical subtleties in this manuscript, including the soundness issue mentioned above. At the same time, it would imply that this paper still has room for improvement.

An additional point I would like to mention is that this paper is not really self-contained, in the sense that several key notions and quantities are not defined or only defined in the Supplementary Materials ($\tilde{U}$ is not explicitly defined at all, the terms "swap time" and "round trip time" are defined in Appendix A.5, $\sigma_p$ in Corollary 1 and Lemma 2 is defined in Appendix A.1). 

All these weaknesses make me to think that another round of revision would be appropriate to properly judge the quality of this paper, whereas there is no such option within the review procedure of ICLR. I therefore cannot recommend acceptance of this paper at least in its current form.

Minor points (page and line numbers refer to the revised version):
- Abstract, line 5: "given sufficient many $P$ chains" would be better phrased as "given $P$ chains", as the big-O notation usually assumes the large-P asymptotic.
- In several places, there are periods after "Figure" and "Table", which are not needed.
- Page 3, line 32: In Lemma 2 there is apparently no such term found as "the second quadratic term". It should appear only after having assumed the equi-acceptance/rejection rates in equation (4), so that the sum becomes proportional to $P$.
- Theorem 1: "the maximal round trip time" should certainly be "the minimal round trip time". / is the ceiling function(. T -> , t)he round trip time
- Table 1: I did not really understand what "non-asymptotic" / "asymptotic" mean, as the big-O notation used here should by definition be asymptotic.
- Corollary 1: the optimal (number of) chains
- Page 4, line 34: The abbreviation SGLD is not defined in this paper.
- Page 4, line 36: similar(ly) to
- Equation (6): The sign of the last term should be "-".
Decision: Reject