Year: 2021

Abstract:
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 modified 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 modified 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES
text: 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).

heading: 4 STRUCTURED SGMCMC
text: 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.

heading: 5 STRUCTURED DROPOUT SGMCMC
text: 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 (Sd-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 Sd-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 rj = 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:
Û (Sd)(θ(t); D̃) ≈ M KE [∑M
i=1 ri ] K∑ k=1 Û(r̃(t,k)θ(t) + (1− r̃(t,k))θ̃(t,k); D̃) (8)
where r̃(t,k) is the kth sample of r̃ for timestep t. A formal justification for Eq. (8) can be found in the Supplement. These energy function approximations lead to the following update step for structured
1K is a hyperparameter that is chosen independent of M ; however, both M and the distribution of r̃ largely influence how small K can be due to how they affect the variance of the gradient of the associated posterior energy function.
Algorithm 2: S-SGMCMC Input: Initial sample θ(0); parameter
partitions θ1, . . . , θM ; step sizes { t}t=0,...,T−1.
Output: q̂(T )(θ) := {θ(t)}t=1,...,T 1 for t = 0 to T − 1 do 2 Sample minibatch D̃(t) ⊂ D 3 for i = 1 to M do 4 Sample θ̃(t)¬i ∼ q̂ (t) ¬i 5 Û (S,t) i = Û([θ (t) i , θ̃ (t) ¬i ]; D̃(t)) 6 end 7 ∇θÛ (S,t) = ∑M i=1∇θÛ (S,t) i 8 θ(t+1) =
SGMCMC_step(θ(t),∇θÛ (S,t), t) 9 end
10 return q̂(T )(θ)
Table 1: IAC and ESS metrics for CIFAR-10, SVHN, and FMNIST with various methods. Subscripts after method names refers to number of equally sized parameter groups, with |θ| meaning every parameter belongs to its own group. Best results are bolded.
CIFAR-10 SVHN FMNIST
Method IAC↓ ESS↑ IAC↓ ESS↑ IAC↓ ESS↑
pSGLD 716 8.01 839 6.82 779 7.09 S-pSGLD2 600 7.44 840 6.80 740 7.55 S-pSGLD4 599 7.4 834 6.83 751 7.45 S-pSGLD8 709 6.41 857 6.67 776 7.24 Sd-pSGLD|θ| 546 8.01 803 7.00 677 8.24 SGHMC 727 7.94 858 6.59 795 6.83 S-SGHMC2 583 7.49 949 5.74 928 5.67 S-SGHMC4 624 7.03 961 5.66 915 5.77 S-SGHMC8 904 4.97 1056 5.30 1142 4.87 Sd-SGHMC|θ| 584 7.7 828 6.56 782 7.08
dropout variant of SGLD (Sd-SGLD):
θ(t+1) = θ(t) − t 2 ∇θÛ (Sd)(θ; D̃) + ξt where ξt ∼ N (0, tI). (9)
The corresponding update rules for the structured dropout variants for pSGLD (Sd-pSGLD) and SGHMC (Sd-SGHMC) are defined in the Supplement. The exact procedure for generating samples of the approximate posterior q̂(t) using structured dropout SGMCMC (Sd-SGMCMC) can also be found in the Supplement.
An example of this method (specifically Sd-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.

heading: 6 EXPERIMENTS
text: 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 classification 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 inefficiency of a MCMC algorithm. IAC is computed as τf = ∑∞ τ=−∞ ρ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−τ ∑N−τ n=1 (fn − µf ) (fn+τ − µf ) and µf = 1 N ∑N n=1 fn. We note that we calculated ĉf (τ) with a fast Fourier transform as it is more computationally efficient than using the direct sum. ESS measures how many independent samples would be equivalent to a chain of correlated samples and is calculated as neff = n1+(n−1)p , where n is the number of samples and p is the autocorrelation. 2 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.

heading: 6.1 DROPOUT RATE & GROUP SIZE INVESTIGATION
text: 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 Sd-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 Sd-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 Sd-pSGLD’s accuracy improves slightly and then remains fairly stable. In the best case, Sd-pSGLD achieves an accuracy of 96.3% with 32 groups and dropout rate of 0.5 which outperforms pSGLD with accuracy of 94.2%. We speculate that the dropout-like behavior is beneficial 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 Sd-pSGLD to lesser and lesser degrees as ρ increases. This behavior is to be expected due to Theorem 2, with Sd-pSGLD exhibiting mixing times more similar to pSGLD when ρ = 0.5 and more similar to S-pSGLD when ρ = 0.1.

heading: 6.2 SYSTEMATIC COMPARISON ON REAL-WORLD DATA
text: 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. (2019). 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 efficiency purposes we limited our scope to models with either fully joint posteriors or fully factorized. As such, for the latter we employed Sd-SGMCMC methods
2We used the TensorFlow implementation for ESS which uses the direct sum for the autocorrelation.
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 SBernoulli-SGMCMC and SUniform-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, SBernoulli-SGMCMC and SUniform-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 findings 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 Sd-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 Sd-SGMCMC with every parameter in a different group mixes the fastest against all other methods. Likewise, for SVHN and FMNIST, Sd-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.

heading: 6.3 EXPLORING PARTITIONING SCHEMES
text: 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.
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), 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 final results are the means with the standard deviations of 5 runs with different seeds. We do not observe any specific 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.

heading: 7 CONCLUSIONS
text: 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 modified energy function. The resulting Markov chain becomes asymptotically decoupled across user-specified groups of parameters, resulting in faster convergence. For better computational efficiency, we proposed Sd-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 efficacy 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. Sd-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.

heading: 8 ETHICS STATEMENT
text: 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.

heading: 9 REPRODUCIBILITY
text: 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.

heading: A THEOREM 1
text: Proof. We begin with some preliminaries from the main text. 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(θ) = − ∑
(x,y)∈D
log p(y|x, θ)− log p(θ). (10)
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:
J(q(θ)) = DKL(q(θ)||p(θ|D)) (11) ≡ Eθ∼q [ log q(θ)
p(θ|D)
] (12)
We then rewrite Eq. 4 as follows:
J(q(θ)) = Eθ∼q [log q(θ)]− Eθ∼q [log p(θ,D)] + C (13) =Eθi∼qi [log qi(θi)] + ∑ i 6=j Eθj∼qj [log qj(θj)]− ∫ log p(θ,D)qi(θi)dθi ∏ i6=j qj(θj)dθj + C (14)
for some i ∈ {1, . . . ,M} where ¬i := {1, . . . ,M} \ {i} and C = log p(D). In order to find the optimal distribution that respects the factorization constraints imposed between parameter groups, we need to minimize this functional over q — or rather every qi. This is done by taking the functional derivative of J with respect to qi, setting it equal to zero, and solving for qi:
δJ(q(θ))
δqi(θi) =
∫ log p(θ,D) ∏ i6=j qj(θj)dθj − 1− log qi(θi) := 0 (15)
=⇒ log qi(θi) = Eθ̃¬i∼q¬i [ log p(θi, θ̃¬i,D) ] − 1 (16)
=⇒ qi(θi) ∝ exp { Eθ̃¬i∼q¬i [ log p(θi, θ̃¬i,D) ]} . (17)
By defining the energy U (S)i (θi) = −Eθ̃¬i∼q¬i [ log p(θi, θ̃¬i,D) ] , we realize that by minimizing
the KL-divergence in Eq. 4, the approximate posterior distribution q = ∏M i=1 qi takes the form of a
Boltzmann distribution as in Eq. 1 with U (S)(θ) = ∑M i=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 DERIVING U (Sd)
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 first noting that U (S) can be represented with a scaled double-expectation:
U (S)(θ) = − M Er∼p(S) [∑M i=1 ri ]Er∼p(S)Eθ̃∼q [log p(rθ + (1− r)θ̃,D)] (18) where p(S)(r) = Cat(r;M−1, . . . ,M−1) and (rθ + (1 − r)θ̃)i is equal to θi if ri = 1 and θ̃i otherwise for i = 1, . . . ,M . Note that this is constructed in this manner specifically so that U (S) remains differentiable with respect to θ. Also note that though the denominator appears superfluous as Er∼p(S) [ ∑M i=1 ri] = 1, it is necessary for certain theoretic properties, as seen in Theorem 2.
By replacing p(S) with a more flexible distribution, we can further generalize and encapsulate different energy functions to sample from. One such choice is p(Sd)(r; ρ) :∝ ∏M i=1 Bern(ri; ρ)1( ∑M i=1 ri > 0) with ρ ∈ (0, 1).3 Substituting p(S) for p(Sd) in Eq. (18) yields a new energy function that we will refer to as U (Sd). 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 (Sd-SGMCMC) with different variants all sharing the same Sd prefix (e.g. Sd-SGHMC).
In practice, the double-expectation in U (Sd) 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 (Sd) 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, configurations with large amounts of parameter groups are typically only feasible when using Sd-SGMCMC as S-SGMCMC would use too much memory and/or compute per sample.

heading: C THEOREM 2
text: Theorem 2. For a given set of parameters θ partitioned into M groups, under minor assumptions (i) U (Sd) → U as ρ → 1 and (ii) U (Sd) → U (S) as ρ → 0. Thus, distributions approximated by Sd-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 [ log p(rθ + (1− r)θ̃,D) ] exists for r ∈ R.
As an aside, this proof assumes that p(Sd)(r; ρ) :∝ ∏M i=1 Bern(ri; ρ)1( ∑M i=1 ri > 0) with ρ ∈ (0, 1); however, the theorem still holds an arbitrary p(Sd) so long as the mean approaches 1 and variance approaches 0 as n→∞.
(i) Let r(n) ∼ p(Sd)(ρ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 finite support R, we find the following:
U (Sd)(θ) = − M Er∼p(Sd) [ ∑M i=1 ri] ∑ r∈R p(Sd)(r; ρn)Eθ̃∼q [ log p(rθ + (1− r)θ̃,D) ] (19)
→ −M M ∑ r∈R 1(∀iri = 1)Eθ̃∼q [log p(θ,D)] as n→∞ (20)
= − log p(θ,D) = U(θ) (21)
(ii) Let r(n) ∼ p(Sd)(ρ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 Supplement). Due to bounded and
3Other choices of distribution that are well justified include any with support over [0, 1]M and with measure 0 over {0}M . Exploring the effects these distributions have are an interesting line of future inquiry.
finite supportR, we find the following:
U (Sd)(θ) = − M Er∼p(Sd) [ ∑M i=1 ri] ∑ r∈R p(Sd)(r; ρn)Eθ̃∼q [ log p(rθ + (1− r)θ̃,D) ] (22)
→ −M 1 ∑ r∈R 1( ∑M i=1 ri = 1) M Eθ̃∼q [ log p(rθ + (1− r)θ̃,D) ] as n→∞ (23)
= − M∑ i=1 Eθ̃∼q[log p([θi, θ̃¬i,D)] = U (S)(θ) (24)
For both Lemmas 1 and 2, let
p(Sd)(r; ρ) = ρ ∑M i=1 ri(1− ρ)M− ∑M i=1 ri
1− (1− ρ)M 1(∀iri ∈ {0, 1})1 ( M∑ i=1 ri > 0 ) (25)
Lemma 1. For r(n) ∼ p(Sd)(ρn), ρn ∈ (0, 1) and n ∈ N, if ρn → 1 as n → ∞ then r(n) → r ∼ δ({1}M ) in distribution as n→∞.
Proof.
p(Sd)(r = {1}M ; ρn) = ρMn (1− ρn)0
1− (1− ρn)M (26)
→ 1 as n→∞ (27) =⇒ r(n) → δ({1}M ) in distribution. (28)
Lemma 2. For r(n) ∼ p(Sd)(ρ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}.
p(Sd)(ri = 1, r¬i = 0; ρn) = ρn(1− ρn)M−1
1− (1− ρn)M (29)
l’Hôspital’s Rule H= (1− ρn)M−1 + ρn(M − 1)(1− ρn)M−2
M(1− ρn)M−1 (30)
→ 1 M as n→∞ (31)
Since the resulting probabilities sum to 1, this implies that r(n) → r ∼ Cat(M−1, . . . ,M−1) in distribution as n→∞.
D DERIVING U (Sd)
To derive U (Sd), we must first start with a shift in perspective on how U (S) is represented. We will rewrite the function in the following way:
U (S)(θ) = − M∑ i=1 Eθ¬i∼q¬i [log p([θi, θ¬i],D)] (32)
= − M Er∼p(S) [ ∑M i=1 ri]
Er∼p(S)Eθ̃∼q [ log p(rθ + (1− r)θ̃,D) ] (33)
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.4
We note that changing the distribution of r leads to different energy functions to sample from. One such choice is to have p(Sd)(r; ρ) ∝ ρ ∑M i=1 ri(1 − ρ)M− ∑M i=1 ri1(∀iri ∈ {0, 1})1( ∑M i=1 ri > 0) for ρ ∈ (0, 1). Note that this is identical to ri iid∼ Bernoulli(ρ) conditional to ∑M i=1 ri > 0. Let the support of p(Sd) be denoted asR = {0, 1}M \ {0}M . This leads to the following energy function:
U (Sd)(θ) = − M Er∼p(Sd) [ ∑M i=1 ri]
Er∼p(Sd)Eθ̃∼q [ log p(rθ + (1− r)θ̃,D) ] . (34)
In practice, a few approximations are made to compute the corresponding U (Sd). Firstly, we approximate p(Sd) 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:
U (Sd)(θ(t); D̃) = − 1 Kρ K∑ k=1 U(r(t,k)θ(t) + (1− r(t,k))θ̃(t,k); D̃) (35)
E ALGORITHM FOR Sd-SGMCMC
The procedure for Sd-SGMCMC can be seen in Algorithm 3.
Algorithm 3: Sd-SGMCMC Input: Initial sample θ(0); parameter partitions θ1, . . . , θM ; data set D; initial auxiliary statistics ξ(0); step sizes { t}t=1,...,T ; masking distribution p(Sd); dropout iterations K. Output: q̂(T )(θ) := {θ(t)}t=1,...,T
1 for t = 0 to T − 1 do 2 Sample minibatch D̃(t) ⊂ D 3 for k = 1 to K do 4 Sample masks r(t,k)1 , . . . , r (t,k) M ∼ p(Sd) 5 Sample θ̃(t,k) ∼ q̂(t)
6 θ(t,k) = [r (t,k) i θ (t) i + (1− r (t,k) i )θ̃ (t,k) i ]i=1,...,M 7 Û (Sd,t) k = Û(θ
(t,k); D̃(t)) 8 end 9 ∇θÛ (Sd,t) = MKE r∼p(Sd) [ ∑M i=1 ri] ∑K k=1∇θÛ (Sd,t) k
10 θ(t+1), ξ(t+1) = SGMCMC_step(θ(t),∇θÛ (Sd,t), ξ(t), t) 11 end 12 return q̂(T )(θ)
4rθ + (1 − r)θ̃ is a slight abuse of notation that is meant to represent masking out θi when ri = 0 and masking out θ̃i when ri = 1.

heading: F SGMCMC UPDATE RULES
text: The update rules for SGLD, pSGLD, and SGHMC are defined as follows:
SGLD θ(t+1) = θ(t) − t 2 ∇θÛ(θ(t)) +N (0, tI) (36)
pSGLD θ(t+1) = θ(t) − t 2
[ R(θ(t))∇θÛ(θ(t)) +
∑ θ ∇θR(θ(t))
] +N (0, tR(θ(t))) (37)
SGHMC θ(t+1) = θ(t) + tM−1m(t+1) (38)
m(t+1) = (1− γ tM−1)m(t) − t∇θÛ(θ(t)) +N (0, 2γ − tV̂ (θ(t))) (39)
where 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.5
The update rules for the S-SGMCMC variants are similarly defined as Eqs. 36-39 but all instances of Û(θ(t)) are replaced with Û (S)(θ(t)). Likewise, replacing with Û (Sd)(θ(t)) yields the Sd-SGMCMC variants.

heading: G ABLATION STUDY
text: This subsection aims to further explore the capabilities of the proposed methodology. More specifically, 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 with 50 neurons 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(3 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. (2020); E & Cho (2020), the concrete compressive strength datasetYeh (1998), and the airfoil self-noise datasetBrooks et al. (1989). 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 final learning rates that we used are presented in Table 3.
5Note that we abuse notation in Eqs. 36-39 where the addition ofN (µ,Σ) denotes the addition of a normally distributed random variable with mean µ and covariance Σ.

heading: H DETAILS ON EXPERIMENTS
text: H.1 QUALITATIVE REGRESSION EXPERIMENTS
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 {~xi = [xi,1, xi,2, xi,3]T }i=1,...,100 are used to sample response values yi ∼ N (~wT~xi+b, σ2) where ~w = [w1, w2, w3]T = [1.5,−0.8, 1.3]T , b = 0.5, and σ2 = 1. More details on the generation process for ~x can be found in the Supplement.
We choose to fit a linear regression model of the same form as the generation process. σ2 is assumed to be known. Thus, θ = [w1, w2, w3, 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 Sd-SGLD. For the latter two sampling schemes, two parameter partitions are tested: (i) two groups of parameters where θ1 = [w1, w2] and θ2 = [w3, b] and (ii) four groups of parameters where θ1 = w1, θ2 = w2, θ3 = w3, and θ4 = b. For Sd-SGLD, ρ = 0.5 and K = 4 was used.
The resulting posterior distributions for (w1, w2) and (w1, w3) from all five scenarios, with SGLD in the leftmost column as our baseline, can be seen in Fig. 1. We observe that, as expected, correlations between (w1, w2) 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 finding 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 Sd-SGLD adequately approximates S-SGLD with just slightly higher variances and very minor correlations between parameter groups being exhibited.
H.2 REAL-WORLD DATA EXPERIMENTS
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 Sd-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 of 500 and a learning rate of 1e-2. For Sd-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 Sd-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 Sd-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 significant time in training. Regarding the parameter partitioning, for Sd-SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16. For Sd-pSGLD, pSGLD, Sd-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.
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 Sd-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 Sd-SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16. For Sd-pSGLD, pSGLD, Sd-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.
SVHN. We also ran all of the experiments for 300 epochs with a batch size of 128. Here for Sd-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 Sd-SGMCMC, we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16. For Sd-pSGLD, pSGLD, Sd-SGHMC and SGHMC we tested their performances with
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 find 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.

Review summaries:
The author proposed a framework to incorporate the independence structure into the posterior inference for faster-mixing speed. To achieve that, the author designed two specific algorithms called S-SGMCMC and S_d-SGMCMC, respectively. Specifically, S-SGMCMC consisted of the following steps. First, the target random variables ($\theta$) are gathered into mutually independent groups. Then, a modified energy function is derived by minimizing the KL divergence between the posterior $q$ and target $p(\theta|D)$. The last step is to apply a standard SGMCMC method to draw samples from the resulting modified energy function. Further, the author also built a connection between dropout and the modified energy function, which results in a structure dropout SGMCMC (S_d-SGMCMC) with better scalability.  The author claimed the resulting algorithms achieved faster mixing speed and better classification accuracy when applied to real-world classification tasks.  

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. 

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.