Year: 2021

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 ) to O(P ) given sufficient many P chains. However, such an innovation largely disappears in big data problems due to the limited chains and extremely few bias-corrected swaps. To handle this issue, we generalize the DEO scheme to promote the non-reversibility and obtain an appealing communication cost O(P logP ) based on the optimal window size. In addition, we also analyze the bias when we adopt stochastic gradient descent (SGD) with large and constant learning rates as exploration kernels. Such a user-friendly nature enables us to conduct large-scale uncertainty approximation tasks without much tuning costs.

heading: 1 INTRODUCTION
text: Langevin diffusion is a standard sampling algorithm that follows a stochastic differential equation
dβt = −∇U(βt)dt+ √ 2τdWt,
where βt ∈ Rd,∇U(·) is the gradient of the energy function U(·),Wt ∈ Rd is a Brownian motion, and τ is the temperature. The diffusion process converges to a stationary distribution π(β) ∝ e− U(β) τ and setting τ = 1 yields a Bayesian posterior. When U(·) is convex, the rapid convergence has been widely studied in Durmus & Moulines (2016); Dalalyan (2017); however, when U(·) is non-convex, a slow mixing rate is inevitable (Raginsky et al., 2017). To accelerate the simulation, replica exchange Langevin diffusion (reLD) proposes to include a high-temperature particle β(P )t , where P ∈ N+\{1}, for exploration. Meanwhile, a low-temperature particle β(1)t is presented for exploitation:
dβ (P ) t = −∇U(β (P ) t )dt+
√ 2τ (P )dW
(P ) t ,
dβ (1) t = −∇U(β (1) t )dt+
√ 2τ (1)dW
(1) t ,
(1)
where τ (P ) > τ (1) and W (P )t is independent of W (1) t . To promote more explorations for the low-temperature particle, the particles at the position (β(1), β(P )) ∈ R2d swap with a probability
aS(β(1), β(P )) = a · ( 1 ∧ e ( 1 τ(1) − 1 τ(P ) )( U(β(1))−U(β(P )) )) , (2)
where a ∈ (0,∞) is the swap intensity. In specific, the conditional swap rate at time t follows that
P(βt+dt = (β(P ), β(1))|βt = (β(1), β(P ))) = aS(β(1), β(P ))dt, P(βt+dt = (β(1), β(P ))|βt = (β(1), β(P ))) = 1− aS(β(1), β(P ))dt.
In the longtime limit, the Markov jump process converges to the joint distribution π(β(1),β(P )) ∝
e −U(β
(1))
τ(1) −U(β
(P ))
τ(P ) . For convenience, we refer to the marginal distribution π(1)(β) ∝ e− U(β) τ(1) and
π(P )(β) ∝ e− U(β) τ(P ) as the target distribution and reference distribution, respectively.

heading: 2 PRELIMINARIES
text: Achieving sufficient explorations requires a large τ (P ), which leads to limited accelerations due to a small overlap between π(1) and π(P ). To tackle this issue, one can bring in multiple particles with temperatures (τ (2), · · · , τ (P−1)), where τ (1) < τ (2) < · · · < τ (P ), to hollow out “tunnels”. To maintain feasibility, numerous schemes are presented to select candidate pairs to attempt the swaps.
APE The all-pairs exchange (APE) attempts to swap arbitrary pair of chains (Brenner et al., 2007; Lingenheil et al., 2009), however, such a method requires a swap time (see definition in section A.5) of O(P 3) and may not be user-friendly in practice.
ADJ In addition to swap arbitrary pairs, one can also swap adjacent (ADJ) pairs iteratively from (1, 2), (2, 3), to (P − 1, P ) under the Metropolis rule. Despite the convenience, the sequential nature requires to wait for exchange information from previous exchanges, which only works well with a small number of chains and has greatly limited its extension to a multi-core or distributed context.
SEO The stochastic even-odd (SEO) scheme first divides the adjacent pairs {(p − 1, p)|p = 2, · · · , P} into E and O, where E and O denote even and odd pairs of forms (2p − 1, 2p) and (2p, 2p + 1), respectively. Then, SEO randomly picks E or O pairs with an equal chance in each iteration to attempt the swaps. Notably, it can be conducted simultaneously without waiting from other chains. The scheme yields a reversible process (see Figure 1(a)), however, the gains in overcoming the sequential obstacle don’t offset the O(P 2) round trip time and SEO is still not effective enough.
DEO The deterministic even-odd (DEO) scheme instead attempts to swap even (E) pairs at even (E) iterations and odd (O) pairs at odd (O) iterations alternatingly† (Okabe et al., 2001). The asymmetric manner was later interpreted as a non-reversible PT (Syed et al., 2021) and an ideal index process follows a periodic orbit, as shown in Figure 1(b). With a large swap rate, Figure 1(c) shows how the scheme yields an almost straight path and a linear round trip time can be expected.
Equi-acceptance The power of PT hinges on maximizing the number of round trips, which is equivalent to minimizing ∑P−1 p=1 1 1−rp (Nadler & Hansmann, 2007b), where rp denotes the rejection
rate for the chain pair (p, p + 1). Moreover, ∑P−1 p=1 rp converges to a fixed barrier Λ as P → ∞ (Predescu et al., 2004; Syed et al., 2021). Applying Lagrange multiplies to the constrained optimization problem leads to r1 = r2 = · · · = rP−1 := r, where r is the equi-rejection rate. In general, a quadratic round trip time is required for ADJ and SEO due to the reversible indexes. By contrast, DEO only yields a linear round trip time in terms of P as P →∞ Syed et al. (2021).

heading: 3 OPTIMAL NON-REVERSIBLE SCHEME FOR PARALLEL TEMPERING
text: 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 of wisdom is to avoid frequent swaps (Dupuis et al., 2012) and to keep the average acceptance rate from 20% to 40% (Kone & Kofke, 2005; Lingenheil et al., 2009; Atchadé et al., 2011).
† E shown in iterations means even iterations; otherwise, it denotes even pairs for chain indexes. The same logic applies toO.
Most importantly, the acceptance rates are severely reduced in big data due to the bias-corrected swaps associated with stochastic energies (Deng et al., 2020), see details in section A.1. As such, maintaining low rejection rates becomes quite challenging and the issue of quadratic costs still exists.

heading: 3.1 GENERALIZED DEO SCHEME
text: As illustrated in Figure.2(e), the blue particle has a larger chance of (1 − r2)2 to move upward 2 steps given W = 2 instead of (1− r)2 when W = 1, although the window number is also halved. Such a trade-off inspires us to analyze the expected round trip time based on the window of size W . 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.

