Structured Stochastic Gradient MCMC
Authors: Antonios Alexos
Alex James Boyd
Stephan Mandt
Venue: ICLR 2022 Submitted (2022)
Keywords: Approximate MCMC
Langevin Dynamics
Stochastic Gradient MCMC

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.

Under review as a conference paper at ICLR 2022
STRUCTURED STOCHASTIC GRADIENT MCMC
Anonymous authors
Paper under double-blind review
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 tradeoff's in these models, variational inference (VI) is often the preferable option. Unfortunately, VI makes strong assumptions both the factorization and functional form of the posterior. In this work, we propose a new nor-maraetttic 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 type algorithm that operates 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 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 ResNet-20 architecture, we obtain better predictive likelihoods and faster mixing time than full SGMCMC.
INTRODUCTION
There has been much recent interest in deep Bayesian neural networks (BNN) due to their reliable confidence estimates and generalization properties (Wilson & Izmailov 2020; Jospin et al. 2020 Cardelli et al.] (019). 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 & Tch 2011
Patterson & Teh 2013). These algorithms owe their scalability to approximating gradients via mini-bat
The main downside of using SGMCMC algorithms is their slow mixing rates in high dimensions. An often faster alternative is variational inference (VT) algorithms that approximate the posterior with simpler (lypically factorized) distribution. This formulation results in an optimization problem that can be solved more efficiently using stochastic optimization (Blei al. 2017 Zhang al. 2018).
One downside of VI approximations is their selid distributional assumptions. A typical choice is to approximate the Bayesian posterior by product of univariate Gaussian distributions. These distributional assumptions are frequently over-simplistic in high-dimensionall models, where the posterior can be highly multi-modal and possibly heavy-tailed. Another downside is that the varia- tional 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 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 factoriza- tion constraints by assuming functional view on variational inference. We show how sample from

---

Under review as a conference paper at ICLR 2022
this optimal distribution by rumning SGMCMC on modified energy function. This energy function obtained by marginalizing the model's joint distribution over previously generated samples from the Markov chain, leading to 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 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 new approximate MCMC scheme running SGMCMC on a modified energy function, trading accuracy for speed. This setup effectively allows sampling from fully joint posterior, completely factorized posterior, and any in- 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 "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 Resnet-20 architectures CIFAR-10, 10,, Fashion MNIST, and SVHN in terms of both runtime and final accuracy.
Our paper is structured as follows: Section presents the related work to our proposal, Section introduces preliminaries regarding the energy function and the stochastic gradient updates, Sections 55d derive our proposed methods, Section details experiments and their results, and Section contains our concluding thoughts.
2 RELATED WORK
Our work connects both to (stochastic) variational inference (Bishop 2006 Hoffman al. 2013
Ranganath et al. 2014, Blei et al.] 2017: Zhang et al.) 2018) and scalable MCMC (Welling &
Teh 201 Chen et al. 2014 Ma al. 2017 /hang et al. 2020 Leimkuhler al. 2019: Wenzel et al. 2020 Izmailov et al. 2021). For space limitations, we focus on the most related work at the
intersection both topics.
Among the earliest works to hybridize both approaches was (de Freitas al. 2001 who constructed a variational proposal distribution in the Metropolos-Hasting 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 method train MCMC kernels with NNs, and Wang et al. (2018): Gong et al. (2018) who leveraged meta learning schemes 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. (016) consider SGD as a type of variational inference, but their approaches did not attempt to close the gap exact MCMC. Other works aim explicitly interpolating between both methods. Domke (017)) proposes 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 variational bound. Recently, Hoffman & Mal (2020) interpreted (parametric) VI as approximate Langevin dynamics and showed that both algorithms have similar transient dynamics.
2

---

