Interacting Contour Stochastic Gradient Langevin Dynamics

1 INTRODUCTION . Stochastic gradient Langevin dynamics ( SGLD ) ( Welling & Teh , 2011 ) has achieved great successes in simulations of high-dimensional systems for big data problems . It , however , yields only a fast mixing rate when the energy landscape is simple , e.g. , local energy wells are shallow and not well separated . To improve its convergence for the problems with complex energy landscapes , various strategies have been proposed , such as momentum augmentation ( Chen et al. , 2014 ; Ding et al. , 2014 ) , Hessian approximation ( Ahn et al. , 2012 ; Li et al. , 2016 ) , high-order numerical schemes ( Chen et al. , 2015 ; Li et al. , 2019b ) , and cyclical learning rates ( Izmailov et al. , 2018 ; Maddox et al. , 2019 ; Zhang et al. , 2020b ) . In spite of their asymptotic properties in Bayesian inference ( Vollmer et al. , 2016 ) and non-convex optimization ( Zhang et al. , 2017 ) , it is still difficult to achieve compelling empirical results for pathologically complex deep neural networks ( DNNs ) . To simulate from distributions with complex energy landscapes , e.g. , those with a multitude of modes well separated by high energy barriers , an emerging trend is to run multiple chains , where interactions between different chains can potentially accelerate the convergence of the simulation . For example , Song et al . ( 2014 ) and Futami et al . ( 2020 ) showed theoretical advantages of appropriate interactions in ensemble/population simulations . Other multiple chain methods include particle-based nonlinear Markov ( Vlasov ) processes ( Liu & Wang , 2016 ; Zhang et al. , 2020a ) and replica exchange methods ( also known as parallel tempering ) ( Deng et al. , 2021a ) . However , the particle-based methods result in an expensive kernel matrix computation given a large number of particles ( Liu & Wang , 2016 ) ; similarly , naı̈vely extending replica exchange methods to population chains leads to a long waiting time to swap between non-neighboring chains ( Syed et al. , 2021 ) . Therefore , how to conduct interactions between different chains , while maintaining the scalability of the algorithm , is the key to the success of the parallel stochastic gradient MCMC algorithms . In this paper , we propose an interacting contour stochastic gradient Langevin dynamics ( ICSGLD ) sampler , a pleasingly parallel extension of contour stochastic gradient Langevin dynamics ( CSGLD ) ( Deng et al. , 2020b ) with efficient interactions . The proposed algorithm requires minimal communication cost in that each chain shares with others the marginal energy likelihood estimate only . As a result , the interacting mechanism improves the convergence of the simulation , while the minimal communication mode between different chains enables the proposed algorithm to be naturally adapted to distributed computing with little overhead . For the single-chain CSGLD algorithm , despite its theoretical advantages as shown in Deng et al . ( 2020b ) , estimation of the marginal energy likelihood remains challenging for big data problems with a wide energy range , jeopardizing the empirical performance of the class of importance sampling methods ( Gordon et al. , 1993 ; Doucet et al. , 2001 ; Wang & Landau , 2001 ; Liang et al. , 2007 ; Andrieu et al. , 2010 ; Deng et al. , 2020b ) in big data applications . To resolve this issue , we resort to a novel interacting random-field function based on multiple chains for an ideal variance reduction and a more robust estimation . As such , we can greatly facilitate the estimation of the marginal energy likelihood so as to accelerate the simulations of notoriously complex distributions . To summarize , the algorithm has three main contributions : • We propose a scalable interacting importance sampling method for big data problems with the minimal communication cost . A novel random-field function is derived to tackle the incompatibility issue of the class of importance sampling methods in big data problems . • Theoretically , we study the local stability of a non-linear mean-field system and justify regularity properties of the solution of Poisson ’ s equation . We also prove the asymptotic normality for the stochastic approximation process in mini-batch settings and show that ICSGLD is asymptotically more efficient than the single-chain CSGLD with an equivalent computational budget . • Our proposed algorithm achieves appealing mode explorations using a fixed learning rate on the MNIST dataset and obtains remarkable performance in large-scale uncertainty estimation tasks . 2 PRELIMINARIES . 2.1 STOCHASTIC GRADIENT LANGEVIN DYNAMICS . A standard sampling algorithm for big data problems is SGLD ( Welling & Teh , 2011 ) , which is a numerical scheme of a stochastic differential equation in mini-batch settings : xk+1 = xk − k N n ∇xŨ ( xk ) + √ 2τ kwk , ( 1 ) where xk ∈ X ∈ Rd , k is the learning rate at iteration k , N denotes the number of total data points , τ is the temperature , andwk is a standard Gaussian vector of dimension d. In particular , Nn∇xŨ ( x ) is an unbiased stochastic gradient estimator based on a mini-batch data B of size n and Nn Ũ ( x ) is the unbiased energy estimator for the exact energy function U ( x ) . Under mild conditions on U , xk+1 is known to converge weakly to a unique invariant distribution π ( x ) ∝ e− U ( x ) τ as k → 0 . 2.2 CONTOUR STOCHASTIC GRADIENT LANGEVIN DYNAMICS . Despite its theoretical guarantees , SGLD can converge exponentially slow when U ( x ) is non-convex and exhibits high energy barriers . To remedy this issue , CSGLD ( Deng et al. , 2020b ) exploits the flat histogram idea and proposes to simulate from a flattened density with much lower energy barriers $ Ψθ ( x ) ∝ π ( x ) /Ψ ζ θ ( U ( x ) ) , ( 2 ) where ζ is a hyperparameter , Ψθ ( u ) = ∑m i=1 ( θ ( i− 1 ) e ( log θ ( i ) −log θ ( i−1 ) ) u−ui−1 ∆u ) 1ui−1 < u≤ui . In particular , { ui } mi=0 determines the partition { Xi } mi=1 of X such that Xi = { x : ui−1 < U ( x ) ≤ ui } , where −∞ = u0 < u1 < · · · < um−1 < um = ∞ . For practical purposes , we assume ui+1 − ui = ∆u for i = 1 , · · · , m− 2 . In addition , θ = ( θ ( 1 ) , θ ( 2 ) , . . . , θ ( m ) ) is the self-adapting parameter in the space Θ = { ( θ ( 1 ) , · · · , θ ( m ) ) ∣∣0 < θ ( 1 ) , · · · , θ ( m ) < 1 & ∑mi=1 θ ( i ) = 1 } . Ideally , setting ζ = 1 and θ ( i ) = θ∞ ( i ) , where θ∞ ( i ) = ∫ Xi π ( x ) dx for i ∈ { 1 , 2 , · · · , m } , enables CSGLD to achieve a “ random walk ” in the space of energy and to penalize the over-visited partition ( Wang & Landau , 2001 ; Liang et al. , 2007 ; Fort et al. , 2011 ; 2015 ) . However , the optimal values of θ∞ is unknown a priori . To tackle this issue , CSGLD proposes the following procedure to adaptively estimate θ via stochastic approximation ( SA ) ( Robbins & Monro , 1951 ; Benveniste et al. , 1990 ) : ( 1 ) Sample xk+1 = xk + k Nn∇xŨΨθk ( xk ) + √ 2τ kwk , ( 2 ) Optimize θk+1 = θk + ωk+1H̃ ( θk , xk+1 ) , Algorithm 1 Interacting contour stochastic gradient Langevin dynamics algorithm ( ICSGLD ) . { Xi } mi=1 is pre-defined partition and ζ is a hyperparameter . The update rule in distributed-memory settings and discussions of hyperparameters is detailed in section B.1.1 in the supplementary material . [ 1 . ] ( Data subsampling ) Draw a mini-batch data Bk from D , and compute stochastic gradients ∇xŨ ( x ( p ) k ) and energies Ũ ( x ( p ) k ) for each x ( p ) , where p ∈ { 1 , 2 , · · · , P } , |Bk| = n , and |D| = N . [ 2 . ] ( Parallel simulation ) Sample x ⊗ P k+1 : = ( x ( 1 ) k+1 , x ( 2 ) k+1 , · · · , x ( P ) k+1 ) > based on SGLD and θk x ⊗ P k+1 = x ⊗ P k + k N n ∇xŨΨθk ( x ⊗ P k ) + √ 2τ kw ⊗ P k , ( 4 ) where k is the learning rate , τ is the temperature , w ⊗ P k denotes P independent standard Gaussian vectors , ∇xŨΨθ ( x ⊗ P ) = ( ∇xŨΨθ ( x ( 1 ) ) , ∇xŨΨθ ( x ( 2 ) ) , · · · , ∇xŨΨθ ( x ( P ) ) ) > , and ∇xŨΨθ ( x ) = [ 1 + ζτ∆u ( log θ ( JŨ ( x ) ) − log θ ( ( JŨ ( x ) − 1 ) ∨ 1 ) ) ] ∇xŨ ( x ) for any x ∈ X . [ 3 . ] ( Stochastic approximation ) Update the self-adapting parameter θ ( i ) for i ∈ { 1 , 2 , · · · , m } θk+1 ( i ) = θk ( i ) + ωk+1 1 P P∑ p=1 θk ( JŨ ( x ( p ) k+1 ) ) ( 1 i=JŨ ( x ( p ) k+1 ) − θk ( i ) ) , ( 5 ) where 1A is an indicator function that takes value 1 if the event A appears and equals 0 otherwise . where∇xŨΨθ ( · ) is a stochastic gradient function of $ Ψθ ( · ) to be detailed in Algorithm 1 . H̃ ( θ , x ) : = ( H̃1 ( θ , x ) , · · · , H̃m ( θ , x ) ) is random-field function where each entry follows H̃i ( θ , x ) = θζ ( JŨ ( x ) ) ( 1i=JŨ ( x ) − θ ( i ) ) , where JŨ ( x ) = m∑ i=1 i1ui−1 < Nn Ũ ( x ) ≤ui . ( 3 ) Theoretically , CSGLD converges to a sampling-optimization equilibrium in the sense that θk approaches to a fixed point θ∞ and the samples are drawn from the flattened density $ Ψθ∞ ( x ) . Notably , the mean-field system is globally stable with a unique stable equilibrium point in a small neighborhood of θ∞ . Moreover , such an appealing property holds even when U ( x ) is non-convex . 3 INTERACTING CONTOUR STOCHASTIC GRADIENT LANGEVIN DYNAMICS . The major goal of interacting CSGLD ( ICSGLD ) is to improve the efficiency of CSGLD . In particular , the self-adapting parameter θ is crucial for ensuring the sampler to escape from the local traps and traverse the whole energy landscape , and how to reduce the variability of θk ’ s is the key to the success of such a dynamic importance sampling algorithm . To this end , we propose an efficient variance reduction scheme via interacting parallel systems to improve the accuracy of θk ’ s .

Summary: This paper samples from distributions with complex energy landscapes by the powerful contour stochastic gradient Langevin dynamics (CSGLD) methods. In order to reduce the variance of random-field function, this paper proposes a simple but effective method, that combines  random-field functions from multiple chains. Since the parameter $\theta$ typically has small size, the communication overhead is marginal. This paper show that compared to CSGLD which only has one chain, the variance ICSGLD with P chains and same computation budget is asymptotically more efficient. The authors evaluate ICSGLD with multiple experiments.