Structured Stochastic Gradient MCMC

1 INTRODUCTION . There has been much recent interest in deep Bayesian neural networks ( BNN ) due to their reliable confidence estimates and generalization properties ( Wilson & Izmailov , 2020 ; Jospin et al. , 2020 ; Cardelli et al. , 2019 ) . BNNs rely on ensemble averages over model parameters typically obtained from Markov chain Monte Carlo ( MCMC ) algorithms , which contrasts to regular neural networks that depend on a single set of parameters . The sheer size of these models requires scalable MCMC approaches based on inexpensive stochastic gradients , of which stochastic gradient Markov chain Monte Carlo ( SGMCMC ) algorithms are the gold standard ( Li et al. , 2016 ; Welling & Teh , 2011 ; Patterson & Teh , 2013 ) . These algorithms owe their scalability to approximating gradients via mini-batching . The main downside of using SGMCMC algorithms is their slow mixing rates in high dimensions . An often faster alternative is variational inference ( VI ) algorithms that approximate the posterior with a simpler ( typically factorized ) distribution . This formulation results in an optimization problem that can be solved more efficiently using stochastic optimization ( Blei et al. , 2017 ; Zhang et al. , 2018 ) . One downside of VI approximations is their solid distributional assumptions . A typical choice is to approximate the Bayesian posterior by a product of univariate Gaussian distributions . These distributional assumptions are frequently over-simplistic in high-dimensional models , where the posterior can be highly multi-modal and possibly heavy-tailed . Another downside is that the variational approximation typically underestimates the posterior variance , leading to poorly calibrated uncertainties and overfitting ( Ormerod & Wand , 2010 ; Giordano et al. , 2015 ; Zhang et al. , 2018 ) . In this work , we derive a fundamentally new SGMCMC approach that takes inspiration from structured VI . While our approach remains a sampling algorithm resembling SGMCMC , we speed up the mixing time by systematically breaking posterior correlations . The resulting algorithm furthermore allows users to specify which posterior correlations to keep and which ones to break . It makes no assumptions on the functional form of the approximate posterior . We call our approach structured SGMCMC since it relies on a structured ( i.e. , only partially factorized ) variational approximation of the posterior ( Wainwright & Jordan , 2008 ) . In more detail , we derive the optimal variational distribution for a given posterior subject to factorization constraints by assuming a functional view on variational inference . We show how to sample from this optimal distribution by running SGMCMC on a modified energy function . This energy function is obtained by marginalizing the model ’ s joint distribution over previously generated samples from the Markov chain , leading to an approximate factorization over user-specified parameter groups . Further , we provide a more robust and computationally efficient approximation to the procedure that allows for interpolation between regular SGMCMC and our structured SGMCMC by taking inspiration from dropout techniques . Both methods are compatible with any Markovian SGMCMC algorithm , including Langevin dynamics and stochastic gradient Hamiltonian Monte Carlo . In sum , our contributions are as follows : • We propose a new approximate MCMC scheme running SGMCMC on a modified energy function , trading accuracy for speed . This setup effectively allows sampling from a fully joint posterior , a completely factorized posterior , and any in-between . • We prove mathematically that the resulting scheme asymptotically generates samples from the best possible posterior approximation subject to user-specified factorization constraints between groups of parameters . • We extend this scheme further by making it more scalable with a dropout-inspired approximation . This new scheme has a hyperparameter that enables a smooth interpolation between full SGMCMC and a `` mean-field '' version where all posterior correlations are broken . • We show in both small and large scale experiments that our method well approximates posterior marginals and gives improved results over SGMCMC on Resnet-20 architectures on CIFAR-10 , Fashion MNIST , and SVHN in terms of both runtime and final accuracy . Our paper is structured as follows : Section 2 presents the related work to our proposal , Section 3 introduces preliminaries regarding the energy function and the stochastic gradient updates , Sections 4 and 5 derive our proposed methods , Section 6 details experiments and their results , and Section 7 contains our concluding thoughts . 2 RELATED WORK . Our work connects both to ( stochastic ) variational inference ( Bishop , 2006 ; Hoffman et al. , 2013 ; Ranganath et al. , 2014 ; Blei et al. , 2017 ; Zhang et al. , 2018 ) and scalable MCMC ( Welling & Teh , 2011 ; Chen et al. , 2014 ; Ma et al. , 2017 ; Zhang et al. , 2020 ; Leimkuhler et al. , 2019 ; Wenzel et al. , 2020 ; Izmailov et al. , 2021 ) . For space limitations , we focus on the most related work at the intersection of both topics . Among the earliest works to hybridize both approaches was ( de Freitas et al. , 2001 ) who constructed a variational proposal distribution in the Metropolos-Hastings step of MCMC . An improved approach to that was introduced in ( Habib & Barber , 2018 ) , where by introducing low-dimensional auxiliary variables they fit a more accurate approximating distribution . Other related advances to MCMC methods were proposed by Levy et al . ( 2017 ) who developed a method to train MCMC kernels with NNs , and Wang et al . ( 2018 ) ; Gong et al . ( 2018 ) who leveraged meta learning schemes in SGMCMC methods . Most recent work focuses on connections between VI and stochastic gradient-based MCMC , or between VI and stochastic gradient descent ( SGD ) . For example , Mandt et al . ( 2016 ; 2017 ) and Duvenaud et al . ( 2016 ) consider SGD as a type of variational inference , but their approaches did not attempt to close the gap to exact MCMC . Other works aim at explicitly interpolating between both methods . Domke ( 2017 ) proposes a divergence bound for hybridizing VI and MCMC , essentially by running Langevin dynamics on a tempered evidence lower bound ( ELBO ) . Salimans et al . ( 2015 ) embody MCMC steps into the variational inference approximation . Ahn et al . ( 2012 ) improve stochastic gradient Langevin dynamics by leveraging the central limit theorem and using the estimated inverse Fisher information matrix to sample from the approximate posterior distribution . Rezende & Mohamed ( 2015 ) interpreted the path of an MCMC algorithm as a variational distribution , and then fitting parameters to tighten a variational bound . Recently , Hoffman & Ma ( 2020 ) interpreted ( parametric ) VI as approximate Langevin dynamics and showed that both algorithms have similar transient dynamics . In contrast to all these approaches , our method is inspired by coordinate ascent variational inference ( Bishop , 2006 ) but uses Langevin updates to generate samples from the target distribution that respects an imposed independence structure . 3 PRELIMINARIES . Variational inference ( VI ) approaches differ from MCMC in two regards : ( 1 ) they impose a structured ( e.g. , fully-factorized ) approximation of the posterior for tractability , and ( 2 ) they often make parametric assumptions . Is it possible to construct a modified scheme that only relies on the assumption ( 1 ) , inheriting the non-parametric nature of MCMC while breaking posterior correlations in a controlled manner ? As follows , we will show how such a scheme can be realized . We will first derive a modified energy function for Langevin dynamics that we can sample from and then prove that its negative exponential results in the optimal posterior approximation subject to specified factorization constraints . Running SGMCMC algorithms on this energy function will consequently generate samples from this distribution . Before we explain our new method , we introduce the setup and common notation . Given data D = { ( xi , yi ) } i=1 , ... , N , parameters θ , a proper prior distribution p ( θ ) , and a likelihood p ( D|θ ) =∏N i=1 p ( yi|xi , θ ) , suppose we are interested in the corresponding posterior distribution p ( θ|D ) ∝ p ( D|θ ) p ( θ ) . A convenient representation of the posterior is as a Boltzmann distribution : p ( θ|D ) ∝ exp { −U ( θ ) } where U ( θ ) = − log p ( θ , D ) = − ∑ ( x , y ) ∈D log p ( y|x , θ ) − log p ( θ ) . ( 1 ) U is typically referred to as the posterior energy function . Note that the posterior distribution is typically intractable due to the normalizing constant . A popular approach for approximating the entire posterior distribution is by deploying Markov chain Monte Carlo ( MCMC ) algorithms . These methods work by producing an empirical distribution of samples in parameter space , often times through the use of a random walk . While being very accurate and having asymptotic guarantees , these methods are known to not scale well with respect to both data and parameters ( Brooks et al. , 2011 ; Geyer , 1992 ) . Stochastic gradient MCMC ( SGMCMC ) is a class of scalable MCMC algorithms that can produce posterior samples through gradients on minibatches of data . These algorithms are largely derived from discretized approximations of continuous-time diffusion processes . Examples of these algorithms include stochastic gradient Langevin dynamics ( SGLD ) ( Welling & Teh , 2011 ) , preconditioned SGLD ( pSGLD ) ( Li et al. , 2016 ) , and stochastic gradient Hamiltonian Monte Carlo ( SGHMC ) ( Chen et al. , 2014 ) . As alluded to , the basis of SGMCMC algorithms is using a sampled minibatch of data D̃ from D to produce an differentiable , unbiased estimate of the posterior energy function : U ( θ ) ≈ Û ( θ ; D̃ ) = − N |D̃| ∑ ( x , y ) ∈D̃ log p ( y|x , θ ) − log p ( θ ) . ( 2 ) Once Û is defined , it is fairly straight forward to generate new samples from the posterior distribution . For instance , the SGLD update is θ ( t+1 ) = θ ( t ) − t 2 ∇θÛ ( θ ( t ) ; D̃t ) + ξt where ξt ∼ N ( 0 , tI ) . ( 3 ) Similar rules for pSGLD and SGHMC can be found in the Supplement . All of these update rules produce a chain of samples up to time step t that ultimately form an empirical distribution p̂ ( t ) ( θ|D ) . Should the algorithms converge , then limt→∞ p̂ ( t ) ( θ|D ) = p ( θ|D ) . 4 STRUCTURED SGMCMC . By design , SGMCMC methods produce a fully joint posterior distribution over parameters θ . For models with a large number of parameters , this can lead to various complications due to the curse of dimensionality . This is typically observed with slow convergence times and potentially unexplored parameter spaces . A viable solution is to break dependencies in the posterior distribution by leveraging ideas commonly used in variational inference ( VI ) . This would reduce the number of various potential posterior correlations that the model would need to capture while sampling . To achieve partial factorization , we must first partition θ into M > 1 distinct , mutually independent groups : θ1 , . . . , θM . This partitioning structure is assumed to be known a priori . We will denote the distribution that respects this partitioning structure as q ( θ ) = ∏M i=1 qi ( θi ) . Similar to VI , we would like this distribution q ( θ ) to best approximate the true posterior distribution p ( θ|D ) according to some criteria , such as KL-divergence . This leads to a natural objective function to minimize : J ( q ( θ ) ) = DKL ( q ( θ ) ||p ( θ|D ) ) ≡ Eθ∼q [ log q ( θ ) p ( θ|D ) ] ( 4 ) The following Theorem 1 proves that there is a unique solution to the non-parametric KL minimization problem described in Eq . ( 4 ) . To describe it , we compose θ = { θi , θ̃¬i } for any i where θ̃ ∼ q and define a structured energy function : U ( S ) ( θ ) = M∑ i=1 U ( S ) i ( θi ) , with U ( S ) i ( θi ) : = Eθ̃∼qU ( { θi , θ̃¬i } ) : = −Eθ̃∼q log p ( θi , θ̃¬i , D ) . ( 5 ) That is , we first define the marginals U ( S ) i ( θi ) , where we marginalize U ( θ ) with respect to all q ( θ ) -factors except qi ( θi ) , and then sum up these marginals to define U ( S ) ( θ ) . A similar partial marginalization procedure is carried out for conjugate exponential family distributions in coordinate ascent VI ( Bishop , 2006 ) . Having a well-defined energy function U ( S ) allows us to use standard SGMCMC methods to approximate the posterior q ( θ ) with samples . This serves as the basis for our proposed algorithm that actually approximates this distribution q ( θ ) , which will be discussed shortly . Theorem 1 . The unique solution to the KL minimization problem given in Eq . 4 is given by the Boltzmann distribution q ( θ ) ∝ exp { − ∑M i=1 U ( S ) i ( θi ) } . Please refer to the Supplement for the proof . In an ideal world , we would be able to use the findings of Theorem 1 directly in conjunction with algorithms like Langevin dynamics and Hamiltonian Monte Carlo to produce empirical distributions for q using U ( S ) ( Liu et al. , 2019 ) . However , this is intractable for two reasons : ( 1 ) these algorithms generally work only well with small amounts of data , and ( 2 ) more importantly , the marginals U ( S ) i ( θi ) do not have a closed-form solution but need to be approximated via samples from q. Luckily , since SGMCMC methods only need access to noisy estimates of U ( S ) , we can run these algorithms on a stochastic estimate of Eq . ( 5 ) , U ( S ) ( θ ) ≈ Û ( S ) ( θ ; D̃ ) = M∑ i=1 Eθ̃∼qÛ ( { θi , θ̃¬i } ; D̃ ) , ( 6 ) where Û ( · ) is defined in Eq . ( 2 ) . In practice , at timestep t for i = 1 , . . . , M we estimate Eθ̃∼qÛ ( { θi , θ̃¬i } ; D̃t ) with a Monte Carlo approximation . In place of θ̃ , we use a single sample of θ̃ ( t ) taken from the current approximate distribution q̂ ( t ) which is composed of samples from previous timesteps ( i.e. , a uniform distribution over { θ ( 1 ) , θ ( 2 ) , . . . , θ ( t ) } ) . This leads to the following update step for structured SGLD ( S-SGLD ) : θ ( t+1 ) = θ ( t ) − t 2 ∇θÛ ( S ) ( θ ; D̃ ) + ξt where ξt ∼ N ( 0 , tI ) . ( 7 ) Similar rules for structured variants of pSGLD ( S-pSGLD ) and SGHMC ( S-SGHMC ) can be found in the Supplement . Additionally , the full procedure for structured SGMCMC ( S-SGMCMC ) can be seen in Algorithm 2 . Remark Since ∇θÛ ( S ) is an unbiased estimator for U ( S ) , we are guaranteed to converge to q from sampling with S-SGMCMC with sufficiently decreasing learning rates so long as we are in a stationary state . While it is unlikely to have the procedure initialize to a stationary state , we observe in practice that our scheme both tends to converge towards and remain in a stationary state . A general proof of convergence is outside the scope of this work and is left to follow-up research . An example of S-SGMCMC can be seen in Fig . 1 ( a-b ) , which features the approximate posterior distributions of a linear regression model with three coefficients and with various independence structures imposed with S-SGLD : ( a ) joint dependence between w1 , w2 , and w3 ; ( b-left ) dependence between w1 and w2 but independence between w3 and the other coefficients ; ( b-right ) fully factorized . Of note is that the bivariate posterior distributions appear to respect the imposed independence structure . Interestingly , it also appears that the variance shrinks as we induce these factorizations which is a commonly seen artifact when using VI .

Summary: This work proposes using a structured variational approximation for stochastic gradient Markov Chain Monte Carlo. This allows someone to choose a factorization for the variational distribution (which factorization is best is unclear; several are studied). Analogously to coordinate ascent variational inference, the authors show that the best approximation is the Boltzmann energy function marginalized over the complements of every parameter group.   However, this structured approximation is computationally expensive and requires the same number of evaluations of the approximation as there are parameter groups. This computational burden is alleviated by a dropout scheme, where instead of sampling from every parameter groups, parameters are masked using a dropout distribution, and the number of stochastic masks is a hyperparameter that controls regularization and fidelity to the structure imposed in the factorization.   Experiments show that this is a viable way to impose structure on a variational distribution, and that mixing times are improved.