Under review as a conference paper at ICLR 2022
In contrast to all these approaches, 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.
PRELIMINARIES
Variational inference (VI) approaches differ from MCMC in two regards: (1) they impose structured (e.g., fully-factorized) approximation of the posterior for tractability, and (2) they often make parametric assumptions. it possible to construct a modified scheme that only relies on the assumption (1), inheriting the parametric nature of MCMC while breaking posterior correlations in controlled manner? As follows, we will show how such a scheme can be realized. We will first derive modified energy function for Langevin dynamics that we can sample from and then prove that negative exponential results in the optimal posterior approximation subject to specified factorization constraints. Running SGMCMC algorithms this energy function will consequently generate samples from this distribution.
Before explain our method, we introduce the setup and common notation. Given data
D = { ys) bi)i.,,...... parameters 0, a proper prior distribution p(0), and a likelihood p(D(0) =
II p(u (r,, 0); suppose we are interested in the corresponding posterior distribution p(0) x
A convenient representation the posterior is as a Boltzmann distribution:
p() &
where U())= log p(0 D)
Σ
logp(y/z,00 p(0). (1)
(a.y)ED
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 of 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 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 al. 2014).
As alluded to, the basis of SGMCMC algorithms is using sampled minibatch of data DD from D to produce an differentiable, unbiased estimate of the posterior energy function:
N
U()) Ú0
Σ
logp(y|æ, 0) log
D
(a.y)e
Once Ú is defined, it is fairly straight forward to generate new samples from the posterior distribution. For instance, the SGLD update is
g(++1) g(t) C
(OL( + & where & ~ N(0, GI).
(3)
2
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 ł that ultimately form an empirical distribution (( (0) Should the algorithms converge, then lim (00 (0(D) = p(O[D).
4 STRUCTURED SGMCMC
By design, SGMCMC methods produce a fully joint posterior distribution over parameters For models with large number of parameters, this can lead to various complications due the curse of
3

---

Under review as a conference paper at ICLR 2022
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 9 into M > 1 distinct, mutually independent groups: 0,,... 0x- This partitioning structure is assumed to be known a priori. We will denote
the distribution that respects this partitioning structure as q(0) = I 9: (0.). Similar to VI, we would like this distribution g(0) to best approximate the true posterior distribution p(0) (D) according to some criteria, such as KL-ddivergence. This leads to natural objective function to minimize:
(()) D(0
J({(O)) = Dxr. (a(0)) = Eggg log
The following Theorem I proves that there is a unique solution to the non-parametric KL minimization problem described in Eq. (4. To describe it, we compose 0 = {0,, 0=i} J for any i where 0 ~ q and define a structured energy function:
(1)) U (O)), with (0.) := Eggg U11,, 0-13) == -Egng log p(0,, P_i, D). (5)
That is, we first define the marginals (0)), where we marginalize U(0) with respect to all ((O)-factors except (0,), and then sum up these marginals to define (0). A similar partial marginalization procedure is carried out for conjugate exponential family distributions in coordinate ascent VI [Bishop 2006). Having well-defined energy function USS) allows us to standard SGMCMC methods to approximate the posterior ((0) with samples. This serves as the basis for our proposed algorithm that actually approximates this distribution (0). which will be discussed shortly.
Theorem 1. The unique solution to the KL minimization problem given in Eq. f is given by the Boltzmann distribution q(0) x exp( Σ (1(5)(()))) (0.)). Please refer to the Supplement for the proof.
In an ideal world, we would be able to use the findings of Theorem directly in conjunction with algorithms like Langevin dynamics and Hamiltonian Monte Carlo produce empirical distributions for q using (SS) (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 (B,) 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. (
(1)) (19 (0, E
(6)
where (.) is defined in Eq. In practice. at timestep for i = 1,..., .....,.. M we estimate Ei- (10,,0=:3: D.) with a Monte Carlo approximation. In place of Ô, we use single sam- ple g(t)) taken from the current approximate distribution g(t) which is composed of samples from previous timesteps (i.e.. a uniform distribution over (0(1) ((2)) ((?),...,(()))) This leads the following update step for structured SGLD (S-SGLD):
( 0t*) (oi(()) (0; D) + & where & ~ AV(0, ez.).
Similar rules for structured variants of PSGLD (S (S-DSGLD) and SGHMC (S-SGHMC) can be found in the Supplement. Additionally, the full procedure for structured SGMCMC (S (S-SGMCMC) can be seen in Algorithm 2
Remark Since (S) is unbiased estimator for (SS), we are guaranteed to converge to from sampling with S- S-SGMMMC with sufficiently decreasing learning rates so long as we are in a stationary state. While it is unlikely to have the procedure initialize to stationary state, 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.
4

---

Under review as a conference paper at ICLR 2022
(a) SGLD (Wi, W2, wa)
(b) S-SGLD
(c) S SGLD
(wi, W2) 1 W0
T W2
1 Wg
(wh, we)
W/
W2 1 w3
W2
3
MAP
W1
W1
W1
W1
W1
Figure I: Sampled posterior distributions between & w2 (top row) and wi & w3 (bottom row) for linear regression model. From left right, (a) the first column has posterior distributions sampled with SGLD, (b) the second and third columns are sampled from S- S-SGLD, and (c) the fourth and fifth are sampled from S.LSGL The imposed parameter structure for each are shown above the plots.
An example of S-SGMCMC can be seen Fig. Ta-b). which features the approximate posterior distributions of linear regression model with three coefficients and with various independence structures imposed with S-SGLD: (a) joint dependence between W1, W22, and wy; (b-left) dependence between wi 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 commonly seen artifact when using VI.
5 STRUCTURED DROPOUT SGMCMC
While S-SGMCMC can successfully break dependencies between parameter groups, does suffer computationally due each parameter update scaling linearly with respect to M. This a single new sample of 0, the model's forward pass needs to be computed M different times on the same batch of data 2,, which can quickly become prohibitively expensive for deep models when M large. Ideally, we would prefer 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-MSG)C 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 SGMGCMC can be found the Supplement.
that for
The main motivation for this technique can be seen by recognizing that the composition (0 ((0) from Eq.
sum of masked values r((() + - -r)((()) where (it) ~ (1) and ry ( 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 i Doing so results in a slightly different energy function and minibatch loss to optimize:
6 can be rewritten
(g(t))
M
((plthh))())) + (lt-))t)) ((t. ())
(8)
KE
Ei
M
ra
k=
where ((t.A) the thth sample for timestep formal justification for Eq. 8) can be found the Supplement. These energy function approximations lead to the following update step for structured
K is hyperparameter that is chosen independent of M: however, both M and the distribution of F largely influence how small K can be due to how they affect the variance of the gradient of the associated posterior energy function.

---

Under review as a conference paper at ICLR 2022
Table IAC and ESS metrics for CIFAR-10, SVHN, and FMNIST with various methods. Subscripts after method names refers number of equally sized param- eter groups, with [0] meaning every parameter belongs to own group. Best results are bolded.
Algorithm 2: S-SGMCMC
Input:
Initial sample p(0); parameter partitions ......, a step sizes tEo....TT ---......
Output: (0)
for /. to do
CIFAR-10
SVHN
FMNIST
Sample minibatch DD:0) CD for to .M do
Method
IAC,
E
IAC
ESS
IAC
ESS
Sample ( ~ (1)
PSGLD
716
8.01
839
6.82
779
7.09
t((5.0) U(0!)) 00 D(0)
S-pSLLD2
600
7.44
840
6.80
740
7.55
S-pSGLD4
599
7.4
834
6.83
751
7.45
end
S-pSGLDg
709
6.41
857
6.67
776
7.24
Vg0(8.1) E. Volis.00 ((++1) SGMCMC_step(0%) VoU(S.)),
Sa-pSGLD (0]
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
end
S-SGHMCa
624
7.03
961
5.66
915
5.77
10 return
(0)
SGHMCC
904
4.97
1056
5.30
1142
4.87
S. S [) 584
7.7
828
6.56
782
7.08
dropout variant of SGLD (SLDSGLD):
0f++1) 0(() ct V v ( + where & ~ e(0,e
The corresponding update rules for the structured dropout variants for pSGLD (Sa-pSGLD) and SGHMC SGHMCC defined in the Supplement. The exact procedure for generating samples of the approximate posterior g(() using structured dropout SGMCMC (8,-SGMCMCC can also be found in the Supplement.
An example of this method (specifically SyLSGLD with F; sdd Bernoulli(0.5) and K 4) used linear regression model can be seen in Fig. [c). Of note, we can see that the dropout variant largely respects the independence structure imposed, but maybe not strictly the exact S-SGLD method seen in Fig. [b). Additionally, the posterior variance also seems to have shrunk similarly to S-SGLD when compared against SGLD.
Masking Distribution Should sid
Bernoulli(p). alongside structure that factorizes by activa- tion components, then the method starts to resemble dropout with rate (Srivastava et al. 20144) The main difference being that instead of replacing parameter value with 0 it is replaced with sample from the approximate posterior distribution at time g(4), While Bernoulli distribution for is
natural choice, there are other distributions that can be chosen as well. For instance, A: ~ N/ 1) or
sid Fi ~ Beta( 3) are both viable distributions and can be seen as analog to Gaussian and Beta-dropout respectively (Srivastava al. 2014 Liu al. 2019). Our experiments will largely focus sampling i from Bernoulli and uniform over (0, (equivalent to Beta(0.5, 0.5)) distributions.
6 EXPERIMENTS
Overview In this section we evaluate our proposed approach various models and datasets. Section T. investigates the impact of the variational approximation on the algorithms' mixing and MNIST (LeCun et 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 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 average.
autocorrelation times using fully-connected network architecture
We have also conducted experiments on uncertainty visualization, where we tested the proposed methodology predictive uncertainty estimation by deploying two-layer fully connected network

---

Under review as a conference paper at ICLR 2022
on 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 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 correlatio between samples in chain and, in turn, describe the inefficiency of MCMC algorithm. IAC is computed as Ty = >r== Ps (T(T) where Ps is the normalized autocorrelation function of the stochastic process that generated the chain for f and is calculated as مO( = ( (0): where CT(T) -ーー Ex= (fn (4)) (fu+r My) and μ4 tे4 Ex= fn. We note that we calculated êy(T) with fast Fourier transform 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 Reff I= where n is the number of samples and p is the autocorrelation We note that 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.1 DROPOUT RATE& & GROUP SIZE INVESTIGATION
The aim of this set of experiments 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 Sa-pSGLDD a Bernoulli(p) masking distribution with dropout rates P E {0.1, 0.5) on fully-connected neural network with 2 hidden layers, with 50 hidden units cach, trained and evaluated with MNIST using the standard train and test split. The model has 42,200 parameters in total. For S-PSGLD and Sa-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. a)) for S-pSGLD we observe that as we increase the number of groups the accuracy drops dramatically whereas Sa-pSGLD's accuracy improves slightly and then remains fairly stable. In the best case, Sa-DSGLD achieves an accuracy of 96.3% with 32 groups and dropout rate of 0.5 which outperforms PSGLD with accuracy of 94. 94.2%. We speculate that the dropout-liks 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 ((M).
Fig. 2b-cc) 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 Sa-PSGLD to lesser and lesser degrees as p increases. This behavior is to be expected due to Theorem 2 with SD-DSGLD exhibiting mixing times more similar to PSGLD when p = 0.5 and more similar to S-pSGLD when p 0.1.
6.2 SYSTEMATIC COMPARISON ON REAL-WORLD DATA
The goal of these experiments is to test the proposed methodology 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 CIFARIO 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 SGMMGCCC methods
We used the TensorFlow implementation for ESS which uses the direct sum for the autocorrelation.
7

---

Under review as a conference paper at ICLR 2022
PSGLD
S-pSGLD
SppSSLLD p=
SppPSLLD p 0.3
SypSGLD,, p 0.5
Model Avera ged Accuracy
0.95
300
1500
0.90
ESS
200
1000
0.80
300
0.75
10a
10a
10a
10a
10
10a
10a
104
109
10a
10"
10*
102
10
10*
M Groups (a)
M Groups
M Groups
(b)
(c)
Figure 2:
Comparison between PSGLD, S-PSGLD, and Sa-PSGLD of various numbers of parameter groups M for (a) accuracy, (b) IAC, and (c) ESS. Right-most points the plots represent models that have every parameter belonging to its own parameter group. S-pSGLD methods are were not able to be evaluated at such extreme values of M due to computational scaling.
(a) CIFAR-10
(b) SVHN
(c) FMNIST
Model, Averaged Accuracy
1.0
(ig SGIMC (@) PSGLD
0.9
0.9
0.8
0.8
0.8
0.7
0.7
0.6
1.0
0.9
0.9
SGMCMC
0.8
0.8
0.8
Spennoa SGMCMC
0.6
0.7
0.7
Suniom Suniom-SGMCMG SGMCMC
0
625
1250
1875
2500
625
1250
0
625
1250
Time(s)
Time(s)
Time(s)
Figure 3: Comparison for (a) CIFAR-10, (b) SVHN and (c) FMNIST using ResNet-20 with (i)
PSGLD and (ii) SGHMC sampling algorithms and their proposed variational variants for model averaged accuracy. Grid search was used 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 SBemoulli Seemoall-SGMMMC and S Mandlen-GGMMCCC respectively, with varying between datasets as determined by a hyperparameter search (detailed in the Appendix).
In Fig. we observe how quickly the proposed methods and the baseline SGMCMC methods approach their optimum accuracy over the course of training. As shown, Spemoott SGMCMC and SUniform-S -SMmmmMSGMMCC....MMMMM 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 CIFARIO 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 S,-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 [ For CIFAR-10, we see that S,-SGMCMC with every parameter in a different group mixes the fastest against all other methods.
Likewise, for SVHN and FMNIST, Sa-DSGLD 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 single sample per Monte Carlo estimate.
8

---

Under review as a conference paper at ICLR 2022
Table 2: MSE for various partitioning schemes on the parameters for a 2-layer fully-connected BNN.
Partition Scheme
Wine
Housing
Obesity
Bike
Concrete
Airfoil
0.0454
0.0233
0.0232
0.0242
0.0226
0.0454
Random (M 3)
#0.001
+0.003
=0.005
#0.001
#0.003
#0.001
0.0494
0.0236
0.0274
0.0247
0.0243
0.0494
By Layer
10.001
+0.003
#0.002
#0.001
#0.003
#0.001
0.0496
0.0233
0.0262
0.0247
0.0238
0.0496
By Neurons
#0.001
+0.003
#0.003
#0.001
#0.003
#0.001
0.0478
0.0236
0.0227
0.025
0.0238
0.0478
Fully-Factorized
+0.002
+0.002
+0.002
#0.001
#0.001
#0.002
6.3 EXPLORING PARTITIONING SCHEMES
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 bike- sharing 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 simple fully connected network with two layers with 50 neurons each, and we use SGLD as an optimizer. As 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 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 valid choice 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.
7 CONCLUSIONS
In an attempt to hybridize MCMC and VI, we proposed S-SGMCMC: 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 Sz-SGMCMC. 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 the chains, and result in similar or better posterior predictive accuracies compared to SGMCMC 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.
variety of (deep) models. Our experimental evaluations
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 disk. Additionally, S-SGMCMC scales poorly with respect to the number of parameter groups. SJ-SGMCMC manages to break this dependency; however, still requires slightly more compute than SGMCMC per sample, but is comparable in wall clock time. Possible future work could focus on more theoretical analyses of S S-SGMCMC, such as formal proofs of convergence.

---

Under review as a conference paper at ICLR 2022
8 ETHICS STATEMENT
The main focus of our work is to train models faster by decreasing the convergence time of their training phase. In this scope are not aware any ethical concerns of our research.
REPRODUCIBILITY
For this work we have made sure 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.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et Tensorflow: system for large-scale machine learning. In 12th Symposium on Operating Systems Design and Implementation 16), 2016.
Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring. In Proceedings of the International Conference on Machine Learning. 2012.
Christopher M Bishop. Pattern recognition and machine learning. springer, 2006.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 2017.
Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of murkov chuin monte carlo. 2011.
Thomas Brooks, D Stuart Pope, and Michael A Marcolini. Airfoil Self-Noise and Prediction. 1989.
Luca Cardelli, Marta Kwiatkowska, Luca Laurenti, Nicola Paoletti, Andrea Patane, and Matthew Wicker. Statistical guarantees for the robustness of bayesian neural networks. Proceedings of the International Joint Conference Artificial Intelligence, 2019.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In Proceedings of the International Conference on Machine Learning, 2014.
Paulo Cortez, António Cerdeira, Fernando Almeida, Telmo Matos, and José Reis. Modeling wine preferences by data mining from physicochemical properties. Decision support systems, 2009.
Thomas M. Cover and Joy A. Thomas. Entropy, Relative Entropy and Mutual Information. John Wiley Sons, Ltd, 2001.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEFE/CVF Conference Computer Vision and Pattern Recognition, 2019.
Nando de Freitas, Pedro Hojen-Sprensen, Michael Jordan, and Stuart Russell. Variational meme. In Proceedings of the Conference on Uncertainty in Artificial Intelligence, 2001.
Joshua Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv preprint arXiv:1711. 2017.
Justin Domke. A divergence bound for hybrids of meme and variational inference and an application to langevin dynamics and sgyi. Proceedings of International Conference on Machine Learning, 2017.
10

---

Under review as a conference paper at ICLR 2022
David Duvenaud, Dougal Maclaurin, and Ryan Adams. Early stopping nonparametric variational inference. In Proceedings of International Conference of Artificial Intelligence and Statistics, 2016.
Sathishkumar V E and Yongyun Cho. A rule-based model for seoul bike sharing demand prediction using weather data. European Journal of Remote Sensing, 2020.
Sathishkumar V E, Jangwoo Park, and Yongyun Cho. Using data mining techniques for bike sharing demand prediction in metropolitan city. Computer Communications, 2020.
Charles J Geyer. Practical markov chain monte carlo. Statistical science, 1992.
Ryan Giordano, Tamara Broderick, and Michael Jordan. Linear response methods for accurate covariance estimates from mean field variational bayes. 2015.
Wenbo Gong, Yingzhen Li, and José Miguel Hernández-Lobato. Meta-learning for stochastic gradient meme. International Conference on Learning Representations, 2018.
Jonathan Goodman and Jonathan Weare. Ensemble samplers with affine invariance. Communications in applied mathematics and compatational science, 2010.
Raza Habib and David Barber. Auxiliary variational memc. In International Conference on Learning Representutions, 2018.
David Harrison Jr and Daniel L Rubinfeld. Hedonic housing prices and the demand for clean air. Journal of environmental economics and management, 1978.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016.
Matthew Hoffman and Yian Ma. Black-box variational inference as a parametric approximation to langevin dynamics. In Proceedings of the International Conference on Machine Learning, 2020.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Wilson. What are bayesian neural network posteriors really like? 2021.
Laurentt Valentin Jospin, Wray Buntine, Farid Boussaid, Hamid Laga, and Mohammed Bennaoum. Hands-on bayesian neural networks-a tutorial for deep learning users. ACM Comput, Surv, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online], 2010.
Benedict Leimkuhler, Charles Matthews, and Tiffany Vlaar. Partitioned integrators for thermodynamic parameterization neural networks. 2019.
Daniel Levy, Matthew D Hoffman,, and Jascha Sohl-Dickstein, Generalizing hamiltonian monte carlo with neural networks. International Conference Learning Representations, 2017.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned stochastic gradient langevin dynamics for deep neural networks, In Proceedings of the AAAI Conference Artificial Intelligence, 2016.
Lei Liu, Yuhao Luo, Xu Shen, Mingzhai Sun, and Bin Li. | -dropout: A unified dropout. IEEE Access, 2019.
Yi-An Ma, Nicholas Foti, and Emily B Fox. Stochastic gradient meme methods for hidden markoy models. In Proceedings the International Conference Machine Learning, 2017.
11

---

Under review as a conference paper at ICLR 2022
Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. In Proceedings of the International Conference on Machine Learning, 2016.
Stephan Mandt, Matthew D Hoffman, David Blei, et al. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 2017.
Eric Nalisnick, José Miguel Hernández-Lobato and Padhraic Smyth. Dropout as structured shrinkage prior. In Proceedings of International Conference on Machine Learning, 2019.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. NeurlPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011.
John T Ormerod and Matt P Wand. Explaining variational approximations. The American Statistician, 2010.
Fabio Mendoza Palechor and Alexis de la Hoz Manotas. Dataset for estimation of obesity levels based on cating habits and physical condition individuals from colombia, peru and mexico. Data in brief. 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, deep
Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, learning library. In Advances in Neural Information Processing Systems. 2019.
Sam Patterson and Yee Whye Teh. Stochastic gradient riemannian langevin dynamies the probability simplex. In Advances in Neural Information Processing Systems, 2013.
Rajesh Rangamath, Sean Gerrish, and David Blei. Black box variational inference. In Proceedings of International Conference of Artificial Intelligence and Statistics, 2014.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of International Conference on Machine Learning, 2015.
Tim Salimans. Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In Proceedings of the International Conference Machine Learning, 2015.
Alan Sokal. Monte carlo methods in statistical mechanics: Foundations and new algorithms. In Functional integration. 1997.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. 2014.
Martin J Wainwright and Michael Irwin Jordan. Graphical models, exponential families, and variational inference. 2008.
Tongzhou Wang, YI WU, Dave Moore, and Stuart J Russell. Meta-learning meme proposals. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Basa-Bi and R. Garnett (eds.), Advances in Neural Information Processing Systems, 2018.
Max Welling and Yee W Teh. Bayesian learning via stochastie gradient langevin dynamics. In Proceedings of the International Conference on Machine Learning, 2011.
Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the bayes posterior in deep neural networks really? In Proceedings of the International Conference on Machine Learning, 2020.
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mmist: machine learning algorithms. 2017.
mnist: a novel image dataset for benchmarking
12

---

Under review as a conference paper at ICLR 2022
I-C Yeh. Modeling of strength of high- peth--eeffrmmece concrete using artificial neural networks. Cement and Concrete research, 1998.
Cheng Zhang. Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Transactions Pattern Analysis and Machine Intelligence, 2018.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. 2020.
A THEOREM 1
Proof. We begin with some preliminaries from the main text. Given data D {( ((tr::)) Wi))e=
parameters proper prior distribution p(0), and a likelihood p(D)0) -I p(y, (z, 0), suppose
we are interested in the corresponding posterior distribution p() x p(D())p(0). A convenient
representation of the posterior is as a Boltzmann distribution:
p(0 () x exp(-U(()) where (()) =
Σ (Z.J)
logp(y/z, 0) - log)(0).
(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(O)) DKI (p())))((((
(11)
((0)
Eagn log p(01D)
(12)
We then rewrite Eq. ass follows:
J(q(0)) Egng [log [0)) Eeg [log (0, + +C
(13)
=Bar~q: (log qe (0:) +
Σ
Raz ~gz [log (0₁)]
logr(0, ()a, (0:)(0; II (0,) + C (14) {
for some {1, (1,...,M. M) where - (l...., 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 q or rather every g. This is done by taking the functional derivative of J with respect to setting it equal to zero, and solving for g::
{J(((0))
|log D) II (,,00)d 1 log (0,) igj
(15)
(0,)
log 9: (0.) = Biling [ogg(0,, 0 i, b.D) -
(16)
(0+) exp 1B logp(0,, θ-i, D)}}
(17)
By defining the energy (()) -Bg [og 0_, D) realize that by minimizing the KL-divergence in Eq. the approximate posterior distribution g = M Boltzmann distribution as Eq. with (15) = -M (0.).
qe takes the form of a
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 200T). This implies that the stationary point of the KL is indeed global optimum and unique.
B DERIVING (Sa)
With just a slight shift in perspective, it is actually possible to further generalize U(S) (and conse- quently S-SGMCMC) to produce a broader class of approximate sampling algorithms. This is done
13

---

Under review as a conference paper at ICLR 2022
by first noting that USS) can be represented with scaled double-expectation:
())))
M
E Brmples Egng [logp(r0 + (1-r)0, D)]
(18)
E Brup(s;
2i=1 rs
where (r) Cat(r; M- ,..., .......M)) and (r0 (1 r)))); is equal 02 if r: 1 and 0, otherwise for :,..,: M. Note that this is constructed in this manner specifically so that U USS) remains differentiable with respect to 0. Also note that though the denominator appears superfluous
as (ypp)()) 55)
Σ> -==1 ri] = 1, it is necessary for certain theoretic properties, as seen in Theorem
By replacing (S) with more flexible distribution, energy functions to sample from. One such choice p(aa) (r: p) :x< I Bern(r;; ( 7; > 0)
further generalize and encapsulate different M
with p E (0, 1) Substituting ( for p ((a) in Eq. 18 yields a new energy function that we will
refer to as (Sa), We note that this choice in distribution leads to a dropout-like behavior (Nalisnick ct al. 2019 Srivastava et al. 2014 where the composition of model parameters as r0 + (1 - r)0 leads to each parameter group A2 having probability of approximately p to be used in prediction and a (1 p) probability of being replaced by 0, from the approximate posterior (in traditional dropout, 0, would instead be replaced with 0). Likewise, we will denote methods that use this energy function for sampling as structured dropout SGMCMC (S-MSGMC with different variants all sharing the same S. prefix (e.g. S,-SGHMC).
In practice, the double-cxpectatiin in (Saa 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 (Sa) in this way, computing 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 computing resources. As such, configurations with large amounts of parameter groups are typically only feasible when using S SGMCMC as S-SGMCMC would use too much memory and/or compute per sample.
C THEOREM 2
Theorem 2. For given set of parameters 0 partitioned into M groups, under minor assumptions (i) U(Sa) + U as + and (ūi) U(S2) + US) as - 0. Thus, distributions approximated by Sa-SGMCMC lie continuum with those generated by S-SGMCMC at one extreme and with those from SGMCMC at the other.
Proof. Assume an arbitrary 0, D, n E N. and that Bagg log + D) exists for r E R. M
As an aside, this proof assumes that ((S)) p) xx however, the theorem still holds arbitrary p(Sa) long the mean approaches 1 and variance approaches as + 00.
i= Bern (ommrrp))1([2 1) > 0) with E (0, 1);
(i) Let (n) ~ pl ((5))) (pn) where Vapn E (0, and Pn + 1. It follows that (n) + {1) M as n + oo in distribution (see Lemma in Supplement). Due to bounded and finite support R, find the following:
(0)
M
Σ PER
p (8a) (r: Bang [logp(ro (1- D)]
(19)
Emppi (8,)
E
Σ M0
Ei=1 ri]
M
I(ViT; = 1))EE [0g as n 00
(20)
FER
= log p(0, D) = U(6)
(21)
(ii) Let (n)
(Sa)) where VnPr E (0. and Pn - follows that (m) +
Cat(M-1 ,..., M-1) + do in distribution (see Lemma 2 iin Supplement). Due to bounded and
*Other choices of distribution that are well justified include any with support over (o, 1] and with measure 0 over {0) M. Exploring the effects these distributions have are interesting line of future inquiry.
14

---

Under review as a conference paper at ICLR 2022
finite support R, find the following:
UT(a) (0)
M
Σ
p( (n: (1-11))) r)8, D)
(22)
E
(
1 i= Til
rER
M
1=1
logg(ro +1 --),)))))) + 00
(23)
M
YER
M
BBng (0gg((0;,0000 =
(24)
For both Lemmas 1 [and? 2
let
p (r:pp p Po p) - P
M
1(ver {0, 13)1
Σ Em>0
(25)
1- (1 p)dd
Lemma 1. For ((n) plsa (pr). Pn E (0,1) and n E N. Pn - then (()) 6((1)²) in distribution as n - co.
-
Proof.
p (50) (r = : = Pn MM (1 Pm)
(26)
1- (1 - Pn)
1 ass n DO
(27)
(()) - 6({1)M) in distribution.
(28)
Lemma 2. For r(m) ~ (pm), Pa E (0,1) and n E N. Pr + n then p())
r
Car(M- .....M 1) in distribution as n + co.
Proof. Leti ML..
Pa(1 - Pma) M1
(5a) = 1,r =1= p:
(29)
1 Pm)
(1 Pn) M -1
+ PM.M 1)(1 M-2
I' ''Hōspitll's Rule
(30)
M(1 Pm pm) M-
as n + 00
(31)
M
Since the resulting probabilities sum to I, this implies that (n) + r CCa(MM ,..,, - in
distribution -
D DERIVING (Sa)
To derive (Sa), we must first start with a shift in perspective on how USS) represented. We will rewrite the function in the following way:
((8)
Ee eag (0ogp((000) 0=), D)]
(32)
M
E Eppptes E Bing [ogg(ro (1-))00+10))))))
(33)
16
> i=
ri]
15

---

Under review as a conference paper at ICLR 2022
where p(5) is Middimensional categorical distribution with uniform weights M and p(r0 + (1 r)0, D) is the joint probability of parameters taking values of r0 + -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( p) & pا (1 p) " (V,r; (0,
0)
for p E (0, 1). Note that this is identical to rj sid Bernoulli(p) conditional to Ei ri > 0. Let the support of (Sa) be denoted as R {0, 1JM {0)M, This leads to the following energy function:
(1)
M
E Brapt 1 Bang [lgg(rr + (1-r)),D))
(34)
E
(5
In practice, a few approximations are made to compute the corresponding U(Sa), Firstly, we approximate ((Sa) with an M M-dimensional Bernoullip)) (p) distribution as the difference is minute when Mp is large. Secondly, the outer expectation in Eq. (34) is approximated with Monte Carlo estimate of K samples. The inner expectation is also approximated with Monte Carlo estimate using the latest approximate posterior q(t), However, just like for S-SGMCMC, only single sample is used. This further leads to:
K
(O(1)
1 Kp
((r(t,)))() + (1- (()-)(()))))) D)
(35)
E ALGORITHM FOR SGMGMMCC
The procedure for S.-SGMCMC can be seen in Algorithm
Algorithm 3: S SMMMMCCM
Input: Initial sample (()); parameter partitions O,...., Oas: data set D; initial auxiliary statistics ( step sizes (t masking distribution p(Sa); dropout iterations K.
Output: (6) = {B(+) I...,
for 0
do
Sample minibatch Dit) C
for k 1 to K do
Sample masks rl (1,k)
...., rM plSa)
(1,h)
Sample (() ~ ((())
(lt.k) = (r(((())(() + (1 (1.kk ( (() I
((Sa.t) (0(t,)) D((1)
end
M (Sait) KE (B DM ri] Σk vott(S.))
((+1))e(te ((t-1) SGMCMC CGMCMC_step((!9 (o((Sa)) ((0)))
10
end
12 return (()
$r0 (1 r) is slight abuse of notation that is meant to represent masking out when r: and masking out when 1.
16

---

Under review as a conference paper at ICLR 2022
F SGMCMC UPDATE RULES
The update rules for SGLD, pSGL.D, and SGHMC defined as follows:
SGLD g((+1) g(())
)
+ +N(00
(36)
2
PSGLD 0(++1) g(t) 4
((O(t))
Σ VOR(OLO) +x((0,e
(37)
2
SGHMC ((++1) = g(() + c, m m (1 -G,M- - Voti(())) +0 ((00,))-cV(((())
(38)
(39)
where E, is the step size at time step t, R(-) and M are preconditioners, 0 friction term, and (-) is an estimate of the covariance induced by the stochastic gradient
The update rules for the S-SGMCMC variants are similarly defined as Eqs. 36 39 but all instances of Ü (((()) are replaced with (O((). Likewise, replacing with (1Sa) (g(t) yields the S3-SGMCMC variants.
G ABLATION STUDY
This subsection aims further explore the capabilities of the proposed methodology. More specifi- cally, we for www.e//aaayyylllllooeeeeeeeeeeee Comnected-Newwork-ee experiment with various parameter partitions.
Parameter Partitions. We tested our proposal with four partitioning schemes on 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, I 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 datse(Cortez al. (2009), the Boston housing datase(Harrison Jr & Rubinfeld 1978), the obesity levels Paleecor & de la Hoz Manotas (2199,, the Seoul bike-sharing datase EE et al. (2020): E Cho (2020). , the concrete compressive strength datase Ych (1998), and the airfoil self-noise Brooke et al. (l.((19))) (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 every partitioning scheme we used the learning rates: 1--,,le-,11e-,,1---- 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
with early stoppage. For every dataset and
for hyperparameter
Partition Scheme
Wine
Housing
Obesity
Bike
Concrete
Airfoil
Random (M 3)
le-3
le-5
le-5
le-4
le-5
lc-4
By Layer
1e-3
le-5
le-4
le-4
le-5
le-4
By Neurons
le-3
le-5
le-4
le-4
le-5
1e-5
Fully-Factorized
le-5
le-4
le-5
le-4
le-4
le-3
Table 3: Best Learning Rates for various partitioning schemes on multiple regression datasets.
where the addition of ((,, distributed random variable with mean / and covariance F.
$ NNoe that abuse notation in Eqs. 36 39
(j,,) E) denotes the addition of a normally
17

---

Under review as a conference paper at ICLR 2022
H DETAILS ON EXPERIMENTS
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, consider a regression task where 100 randomly sampled three-dimensional covariates { = [Ta,1 Zi,2, Ti,3] 3"] ii=1...,1000 100 are used to sample response values y/8 ~ >N 2,+b, o²) where प = [w1 W2, Wy] = [1.5. -0.8, 1.3)", b 0.3, and 0² 1. More details on the generation process for can be found in the Supplement.
We choose to fit linear regression model of the same form as the generation process. o² is assumed to be known. Thus, 0 = [1, W22 W3, b). 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 MM ~ [0.52, 0.31, 0.47, 0.84].
The approximated posterior distributions for 0 are found using SGLD, S-SGLD, and S& For the latter two sampling schemes, two parameter partitions are tested: (i) two groups of parameters where 0, = [w', w2] and 02 = (w3, bj and (ii) four groups of parameters where 01 = 11, 0g = W2, 03 = W3, and 6. = b. For Sa-SGLD, p 0.5 and K 4 was used.
The resulting posterior distributions for (W1, w2) and (W1, W's) from all five scenarios, with SGLD in the leftmost column as our baseline, can be seen in Fig. [ We observe that, expected, correlations between (Wi, 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 commonly reported finding for VI techniques and is interesting note that our non-parametric methods appear to exhibit this behavior as well. Finally, it appears that the Su-SGLD adequately approximates S-SGLD with just slightly higher variances and very minor correlations between parameter groups being exhibited.
H.2 HEAL-WORLD. DATA EXPERIMENTS
Framework details. In this subsection, we provide more detailed results for our experiments and grid search for FMNIST, CIFARIO, and SVHN. We note that all the code apart from the metrics was written in PyTorch (Paszke et al. 2019). Regarding the metries, ESS was adopted from the TensorFlow probability library (Dillon et al. 2017 Abadi 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 SSGD and S S-pSGLD. We denoted previously that K p the number of groups. So every parameter will go to the i mod Kpp group where i is the parameter index. If, for instance, Kp then parameter 1 will go to group parameter 2 will go to group parameter 9 will go 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 learning rate of 1e-2. For St-DSGLD, the K is set to 300, which is the forward passes that the model does within epoch. For the grouping of the parameters, for Sa-DSGLD we used group sizes of 2,4,8,32, 128.5 sizes of 2,8,32, 2000000044499.88 4096,819999 01996,8119.111 112,666449,,6,84..333 16344,2776 and 42200; and for S-pSGLD we used groups and 8192.
FashionMNIST. We ran all experiments for 300 epochs with batch size of 500. For SD-SGHMC the K is set to 2, which the forward passes that the model does within epoch. We observed with experimenting with K that we do not need to set K very high, and even 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 S,MSMMCM we put every parameter in a different group, and for S-SGMCMC we used groups of 2.4,8, and 16. For Sa-pSGLD, PSGLD, SGHMC and SGHMC tested their performances with learning rates of 1e-2,te-3.1e-411e le-5. For SSLLD we used a learning rate of le-3 and for S-SGHMC a learning rate of le-2.
18

---

Under review as a conference paper at ICLR 2022
Table 4: Evaluation Metrics on FashionMNIST with pSGLD, Sa-DSGLD and S-PSGLD
Method
dropout
LR
IAC
ESS
Accuracy
Sa-pSGLD [0)
0.1
le-05
1018
5.63
0.918
Sa-PSGi.D
0.1
le-04
808
0.925
Sa-pSGLD
0.1
le-03
754
7.48
0.924
Sa-pSGLD
0.1
le-02
723
8.05
0.911
Sa-pSGLD
0.5
le-05
778
7.08
0.923
Sa-PSGL.D
0.5
le-04
777
7.15
0.923
Sa-pSGLD
0.5
le-03
737
7.57
0.924
Sa-pSGLD
0.5
le-02
677
8.24
0.91
PSGLD
le-5
779
7.09
0.924
PSGLD
le-4
774
7.16
0.911
pSGLD
le-3
770
7.26
0.809
pSGLD
le-2
745
7.48
0.724
S-pSGLD,
le-3
740
7.55
0.918
S-pSGLD,
le-3
751
7.45
0.919
S-pSGLDs
le-3
776
7.24
0.919
S-pSGLD11
le-3
855
6.64
0.916
Table 5: Evaluation Metrics on FashionMNIST with SGHMC, S S,HSGHMC and S-SGHMC
Method
dropout
LR
IAC
ESS
Accuracy
S.-SGHMC
0.1
le-05
782
7.08
0.412
S4-GGHMC
0.1
le-04
888
6.41
0.796
S,SGHMCC
le-03
793
6.98
0.92
S4SGHMC
0.1
le-02
1113
5.06
0.922
S.SGHHMC
0.5
1e-05
790
6.93
0.207
S,-SGHMC
0.5
le-04
789
6.9
0.758
S.-SGHMC
0.5
le-03
796
6.81
0.0.92
S4-SGHMC
0.5
le-02
923
5.7
0.927
SGHMC
le-5
791
6.93
0.206
SGHMC
le4
789
6.9
0.751
SGHMC
le-3
795
6.83
0.92
SGHMC
le-2
920
5.72
0.928
S-SGHMC,
Ic-2
928
5.67
0.928
S-SGHMC,
le-2
915
5.77
0.927
S-SGHMCa
le-2
1142
4.87
0.919
S-SGHMC16
le-2
1121
4.92
0.906
19

---

Under review as a conference paper at ICLR 2022
The setup is similar to the one used in FashionMNIST for 300 epochs with a batch size of 128. For S2-SGHMC, the K is set to 2, which K the forward passes that the model does within 1 epoch. Regarding the parameter partitioning, for SMMMGCC we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4,8, and 16. For Sa-pSGLDD, PSGLD, S5-SGMMC and SGHMC we tested their performances with learning rates of 1--2.le-..le-,Jee- For S-pSGLD, we used learning rate of le-3, and for S-SGHMC, learning rate of le-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 and and
CIFARIO.
ran all experiments
Table 6: Evaluation Metries on CIFARIO with pSGLD, Sa-DSGLD and S-pSGLD
Method
dropout
LR
IAC
ESS
Accuracy
Sa-pSGLD
0.1
le-02
623
7.23
0.191
Sa-pSGLD
0.1
le-03
572
7.6
0.896
Sa-pSGLD
0.1
1e-04
692
6.45
0.921
Sa-PSGLD
0.1
le-05
922
4.88
0.922
Sa-pSGLD
0.5
le-02
546
8.01
0.768
Sa-pSGLD
0.5
le-03
582
7.88
0.918
SappSGLLD [
0.5
1e-04
691
6.85
0.926
Sa-DSGLD [1
0.5
le-05
620
7.22
0.927
PSGLD
le-2
716
8.01
0.666
PSGLD
le-3
740
7.87
0.866
PSGLD
le-4
780
7.41
0.914
PSGLD
le-5
831
6.89
0.926
S-pSGLD,
le-3
600
7.44
0.894
S-pSGLD4
le-3
599
7.4
0.905
S-pSLLDs
le-3
709
6.41
0.881
S-pSGLD16
1e-3
767
5.93
0.836
Table Evaluation Metrics on CIFARIO with SGHMC, S,-SGHMC and S-SGHMC
Method
dropout
LR
IAC
ESS
Accuracy
Sa-DSGLDD [
0.1
le-02
608
7.16
0.91
Sa-pSGLD
0.1
le-03
975
4.6
0.922
Sa-pSGLD
0.1
le-04
654
6.63
0.869
SyrpSGLD
0.1
Ie-05
652
6.65
0.724
SapDGLLD
0.5
1e-02
584
7.7
0.918
Sa-PSLLD
0.5
le-03
751
6.23
0.927
Sa-DSGLD
0.5
le-04
679
6.73
0.886
S SpppSGLD
0.5
le-05
772
6.01
0.778
PSGLD
le-2
727
7.94
0.86
PSGLD
le-3
832
6.84
0.926
pSGLD
le-4
862
6.57
0.885
PSGLD
le-5
858
6.6
0.746
S-pSGLD2
le-3
583
7.49
0.913
S-pSGLD,
le-3
624
7.03
0.919
S-pSGLDs
le-3
904
4.97
0.908
S-pSGLD16
le-3
822
5.47
0.774
SVHN. We also ran all of the experiments for 300 epochs with batch size of 128. Here for SHHGG the K is set to 2, which is the forward passes that the model does within I epoch. We note that K here is less than on CIFAR1O and FashionMNIST, but as we mentioned before, this does not make a difference for results, as have tested. Regarding the parameter partitioning, for SGMMGM we put every parameter in a different group, and for S-SGMCMC we used groups of 2,4.8, and 16. For Sp-DDGL pSGLD, SGHGG and SGHMC we tested their performances with
20

---

Under review as a conference paper at ICLR 2022
0.9
0.8
0.7
Model Averaged Accuracy
SurssaaL pro.., Lit-te2
SasSOLDw
0.6
pacils aoos La
supartt anow, tneTTw
0.5
Lit-
pacun.
poOLD. Li-
gyooni 13-1ea
Supaaa,, Lif-ted
0.4
Segereann 12-
Sycotn LBwled
0.3
0.2
0
2000
4000
6000
8000 Time(s)
10000
12000
14000 16000
Figure 4: Model comparison with different hyperparameters for the PSGLD versions of our baselines and proposed algorithms. We evaluated expected accuracy over iterations CIFAR-10. 10. Subscript "d" denotes the dropout version of our approach. I.R refers to the learning rate, and denotes the dropout rate.
21

---

Under review as a conference paper at ICLR 2022
0.9
0.8
Model Averaged Accuracy
0.6
0.5
....
5, SCHMG
p-0.1. LR-1e2
5, S.NNNME
1R:
p-3.1
9, S0HMM pm0.1.
185
0.4
SCHMG
p-8.5.
5, S000M pa0.5. 1R-1e3
52.901MF pm0.s,
4y 9p-90t1Maee pmo.s, Lit-les
9C:3MC 18-1e
0.3
12t-1ed
8-S01MM,, 12-1e>2
#-00MM Lit te-2
SCSMM, 1R te
0.2
5-ScHrMC,6, 1R=1w2
2000
4000
6000
8000 10000 12000 14000 16000 Time(s)
Figure 5: Model comparison with different hyperparameters for the SGHMC versions of our baselines and proposed algorithms. We evaluated expected accuracy over iterations CIFAR-10 10. Subscript "d" denotes the dropout version of our approach. I.R refers to the learning rate, and denotes the dropout rate.
22

---

Under review as a conference paper at ICLR 2022
learning rates of for S-SGHMC, a learning rate of le-2. Same as CIFARIO, we conducted 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 Figure in the main paper. The strategy that we followed is the same in CIFARIO and is presented in Figs. 6nd 7
For S-pSGLD used learning rate of le-4, and
0.9
M MA
0.8
Supiona maa. Lit SepsOLD Lit-te Bee a SctEE aw I. -R- BonD
0.7
Model Averaged Accuracy
U.Be te
Suppaaa awoti Libeaeia Se poctin ao LiB Bar SooD ae oin, URR te Supp prnni
0.6
Suppoo amots, Literted
Lit-le
0.5
-- Fe sh pwo.s. LR-
p000. twin
$SCLD
0.4
PSCLL
$SCLLDD 1.8 te
DSCLD, IR- 1ed
S-paCLLD, LR- lod
Sigatii, s go0LDie LRale-4
0.3
0.2
1/1 WWWWWWW M
0.1
0
50
100
150 Epochs
200
250
Figure 6: Model comparison with different hyperparameters for the PSGLD versions of our baselines and proposed algorithms. We evaluated expected accuracy over iterations on SVHN. Subscript "d" denotes the dropout version of our approach. LR refers to the learning rate, and p denotes the dropout rate.
23

---

Under review as a conference paper at ICLR 2022
0.9
0.8
0.7
Model Averaged Accuracy
0.6
S0idMCC pw0.1 18-Jee Je
SCHMM 18- 1
0.5
S,5001MG pmo Te-
Sy-0C1MM pmo.1. ted
9,90MMM p-0.1.
W.SCHHM
2-0:1: 16
0.4
S,5001MG
p00.5, 1w-d
pm0.5. 1.2- two
SyS0HMM0
p-0.5 18-1e3
5, pa0.5
0.3
S,5001M pmo.8. twc
SCNMC Lt--e-
9CHMC 1R-1e-5
9CHMC: 18-1e4
0.2
SCHDC 1R-1ee
6-90rMr, 1if-lee2
S-5001M Lit-lack
0.1
SCNMC, [R-le
6-SCHMM 12m1w?
0
50
100
150 Epochs
200
250
Figure 7: Model comparison with different hyperparameters for the SGHMC versions of our baselines and proposed algorithms. We evaluated expected accuracy over iterations on SVHN. Subscript "d" denotes the dropout version of our approach. LR refers to the learning rate, and p denotes the dropout rate.
24


Reviews:
[{"review_id": "PQkks9wh3_q", "reviewer": "Reviewer_qXUG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is a very well-written, clear paper, with nice experiments to guide intuition. Congratulations! The effort shows. \n\nIf the paper \"attempt[s] to hybridize MCMC and VI\", why not compare to structured variational inference? Is the only difference between this approach and variational inference the additional noise in the gradient update?\n\nIf so, it could be stated in a sentence, which would help me be less confused. After all, Equation (4) is the same objective function, theorem 1 is roughly the same proof as coordinate ascent variational inference (where the parameter groups M are simply coordinates i), then these are more of an 'observation' than a theorem. But the key insight of using the coordinate ascent variational inference for SGMCMC is valuable and the main point of thh paper, and calling it a theorem may be helpful for readers. So it's up to the authors for deciding what's best - just wanted to point out that these connections and derivation analogs could be made more clear, which would make the method easier to understand and exposition more straightforward.\n\n- The most confusing sentence for me was belot Eq. 6: \"in practice, at timestep t...\". Maybe describe how \\hat q is composed of samples from previous timesteps. \n\n- small wording choices: \"completely joint\", \"well-approximates\", \"solid distributional assumptions\"\n\n- in equation 8, rho_i is not defined\n\nI also appreciate the authors including clean code as a supplementary zip file for reproducibility.\n", "summary_of_the_paper": "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. \n\nHowever, 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. \n\nExperiments show that this is a viable way to impose structure on a variational distribution, and that mixing times are improved.", "summary_of_the_review": "This is a clearly written paper with a tightly scoped contribution. I advocate for acceptance, and looking forward to follow up work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "O5jweuJO3iR", "reviewer": "Reviewer_NCgf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I enjoyed reading the paper and I think the proposed method has merit, but I also do have some concerns.\n\n1)\tThe paper does not have proof of convergence for the proposed algorithm, a proof of convergence would have made the paper stronger. \n\n2)\tAlthough the paper deals with multivariate models and a potentially large number of parameters, the evaluation metrics are univariate. For example, the ESS in table 1. The authors should report the results for multivariate ESS, especially considering that on, for example, CIFAR-10 the results are only marginally better than SGHMC.  \n\n3)\tPlease note, that in the recent paper [1] it was shown that ESS and “other standard MCMC metrics which do not account for sample bias are not appropriate diagnostic tools for SGMCMC”. This paper [1] proposes to use the kernel Stein discrepancy metric for the assessment of SGMCMC methods. I would recommend the authors include it in their analysis. \n\n4)\tOn page 4 the authors assume that “This partitioning structure is assumed to be known a priori.” when talking about factorization of the parameters into mutually independent groups. While it can be naturally assumed in some models, is this a straightforward assumption of BNN, since the proposed algorithm is in particular motivated by application to BNN?\n\n5)\tOn page 4 the authors mention: “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. “ I would encourage the authors to directly refer to a figure or table which illustrates this claim. \n\n6)\tOn page 3 it states: “The gold standard 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 through a random walk in parameter space. “  \nI would argue that this is not a correct statement as not all MCMC methods deploy random walk behavior (for example, not HMC or NUTS). \n\n7)\tIn the appendix on page 17 you mention “We observe that as we break dependencies we capture similar uncertainty intervals. ”. To me it looks like the confidence intervals actually change quite a bit, especially left/right of figure 4a and left/right of figures 4b and 4c, the confidence intervals become quite a bit more narrow.  Can this again be attributed to VI behavior? \n\n8)\tIn Figure 6 the average accuracy for the last model in the list is quite surprising. Do you have any intuition about that? \n\n9)\tPlease, check your references, both formatting, and whether some of the papers you cite have been published in the meantime. \n\n10) On page 3 the sentence \" The answer is affirmative and will be answered as follows\" probably can be re-written in a better way.  \n\n[1] Nemeth, Christopher, and Paul Fearnhead. \"Stochastic gradient Markov chain Monte Carlo.\" Journal of the American Statistical Association 116.533 (2021): 433-450.\n\nUPD: I increased my score for correctness to 3 after the authors reply and overall score to 5. ", "summary_of_the_paper": "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. ", "summary_of_the_review": "The paper takes a relevant direction in research trying to make inference more scalable for BNN (and other models) by combining MCMC and VI in a stochastic gradient MCMC framework. The paper, however, does not provide theoretical guarantees of convergence while empirical evaluation metrics were not chosen optimally (only univariate metrics which are also not optimal for SGMCMC methods as shownin [1]). Consequently, I would encourage the authors to address the critical points and resubmit the paper in the future. \n\n[1] Nemeth, Christopher, and Paul Fearnhead. \"Stochastic gradient Markov chain Monte Carlo.\" Journal of the American Statistical Association 116.533 (2021): 433-450.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Evlobf5AJvr", "reviewer": "Reviewer_AW5Q", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strength\nPersonally, I find the proposed method interesting. The paper is written clearly and easy to follow, and the overall idea is easy to understand. I have checked the proofs of the theorem and they seem to be correct. To support the claims, the proposed methods were applied to real-world large data set to confirm their advantages compared to the standard full SGMCMC with some ablation studies, although I still think there is room for improvement. \n\n## Weakness and concerns\nAlthough the proposed method is easy to follow, I still find some parts unclear and elaborations are needed. First, the author claims that building the independent structure into the posterior can help improve the mixing speed (bottom of page 1, the first paragraph in the conclusion, etc.). However, I cannot see direct reasoning behind those claims, could you elaborate more on this? In section 6.1, I am not sure I understand the reasoning behind the better performances of S_d-SGMCMC. Can you elaborate more on \"regularizing the model\"?\nOn page 4, why $\\tilde{\\theta}^{(t,i)}$ is composed of samples from previous timesteps? I thought the author mentioned that a **single** sample $\\tilde{\\theta}^{(t,i)}$ from the **current** timestep is used for Monte Carlo approximation. I recommend putting more details in algorithm 1. For example, elaborating how $\\tilde{\\theta}^t_{\\neg i}$ is drawn from $\\hat{q}^t_{\\neg i}$. \nIn addition, why do appendix B and D use the same title? Can you merge them together?\n\nAnother concern is the theoretical soundness of proposed algorithms. For S-SGMCMC, the Monte Carlo estimation of the modified energy function requires previous samples. Does this break the Markov assumption of SGMCMC since it should only depend on the **current** samples? For example, when the sampler does not reach the stationary stage, $\\hat{q}^t$ still evolves with time, thus, the samples from the previous timesteps are not from $\\hat{q}^t$. \nApart from the stationary distribution of S-SGMCMC, I also wonder what is the stationary distribution (if exists) of the dropout version? How different is the stationary distribution of the dropout version compared to the optimal $q$? Any theoretical guarantees on the correctness of the dropout version? \n\nIn terms of the empirical evaluation, I wonder about the uncertainty quantification ability of the proposed algorithms, which is an important metric for SGMCMC methods, especially since the proposed methods also seem to underestimate the posterior variance. In addition, in the abstract, the author mentioned better predictive likelihoods. However, I can only find the accuracy metric in the experiment, which is different from the predictive likelihood. \n\n\n\n", "summary_of_the_paper": "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. \nThe author claimed the resulting algorithms achieved faster mixing speed and better classification accuracy when applied to real-world classification tasks. \n", "summary_of_the_review": "The paper presented an interesting approach to incorporate independent structures into posterior inference. However, there are still some ambiguities in motivation and theoretical soundness. To better support the claims, some metrics or experiments regarding uncertainty quantification should be added or discussed since the proposed methods also seem to underestimate the variance.  \n\n---\nI have read the author's responses. It address some of my concerns. However, for my last concern regarding the uncertainty (which is one of the reasons we consider MCMC over variational inference), the author did not mention it at all. Since this algorithm also underestimates the uncertainty due to minimizing the exclusive KL, I recommend the author to demonstrate that it can provide reasonable uncertainty quantification. Therefore, I will keep my original evaluations. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EnZJyYn8k4m", "reviewer": "Reviewer_zQoE", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. The proposed algorithm is not well motivated: the VI approach aims to obtain a point estimate, SGMCMC aims to draw a sequence of samples from the target posterior, but the proposed algorithm is to simulate a sequence of samples from a modified distribution.\nIt is unclear how well the modified distribution approximates the target distribution and how useful the pseudo-posterior samples are for statistical inference of the target distribution.   \n\n2. The comparison with the existing SGMCMC algorithm is not convincing. It seems that they are ``comparable in wall clock time'', while the baselines pSGLD and SGHMC might not be the state-of-the-art.  ", "summary_of_the_paper": "This paper proposes a hybrid SGMCMC algorithm, a Langevin-type algorithm that operates on a modified energy function based on the variational inference (VI) approach. ", "summary_of_the_review": "This paper proposes a hybrid SGMCMC algorithm, but the underlying theory is not fully developed and the performance of the algorithm is not fully explored. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose a flexible variational posterior approximation, relaxing unrealistic factorization and strong parametric constraints that are standard. There was a mixed reception from reviewers. Overall, the paper is on the borderline. The presentation and empirical investigation could be changed so that the nice contributions in the paper are more easily recognized. Indeed, after rebuttal several reviewers still felt like their concerns were not fully addressed. One reviewer was concerned about the evaluation metrics, and wanted to see Stein discrepancy instead of ESS, and did not feel the ESS was sufficiently motivated (as described in updated comments). Another reviewer felt the uncertainty of the predictive distribution was sufficiently well evaluated. Another reviewer generally satisfied by the response. The decision could go either way, but the paper would probably be more widely appreciated by a significant revision, carefully taking into account the questions of the reviewers. The authors are encouraged to accommodate reviewer questions in future versions of the paper.
Decision: Reject