heading: 3.2 ANALYSIS OF ROUND TRIP TIME
text: To bring sufficient interactions between the reference distribution π(P ) and the target distribution π(1), we expect to minimize the expected round trip time T (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 B1 and B2 in section B, for P (P ≥ 2) chains with window size W (W ≥ 1) and rejection rates {rp}P−1p=1 , we have
E[T ] = 2WP + 2WP P−1∑ p=1 rWp 1− rWp . (4)
The proof in section B.1 shows that E[T ] increases as we adopt larger number of chains P and rejection rates {rp}P−1p=1 . In such a case, the round trip rate PE[T ] is also maximized by the key renewal theorem. In particular, applying W = 1 recovers the vanilla DEO scheme.

heading: 3.3 ANALYSIS OF OPTIMAL WINDOW SIZE AND ROUND TRIP TIME
text: By Lemma 1, we observe a potential to remove the second quadratic term given an appropriate W . Such a fact motivates us to study the optimal window size W to achieve the best efficiency. Under
§The generalized DEO with the optimal window size is denoted by DEO? and will be studied in section 3.3.
the equi-acceptance settings, by treating the window size W as a continuous variable and taking the derivative of E[T ] with respect to W , we have
∂
∂W E[T ] =
2P (1− rW )2 { (1− rW )2 + (P − 1)rW (1− rW +W log r) } , (5)
where r is the equi-rejection rate for adjacent chains. Define x := rW ∈ (0, 1), where W = logr(x) = log x log r . The following analysis hinges on the study of the solution g(x) = (1− x)
2 + (P − 1)x(1− x+ log(x)) = 0. By analyzing the growth of derivatives and boundary values, we can easily identify the uniqueness of the solution. Then, we proceed to verify that 1P logP yields an asymptotic
approximation such that g( 1P logP ) = − log(logP ) logP +O ( 1 logP ) → 0 as P →∞. In the end, we have
Theorem 1. Under Assumptions B1 and B2 based on equi-acceptance settings, if P = 2, 3, the maximal round trip time is achieved when W = 1. If P ≥ 4, with the optimal window size W? ≈ ⌈ logP+log logP − log r ⌉ , where d·e is the ceiling function. The round trip time follows O(P logP− log r ).
The above result yields a remarkable round trip time of O(P logP ) by setting the optimal window size W?. By contrast, the vanilla DEO scheme only leads to a longer time of O(P 2) §. 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.

heading: 3.4 DISCUSSIONS ON THE OPTIMAL NUMBER OF CHAINS
text: Note that in practice given P parallel chains, a large P leads to a 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 P? > minp σp 3τ(p) log( τ (P ) τ(1) ), where σp is defined in Eq.(15).
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 as many chains as possible to yield a small enough equi-rejection rate r to maintain the non-reversibility.
Cutoff phenomenon On the one hand, when we only afford at most P chains, where P < P?, a large equi-rejection rate r is inevitable and DEO? is preferred over DEO; on the other hand, the rejection rate r goes to 0 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 CIFAR100 example, which is hard to achieve due to the limited computational budget and further motivates us to adopt finite chains with a target swap rate S to balance between acceleration and accuracy.

heading: 4 USER-FRIENDLY APPROXIMATE EXPLORATIONS IN BIG DATA
text: 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 a large rejection rate r, − log(r) = 1− r, which means 1− log(r) = O( r 1−r ).
Efficient explorations not only require a high temperature but also prefer a large learning rate. Such a demand inspires us to consider SGD with a constant learning rate η as the exploration component
βk+1 = βk − η(∇U(βk) + ε(βk)) = βk − η∇U(βk) + √ 2η (η
2
) ε(βk), (6)
where ε(βk) ∈ Rd is the stochastic gradient noise. Under mild normality assumptions on ε (Mandt et al., 2017; Chen et al., 2020), βk converges approximately to an invariant distribution, where the underlying temperature linearly depends on the learning rate η. Motivated by this fact, we propose an approximate transition kernel Tη with P parallel SGD runs based on different learning rates
Exploration:  β (P ) k+1 = β (P ) k − η(P )∇Ũ(β (P ) k ),
· · · β
(2) k+1 = β (2) k − η(2)∇Ũ(β (2) k ),
Exploitation: β(1)k+1 = β (1) k − η (1)∇Ũ(β(1)k ) + optional︷︸︸︷ Ξk ,
(7)
where η(1) < η(2) < · · · < η(P ), Ξk ∼ N (0, 2η(1)τ (1)), and τ (1) is the target temperature. Since there exists an optimal learning rate for SGD to estimate the desired distribution through Laplace approximation (Mandt et al., 2017), the exploitation kernel can be also replaced with SGD based on constant learning rates if the accuracy demand is not high. Regarding the validity of adopting different learning rates for parallel tempering, we leave discussions to section A.2.

heading: 4.1 APPROXIMATION ANALYSIS
text: 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 Eq.(2). To tackle this issue, we utilize the randomness in stochastic energies and propose a deterministic swap condition for the approximate kernel Tη in Eq.(7) such that
Deterministic swap condition: (β(p),β(p+1))→ (β(p+1),β(p)) if Ũ(β(p+1)) + C < Ũ(β(p)), (8) where p ∈ {1, 2, · · · , P − 1}, C > 0 is a correction buffer to approximate the Metropolis rule Eq.(2). Lemma 2. Assume the energy normality assumption (C1), then for any fixed ∂Up := U(β(p)) − U(β(p+1)), there exists an optimal C? ∈ (0, ( 1τ(p) − 1 τ(p+1)
)σ2p] that perfectly approximates the random event S̃(β(p),β(p+1)) > u, where σp defined in Eq.(15) and u ∼ Unif [0, 1].
The proof is postponed in section C.1, which paves the way for the guarantee that a deterministic swap condition may replace the Metropolis rule Eq.(2) for approximation tasks. In addition, the normality assumption can be naturally extended to the asymptotic normality assumption (Quiroz et al., 2019; Deng et al., 2021) given large enough batch sizes.
Admittedly, the approximation error still exists for different ∂Up. By the mean-value theorem, there exists a tunable C to optimize the overall approximation. Further invoking the central limit theorem such that ε(·) in Eq.(6) approximates a Gaussian distribution with a fixed covariance, we can expect a bounded approximation error for the SGD-based exploration kernels (Mandt et al., 2017). Theorem 2. Consider the exact transition kernel T and the proposed approximate kernel Tη , which yield stationary distributions π and πη, respectively. Under smoothness (C2) and dissipativity assumptions (C3) (Mattingly et al., 2002; Raginsky et al., 2017; Xu et al., 2018), T satisfies the geometric ergodicity such that there is a contraction constant ρ ∈ [0, 1) for any distribution µ:
‖µT − π‖TV ≤ ρ‖µ− π‖TV, where ‖ · ‖TV is the total variation (TV) distance. Moreover, assume that ε(·) ∼ N (0,M) for some positive definite matrixM (C4) (Mandt et al., 2017), then there is a uniform upper bound of the one step error between T and Tη such that
‖µT − µTη‖TV ≤ ∆max,∀µ,
where ∆max ≥ 0 is a constant. Eventually, the TV distance between π and πη is bounded by
‖π − πη‖TV ≤ ∆max 1− ρ .
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 a well-controlled approximations by solely tuning η = (η(1), · · · , η(P )) and C. In addition, our proposed algorithm for uncertainty approximation is highly related to non-convex optimization. For the detailed discussions, we refer interested readers to section A.4.

heading: 4.2 EQUI-ACCEPTANCE PARALLEL TEMPERING ON OPTIMIZED PATHS
text: Stochastic approximation (SA) is a standard method to achieve equi-acceptance (Atchadé et al., 2011; Miasojedow et al., 2013), however, implementing this idea with fixed η(1) and η(P ) is rather non-trivial. Motivated by the linear relation between learning rate and temperature, we propose to adaptively optimize the learning rates to achieve equi-acceptance in a user-friendly manner. Further by the geometric temperature spacing commonly adopted by practitioners (Kofke, 2002; Earl & Deem, 2005; Syed et al., 2021), we adopt the following scheme on a logarithmic scale such that
∂ log(υ (p) t ) = h (p)(υ (p) t ), (9)
where p ∈ {1, 2, · · · , P − 1}, υ(p)t = η (p+1) t − η (p) t , h (p)(υ (p) t ) :=
∫ H(p)(υ
(p) k ,β)π (p,p+1)(dβ) is the mean-field function, π(p,p+1) is the joint invariant distribution for the p-th and p+ 1-th processes. In particular, H(p)(υ(p)k ,β) = 1Ũ(β(p+1))+C<Ũ(β(p))− S is the random-field function to approximate h(p)(υ
(p) k ) with limited perturbations, υ (p) k † implicitly affects the distribution of the indicator function, and S is the target swap rate. Now consider stochastic approximation of Eq.(9), we have
log(υ (p) k+1) = log(υ (p) k ) + γkH (p)(υ (p) k ,βk), (10)
where γk is the step size. Reformulating Eq.(10), we have
υ (p) k+1 = max(0, υ (p) k )e γkH(υ (p) k ),
where the max operator is conducted explicitly to ensure the sequence of learning rates is nondecreasing. This means that given fixed boundary learning rates (temperatures) η(p−1)k and η (p+1) k , applying η(p) = η(p−1) + υ(p) and η(p) = η(p+1) − υ(p+1) for p ∈ {2, 3, · · · , P − 1} lead to
η (p−1) k + max(0, υ (p) k )e γkH(υ (p) k )︸ ︷︷ ︸
forward sequence
= η (p) k+1 = η (p+1) k −max(0, υ (p+1) k )e γkH(υ (p+1) k )︸ ︷︷ ︸
backward sequence
. (11)
Adaptive learning rates (temperatures) Now given a fixed η(1), the sequence η(2), η(3), · · · , η(P ) can be approximated iteratively via the forward sequence of (11); conversely, given a fixed η(P ), the backward sequence η(P−1), η(P−2), · · · , η(1) can be decided reversely as well. Combining the forward and backward sequences, η(p)k+1 can be approximated via
η (p) k+1 : =
η (p−1) k + η (p+1) k
2 +
max(0, υ (p) k )e γkH(υ (p) k ) −max(0, υ(p+1)k )eγkH(υ (p+1) k )
2 , (12)
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− 1, p) and (p, p+ 1) until an equilibrium is achieved. This is the first attempt to achieve equi-acceptance given two fixed boundary values to our best knowledge. By contrast, Syed et al. (2021) proposed to estimate the barrier Λ to determine the temperatures and it easily fails in big data given a finite number of chains and bias-corrected swaps.
†For convenience, υ(p)t denotes the continuous-time diffusion at time t and υ (p) k represents the discrete
approximations at iteration k.
Algorithm 1 Non-reversible parallel tempering with SGD-based exploration kernels (DEO?-SGD). Input Number of chains P ≥ 3, boundary learning rates η(1) and η(P ), target swap rate S. Input Optimal window size W := ⌈ logP+log logP − log(1−S) ⌉ , total iterations K, and step sizes {γk}Kk=0.
for k = 1 to K do βk+1 ∼ Tη(βk) following Eq.(7) . Exploration / exploitation phase (parallelizable) P = { ∀p ∈ {1, 2, · · · , P} : p mod 2 = b k
W c mod 2
} . . Generalized even/odd iterations
for p = 1, 2 to P − 1 do A(p) := 1
Ũ(β (p+1) k+1 )+Ck<Ũ(β (p) k+1 )
G(p) := 1k modW=0. . Open the gate to allow swaps if p ∈ P and G(p) and A(p) then
Swap: β(p)k+1 and β (p+1) k+1 . . Communication phase (parallelizable)
Freeze: G(p) = 0. . Close the gate to refuse swaps end if if p > 1 then
Update learning rate (temperature) following Eq.(12) end if
end for Adaptive correction buffer: Ck+1 = Ck + γk
( 1 P−1 ∑P−1 p=1 A (p) − S ) .
end for Output Models collected from the target temperature {β(1)k } K k=1.
Adaptive correction buffers In addition, equi-acceptance does not guarantee a convergence to the desired acceptance rate S. To avoid this issue, we propose to adaptively optimize C as follows
Ck+1 = Ck + γk
( 1
P − 1 P−1∑ p=1 1 Ũ(β (p+1) k+1 )+Ck−Ũ(β (p) k+1)<0 − S
) . (13)
As k →∞, 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-specific 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 1. Extensions of SGD with a preconditioner (Li et 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 DEO?-pSGD and DEO?-mSGD, respectively.

heading: 5 EXPERIMENTS

heading: 5.1 SIMULATIONS OF MULTI-MODAL DISTRIBUTIONS
text: We first simulate the proposed algorithm on a distribution π(β) ∝ exp(−U(β)), where β = (β1, β2), U(β) = 0.2(β21 + β 2 2)− 2(cos(2πβ1) + cos(2πβ2)). The heat map is shown in Figure 3(a) with 25 modes of different volumes. To mimic big data scenarios, we can only access stochastic gradient ∇Ũ(β) = ∇U(β) + 2N (0, I2×2) and stochastic energy Ũ(β) = U(β) + 2N (0, I).
We first run DEO?-SGD×P16 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 S for the accelerationaccuracy trade-off. Fig.3 shows that fixing S = 0.2 or 0.3 is too conservative and underestimates the
uncertainty on the corners; S = 0.6 results in too many radical swaps and eventually leads to crude estimations; by contrast, S = 0.4 yields the best uncertainty approximation among the five choices.
Next, we select S = 0.4 and study the round trips. We observe in Fig.4(a) that the vanilla DEO only yields 18 round trips every 1,000 iterations; by contrast, slightly increasing W tends to improve the 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 is adaptively estimated to ensure the average acceptance rates converge to S = 0.4.
We compare the proposed algorithm with parallel SGLD based on 20,000 iterations and 16 chains (SGLD×P16); we fix the learning rate 0.003 and a temperature 1. We also run cycSGLD×T16, which is short for a single long chain based on 16 times of budget and cosine learning rates (Zhang et al., 2020) of 100 cycles. We see in Figure 5(b) that SGLD×P16 has good explorations but fails to quantify the uncertainty. Figure 5(c) shows that cycSGLD×T16 explores most of the modes but overestimates some areas occasionally. Figure 5(d) demonstrates the DEO-SGD with 16 chains (DEO-SGD×P16) estimates the uncertainty of the centering 9 modes well but fails to deal with the rest of the modes. As to DEO?-SGD×P16, the approximation is rather accurate, as shown in Fig.5(e).
We also present the index process for both schemes in Fig.6. We see that the vanilla DEO scheme results in volatile paths and a particle takes quite a long time to complete a round trip; by contrast, DEO? only conducts at most one cheap swap in a window and yields much more deterministic paths.

heading: 5.2 UNCERTAINTY APPROXIMATION AND OPTIMIZATION FOR IMAGE DATA
text: 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 it by DEO?-mSGD×P10. 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-mSGD×P10 with the same setup except that the window size is 1; the standard ensemble mSGD×P10 is also included with a learning rate of 0.005. In addition, we include two baselines based on a single long chain, i.e. we run stochastic gradient Hamiltonian Monte Carlo
(Chen et al., 2014) 5000 epochs with cyclical learning rates and 50 cycles (Zhang et al., 2020) and refer to it as cycSGHMC×T10; we run SWAG×T10 (Maddox et al., 2019) under a similar setup. In particular for DEO?-mSGD×P10, we tune the target swap rate S and find an optimum at S = 0.005. We compare our proposed algorithm with the four baselines and observe in Table.2 that mSGD×P10 can easily obtain competitive results simply through model ensemble (Lakshminarayanan et al., 2017), which outperforms cycSGHMC×T10 and cycSWAG×T10 on ResNet20 and ResNet32 models and perform the worst among the five methods on ResNet56; DEO-mSGD×P10 itself is already a pretty powerful algorithm, however, DEO?-mSGD×P10 consistently outperforms the vanilla alternative.
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 as the training proceeds, the learning rates converge to fixed points and the acceptance rates for different pairs converge to the target swap rate.
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 a novel viewpoint to interpret why cyclic learning rates work well empirically. Nevertheless, the stochastic and parallel manner further improves the margin and eventually leads to the most efficient approximations in this task.

heading: 6 CONCLUSION
text: In this paper, we show how to conduct efficient PT in big data problems. To 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 in a significant acceleration of O( PlogP ) times. For a user-friendly purpose, we propose a deterministic swap condition to interact with SGD-based exploration kernels and provide a theoretical guarantee to control the bias solely depending on the learning rate η and the correction buffer C; we also provide a practical algorithm to adaptively approximate η and C for achieving the optimal efficiency in constrained settings.

heading: A BACKGROUND

heading: A.1 REPLICA EXCHANGE STOCHASTIC GRADIENT LANGEVIN DYNAMICS
text: To approximate the replica exchange Langevin diffusion Eq.(1) based on multiple particles (β
(1) k+1,β (2) k+1, · · · ,β (P ) k+1) in big data scenarios, replica exchange stochastic gradient Langevin dy-
namics (reSGLD) proposes the following numerical scheme:
Exploration:  β (P ) k+1 = β (P ) k − η(P )∇Ũ(β (P ) k ) + √ 2η(P )τ (P )ξ (P ) k ,
· · · β
(2) k+1 = β (2) k − η(2)∇Ũ(β (2) k ) +
√ 2η(2)τ (2)ξ
(2) k ,
Exploitation: β(1)k+1 = β (1) k − η (1)∇Ũ(β(1)k ) + √ 2η(1)τ (1)ξ (1) k ,
(14)
where η(·) is the learning rate, ξ(·)k is a standard d-dimensional Gaussian noise and each subprocess follows a stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011). Further assuming the energy normality assumption C1 in section C.1, there exists a σp such that
Ũ(β(p))− Ũ(β(p+1)) ∼ N (U(β(p))− U(β(p+1)), 2σ2p), for any β(p) ∼ π(p), (15)
where p ∈ {1, 2, · · · , P − 1}, in what follows, Deng et al. (2020) proposed the bias-corrected swap function as follows
aS̃(β (p) k+1,β (p+1) k+1 ) = a ·
( 1 ∧ e ( 1 τ(p) − 1 τ(p+1) )( Ũ(β (p) k+1)−Ũ(β (p+1) k+1 )− ( 1 τ(p) − 1 τ(p+1) ) σ2p )) , (16)
where p ∈ {1, 2, · · · , P − 1}, (
1 τ(p) − 1 τ(p+1) ) σ2p is a correction term to avoid the bias and the swap
intensity a can be then set to 1 min{η(p),η(p+1)} for convenience. Namely, given a particle pair at location (β(p), β(p+1)) in the k-th iteration, the conditional probability of the swap follows that
P(βk+1 = (β(p+1), β(p))|βk = (β(p), β(p+1))) = S̃(β(p), β(p+1)), P(βk+1 = (β(p), β(p+1))|βk = (β(p), β(p+1))) = 1− S̃(β(p), β(p+1)).
A.2 INHOMOGENOUS SWAP INTENSITY VIA DIFFERENT LEARNING RATES
Assume the high-temperature process also applies η(p+1) ≥ η(p) > 0, where p ∈ {1, 2 · · · , P − 1}. At time t, the swap intensity can be interpreted as being 0 in a time interval [t, t+ η(p+1) − η(p)) and being 1
η(p) in [t+ η(p+1)− η(p), t+ η(p+1)). Since the coupled Langevin diffusion process converges
to the same joint distribution regardless of the swaps and the swap intensity a varies in a bounded domain, the convergence of numerical schemes is not much affected except the numerical error.

