S TRUCTURED S TOCHASTIC G RADIENT MCMC
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic gradient Markov Chain Monte Carlo (SGMCMC) is considered the gold standard for Bayesian inference in large-scale models, such as Bayesian neural networks.
Since practitioners face speed versus accuracy tradeoffs in these models, variational inference (VI) is often the preferable option.
Unfortunately,
VI makes strong assumptions on both the factorization and functional form of the posterior.
In this work, we propose a new non-parametric variational approximation that makes no assumptions about the approximate posterior’s functional form and allows practitioners to specify the exact dependencies the algorithm should respect or break.
The approach relies on a new Langevin-type algorithm that operates on a modiﬁed energy function, where parts of the latent variables are averaged over samples from earlier iterations of the Markov chain.
This way, statistical dependencies can be broken in a controlled way, allowing the chain to mix faster.
This scheme can be further modiﬁed in a “dropout” manner, leading to even more scalability.
By implementing the scheme on a ResNet-20 architecture, we obtain better predictive likelihoods and faster mixing time than full SGMCMC.
I NTRODUCTION
There has been much recent interest in deep Bayesian neural networks (BNN) due to their reliable conﬁdence 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 efﬁciently 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 overﬁtting (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 modiﬁed 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-speciﬁed parameter groups.
Further, we provide a more robust and computationally efﬁcient 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 modiﬁed 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-speciﬁed 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-ﬁeld" 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 ﬁnal accuracy.
Our paper is structured as follows:
Section 2 presents the related work to our proposal, Sectionintroduces preliminaries regarding the energy function and the stochastic gradient updates, Sectionsand 5 derive our proposed methods, Section 6 details experiments and their results, and Sectioncontains our concluding thoughts.
R ELATED W ORK
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 ﬁt 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.
; 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.
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 ﬁtting 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.
P RELIMINARIES
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 modiﬁed 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 ﬁrst derive a modiﬁed 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 speciﬁed 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 = { ( x i , y i ) } i =1 ,...,N , parameters θ , a proper prior distribution p ( θ ) , and a likelihood p ( D| θ ) =
(cid:81)
Ni =1 p ( y i | x i , θ ) , 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 )
= − (cid:88)
( x,y ) ∈D log p ( y | x, θ ) − log p ( θ ) .
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:
[Equation]
Once ˆ U is deﬁned, it is fairly straight forward to generate new samples from the posterior distribution.
For instance, the SGLD update is
[Equation]
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 lim t →∞ ˆ p ( t ) ( θ |D )
=
p ( θ |D ) .
S TRUCTURED 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 ﬁrst 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 ( θ ) =
(cid:81) Mi
=1 q i ( θ 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:
[Equation]
The following Theorem 1 proves that there is a unique solution to the non-parametric KL minimization problem described in Eq.
To describe it, we compose θ = { θ i , ˜ θ ¬ i } for any i where ˜ θ ∼ q and deﬁne a structured energy function:
[Equation]
That is, we ﬁrst deﬁne the marginals U ( S )
i ( θ i ) , where we marginalize U ( θ ) with respect to all q ( θ ) -factors except q i ( θ i ) , and then sum up these marginals to deﬁne U ( S )
( θ ) .
A similar partial marginalization procedure is carried out for conjugate exponential family distributions in coordinate ascent VI (Bishop, 2006).
Having a well-deﬁned 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 {− (cid:80) Mi
=1 U ( S ) i ( θ i ) } .
Please refer to the Supplement for the proof.
In an ideal world, we would be able to use the ﬁndings 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.
[Equation]
where ˆ U ( · ) is deﬁned in Eq.
In practice, at timestep t for i = 1 , . . .
, M we estimate
E ˜ θ ∼ q ˆ U ( { θ 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):
[Equation]
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 sufﬁciently 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.
[Figure] [Figure]
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 coefﬁcients and with various independence structures imposed with S-SGLD: (a) joint dependence between w 1 , w 2 , and w 3 ; (b-left) dependence between w 1 and w 2 but independence between w 3 and the other coefﬁcients; (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.
S TRUCTURED D ROPOUT SGMCMC
While S-SGMCMC can successfully break dependencies between parameter groups, it does suffer computationally due to each parameter update scaling linearly with respect to M .
This means that for a single new sample of θ , the model’s forward pass needs to be computed M different times on the same batch of data ˜ D , which can quickly become prohibitively expensive for deep models when M is large.
Ideally, we would prefer a method that both closely resembles the S-SGMCMC procedure and scales independently from the partitioning scheme.
This section presents such a method that achieves this, which we call structured dropout SGMCMC (S d -SGMCMC), as well as an informal motivation and derivation of the method.
More formal details and a theorem proving both SGMCMC and S-SGMCMC are limiting cases for S d -SGMCMC can be found in the Supplement.
The main motivation for this technique can be seen by recognizing that the composition { θ ( t ) i , ˜ θ ( t ) ¬ i } from Eq.
(6) can be rewritten as a sum of masked values rθ ( t )
+ (1 − r )
˜ θ ( t ) where ˜ θ ( t ) ∼ q ( t ) and r j = 1 ( i = j ) for i = 1 , . .
, M .
We can decouple the computational scaling from the number of parameter groups M by replacing the M deterministic masks r ’s with K stochastically sampled masks ˜ r .
1 Doing so results in a slightly different energy function and minibatch loss to optimize:
[Equation]
where ˜ r ( t,k ) is the k th sample of ˜ r for timestep t .
A formal justiﬁcation for Eq.
(8) can be found in the
Supplement.
These energy function approximations lead to the following update step for structured
Input: Initial sample θ (0) ; parameter partitions θ 1 , . .
, θ M ; step sizes { (cid:15) t } t =0 ,...,T − 1 .
Output: ˆ q ( T ) ( θ ) :
= { θ ( t ) } t =1 ,...,T 1 for t = 0 to T − 1 do
Sample minibatch ˜ D ( t )
⊂ D for i = 1 to M do
Sample ˜ θ ( t ) ¬ i ∼ ˆ q ( t ) ¬
i
ˆ U ( S,t ) i
= ˆ U ([ θ ( t ) i , ˜ θ ( t ) ¬ i ]; ˜ D ( t ) ) end ∇ θ ˆ U ( S,t )
= (cid:80) Mi
=1 ∇ θ ˆ U
( S,t ) i θ ( t +1) =
SGMCMC_step ( θ ( t ) , ∇ θ ˆ U ( S,t ) , (cid:15) t ) 9 end 10 return
ˆ q ( T ) ( θ )
[Table]
S-
[Equation]
S-SGHMC 7.
S-SGHMC 4.
S d -SGHMC | θ |7.
dropout variant of SGLD (S d -SGLD):
[Equation]
The corresponding update rules for the structured dropout variants for pSGLD (S d -pSGLD) and
SGHMC (S d -SGHMC) are deﬁned in the Supplement.
The exact procedure for generating samples of the approximate posterior ˆ q ( t ) using structured dropout SGMCMC (S d -SGMCMC) can also be found in the Supplement.
An example of this method (speciﬁcally S d -SGLD with ˜ r i iid ∼ Bernoulli (0 . 5) and K = 4 ) used on a linear regression model can be seen in Fig.
1(c).
Of note, we can see that the dropout variant largely respects the independence structure imposed, but maybe not as strictly as the exact S-SGLD method seen in Fig. 1(b).
Additionally, the posterior variance also seems to have shrunk similarly to S-SGLD when compared against SGLD.
Masking Distribution
Should ˜ r
i iid ∼ Bernoulli ( ρ ) , alongside a structure that factorizes by activation components, then the method starts to resemble dropout with rate ρ (Srivastava et al., 2014).
The main difference being that instead of replacing a parameter value with 0 it is replaced with a sample from the approximate posterior distribution at time t : ˆ q ( t ) .
While a Bernoulli distribution for ˜ r is a natural choice, there are other distributions that can be chosen as well.
For instance, ˜ r i iid ∼ N (0 , 1) or ˜ r i iid ∼ Beta ( α, β ) are both viable distributions and can be seen as analog to Gaussian and Beta-dropout respectively (Srivastava et al., 2014; Liu et al., 2019).
Our experiments will largely focus on sampling ˜ r from Bernoulli and uniform over [0 , 1] (equivalent to Beta (0 . 5 , 0 . 5) ) distributions.
E XPERIMENTS
Overview In this section we evaluate our proposed approach on various models and datasets.
Section 6.1 investigates the impact of the variational approximation on the algorithms’ mixing and autocorrelation times using a fully-connected network architecture on MNIST (LeCun et al., 2010).
Section 6.2 studies our methods with ResNet-20 (He et al., 2016) on CIFAR-10 (Krizhevsky et al., 2009), SVHN (Netzer et al., 2011), and Fashion MNIST (Xiao et al., 2017) and compares them for their accuracy and mixing time.
Our experiments reveal that the chains in our proposed methods mix faster than SGMCMC and achieve either comparable or even higher accuracies on average.
We have also conducted experiments on uncertainty visualization, where we tested the proposed methodology on predictive uncertainty estimation by deploying a two-layer fully connected network on a toy dataset.
The uncertainty experimental setup and results, along with more technical details for the other experiments, can be found in the Appendix.
Metrics The primary predictive metric of interest use to evaluate our proposal is classiﬁcation accuracy.
We take the average of an ensemble of 100 models whose weights are sampled from the past samples of the parameters chains in order to calculate the accuracy.
Additionally, we also monitor the mixing time of the chains of our methods with both integrated autocorrelation time (IAC) (Sokal, 1997; Goodman & Weare, 2010) and effective sample size (ESS) (Geyer, 1992).
IAC measures the correlation between samples in a chain and, in turn, describe the inefﬁciency of a MCMC algorithm.
IAC is computed as τ f = (cid:80) ∞
τ = −∞ ρ f ( τ ) where ρ f is the normalized autocorrelation function of the stochastic process that generated the chain for f and is calculated as ˆ ρ f ( τ ) =
ˆ c f ( τ ) /
ˆ c f (0) ; where ˆ c f ( τ ) =
1 N − τ (cid:80) N − τ n =1 ( f n − µ f ) ( f n + τ − µ f ) and µ f = 1 N (cid:80) Nn =1 f n .
We note that we calculated ˆ c f ( τ ) with a fast Fourier transform as it is more computationally efﬁcient than using the direct sum.
ESS measures how many independent samples would be equivalent to a chain of correlated samples and is calculated as n eﬀ = n 1+( n − 1) p , where n is the number of samples and p is the autocorrelation.
We note that a model with higher ESS and lower IAC have faster mixing time.
Please see the
Appendix for the detailed implementation details and experimental setup for the metrics and our models.
6. D ROPOUT R ATE & G ROUP S IZE I NVESTIGATION
The aim of this set of experiments is to study the effects that the number of independent parameter groups (or alternatively, the amount of allowed posterior correlations) has on accuracy and mixing time when using our proposed methods.
We compare pSGLD, S-pSGLD, and S d -pSGLD a Bernoulli ( ρ ) masking distribution with dropout rates
ρ ∈ { 0 . 1 , 0 . 3 , 0 . 5 } on a fully-connected neural network with 2 hidden layers, with 50 hidden units each, trained and evaluated with MNIST using the standard train and test split.
The model has 42,200 parameters in total.
For S-pSGLD and S d -pSGLD, these parameters are evenly distributed into M groups where M ranges from 4 to 42,200.
Accuracy, IAC, and ESS are reported in Fig.
2 using 100,000 posterior samples after a 150,000 burn in period.
More details on the implementation of the model regarding training and evaluation can be found in the
Appendix.
As shown in Fig.
2(a), for S-pSGLD we observe that as we increase the number of groups the accuracy drops dramatically whereas S d -pSGLD
’s
accuracy improves slightly and then remains fairly stable.
In the best case, S d -pSGLD achieves an accuracy of 96.3% with 32 groups and dropout rate of 0.
which outperforms pSGLD with accuracy of 94.2%.
We speculate that the dropout-like behavior is beneﬁcial for regularizing the model (much like normal dropout), hence the improved accuracy across all dropout rates.
Similarly, a single sample used for the Monte Carlo estimate in S-SGMCMC may not be enough as the number of groups M increase; however, increasing the number of samples in this scenario is infeasible due to S-SGMCMC scaling as O ( M ) .
Fig.
2(b-c) portrays the comparison between number of groups and mixing time metrics IAC and ESS.
As the number of groups gradually increase, we note that S-pSGLD mixes faster, as does S d -pSGLD to lesser and lesser degrees as ρ increases.
This behavior is to be expected due to Theorem 2, with
S d -pSGLD exhibiting mixing times more similar to pSGLD when ρ = 0 .
5 and more similar to
S-pSGLD when ρ = 0 .
6. S YSTEMATIC C OMPARISON ON R EAL -W ORLD D ATA
The goal of these experiments is to test the proposed methodology on larger-scale datasets which mimic real-world data: CIFAR-10, SVHN, and FMNIST.
We evaluate our methods on performance accuracy and on mixing times of the chains.
We employ ResNet-20 for SVHN and FMNIST without any data augmentation to assess our methods.
For CIFAR10 we employ the same data augmentation process as proposed in Cubuk et al.
We evaluate the precision of the methods on accuracy over time and the overall mixing time of them on IAC and ESS with 2 base algorithms: pSGLD and SGHMC.
For efﬁciency purposes we limited our scope to models with either fully joint posteriors or fully factorized.
As such, for the latter we employed S d -SGMCMC methods
[Figure]
[Figure]
their proposed variational variants for model averaged accuracy.
Grid search was used to determine optimal hyperparameters for each method — more details can be found in the Appendix.
as S-SGMCMC would not be feasible with the amount of parameter groups present.
Bernoulli ( ρ ) and uniform masking distributions were investigated and are denoted as S Bernoulli -SGMCMC and
S Uniform -SGMCMC respectively, with ρ varying between datasets as determined by a hyperparameter search (detailed in the Appendix).
In Fig.
3 we observe how quickly the proposed methods and the baseline SGMCMC methods approach their optimum accuracy over the course of training.
As is shown, S Bernoulli -SGMCMC and
S Uniform -SGMCMC appear to achieve optimal accuracy values much faster than SGMCMC on all datasets and with all base sampling schemes.
In some cases, the variational methods achieve better accuracy values than the baseline methods, as seen for CIFAR10 in Fig. 3.
Mixing Time Comparisons We further validated our ﬁndings from Section 6.1 by evaluating the
IAC and ESS on larger datasets using various methods.
Both pSGLD and SGHMC were used as base methods in conjunction with both S-SGMCMC and S d -SGMCMC using a Bernoulli masking distribution.
IAC and ESS were calculated for these methods using the latest 5,000 samples after sampling for 300 epochs; the results of which can be found in Table 1.
For CIFAR-10, we see that
S d -SGMCMC with every parameter in a different group mixes the fastest against all other methods.
Likewise, for SVHN and FMNIST, S d -pSGLD with every parameter belonging to its own group mixes faster than all other methods.
At times it does appear that increasing the number of parameter groups causes slower mixing time for S-SGMCMC.
This could potentially be attributed to large variance in the gradients from using only a single sample per Monte Carlo estimate.
[Table]
6. E XPLORING P ARTITIONING S CHEMES
This part of the study aims to explore the capabilities of the proposed methodology further.
Here we explore different parameter partitioning schemes on regression datasets
[Table]
Here we present the results with different partitions on various regression datasets.
We used 7 different datasets: the wine quality dataset (Cortez et al., 2009), the Boston housing dataset (Harrison Jr &
Rubinfeld, 1978), the obesity levels dataset (Palechor & de la Hoz Manotas, 2019), the Seoul bikesharing dataset (E et al., 2020; E & Cho, 2020), the concrete compressive strength dataset (Yeh, 1998)
[Table]
and the airfoil self-noise dataset (Brooks et al., 1989).
For the evaluation we chose a simple fully connected network with two layers with 50 neurons each, and we use SGLD as an optimizer.
As a performance metric we chose mean squared error (MSE).
We did hyperparameter tuning with different learning rates and the ﬁnal results are the means with the standard deviations of 5 runs with different seeds.
We do not observe any speciﬁc systematic trends on the partitions, apart from the fact that in some cases random performs better.
In that way the use of either random partitioning or the fully-factorized partitioning, where every parameter is in a different group appears to be a valid choice a priori ; especially the latter since we have noted earlier the faster mixing times associated with this partitioning scheme.
More details about the partitioning schemes experiments can be found in the Appendix.
C ONCLUSIONS
In an attempt to hybridize MCMC and VI, we proposed S-SGMCMC: an approach that produces samples from an structured posterior by running SGMCMC on a modiﬁed energy function.
The resulting Markov chain becomes asymptotically decoupled across user-speciﬁed groups of parameters, resulting in faster convergence.
For better computational efﬁciency, we proposed S d -SGMCMC: a further generalization of S-SGMCMC inspired by dropout.
This extension allows interpolating between a SGMCMC algorithm and its corresponding S-SGMCMC method.
Our experimental results demonstrate that the proposed methods impose structure over posterior distributions, increase mixing times of the chains, and result in similar or better posterior predictive accuracies compared to SGMCMC on a variety of (deep) models.
Our experimental evaluations have provided strong empirical evidence for the efﬁcacy of our approach.
We also showed that the proposed approach is compatible with various deep learning architectures, including ResNet-20, and various datasets.
Despite its proven capabilities, our proposed methodology does come with some limitations.
Namely, for quick access our methods require keeping chains of samples on the GPU whereas the baseline
SGMCMC methods can simply save samples to disk.
Additionally, S-SGMCMC scales poorly with respect to the number of parameter groups.
S d -SGMCMC manages to break this dependency; however, it still requires slightly more compute than SGMCMC per sample, but it is comparable in wall clock time.
Possible future work could focus on more theoretical analyses of S-SGMCMC, such as formal proofs of convergence.
E THICS S TATEMENT
The main focus of our work is to train models faster by decreasing the convergence time of their training phase.
In this scope we are not aware of any ethical concerns of our research.
R EPRODUCIBILITY
For this work we have made sure to guarantee reproducibility of the results.
We provide all the technical details of our experiments and their implementations, like the hyperparameters, the data, the frameworks and the experimental setups.
We have used only open source datasets that are easily accessible to the public.
Finally we commit to release the code that we implemented for this work via a public repository.
A T HEOREM
Proof.
We begin with some preliminaries from the main text.
Given data D
= { ( x i , y i ) } i =1 ,...,N , parameters θ , a proper prior distribution p ( θ ) , and a likelihood p ( D| θ ) =
(cid:81)
Ni =1 p ( y i | x i , θ ) , 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 ( θ ) = − (cid:88) ( x,y ) ∈D log p ( y | x, θ ) − log p ( θ ) .
U is typically referred to as the posterior energy function .
Note that the posterior distribution is typically intractable due to the normalizing constant.
We also write the equation for KL divergence from the main text:
[Equation]
[Equation]
We then rewrite Eq.
4 as follows:
[Equation]
[Equation]
q
[Equation]
14) for some i ∈ { 1 , . . .
, M } where ¬ i := { 1 , . . .
, M } \ { i } and C = log p ( D ) .
In order to ﬁnd the optimal distribution that respects the factorization constraints imposed between parameter groups, we need to minimize this functional over q — or rather every q i .
This is done by taking the functional derivative of J with respect to q i , setting it equal to zero, and solving for q i :
[Equation]
[Equation]
By deﬁning the energy U ( S ) i
[Equation]
, we realize that by minimizing the KL-divergence in Eq.
4, the approximate posterior distribution q = (cid:81)
Mi
=1 q i takes the form of a
Boltzmann distribution as in Eq.
1 with U ( S ) ( θ ) =
(cid:80)
Mi
=1 U ( S ) i ( θ i ) .
It remains to be shown that the solution is unique.
To this end, we refer to the convexity of the KL divergence in function space (Cover & Thomas, 2001).
This implies that the stationary point of the
KL is indeed a global optimum and unique.
B D ERIVING U ( S d )
With just a slight shift in perspective, it is actually possible to further generalize U ( S ) (and consequently S-SGMCMC) to produce a broader class of approximate sampling algorithms.
This is done by ﬁrst noting that U ( S ) can be represented with a scaled double-expectation:
[Equation]
where p ( S ) ( r ) =
Cat ( r ; M − 1 , . .
, M − 1 ) and ( rθ + (1 − r )
˜ θ )
i is equal to θ i if r i
= 1 and ˜ θ
i otherwise for i = 1 , . . .
, M .
Note that this is constructed in this manner speciﬁcally so that U ( S ) remains differentiable with respect to θ .
Also note that though the denominator appears superﬂuous as E r ∼ p ( S ) [ (cid:80) Mi
=1 r i ] = 1 , it is necessary for certain theoretic properties, as seen in Theorem 2.
By replacing p ( S ) with a more ﬂexible distribution, we can further generalize and encapsulate different energy functions to sample from.
One such choice is p ( S d ) ( r ; ρ ) : ∝ (cid:81)
Mi =1 Bern ( r i ; ρ ) 1 ( (cid:80) Mi
=1 r i > 0) with ρ ∈ (0 , 1) .
3 Substituting p ( S ) for p ( S d ) in Eq.
(18) yields a new energy function that we will refer to as U ( S d ) .
We note that this choice in distribution leads to a dropout-like behavior (Nalisnick et al., 2019; Srivastava et al., 2014), where the composition of model parameters as rθ + (1 − r )˜ θ leads to each parameter group θ i having a probability of approximately ρ to be used in a prediction and a (1 − ρ ) probability of being replaced by ˜ θ i from the approximate posterior (in traditional dropout, θ i would instead be replaced with 0).
Likewise, we will denote methods that use this energy function for sampling as structured dropout SGMCMC (S d -SGMCMC) with different variants all sharing the same S d preﬁx (e.g. S d -SGHMC).
In practice, the double-expectation in U ( S d ) is jointly approximated using a Monte Carlo estimate with K samples.
This leads to Eq.
(8) in the main paper.
We note that by approximating U ( S d ) in this way, computing a gradient no longer scales on the order of O ( M ) , but rather O ( K ) .
This means that the choice of structure imposed on the posterior distribution remains independent of computing resources.
As such, conﬁgurations with large amounts of parameter groups are typically only feasible when using S d -SGMCMC as S-SGMCMC would use too much memory and/or compute per sample.
C T HEOREM
Theorem 2.
For a given set of parameters θ partitioned into M groups, under minor assumptions (i) U ( S d ) → U as ρ → 1 and (ii) U ( S d ) → U ( S ) as ρ → 0 .
Thus, distributions approximated by
S d -SGMCMC lie on a continuum with those generated by S-SGMCMC at one extreme and with those from SGMCMC at the other.
Proof.
Assume an arbitrary θ , D , n ∈ N , and that E ˜ θ ∼ q (cid:104) log p ( rθ + (1 − r )
˜ θ, D ) (cid:105) exists for r ∈ R .
As an aside, this proof assumes that p ( S d ) ( r ; ρ ) : ∝ (cid:81) Mi =1 Bern ( r i ; ρ ) 1 ( (cid:80) Mi =1 r i > 0) with ρ ∈ (0 , 1) ; however, the theorem still holds an arbitrary p ( S d ) so long as the mean approaches 1 and variance approaches 0 as n → ∞ .
(i) Let r ( n ) ∼ p ( S d ) ( ρ n ) where ∀ n ρ n ∈ (0 , 1) and ρ n → 1 .
It follows that r ( n ) → { 1 } M as n → ∞ in distribution (see Lemma 1 in Supplement).
Due to bounded and ﬁnite support R , we ﬁnd the following:
[Equation]
(ii) Let r ( n ) ∼ p ( S d ) ( ρ n ) where ∀ n ρ n ∈ (0 , 1) and ρ n → 0 .
It follows that r ( n ) → r ∼
Cat ( M − 1 , . . .
, M − 1 ) as n → ∞ in distribution (see Lemma 2 in
Due to bounded and ﬁnite support R , we ﬁnd the following:
[Equation]
For both Lemmas 1 and 2, let
[Equation]
Lemma 1.
For r ( n ) ∼ p ( S d ) ( ρ n ) , ρ n ∈ (0 , 1) and n ∈ N , if ρ n → 1 as n → ∞ then r ( n ) → r ∼ δ ( { 1 } M ) in distribution as n → ∞ .
Proof.
[Equation]
[Equation]
[Equation]
Lemma 2.
For r ( n ) ∼ p ( S d ) ( ρ n ) , ρ n ∈ (0 , 1) and n ∈ N , if ρ n → 0 as n → ∞ then r ( n ) → r ∼
Cat ( M − 1 , . . .
, M − 1 ) in distribution as n → ∞ .
Proof.
Let i ∈ { 1 , . . .
, M } .
[Equation]
Rule
[Equation]
Since the resulting probabilities sum to 1, this implies that r ( n ) → r ∼ Cat ( M − 1 , . . .
, M − 1 ) in distribution as n → ∞ .
D D ERIVING U ( S d )
To derive U ( S d ) , we must ﬁrst start with a shift in perspective on how U ( S ) is represented.
We will rewrite the function in the following way:
[Equation]
where p ( S ) is a M -dimensional categorical distribution with uniform weights M − 1 and p ( rθ + (1 − r )
˜ θ, D ) is the joint probability of parameters taking values of rθ
+ (1 − r )
˜ θ and data D .
We note that changing the distribution of r leads to different energy functions to sample from.
One such choice is to have p ( S d ) ( r ; ρ )
∝ ρ
[Equation]
Mi =1 r i (1 − ρ ) M − (cid:80) Mi
=1 r i 1 ( ∀ i r i ∈ { 0 , 1 } ) 1 ( (cid:80) Mi
=1 r i > 0) for ρ ∈ (0 , 1) .
Note that this is identical to r i iid ∼ Bernoulli ( ρ ) conditional to (cid:80)
Mi =1 r i > 0 .
Let the support of p ( S d ) be denoted as R = { 0 , 1 } M \ { 0 } M .
This leads to the following energy function:
[Equation]
In practice, a few approximations are made to compute the corresponding U ( S d ) .
Firstly, we approximate p ( S d ) with an M -dimensional Bernoulli ( ρ ) distribution as the difference is minute when
Mρ is large.
Secondly, the outer expectation in Eq.
(34) is approximated with a Monte Carlo estimate of K samples.
The inner expectation is also approximated with a Monte Carlo estimate using the latest approximate posterior ˆ q ( t ) .
However, just like for S-SGMCMC, only a single sample is used.
This further leads to:
[Equation]
E A LGORITHM FOR S d -SGMCMC
The procedure for S d -SGMCMC can be seen in Algorithm 3.
Input
Output: ˆ q ( T ) ( θ ) :
= { θ ( t ) } t =1 ,...,T 1 for t = 0 to T − 1 do
Sample minibatch ˜ D ( t ) ⊂ D for k = 1 to K do
Sample masks r ( t,k ), . . .
, r ( t,k ) M ∼ p ( S d )
Sample ˜ θ ( t,k ) ∼ ˆ q ( t ) θ ( t,k ) =
[ r ( t,k ) i θ ( t ) i + (1 − r ( t,k ) i )
˜
θ ( t,k ) i ] i
=1 ,...,M ˆ U ( S d ,t ) k
= ˆ U ( θ ( t,k ) ; ˜ D ( t ) ) end ∇ θ ˆ U ( S d ,t ) =
M K E r ∼ p
( Sd ) [ (cid:80) Mi
=1 r i ] (cid:80)
Kk =1 ∇ θ ˆ U ( S d ,t ) k
θ ( t +1) , ξ ( t +1)
= SGMCMC_step ( θ ( t ) , ∇ θ ˆ U ( S d ,t ) , ξ ( t ) , (cid:15) t ) 11 end 12 return ˆ q ( T ) ( θ )
The update rules for SGLD, pSGLD, and SGHMC are deﬁned as follows: [Equation] where (cid:15) t is the step size at time step t , R ( · ) and M are preconditioners, γ ≥ 0 is a friction term, and ˆ V ( · ) is an estimate of the covariance induced by the stochastic gradient.
The update rules for the S-SGMCMC variants are similarly deﬁned as Eqs. 36-39 but all instances of ˆ U ( θ ( t ) ) are replaced with ˆ U ( S ) ( θ ( t ) ) . Likewise, replacing with ˆ U ( S d ) ( θ ( t ) ) yields the S d -SGMCMC variants.
G A BLATION S TUDY
This subsection aims to further explore the capabilities of the proposed methodology.
More speciﬁcally, we visualize uncertainty for a two-layer Fully Connected Network and experiment with various parameter partitions.
Parameter Partitions.
We tested our proposal with four partitioning schemes on a 2 layer withneurons fully connected network on a regression task.
The partitioning schemes that we used are the following: (a) the parameters are split into 3 groups randomly, (b) the parameters are split by layer( layers, 1 input and 2 hidden), (c) by activating neurons inside the layers and (d) every parameter belongs in each own group.
We used 7 different datasets: the wine quality datsetCortez et al.
(2009), the Boston housing datasetHarrison Jr & Rubinfeld (1978), the obesity levels datasetPalechor & de la
Hoz Manotas (2019), the Seoul bike-sharing datasetE et al.
; E & Cho (2020), the concrete compressive strength datasetYeh (1998), and the airfoil self-noise datasetBrooks et al.
Every dataset was split into 75% training data, 10% validation data, and 15% test data.
We trained the model on training set and validated it in the validation set with an early stoppage.
For every dataset and every partitioning scheme we used the learning rates: 1e-3,1e-4,1e-5,1e-6,1e-7 for hyperparameter tuning.
For each combination of partition and dataset, we chose the learning rate that provides the best accuracy score on the test set.
In this case, as an accuracy score, we used the Mean Squared
Error.
The ﬁnal learning rates that we used are presented in Table 3.
[Table]
Random ( M = 3 ) 1e1e- 1e1e- 1e1e-
By Layer 1e1e- 1e1e- 1e1e-
By Neurons 1e1e- 1e1e- 1e1e-
Fully-Factorized 1e1e- 1e1e- 1e1e- datasets
H D ETAILS ON E XPERIMENTS H. Q UALITATIVE R EGRESSION E XPERIMENTS
First, we aim to showcase qualitative differences in the empirical posterior distributions generated by a baseline SGMCMC algorithm and our proposed variants.
To do so, we consider a regression task where 100 randomly sampled three-dimensional covariates { (cid:126)x i =
[ x i, 1 , x i, 2 , x i, 3 ] T } i =1 ,..., 100 are used to sample response values y i ∼ N ( (cid:126)w T (cid:126)x i + b, σ 2 ) where (cid:126)w =
[ w 1 , w 2 , w 3 ] T = [1 . 5 , − 0 .
8 , 1 . 3] T , b = 0 .
5 , and σ 2 = 1 .
More details on the generation process for (cid:126)x can be found in the Supplement.
We choose to ﬁt a linear regression model of the same form as the generation process.
σ 2 is assumed to be known.
Thus, θ =
[ w 1 , w 2 , w 3 , b ] .
A standard normal distribution is used as the prior for each parameter.
Due to conjugacy, the posterior distribution can be calculated analytically.
As such, the
MAP is roughly ˆ θ MAP ≈ [0 . 52 , 0 . 31 , 0 . 47 , 0 . 84] .
The approximated posterior distributions for θ are found using SGLD, S-SGLD, and S d -SGLD.
For the latter two sampling schemes, two parameter partitions are tested: (i) two groups of parameters where θ 1 = [ w 1 , w 2 ] and θ 2 =
[ w 3 , b ] and (ii) four groups of parameters where θ 1 = w 1 , θ 2 = w 2 , θ 3 = w 3 , and θ
= b .
For S d -SGLD, ρ = 0 . 5 and K = 4 was used.
The resulting posterior distributions for ( w 1 , w 2 ) and ( w 1 , w 3 ) from all ﬁve scenarios, with SGLD in the leftmost column as our baseline, can be seen in Fig.
We observe that, as expected, correlations between ( w 1 , w 2 ) still exist when they are allocated to the same parameter group and become apparently independent when assigned to different groups.
We also note that the variance of the distributions shrink as the parameter space is partitioned into smaller groups.
The underestimation of posterior variance is a commonly reported ﬁnding for VI techniques and is interesting to note that our non-parametric methods appear to exhibit this behavior as well.
Finally, it appears that the S d -SGLD adequately approximates S-SGLD with just slightly higher variances and very minor correlations between parameter groups being exhibited.
H. R EAL -W ORLD D ATA E XPERIMENTS
Framework details.
In this subsection, we provide more detailed results for our experiments and a grid search for FMNIST, CIFAR10, and SVHN.
We note that all the code apart from the metrics was written in PyTorch (Paszke et al., 2019).
Regarding the metrics, ESS was adopted from the
TensorFlow probability library (Dillon et al., 2017; Abadi et al., 2016) and IAC was calculated in python.
For all the experiments, we used a seed of 2.
Moreover, we note that we grouped the parameters in an ordered way for S d -pSGLD and S-pSGLD.
We denoted previously that Kρ is the number of groups.
So every parameter will go to the i mod Kρ group where i is the parameter index.
If, for instance, Kρ is 8 then parameter 1 will go to group 1, parameter 2 will go to group 2, parameter 9 will go to group 1, etc.
If Kρ is the same as the number of parameters, every parameter will go into its own group.
MNIST.
Regarding MNIST, we ran all the experiments for 500 epochs with a batch size ofand a learning rate of 1e-2.
For S d -pSGLD, the K is set to 300, which is the forward passes that the model does within 1 epoch.
For the grouping of the parameters, for S d -pSGLD we used group sizes of 2,4,8,32,128,512,2048,4096,8192,16384,32768 and 42200; and for S-pSGLD we used groups sizes of 2,8,32,128,512,2048,4096 and 8192.
FashionMNIST.
We ran all experiments for 300 epochs with a batch size of 500.
For S d -SGHMC the K is set to 2, which is the forward passes that the model does within 1 epoch.
We observed with experimenting with K that we do not need to set K very high, and even a small number like 2 that we used here is enough to produce the same results as with an K of 200 or 300.
In this way, we save signiﬁcant time in training.
Regarding the parameter partitioning, for S d -SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16.
For
S d -pSGLD, pSGLD, S d -SGHMC and SGHMC we tested their performances with learning rates of 1e-2,1e-3,1e-4,1e-5.
For S-pSGLD we used a learning rate of 1e-3 and for S-SGHMC a learning rate of 1e-2.
[Table]
[Table]
CIFAR10.
The setup is similar to the one we used in FashionMNIST as we ran all experiments for 300 epochs with a batch size of 128.
For S d -SGHMC, the K is set to 2, which K is the forward passes that the model does within 1 epoch.
Regarding the parameter partitioning, for S d -SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16.
For S d -pSGLD, pSGLD, S d -SGHMC and SGHMC we tested their performances with learning rates of 1e-2,1e-3,1e-4,1e-5.
For S-pSGLD, we used a learning rate of 1e-3, and for S-SGHMC, a learning rate of 1e-2.
We focused our strategy on evaluating the accuracy of the different combinations of hyperparameters with the proposed methods, as can be seen in Figs.
4 and 5.
Quantitative results on
IAC, ESS and maximum accuracy are depicted in Tables 6 and 7.
[Table] [Table]
1e6.
pSGLD1e-  7. pSGLD1e-  6. pSGLD1e-  6. pSGLD1e-  6.
S-pSGLD1e-  7.
S-pSGLD1e-  7.
S-pSGLD1e-  4.
S-pSGLD1e-  5.
SVHN.
We also ran all of the experiments for 300 epochs with a batch size of 128.
Here for
S d -SGHMC, the K is set to 2, which is the forward passes that the model does within 1 epoch.
We note that K here is less than on CIFAR10 and FashionMNIST, but as we mentioned before, this does not make a difference for our results, as we have tested.
Regarding the parameter partitioning, for
S d -SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16.
For S d -pSGLD, pSGLD, S d -SGHMC and SGHMC we tested their performances with
[Figure]
M o d e l A v e r a g
e d
A cc u r a c y
[Figure]
[Figure]
M o d e l A v e r a g
e d
A cc u r a c y
[Figure]
learning rates of 1e-1,1e-2,1e-3,1e-4,1e-5,1e-6.
For S-pSGLD we used a learning rate of 1e-4, and for S-SGHMC, a learning rate of 1e-2.
Same as in CIFAR10, we conducted a grid search for learning rate, dropout rate, and optimizers to ﬁnd the best performing models and test them for their accuracy.
We can observe these results in Figure 3 in the main paper.
The strategy that we followed is the same as in CIFAR10 and is presented in Figs.
6 and 7.
[Figure]
M o d e l A v e r a g
e d
A cc u r a c y
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
M o d e l A v e r a g
e d
A cc u r a c y
[Figure]