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 addresses the problem of sampling from distributions with complex energy landscapes. The authors propose an extension of the contour Stochastic Gradient Langevin dynamics (CSGLD) sampler to efficiently simulate from big-data distributions. The extension is called “interacting” since $P$ different CSGLD are updated simultaneously and the random-field function is obtained as a Monte-Carlo average over the random field functions of these multiple chains run in parallel. This procedure allows to reduce the variance of the self-adapting parameter $\theta$ improving the marginal energy likelihood estimate, and is well-suited for distributed computing. The different chains only share the low-dimensional vector $\theta$ and do not share the model parameters, that are typically high dimensional. The authors prove that the proposed algorithm is asymptotically more efficient than its single-chain counterpart with the same computational budget. This theorem holds for a step size that decreases polynomially in time. Finally, they compare their algorithm to alternative methods on different tasks (e.g., mode exploration on MNIST dataset, uncertainty estimation) achieving competitive performances.