heading: A.3 ACCURACY-ACCELERATION TRADE-OFF
text: Despite the exponential acceleration potential, the average bias-corrected swap rate is significantly reduced such that E[S̃] = O ( Se − ( 1 τ(p) − 1 τ(p+1) )2 σ2p 8 )
(Deng et al., 2021), where S is the swap function following Eq.(2) in full-batch settings. Including more parallel chains are promising to alleviate this issue, however, it is often inevitable to sacrifice some accuracy to obtain more swaps and accelerations, as discussed in section 3.4. As such, it inspires us to devise the user-friendly SGD-based approximate exploration kernels.

heading: A.4 CONNECTION TO NON-CONVEX OPTIMIZATION
text: Parallel to our work, a similar SGLD×SGD framework was proposed by Dong & Tong (2021) for non-convex optimization, where SGLD and SGD work as exploration and exploitation kernels,
respectively. By contrast, our algorithm performs exactly in the opposite for uncertainty approximation because SGLD is theoretically more appealing for the exploitations based on small learning rates instead of explorations, while the widely-adopted SGDs are quite attractive in exploration due to its user-friendly nature and ability in exploring wide optima.
If we manipulate the scheme to propose an exact swap in each window instead of at most one in the current version, the algorithm shows a better potential in non-convex optimization. In particular, a larger window size corresponds to a slower decay of temperatures in simulated annealing (SAA) (Mangoubi & Vishnoi, 2018). Such a mechanism yields a larger hitting probability to move into a sub-level set with lower energies (losses) and a better chance to hit the global optima. Nevertheless, the manipulated algorithm possess the natural of parallelism in cyclical fashions.

