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: The paper proposes a new hybrid method between MCMC and VI. The main idea of the paper is construction of a new energy function which allows to speed up the sampling significantly in comparison to the SGMCMC (stochastic gradient MCMC). An additional modification of the proposed algorithm is done by adopting a drop-out inspired approximation which allows for even better scalability.