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: The paper proposes an MCMC algorithm containing two key ideas: the approximation of the target density with a simpler function (proposed by [Deng, 2020]) and the parallel simulation of many chains that allow for better capturing the global properties of the target density. To be more precise, the main contribution of the paper is the extension of the algorithm by [Deng, 2020] to the case of multiple chains.   The proposed extension goes as follows. The original algorithm by [Deng, 2020] approximates the target density with piecewise continuous functions using the level sets (slicing the values of the energy function). These functions are parameterized by vector $\theta$, which needs to be defined. In practice, $\theta$ is estimated using the samples from the chain. The current paper then proposes to run many chains in parallel and use all of the available samples for the estimation.  The efficiency of the proposed technique is analyzed both theoretically and empirically. For theoretical analysis, the authors provide convergence speed and asymptotic distribution of the residuals $\theta - \theta^*$. Overall, the analysis favors the proposed method over the original algorithm. The empirical comparison is made on several tasks. Except for the toy task, the authors consider learning policy for contextual bandit problem and sampling from the posterior distribution of a neural network.  [Deng, 2020]: Deng, Wei, Guang Lin, and Faming Liang. "A Contour Stochastic Gradient Langevin Dynamics Algorithm for Simulations of Multi-modal Distributions." arXiv preprint arXiv:2010.09800 (2020).