heading: A.5 OTHERS
text: Swap time refers to the communication time to conduct a swap in each attempt. For example, chain pair (p, p+ 1) of ADJ requires to wait for the completion of chain pairs (1, 2), (2, 3), · · · , (p− 1, p) to attempt the swap and leads to swap time of O(P ); however, SEO, DEO, and DEO? don’t have this issue because in each iteration, only even or odd chain pairs are attempted to swap, hence the swap time is O(1).
Round trip time refers to the time (stochastic variable) used in a round trip. A round trip is completed when a particle in the p-th chain, where p ∈ [P ] := {1, 2, · · · , P}, hits the index boundary index at both 1 and P and returns back to its original index p.

heading: B ANALYSIS OF ROUND TRIPS
text: 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 β(p) ∼ π(p) for any p ∈ [P ];
• (B2) Weak independence: For any β̄(p) simulated from the p-th chain conditional on β(j), U(β̄(j)) and U(β(j)) are independent.

heading: B.1 ANALYSIS OF ROUND TRIP TIME
text: Proof of Lemma 1. For t ∈ N, define Zt ∈ [P ] = {1, 2, · · · , P} as the index of the chain a particle arrives after t windows. Define δt ∈ {1,−1} to indicate the direction of the swap a particle intends to make during the t-th window; i.e., the swap is between Zt and Zt + 1 if δt = 1 and is between Zt and Zt − 1 if δt = −1. Define U := min{t ≥ 0 : Zt = P, δt = −1} and V := min{t ≥ 0 : Zt = 1, δt = 1}. Define rp := P[ reject the swap between Chain p and Chain p+1 for one time] . Define up,δ := E[U |Z0 = p, δ0 = δ] for δ ∈ {1,−1} and vp,δ := E[V |Z0 = p, δ0 = δ]. Then the expectation of round trip time T is
E[T ] = W (u1,1 + uP,−1). (17)
By the Markov property, for up,δ , we have
up,1 = r W p (up,−1 + 1) + (1− rWp )(up+1,1 + 1) (18)
up,−1 = r W p−1(up,1 + 1) + (1− rWp−1)(up−1,−1 + 1), (19)
where rWp denotes the rejection probability of a particle in a window of size W at the p-th chain.
According to Eq.(18) and Eq.(19), we have
up+1,1 − up,1 = rWp (up+1,1 − up,−1)− 1 (20) up,−1 − up−1,−1 = rWp−1(up,1 − up−1,−1) + 1. (21)
Define αp = up,1 − up−1,−1. Then by definition, Eq.(20), and Eq.(21), we have αp+1 − αp = (up+1,1 − up,−1)− (up,1 − up−1,−1)
= (up+1,1 − up,1)− (up,−1 − up−1,−1) = rWp αp+1 − rWp−1αp − 2,
which implies that ap+1 − ap = −2, (22)
for ap := (1− rWp−1)αp. Thus, by Eq.(22), we have
ap = a2 − 2(p− 2). (23) By definition, u1,−1 = u1,1 + 1. According to Eq.(20), for p = 1, we have
a2 = (1− rW1 )(u2,1 − u1,−1) = (1− rW1 ) [ u1,1 − u1,−1 + rW1 (u2,1 − u1,−1)− 1 ] = −2(1− rW1 ) + rW1 a2.
(24)
Since rp ∈ (0, 1) for 1 ≤ p ≤ P , Eq.(24) implies a2 = −2 which together with Eq.(23) implies
(1− rWp−1)αp = ap = −2(p− 1), (25) and therefore
rWp−1αp = −2(p− 1) rWp−1
1− rWp−1 . (26)
According to Eq.(21) and Eq.(26), we have
uP,1 − u1,1 = P−1∑ p=1 rWp (up+1,1 − up,−1)− (P − 1) (27)
= P−1∑ p=1 rWp αp+1 − (P − 1) (28) =− 2 P−1∑ p=1 rWp 1− rWp p− (P − 1). (29)
Since uP,1 = 1, we have
u1,1 = uP,1 + 2 P−1∑ p=1 rWp 1− rWp p+ (P − 1)
= P + 2 P−1∑ p=1 rWp 1− rWp p.
(30)
Similarly, for vp,δ , we also have
vp,1 = r W p (vp,−1 + 1) + (1− rWp )(vp+1,1 + 1)
vp,−1 = r W p−1(vp,1 + 1) + (1− rWp−1)(vp−1,−1 + 1).
With the same analysis, we have
bp+1 − bp = −2 (31) vp,−1 − vp−1,−1 = rWp−1(vp,1 − vp−1,−1) + 1, (32)
where bp := (1− rWp−1)βp and βp := vp,1 − vp−1,−1. According to Eq.(32), we have
vP,−1 − v1,−1 = P−1∑ p=1 rWp βp+1 + (P − 1). (33)
By definition, vP,1 = vP,−1 + 1. According to Eq.(32), for p = P , we have
bP = (1− rWP−1)(vP,1 − vP−1,−1) = (1− rWP−1) [ vP,1 − vP,−1 + rWP−1(vP,1 − vP−1,−1) + 1 ] = rWP−1bP + 2(1− rWP−1).
(34)
Since rp ∈ (0, 1) for 1 ≤ p ≤ P , Eq.(34) implies bP = 2. Then according to Eq.(31), we have bp = 2(P − p+ 1),
and therefore
rWp−1βp = 2 rWp−1
1− rWp−1 (P − p+ 1). (35)
By Eq.(33) and Eq.(35), we have
vP,−1 − v1,−1 = 2 P−1∑ p=1 rWp 1− rWp (P − p) + (P − 1).
Since v1,−1 = 1, we have
vP,−1 = 2 P−1∑ p=1 rWp 1− rWp (P − p) + P. (36)
According to Eq.(17), Eq.(30) and Eq.(36), we have
E[T ] = W (u1,1 + uP,−1) = 2WP + 2WP P−1∑ p=1 rWp 1− rWp .
The proof is a generalization of Theorem 1 (Syed et al., 2021); when W = 1, the generalized DEO scheme recovers the DEO scheme. For the self-consistency of our analysis, we present it here anyway.

