Non-reversible Parallel Tempering for Uncertainty Approximation in Deep Learning

1 INTRODUCTION . 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 . 2 PRELIMINARIES . 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 ) . 3 OPTIMAL 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 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 . 3.1 GENERALIZED DEO SCHEME . 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 . 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 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 . 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 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 .

Summary: 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.