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