heading: B.2 ANALYSIS OF OPTIMAL WINDOW SIZE
text: Proof of Theorem 1. By treating W ≥ 1 as a continuous variable and taking the derivative of E[T ] with respect to W and we can get
∂
∂W E[T ] = 2P
[ 1 +
P−1∑ p=1 rWp (1− rWp ) +W P−1∑ p=1 rWp log rp (1− rWp )2
]
= 2P { 1 +
P−1∑ p=1 rWp − r2Wp +WrWp log rp (1− rWp )2
}
= 2P { 1 +
P−1∑ p=1 rWp 1− rWp +W log rp (1− rWp )2
} .
(37)
Assume that rp = r ∈ (0, 1) for 1 ≤ p ≤ P . Then we have
E[T ] = 2WP + 2WP (P − 1) r W
1− rW . (38)
∂
∂W E[T ] =
2P (1− rW )2 { (1− rW )2 + (P − 1)rW (1− rW +W log r) } . (39)
Define x := rW ∈ (0, 1). Hence W = logr(x) = log x log r and
∂
∂W E[T ] = f(x) :=
2P (1− x)2 { (1− x)2 + (P − 1)x(1− x+ log x) }
(40)
Thus it suffices to analyze the sign of the function g(x) := (1− x)2 + (P − 1)x(1− x+ log(x)) for x ∈ (0, 1). For g(x), we have
g′(x) = (4− 2P )(x− 1) + (P − 1) log(x)
and g′′(x) = 4− 2P + P − 1
x .
Thus, limx→0+ g′(x) = −∞, limx→1 g′(x) = g′(1) = 0 and g′′(x) is monotonically decreasing for x > 0.
1. For P > 2, we know that g′′(x) > 0 when 0 < x < P−12(P−2) and g ′′(x) < 0 when
x < P−12(P−2) . Therefore, g ′(x) is maximized at P−12(P−2) .
(a) If P = 3, by log(1+y) < y for y > 0, we have g′(x) = 2 ( log(1+(x−1))−(x−1) ) <
0 for any x ∈ (0, 1). Thus g(x) > g(1) = 0 for x ∈ (0, 1) and therefore ∂∂W E[T ] > 0 for W ∈ N+. E[T ] is globally minimized at W = 1.
(b) If P > 3, we have the following lemma and the proof is postponed in section B.2.1. Lemma 3 (Uniqueness of the solution). For P > 3, there exists a unique solution x∗ ∈ (0, 1) such that g(x) > 0 for ∀x ∈ (0, x∗) and g(x) < 0 for ∀x ∈ (x∗, 1). Moreover, W ∗ = logr(x ∗) is the globally minimizer for the round trip time.
2. For P = 2, we know that g′′(x) > 0 for x ∈ (0, 1). Thus g′(x) < g(1) = 0 for x ∈ (0, 1) and therefore g(x) > g′(1) = 0 for x ∈ (0, 1). Then according to Eq.(37), ∂∂W E[T ] > 0 for W ∈ N+. E[T ] is globally minimized at W = 1.
In what follows, we proceed to prove that 1P logP is a good approximation to x ∗. In fact, we have
g
( 1
P logP
) = ( 1− 1
P logP
)2 +
P − 1 P logP
( 1− 1
P logP − logP − log(logP ) ) = 1 + o ( 1
P
) − 1 + 1
logP − log(logP ) logP +O
( 1
P ) = − log(logP )
logP +O
( 1
logP ) Thus, limP→∞ g( 1P logP ) = 0.
For x = 1P logP , we have W = logr ( 1 P logP ) = logP+log logP− log r . Then according to Eq.(38), for P ≥ 4,
E[T ] = 2P logP + log logP
− log r
[ 1 + (P − 1) 1 P logP
1− 1P logP
]
= ( 1 + 1
logP
1
1− 1P logP
) 2
− log r (P logP + P log logP )
≤ ( 1 + 1
log 4
1
1− 14 log 4
) 2
− log r (P logP + P log logP )
< 4
− log r (P logP + P log logP )
(41)
In conclusion, for P = 2, 3, the maximum round trip rate is achieved when the window size W = 1. For P ≥ 4, with the window size W = logP+log logP− log r , the round trip rate is at least Ω ( − log r logP ) .
Remark: Given finite chains with a large rejection rate, the round trip time is only of order O(P logP ) by setting the optimal window size W ≈ ⌈ logP+log logP − log r ⌉ . By contrast, the vanilla DEO
scheme with a window of size 1 yields a much longer time of O(P 2), where 1− log(r) = O( r 1−r ) based on Taylor expansion and a large r 0.

heading: B.2.1 TECHNICAL LEMMA
text: Proof of Lemma 3. To help illustrate the analysis below, we plot the graphs of g′(x) and g(x) for x ∈ (0, 1) and P = 5 in Figure 9. For P > 3, since g′(x) is maximized at x = P−12(P−2) ∈ (0, 1) with g′′(x) > 0 when 0 < x < P−12(P−2) and g ′′(x) < 0 when P−12(P−2) < x < 1, limx→0+ g ′(x) = −∞, and g′(1) = 0, we know that g′( P−12(P−2) ) > 0 and there exists x0 ∈ (0, P−1
2(P−2) ) such that g′(x) < 0 (i.e., g′(x) is monotonically decreasing) for any x ∈ (0, x0) and g′(x) > 0 (i.e., g′(x) is monotonically increasing) for any x ∈ (x0, 1). Then g(x) on (0, 1) is globally minimized at x = x0. Moreover, limx→0+ g(x) = 1 and g(1) = 0. Thus, g(x0) < g(1) = 0 and there exists x∗ ∈ (0, x0) ( (0, 1) such that g(x) > 0 if x ∈ (0, x∗) and g(x) < 0 if x ∈ (x∗, 1).
Meanwhile, by Eq.(37) and the definition of x, we know that the sign of ∂∂W E[T ] is the same with that of g(x) for W = logr(x). Thus, ∂ ∂W E[T ] < 0 when W < W ∗ := logr(x ∗) and ∂∂W E[T ] > 0 when W > W ∗, which implies that E[T ] is globally minimized at W ∗ = logr(x∗) with some x∗ ∈ (0, 1).

heading: B.3 DISCUSSIONS ON THE OPTIMAL NUMBER OF CHAINS IN BIG DATA
text: To shed light on the suggestions of the optimal number of chains, we lay out two additional assumptions:
• (B3) Integrability: U3 is integrable with respect to π(1) and π(P ).
• (B4) Approximate geometric spacing: Assume that τ (p+1) τ(p) ≡ ( τ(P ) τ(1) ) 1 P−1 +o( 1P ) for p ∈ [P −1].
Assumption B3 is a standard assumption in Syed et al. (2021) to establish the convergence of the summation of the rejection rates. Assumption B4 allows extra perturbations to the geometric spacing assumption (Kofke, 2002) and is empirically verified in the CIFAR100 example in Figure 7(c).
Proof of corollary 1. Before we start the proof, we denote by r̃ and s̃ the average rejection and acceptance rates in full-batch settings, respectively. According to section 5.1 in Syed et al. (2021), Assumptions B1, B2 and B3 lead to
r̃ = Λ
P +O
( 1
P 3
) , (42)
where Λ is the barrier, which becomes larger if π(1) and π(P ) have a smaller probability overlap (Predescu et al., 2004; Syed et al., 2021). Then acceptance rate in full-batch settings is
s̃ = 1− r̃ = 1− Λ P +O
( 1
P 3
) . (43)
By Lemma D4 (Deng et al., 2021) and Eq.(15), a noisy energy estimator with variance (τ(p)−τ(p+1))2
τ(p)2τ(p+1)2 σ2p yields an average swap rate of O
( s̃ exp ( − (τ (p)−τ(p+1))2
8τ(p)2τ(p+1)2
) σ2p ) . In what follows,
the average rejection rate in mini-batch settings becomes r = 1−O ( s̃ exp ( − (τ (p) − τ (p+1))2
8τ (p) 2 τ (p+1) 2
) σ2p ) , (44)
where p ∈ [P ] and more accurate estimates is studied in proposition 2.4 (Gelman et al., 1997).
Define ∆τ = τ (P )
τ(1) . By assumption B4 and Taylor’s theorem, we have(
τ (p+1) − τ (p)
τ (p)
)2 := ( ∆ 1 P−1 τ − 1 + o ( 1
P
))2 = ( log ∆τ P − 1 )2 + o ( 1 P 2 ) . (45)
Denote γp = σp
τ(p+1) . It follows that
exp ( − (τ (p) − τ (p+1))2
8τ (p) 2 τ (p+1)
2 σ 2 p ) = exp ( − ( (log ∆τ ) 2
8(P − 1)2 + o
( 1
P 2
)) σ2p
τ (p) 2
)
= exp ( − (γp log ∆τ ) 2
8(P − 1)2 + o
( 1
P 2
)) .
(46)
Plugging it to Eq.(44), we have r = 1−O ((
1− Λ P +O
( 1
P 3
)) exp ( − (γp log ∆τ ) 2
8(P − 1)2 + o
( 1
P 2
))) . (47)
Recall in Eq.(41), the round trip time follows that E[T ] = O ( P logP
− log r
) . (48)
By log(1− t) ≤ −t for t ∈ [0, 1) and 11−t = 1 + t+O(t 2) for a small t, we have
− 1 log r = − 1 log { 1−O (( 1− ΛP +O ( 1 P 3 )) exp ( − (γp log ∆τ ) 2 8(P−1)2 + o ( 1 P 2 )))} ≤ 1
O (( 1− ΛP +O ( 1 P 3 )) exp ( − (γp log ∆τ ) 2 8(P−1)2 + o ( 1 P 2 ))) = O (( 1 + Λ
P +O
( 1
P 3
)) exp ( (γp log ∆τ ) 2
8(P − 1)2 + o
( 1
P 2
))) .
(49)
Combining Eq.(49) and Eq.(48), we have E[T ] =O ( P logP ( 1 + Λ
P +O
( 1
P 3
)) exp ( µ2p
(P − 1)2 +O
( 1
P 2
)))
=O ( P logP exp ( µ2p P 2 )) ,
(50)
where µp := γp log ∆τ
2 √ 2 . The optimal P to minimize E[T ] = O
( P logP exp ( µ2p P 2 )) is equivalent to
the minimizer of hp(x) = log { x log x exp ( µ2p x2 )} = log x+ log log x+ µ2p x2 for x ≥ 4. Then
h′p(x) = 1
x +
1
x log x − 2µ2p x3 = x2(1 + 1log x )− 2µ 2 p x3 , (51)
Define s(x) = x2(1 + 1log x ). Then we have
s′(x) = 2x+ 2x log x− x
(log x)2 > 0, (52)
for x ≥ 4. Thus s(x) increases with x for x ≥ 4. Obviously, limx→∞ s(x) = ∞. Since µp 1, we know s(4) < 0. Thus, there is a single zero point P? ∈ (4,∞) of s(x) such that h′p(x) < 0 if x ∈ [4, P?) and h′p(x) > 0 if x > P?. Therefore, hp(x) is minimized at P? with
P 2?
( 1 + 1
logP?
) ∈ (
2 min p µ2p, 2 max p µ2p
) .
For any P? ≥ 2, we have √ 2/ (1 + 1/ logP?) ∈ (1.1, √
2). Thus, we have the optimal number of chains in mini-batch settings following that
P? ∈ (1.1 min p µp, √ 2 max p
µp) ∈ (
min p σp 3τ (p+1) log ∆τ ,max p σp 2τ (p+1) log ∆τ
) . (53)

heading: B.4 EMPIRICAL JUSTIFICATION OF THE OPTIMAL NUMBER OF CHAINS
text: To obtain an estimate of the optimal number of chains for PT in big data problems, we study the standard deviation of the noisy energy estimators on CIFAR100 via ResNet20 models. We first pre-train a model and then try different learning rates to check the corresponding standard deviation of energy estimators. The reason we abandon the temperature variable is that the widely used data augmentation has drastically affected the estimation of the temperature (Wenzel et al., 2020). Motivated by the linear relation between the learning rate and the temperature in Eq.(6) and the cold posterior affect with the temperature much smaller than 1 (Aitchison, 2021), applying Figure 10 and Eq.(53) concludes that P? is achieved in an order of thousands in the CIFAR100 example.
The conclusion is different from Syed et al. (2021) since big data problems require a much smaller swap rate to maintain the unbiasedness of the swaps. On the one hand, 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 in terms of limited memory and computational budget.

heading: C APPROXIMATION ERROR FOR SGD-BASED EXPLORATION KERNELS
text: This section proposes to analyze the approximation error when we adopt SGDs as the approximate exploration kernels. Since it is independent of section B, different assumptions may be required.
The first subsection shows that for any Gaussian-distributed energy difference estimator Ũ(β(p))− Ũ(β(p+1)), there exists an optimal C? such that the deterministic swap condition Ũ(β(p+1)) +C? < Ũ(β(p)) perfectly approximates the random event S̃(β(p),β(p+1)) > u, where u ∼ Unif[0, 1].

heading: C.1 GAP BETWEEN ACCEPTANCE RATES
text: To approximate the error when we adopt the deterministic swap condition Eq.(8), we first assume the following condition
• (C1) Energy normality: The stochastic energy estimator for each chain follows a normal distribution with a fixed variance.
The above assumption has been widely assumed by Ceperley & Dewing (1999); Bardenet et al. (2017); Seita et al. (2017), which can be easily extended to the asymptotic normality assumption depending on large enough batch sizes (Quiroz et al., 2019; Deng et al., 2021).
Proof of Lemma 2. Denote ∂Tp = 1τ(p) − 1 τ(p+1) > 0 for any p ∈ [P − 1], the energy difference ∂Up = U(β (p)) − U(β(p+1)). By the energy normality assumption C1 and Eq.(15), the energy difference estimator follows that ∂Ũp = ∂Up(·) + √
2σpξ, where ξ ∼ N (0, 1). Recall from Eq.(16) that the bias-corrected acceptance rate follows that
S̃(β (p) t ,β (p+1) t ) = 1 ∧ e ∂Tp
( Ũ(β
(p) t )−Ũ(β (p+1) t ) ) −∂T 2pσ 2 p
= 1 ∧ e∂Tp∂Up+ √ 2∂Tpσpξ−∂T 2pσ 2 p ,
(54)
(i) For a uniformly distributed variable µ ∼ Unif[0, 1], the base swap function satisfies that
P ( S̃(β
(p) t ,β (p+1) t ) > µ|∂Up
) = ∫ 1 0 P(S̃(β(p)t ,β (p+1) t ) > u|∂Up)du
= ∫ 1 0 P(e∂Tp∂Up+ √ 2∂Tpσpξ−∂T 2pσ 2 p > u|∂Up)du
= ∫ 1 0 P ( ξ > −∂Up + ∂Tpσ2p + log u ∂Tp√ 2σp ) du
(55)
(ii) For the deterministic swaps based on Eq.(8), the approximate swap function follows that
P ( Ũ(β
(p+1) k ) + C < Ũ(β (p) k )|∂Up ) = P(C < ∂Up + √ 2σpξ|∂Up) = P ( ξ > −∂Up + C√
2σp
) .
(56)
Denote by D(u,C) = ∣∣∣∣P(ξ > −∂Up+∂Tpσ2p+ log u∂Tp√2σp ) − P ( ξ > −∂Up+C√ 2σp )∣∣∣∣ ∈ [0, 2). Combining (55) and (56), we can easily derive that
∆(C) = P ( S̃(β
(p) t ,β (p+1) t ) > µ|∂Up
) − P ( Ũ(β
(p+1) k ) + C < Ũ(β (p) k )|∂Up
)
= ∫ e−∂T2pσ2p+∂TpC 0 P ( ξ > −∂Up + ∂Tpσ2p + log u ∂Tp√ 2σp ) − P ( ξ > −∂Up + C√ 2σp ) ︸ ︷︷ ︸
D(u,C) for small enough u
du
+ ∫ 1 e −∂T2pσ 2 p+∂TpC P ( ξ > −∂Up + ∂Tpσ2p + log u ∂Tp√ 2σp ) − P ( ξ > −∂Up + C√ 2σp ) ︸ ︷︷ ︸
−D(u,C) for large enough u
du
= ∫ e−∂T2pσ2p+∂TpC 0 D(u,C)du− ∫ 1 e −∂T2pσ 2 p+∂TpC D(u,C)du
Clearly, ∆(C) is continuous in C and it is straightforward to verify that
∆(∂Tpσ 2 p) = ∫ e−∂T2pσ2p+∂T2pσ2p 0 D(u, ∂Tpσ 2 p)du− ∫ 1 e −∂T2pσ 2 p+∂T 2 pσ 2 p D(u, ∂Tpσ 2 p)du
= ∫ 1 0 D(u, ∂Tpσ 2 p)du ≥ 0
Similarly, ∆(−∞) = − ∫ 1
0 D(u,C)du ≤ 0. Moreover, the physical construction suggests that the
threshold C should be strictly positive to avoid radical swap attempts, which implies that there is an optimal solution C? ∈ (0, ∂Tpσ2p] that solves ∆(C?) = 0.
Remark 1. Note that the optimal C? may lead to few swaps given a 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 E[S̃(β(p)t ,β (p+1) t )] = O(e
− ∂T2pσ 2 p
8 ). This
suggests that for a well-approximated threshold C, the error is at most E[∆(C)] = O(e− ∂T2pσ 2 p 8 ).

heading: C.2 APPROXIMATE UPPER BOUND
text: The second subsection completes the proof regarding the numerical approximation of SGD-based exploration kernels. Nevertheless, we still require some standard assumptions.
• (C2) Smoothness: The function U(·) is C-smooth if there exists a positive Lipschitz constant C such that ‖∇U(x)−∇U(y)‖2 ≤ C‖x− y‖2 for every x, y ∈ Rd.
• (C3) Dissipativity: The function U(·) is (a, b)-dissipative if there exist positive constants a and b such that 〈x,∇U(x)〉 ≥ a‖x‖2 − b.
• (C4) Gradient normality: The stochastic gradient noise ε(·) in Eq.(6) follows a Normal distribution N (0,M) with a fixed positive definite matrixM.
The assumptions C2 and C3 are standard to show the geometric ergodicity of Langevin diffusion and the diffusion approximations for non-convex functions (Mattingly et al., 2002; Raginsky et al., 2017; Xu et al., 2018). The assumption C4 is directly motivated by Mandt et al. (2017) to track the Ornstein Uhlenbeck (OU) process with a Gaussian invariant measure.
Proof of Theorem 2. The transition kernel T of the continuous-time replica exchange Langevin diffusion (reLD) follows that T = ∏ p T (p), where each sub-kernel T (p) draws a candidate θ ∈ Rd with the following probability
T (p)(θ|β(p),β(q)) = (1− S(β(p)t ,β (q) t ))Qp(θ|β (p)) + S(β (p) t ,β (q) t )Qq(θ|β (q)),
where p, q ∈ {1, 2, · · · , P}, |p− q| = 1 and q is selected based on the DEO? scheme (3), Qp(θ|β) is the proposal distribution in the p-th chain that models the probability to draw the parameter θ via Langevin diffusion conditioned on β, and S(β(p)t ,β (q) t ) follows the swap function in Eq.(2).
Now we consider the following approximations:
reLD I=⇒ reSGLD II=⇒ D̂EO?-SGLD III =⇒ DEO?-SGLD IV =⇒ DEO?-SGD,
where D̂EO?-SGLD resembles the DEO? scheme except that there are no restrictions on conducting at most one swap.
I: Numerical approximation via SGLD By Lemma 5, we show that approximating replica exchange Langevin diffusion via reSGLD yields a weak error only depending on the learning rate; similar results depending on the learning rate and the noise in gradient and energy noise have been studied in Lemma 1 of (Deng et al., 2020).
II: Restrictions on at most one swap Adopting at most one swap in the DEO? scheme includes a 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 D̂EO?-SGLD differs from DEO? in that DEO?-SGLD freezes at most W − 1 attempts of swaps in each neighboring chains in each window, while D̂EO?-SGLD keeps all of these positions open. Recall that the bias-corrected swap rate follows that S̃ = O ( Se−O(σ 2) )
by invoking Eq.(16). Following a 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)︸ ︷︷ ︸
=0
=
O(e−O(σ 2)). In other words, the restrictions on at most one swap only lead to a 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 a 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̃(β(p)t ,β (p+1) t ) > µ through Ũ(β (p) k ) + C? < Ũ(β (p+1) k ) in the average sense for any β (p) t and β (p+1) t .
IV: Laplace approximation via SGD Assumption C4 guarantees that there exists an optimal learning rate η to conduct Laplace approximation for the target posterior (Mandt et al., 2017). Further, by the Bernstein-von Mises Theorem, the approximation error goes to 0 as the number of total data points goes to infinity.
Combining the above approximations, there exists a finite constant ∆max ≥ 0 depending on the learning rates η, the variance of noisy energy estimators, and the choice of correction terms such that the total approximation error of the SGD exploration kernels with deterministic swap condition is upper bounded by ∆max. For any joint probability density µ, the distance between the distributions generated by one step of the exact transition kernel T and the approximate transition kernel Tη in
Eq.(7) is upper bounded by∫ β dΩ(β) ∣∣∣∣µT (β)− µTη(β)∣∣∣∣ ≤ ∫ β dΩ(β) ∣∣∣∣ ∫ φ,ψ dµ(φ,ψ)∆max ( Q(β|φ)−Q(β|ψ)
)∣∣∣∣ ≤ ∆max ∫ dΩ(β) ∫ φ,ψ dµ(φ,ψ) ( Q(β|φ) +Q(β|ψ)
) ≤ 2∆max,
where Q(·|·) = ∏ pQp(·|·). Thus, the one-step total variation distance between the exact kernel and the proposed approximate kernel is upper bounded by
‖µT (β)− µTη(β)‖TV = 1
2 ∫ β dΩ(β) ∣∣∣∣µT (β)− µTη(β)∣∣∣∣ = ∆max. (57) The following is a restatement of Lemma 3 in the supplementary file of Korattikara et al. (2014).
Lemma 4. Consider two transition kernels T and Tη , which yield stationary distributions π and πη , respectively. If T follows a contraction such that there is a constant ρ ∈ [0, 1) for all probability distributions µ:
‖µT − π‖TV ≤ ρ‖µ− π‖TV,
and the uniform upper bound of the one step error between T and Tη follows that ‖µT − µTη‖TV ≤ ∆,∀µ,
where ∆ ≥ 0 is a constant. Then the total variation distance between π and πη is bounded by
‖π − πη‖TV ≤ ∆
1− ρ .
Under the smoothness assumption C2 and the dissipative assumption C3, the geometric ergodicity of the continuous-time replica exchange Langevin diffusion to the invariant measure π has been established (Chen et al., 2019). Combining the exponential convergence of the KL divergence (Deng et al., 2020) and the Pinsker’s inequality (Csiszár & Körner, 2011), we have that for any probability density µ, there exists a contraction parameter ρ ∈ (0, 1) such that
‖µT − π‖TV ≤ ρ‖µ− π‖TV, where ρ depends on the spectral gap established in Raginsky et al. (2017) and the swap rate.
Applying Lemma 4, the total variation distance can be upper bounded by
‖π − πη‖TV ≤ ∆max 1− ρ ,
which concludes the proof of Theorem 2.

heading: C.2.1 BIAS ANALYSIS FOR GENERALIZED WINDOWS IN BIG DATA PROBLEMS.
text: Following the empirical setup in section 5.1 (Deng et 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 as 1 and 10, respectively. The swap condition follows from Eq.(16). We run the algorithm 3,000,000 iterations and also include a 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 at sd = 3, where the blue curves show remarkable approximations.
†Further increasing sd (or the variance) may yield fewer swaps and less accelerations.

heading: C.3 WEAK ERROR
text: In this section, we show the weak error (Lemma 5 below) between the Replica exchange Langevin diffusion (β(1)t ,β (2) t , · · · ,β (P ) t ) and the Replica exchange stochastic gradient Langevin diffusion (β̃ (1) t , β̃ (2) t , · · · , β̃ (P ) t ). We denote P := PW ×N, where PW is the infinite dimensional Wiener measure and N is the Poisson measure independent of PW and has some constant jump intensity. In our general framework below, the jump process α is introduced by swapping the diffusion matrix of the 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 [iη, (i+1)η], for i ∈ N+. Without loss of generality, the replica exchange Langevin diffusion (reLD) is reformulated as below. For any fixed learning rate η > 0, we define dβt = −∇U(βt)dt+ Σ(αt)dWt,P (α(t) = j|α(t− dt) = l,β(bt/ηcη) = β) = aS(β)η1{t=bt/ηcη} + o(dt), for l 6= j, (58) where ∇U(β) := ( ∇U(β(1)),∇U(β(2)), · · · ,∇U(β(P )) )T , and 1t=bt/ηcη is the indicator func-
tion, i.e. for every t = iη with i ∈ N+, given β(iη) = β, we have P (α(t) = j|α(t− dt) = l) = rS(β)η, where S(β) is defined as min{1, S(β(1),β(2))} and S(β(1),β(2)) is defined in (2). In this case, the Markov Chain α(t) is a constant on the time interval [bt/ηcη, bt/ηcη + η) with some state in the finite-state spaceM = {0, 1, · · · , 2P−1 − 1} and the generator matrix Q follows, Q11 = −aS(β(1),(2))ηδ(t− bt/ηcη), Q12 = aS(β(1),(2))ηδ(t− bt/ηcη) Qi(i−1) = aS(β
(i−1),(i))ηδ(t− bt/ηcη), Qi(i+1) = aS(β(i),(i+1))ηδ(t− bt/ηcη) Qii = −aS(β(i−1),(i))ηδ(t− bt/ηcη)− aS(β(i),(i+1))ηδ(t− bt/ηcη) QP (P−1) = aS(β (P−1),(P ))ηδ(t− bt/ηcη), QPP = −aS(β(P−1),(P ))ηδ(t− bt/ηcη),
where 1 < i < P , and all the other terms in Q remain zero. We denote δ(·) as a Dirac delta function. Furthermore, for general adjacent swap with odd iterations (denoted as QO), and even iterations (denoted as QE), the generator matrices QO and QE ∈ R(Pd)×(Pd) have the following form,
QO =  Q11 Q12 0 0 · · · Q21 −Q21 0 0 0 0 0 −Q34 Q34 0 0 0 Q43 −Q43 0 0 0 0 0 · · ·
· · · · · ·
 , QE =  0 0 0 0 0 0 −Q23 Q23 0 0 · · · 0 Q32 −Q32 0 0 0 0 0 0 −Q45 Q45 · · · 0 0 0 Q54 −Q54 0
· · · · · ·
 ,
For the odd iteration case, the diffusion matrix for the first two chains β(1) and β(2) has the following form according to the state of Σ(αt), i.e. (Σ(0),Σ(1)) := √2τ (1)Id 0 · · ·0 √2τ (2)Id · · ·
· · · · · · · · ·  (Pd)×(Pd) , √2τ (2)Id 0 · · ·0 √2τ (1)Id · · · · · · · · · · · ·  (Pd)×(Pd) , where the rest of the diffusion matrix remain unchanged. Based on the observation of the transition matrix of QO, and QE , there is no interaction between every pairs (each pair has two chains) in
each iteration, we will prove the weak error for only two chains. We can then add different square blocks together as shown in QO and QE . From now on, we restrict to the two chain case with β = (β(1),β(2)). Next, we denote β̃ as the Replica exchange stochastic gradient Langevin diffusion, for the same learning rate η > 0 as above, we have dβ̃ηt = −∇Ũ(β̃ η bt/ηcη)dt+ Σ(α̃bt/ηcη)dWt, P ( α̃(t) = j|α̃(t− dt) = l, β̃(bt/ηcη) = β̃ ) = S̃(β̃)η1{t=bt/ηcη} + o(dt), for l 6= j, (59)
where∇Ũ(β) := ( ∇Ũ(β(1)) ∇Ũ(β(2)) ) and S̃(β̃) = min{1, S̃η,m,n(β̃(1), β̃(2))} and S̃η,m,n(β̃(1), β̃(2)) is
shown in Eq.(16). The distribution of process (β̃t)0≤t≤T is denoted as µT := PG̃×NS̃ , where α̃ is a Poisson process with jump intensity aS̃(β̃)ηδ(t− bt/ηcη) on the time interval [bt/ηcη, bt/ηcη+ η). For any open set U ⊂ R2d, which is contained in a large enough open ball with radius R centered at β(0), we define stopping time
τU = inf{t ≥ 0|(β(t), α(t)) /∈ U ×M)}. (60)
Lemma 5. Assume the smoothness assumption C2 holds. For any large enough open set U , we have
|Ex,i[h(β(T ∧ τU ), α(T ∧ τU ))]− Ex,i[h(β̃(T ∧ τU ), α̃(T ∧ τU ))]| = O(η).
for any continuous differentiable and polynomial growth function h.
Proof. We denote βt = (β (1) t ,β (2) t ), which follows the dynamics below,
dβt = −∇U(βt)dt+ Σ(αt)dWt Denote u(βt = x, t, i) = Ex,i[h(βT∧τU , αT∧τU )|(βt, αt) = (x, i)], i ∈M. For t < T ∧ τU , applying Feynman-Kac formula (see e.g. Baran et al. (2013)), we have{
∂tu(x, t, i) + Lu(x, t, i) = 0, Lu(x, t, i) = tr(ΣΣT∇2u(x, t, i))− 〈∇u(x, t, i),∇U〉+Q(x, t)u(x, t, ·)(i) (61)
Applying generalized Itô formula, for each i ∈M we have h(βt, αt)− h(β0, α0) = ∫ t
0
Lg(βs, αs)ds+M1(t) +M2(t), (62)
where M1 +M2 is a local martingale. See details in (Yin & Zhu, 2010; Skorokhod, 2009). We thus get, for kη ≤ t < (k + 1)η < T ∧ τU ,
Ex,i[h(βt, αt)]− h(x, i) = Ex,i ∫ t
0
Lh(βs, αs)ds.
For 0 < t1 < t2 < · · · , tN = T ∧ τU , with tk = kη, we next derive the Itô formula for h(β̃t, αt).
Note that β and β̃ are defined by using the same P-Brownian motionW , but with two different jump intensity on the time interval [bt/ηcη, bt/ηcη + η). We have the following approximation,
β̃tk − β̃tk−1 = −∇Ũ(β̃tk−1)η + Σ(α̃tk−1)N (0, η).
For stochastic noise ξtk−1 , with E[ξtk−1 ] = 0, we rewrite the above approximation as below,
β̃tk − β̃tk−1 = −(∇U(β̃tk−1) + ξtk−1)η + Σ(α̃tk−1)N (0, η).
For notation convenience, for tk−1 ≤ t < tk, we keep the following convention,
Σ̃(β̃t) := Σ(β̃ η bt/ηcη) = Σ(β̃ η tk−1 ), −∇Ũ(β̃t) := −∇Ũ(Σ(β̃ηbt/ηcη)) = −∇Ũ(β̃tk−1).
For any tk−1 ≤ t < tk, we apply Itô formula for u(β̃t, t, i),
du(β̃t, t, i) = (∂tu+ L̃u(β̃t, i)dt+ Σ(α̃bt/ηcη) ∂u
∂x (β̃t)dWt + dM̃1(t) + dM̃2(t), (63)
= [(L̃ − L)u(β̃t, t, i)]dt+ Σ(α̃bt/ηcη) ∂u
∂x (β̃t)dWt + dM̃1(t) + dM̃2(t),
where the last equality follows from (61), and for tk−1 ≤ t < tk,
L̃u(β̃t, t, i) = tr(Σ̃Σ̃T∇2u(β̃t, t, i))− 〈∇Ũ(β̃t),∇u〉+ Q̃(β̃t, t)u(β̃t, t, ·)(i).
Evaluate the integral on time interval [0, T ∧ τU ], and take expectation Ex,i, we have Ex,i[u(β̃T∧τU , T ∧ τU , i)− u(β0, 0, i)] = ∫ T∧τU
0
Ex,i[(L̃ − L)u(β̃t, t)]dt
= I1 + I2 + I3.
Where we have I1 := − ∫ T∧τU
0
Ex,i[〈∇Ũ(β̃t),∇u(β̃t, t)〉 − 〈∇U(β̃t),∇u(β̃t, t)〉]dt,
I2 := ∫ T∧τU
0
Ex,i[tr(Σ̃Σ̃T∇2u(β̃t, t))− tr(ΣΣT∇2u(β̃t, t))]dt,
I3 := ∫ T∧τU
0
Ex,i[Q̃(β̃t, t)u(β̃t, t, ·)−Q(β̃t, t)u(β̃t, t, ·)]dt.
In the next step, we introduce a sequence of stopping time based on our definition of process β and β̃. For j ∈ N+, we denote ζ ′js as a stopping times defined by ζj+1 := inf{t > ζj : α(t) 6= α(ζj)} and N(T ) = max{n ∈ N : ζn ≤ T}. According to the definition of process β, we have ζj = lη, for some l ∈ N+. Similarly, we have stopping time ζ̃j+1 := inf{t > ζ̃j : α̃(t) 6= α̃(ζj)} associated with process β̃, and α̃(t) follows the same trajectory of α(t). We thus have the following estimates for the terms I1, I2, and I3.
Estimate of I1: as for I1, −∇U(·) does not depend on α, we get
I1 = − K−1∑ k=0 ∫ tk+1 tk Ex,i[〈(∇U(β̃tk) + ξtk −∇U(β̃t)),∇u(β̃t, t)〉]dt
= − K−1∑ k=0 ∫ tk+1 tk Ex,i[〈(∇U(β̃tk)−∇U(β̃t, t)),∇u(β̃t, t)〉]dt,
where we use the assumption Ex,i[ξtk ] = 0, and independent property of the noise ξtk . For tk ≤ t < tk+1, let
ρ1(β̃t, t) = 〈(∇U(β̃tk)−∇U(β̃t)),∇u(β̃t, t)〉,
applying Itô formula (same as (63) above), taking expectation, and using the fact Q̃ρ1(β̃t, t) = aS̃(β̃
(1) t , β̃ (2) t )[ρ1((β (1) t ,β (2) t ), t)− ρ1((β (2) t ,β (1) t ), t)] = 0, we get
dEx,i[ρ1(β̃t, t)] dt = Ex,i[∂tρ1(β̃t, t) + tr(Σ̃Σ̃T∇2ρ1(β̃t, t))− 〈∇Ũ(β̃t),∇ρ1(β̃t, t)〉] ≤ 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,
Ex,i[ρ1(β̃t, t)] ≤ Ckη,
which implies I1 ≤ TηCmax.
Estimate of I2: for each time interval tk ≤ t < tk+1, falling between two stopping time, the matrix Σ(α(t)) = Σ̃(α̃(t)) remains constant, and tr(ΣΣT∇2u) = tr(Σ̃Σ̃T∇2u) = ∑P j=0 τj∆β̃ju, we have
I2 = 0.
Estimate of I3: we rewrite the time interval [0, T ∧ τU ] as the summation of the stopping times, i.e.∫ T∧τU 0 = ∑N(T∧τU )−1 l=0 ∫ ζl+1 ζl , we thus have
I3 := ∫ T∧τU
0
Ex,i[Q̃(β̃t, t)u(β̃t, t, ·)−Q(β̃t, t)u(β̃t, t, ·)]dt
= N(T∧τU )−1∑ l=0 ∫ ζl+1 ζl Ex,i[Q̃(β̃t, t)u(β̃t, t, ·)−Q(β̃t, t)u(β̃t, t, ·)]dt.
Recall that, Q(x)u(x, ·)(i) = ∑ j∈M,j 6=i Qij(u(x, j)− u(x, i)),
and we denote Q̃ij as the components of the generator matrix Q̃ of β̃. For ζl ≤ t < ζl+1, we denote
ρl3(β̃t, t) := Q̃(β̃t, t)u(β̃t, t, ·)(α̃(τl))−Q(β̃t, t)u(β̃t, t, ·)(α(τl)) = ∑
j∈M,j 6=α(τl)
(Q̃α(τl),j(β̃t)−Qα(τl),j(β̃t))(u(β̃t, t)(α(j)− u(β̃t, t)α(τl))).
In particular, for the two chain case,M = {0, 1}. For any t ≤ T ∧ τU , there exists large enough open ball with radius R, denoted as U ⊂ B(β(0),R), such that U ⊂ B(β(0),R), which implies the uniform bound of u(βt, t) ≤ C, for t ≤ T ∧ τU . We thus get the estimates
|I3| ≤ N(T∧τU )−1∑
l=0
∫ ζl+1 ζl Ex,i[|ρl3(β̃t, t)|]
≤ 2C N(T∧τU )−1∑
l=0
∫ ζl+1 ζl ∑ j∈M,j 6=α(τl) Ex,i[|(Q̃α(τl),j(β̃t)−Qα(τl),j(β̃t))|]
≤ 2Caη N(T∧τU )−1∑
l=0
∫ ζl+1 ζl ∑ j∈M,j 6=α(τl) Ex,i[|(S̃α(τl),j(β̃t)− Sα(τl),j(β̃t))|]
≤ 2CKaη.
where the last inequality follows from the fact that S, S̃ ≤ 1, and N(T ∧ τU ) ≤ K. Combing all the estimates in I1, I2, I3, the proof is thus completed.

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

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. 

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.