Parameter Expanded Stochastic Gradient Markov Chain Monte Carlo
Authors: Hyunsu Kim
Giung Nam
Chulhee Yun
Hongseok Yang
Juho Lee
Venue: ICLR 2025 Poster (2025)
Keywords: SGMCMC
Bayesian Neural Network
Parameter Expansion
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Abstract: Bayesian Neural Networks (BNNs) provide a promising framework for modeling predictive uncertainty and enhancing out-of-distribution robustness (OOD) by estimating the posterior distribution of network parameters. Stochastic Gradient Markov Chain Monte Carlo (SGMCMC) is one of the most powerful methods for scalable posterior sampling in BNNs, achieving efficiency by combining stochastic gradient descent with second-order Langevin dynamics. However, SGMCMC often suffers from limited sample diversity in practice, which affects uncertainty estimation and model performance. We propose a simple yet effective approach to enhance sample diversity in SGMCMC without the need for tempering or running multiple chains. Our approach reparameterizes the neural network by decomposing each of its weight matrices into a product of matrices, resulting in a sampling trajectory that better explores the target parameter space. This approach produces a more diverse set of samples, allowing faster mixing within the same computational budget. Notably, our sampler achieves these improvements without increasing the inference cost compared to the standard SGMCMC. Extensive experiments on image classification tasks, including OOD robustness, diversity, loss surface analyses, and a comparative study with Hamiltonian Monte Carlo, demonstrate the superiority of the proposed approach.

Published as a conference paper at ICLR 2025
PARAMETER EXPANDED STOCHASTIC GRADIENT MARKOV CHAIN MONTE CARLO
Hyunsu Kim, Giung Nam, Chulhee Yun, Hongseok Yang, & Juho Lee
KAIST, South Korea (kim. f.m.hynsu, g ung, chu chel yun, wwwgshok........o.ooooo.ooo... yang, juholee)kkaist.
ABSTRACT
Bayesian Neural Networks (BNNs) provide a promising framework for model- ing predictive uncertainty and enhancing out-of-distribution robustness (OOD) by estimating the posterior distribution of network parameters. Stochastic Gradient Markov Chain Monte Carlo (SGMCMC) one of the most powerful methods for scalable posterior sampling in BNNs, achieving efficiency by combining stochas- tic gradient descent with second-order Langevin dynamics. However. SGMCMC often suffers from limited sample diversity in practice, which affects uncertainty estimation and model performance. We propose a simple yet effective approach to enhance sample diversity in SGMCMC without the need for tempering run- ning multiple chains. Our approach reparameterizes the neural network by de- composing each of its weight matrices into product of matrices, resulting in a sampling trajectory that better explores the target parameter space. This approach produces more diverse set of samples, allowing faster mixing within the same computational budget. Notably, our sampler achieves these improvements with- out increasing the inference cost compared to the standard SGMCMC. Extensive experiments on image classification tasks, including OOD robustness, diversity, loss surface analyses, and a comparative study with Hamiltonian Monte Carlo, demonstrate the superiority of the proposed approach.
INTRODUCTION
Bayesian Neural Networks (BNNs) provide promising framework for achieving predictive uncer- tainty and -it-o-i-disiibtiionn (OOD) robustness (Goan & Fookes, 2020). Instead of using gradient- descent optimization algorithm typical for neural networks (NNs), BNNs are trained by estimating the posterior distribution p(O where represents the BNN parameters and the given dataset. Variational methods (Liu & Wang, 2016; David M. Blei & McAuliffe, 2017) and sampling methods are commonly used to estimate the posterior. A representative sampling method is Stochastic Gra- dient Markov Chain Monte Carlo (SGMCMC; Welling & Teh, 2011; Chen al., 2014; Ma al., 2015), which has been highly successful for large-scale Bayesian inference by leveraging both the exploitation of stochastic gradient descent (SGD) and the exploration of Hamiltonian Monte Carlo (HMC; Duane et al., 1987). Based on an appropriate stochastic differential equation, SGMCMC approximately draws samples from the posterior, while avoiding the computation of the intractable posterior density p(0) 'D) and instead using its unnormalized tractable counterpart p( (())
However, despite its powerful performance, SGMCMC suffers from the issue of low sample di- versity. Drawing diverse samples of the BNN parameters 0 is important because it usually leads to the diversity of the sampled functions fo of the BNN, which in turn induces an improved ap- proximation of the likelihood pDD)0). As a result, SGMCMC often fails to replace less less-principled alternatives in practice, in particular, deep ensemble (DE; Lakshminarayanan et al., 2017), which generates multiple samples of the BNN parameters 0 simply by training neural network multiple times with different random initializations and SGD. Although DE is technically not principled Bayesian method, it partially approximates the posterior with high sample diversity, leading good uncertainty estimation (Ovadia et al., 2019; Ashukha et al., 2020). When training time is not a concern, DE usually outperforms SGMCMC in terms of both accuracy and uncertainty estimation.

---

Published as a conference paper at ICLR 2025
Existing approaches for addressing the sample-diversity issue of SGMCMC mostly focus on mod- ifying the dynamics of SGMCMC directly, e.g.. by adjusting the step size schedule (Zhang et al,, 2020), introducing preconditioning into the dynamics (Ma et al., 2015; Gong et al., 2019; Kim et al., 2024), or running multiple SGMCMC chains to explore different regions of the loss surface (Gallego & Insua, 2020; Deng al., 2020). However, such direct modification of the dynamics commonly requires extra approximations, such as the estimation of the Fisher information, bi-level optimiza- tion learn appropriate preconditioning, or high computational resources, such as a large amount of memory due to the use of multiple chains.
In this paper. we propose simple yet effective approach for increasing the sample diversity of SGMCMC without requiring explicit preconditioning or multiple chains. Our approach modifies the dynamics of SGMCMC indirectly by reparameterizing the BNN parameters. the original parameter matrices of the BNN are decomposed into the products of matrices of new parameters. Specifically, for a given multilayer perceptron (MLP), when W Rm=x a parameter matrix of an MLP layer, approach reparameterizes W as the following matrix product:
the approach,
W- PVQ,
(1)
for new parameter matrices V R P Rmxm and Q E R for the same layer. We call the approach as Parameter Expanded SGMCMC (PX-SGMCMC). In the paper, we provide theoret- ical and empirical evidence that our reparametrization alters the dynamics of SGMCMC such that PX-SGM PX-SGMCMC explores the target potential energy surface better than the original SGMCMC. The modified dynamics introduce preconding on the gradient of the potential energy and causes an effect of increasing step size implicitly. While simply growing the step size often hinders the con- vergence of gradient updates, the implicit step size scaling caused by the precondintioning improves the convergence. Although PX-SGMCMC needs more BNN parameters during training, its infer- ence cost remains the same as SGMCMC because the matrices P, V, Q in Eq. 1 can be reassembled into the single weight matrix W for inference.
We evaluate the performance of PX SGMCMC on various image classification tasks, with residual networks (He al., 2016), measuring both in-distribution and OOD performance. Furthermore, we assess sample diversity various ways, such as measuring ensemble ambiguity, comparing PX-SGMCMC with HMC (which is considered an oracle method), and visualizing the sampling trajectories over the loss surface. Our evaluation shows that PX-SGMCMC outperforms SGMCMC and other baselines by significant margin, producing more diverse function samples and achieving better uncertainty estimation and OOD robustness than these baselines.
2
PRELIMINARIES
2.1 NOTATION
We begin by presenting the mathematical formulation of neural networks. Specifically, multilayer perceptron (MLP; Rosenblatt, 1958) with L layers transforms inputs (o) to outputs y (1.) through the following transformations:
a (w
6"),, for -, L- and (L) (1 + (1)
(2)
where (1) denotes the feature at the (-th layer, W(() and (()) respectively are the weight and bias parameters at the 1-th layer, and (.) indicates the activation function applied element-wise
2.2 BAYESIAN INFERENCE WITH STOCHASTIC GRADIENT MCMC
Bayesian model averaging. In Bayesian inference, our goal is not to find a single best estimate of the parameters, such as the maximum a posteriori estimate, but instead to sample from the poste- rior distribution DD) of the parameters 0 given the observed data D. The prediction for new datapoint r is then given by Bayesian model averaging (BMA),
plylz, D) = / p(ulæ, ())d
(3)
which can be approximated by Monte Carlo integration ply|æ, D) xe EM p(y|a, 0mM using finite posterior samples 0i....;0x. p(D)D). This Monte-Carlo integration commonly used
2

---

Published as a conference paper at ICLR 2025
in Bayesian deep learning with posterior samples generated by sampling method, as the posterior distribution of the neural network parameters cannot be expressed in closed form in practice.
Langevin dynamies for posterior simulation. Due to the intractability of the posterior p( p())D) for the neural network parameters 0, we often work with its unnormalized form. Specifically, we typically introduce the potential energy, defined the negative of the unnormalized hog-posterior:
( = - log pl 0) - log p(0).
(4)
Simulating the following Langevin dynamics over the neural network parameters O,
d0 -0MM-rraa MM'rr(, dr -rdd (OU( + /
(5)
yields a trajectory distributed according to exp (-U(0) where setting T = provides posterior samples for computing the BMA integral (Eq. 3). Here. M is the mass, 7 is the damping constant, W represents a standard Wiener process, and '" is the temperature.
Langevin dynamics with stochastic gradients. Computing the gradient (Q) over the entire dataset D becomes intractable as the dataset size increases. Inspired by stochastic gradient meth- ods (Robbins & Monro, 1951), SGMCMC introduces a noisy estimate of the potential energy:
(6) where the log-likelihood is computed only for a mini-batch of data B C D, replacing the full- data gradient with the method, with Stochastic Gradient Langevin Dynamics (SGLD; Welling & Teh, 201 1) and Stochastic Gradient Hamiltonian Monte Carlo (SGHMC; Chen et al., 2014) being two representatives:
(0) ((DI/ (B) log p(B]0) log)(0),
gradient. practice. simulations rely on the semi-implicit Euler
(SGLD) 0 cM Vo (0) + ((0,2(MM)
(7)
(SGHMC) e 0 Mr r (1- (1 - (Vo -x0
(8)
3 PARAMETER EXPANDED SGMCMC
3.1 REPARAMETTIIZATION
Our solution for the sample-diversic problem of SGMCMC is motivated by the intriguing prior sults on deep linear neural networks (DLNNs) (Arora et al., 2018; 2019a; He et al., 2024), which are just MLPs with linear or even the identity activation function (i.e., (1) = l). Although DLNNs equivalent to linear models in terms of expressiveness, their training trajectories during (stochastic) gradient descent are very different from those of the corresponding linear models. Existing results show that, under proper assumptions, DLNNS exhibit faster convergence (Arora et al., 2018; 2019a) and have an implicit bias (distinct from linear models) to converge to solutions that generalize bet- ter (Woodworth et al., 2020; Arora et al., 201 bb; Gunasekar et al., 2018). Observe also that all the layers each DLNN can be reassembled into a single linear layer so that the DLNNs do not incur overhead during inference when compared to the linear models.
Building on these results and observation, we introduce the expanded parametrization (EP) of an T-layer MLP f in Eq. with parameters (W(I), W(L), b(1) ",,,,,,,,,, ((()) as
(1)) (1)
Lid and b(0) a for I = L.....,.
(9)
Q Q for some new parameter matrices P!() and Q." called expanded matrices. Here, represents the number of expanded matrices the left side, and d the right side, and the total number of expanded matrices is c c d 0 (note that when e = = 0, (() and Q11)) are identity matrices). While the expanded matrices p!!)) and Q do not need to be square, their products, ((() and Q must be so in order to ensure that the base matrix ( retains the dimensionality of (!)), In this paper, we square matrices for p!")'s and (")s,, which lets us minimize additional memory overhead while making sure that the reparametrization does not introduce any additional non-global local minima; this is guaranteed when the widths of intermediate layers DLNN are greater than or equal to both the input and output dimensions (Laurent & von Brecht, 2018: 8; Yun et al., 2019).
where (()) A ! ... P and Q
!
L:at
3

---

Published as a conference paper at ICLR 2025
Under our EP, the position variable in the SGLD algorithm is given by
0 = (P!) ..... P o ..... 0 a )
7
(10)
1=1
while in the SGHMC algorithm, the momentum variable, in addition to the position, is defined as
R (1) R (0,SO,, ...... T c()) )
(11)
1=1
where this new momentum variable is related the original momentum variable of SGHMC as in Eq. 9. SP only has S, while R and 'T are introduced by the expanded parameters in the HP setting. We call SGLD and SGHMC under these EPs broadly as Parameter Expanded SGMCMC (PX-SGMCMC) methods. Although the dynamics of such PX-SGMCMC method follows the update formulas Eq. 7, it differs from the dynamics of the corresponding SGMCMC method sig- nificantly. The former is the preconditioned variant of the latter where gradients in the SGMCMC's dynamics, such V Venrie (0orig) for the original position variable Oorig. are replaced by precondi- tioned versions, and this preconditioning produces extraordinary directions of gradient steps in the PX-SGMCMC. In the next section, dive into the new dynamics induced by the preconditioning, and analyze the effect of the preconditioning on the exploration of PX-SGMCMC in terms of the depth of EP and the maximum singular values of matrices in it.
3.2 THEORETICAL ANALYSIS
In this section, we provide theoretical analysis of the EP proposed in Section 3.1. For simplicity,
the following theorem and proof focus the SGLD method described in Eq. 7. To highlight the effect of EP on SGLD, we first need to understand how the combined parameter W evolves under gradient flows when each of its components follows its own gradient flow. The following lemma explains the preconditioning matrix induced by EP. Proofs can be found in Appendix A.
Lemma 3.1 (Dynamics of EP). For an arbitrary function F whose parameter W I'c
W, ... We with its vectorization X = vec (W 1.),),, assume thut the gradient update of each W for i E (I, is defined as the following PDE:
( = Vw -W.F(W ((.))).......))..
(12)
Then, their multiplication X satisfies the following dynamics:
dx(1) - Px VxF(X())
I, (e= 1),
W T W2(t) & I + 1 W (t) (c=2)
where Px(1)
W 2: (t) T (t. @ I W
- (l))))) W ((+ 8 W (1)-1())(15110177
(e 2).
The operator & refers the Kronecker product, and Px(1) denotes the symmetric and positive semi-definite matrix.
This unique gradient flow is known be unattainable through regularization (Arora al., 2018) in the standard parameterization (SP). Furthermore, the singular values or eigenvalues of Px (t) are closely tied to the singular values of the parameters W. In the following, we show that the new dynamics induced by EP promotes greater exploration of the energy surface, which scales with the depth of EP and the maximum singular value across the EP parameters Pt(t), (tt., Qi(t) in Eq. 9.
Theorem 3.2 (Exploration). Assume the following bounds on the expectations of the norms of the gradients, the stochastic gradients, and the Gaussian noise Eq. 7:
VU(W" [ < sh E []] (V))))()))((()))
(t) Ss, E|[[TTM 2 SC.
(13)
where the elements of M corresponding to the expanded parameters P Q are zero. Also assume that the maximum singular value of each parameter matrix in EP is bounded as follows:
sup (tt)= M,
?
V(t) max {mmax (Pi ...., max (P (t) V (t)), O mas (Qi (t))..... MIx (t(())) (14)
4

---

Published as a conference paper at ICLR 2025
Then, due the preconditioning in Lemma 3.1, the Euclidean distance of two SGLD samples at consecutive time steps is upper-bounded by the following term, which depends on the depth c- +++++;
E (() W(t + 1)112] & eL²(c+ d 1)M (h 5) eLC.
(15)
Note that the depth (c+d+I) and the maximum singular value M decrease, the upper bound on the distance of two consercutive samples gets smaller.
Although the bound in Theorem 3.2 is not necessarily the maximum distance between consecutive samples, its dependency on the depth (c+dd11) of our EP suggests that the preconditioning induced by EP may improve the exploration of the SGLD and lead to the generation of more diverse samples. Note that the bound Eq. 15 is also proportional the step size e. However, in practice, using large step size above a certain threshold rather hinders the performance. In our experiments. the performance monotonically improved as we increased the depth of EP, while converging a certain fixed level in the end. This suggests that the preconditioning of our EP induces a form of implicit step-size scaling while maintaining the stability of the gradient-desces steps.
3.3 IMPLEMENTING EP FOR DIFFERENT ARCHITECTURES
Linear layers. We follow the Eq. 9. The depth and width each matrix P,, V. or Q; can be adjusted, but the width should be at least as large that of the corresponding input dimension.
Convolution layers. There are four dimension axes in the standard convolution layer. Let k. c, Co be the sizes of the kernel, the input channel, and the output channel, and W Rkxkxx, be the kernel matrix. If we naively reparameterize W with E ke xke. V E TB and Q ggke,kke, the memory and computation overheads become significant. Thus, we let P and Q operate on the channel Ci, dimensions only, and each kernel dimension is multiplied by the same matrices by defining P E RexXe., V IR 1xkxkxe,Xe, and Q E R2xx That is, in the index notation, our reparameterization is:
abis = Σ Vabur u.d
b, = Σ Pinau.
(16)
In contrast to the naive reparametrization, where the number parameters roughly increases by three folds when the depth of the reparametrization is 3 (i.e., (PVQ) = 3 #Params the above reparametrization requires only #Params = (1+2/k²) 2 #Params(W) that case.
Normalization layers. Normalization layers, such Batch Normalization (loffe & Szegedy, 2015), Layer Normalization (Ba, 2016), and Filter Response Normalization (FRN; Singh & Krishnan, 2020), contain few parameter vectors. For example, FRN used in Izmailov et al. (2021) sists of the scale, bias, and threshold vectors, s E IR ",, E R E RR As in the case of the linear layer, we can simply multiply matrices on the left-hand side.
Σ Ph St
t, IPPth a.t
(17)
The row and column dimensions of P and Q should be at least Co.
RELATED WORK
Linear parameter expansion. Linear parameter expansion techniques are relatively underexplored in deep learning, as they do not inherently increase the expressivity of deep neural networks, particu- larly non-linear architectures. While this approach has shown some benefits in linear networks, is often overlooked in modern deep learning applications. A few notable works, however, have em- ployed parameter expansion in the context of convolutional neural networks. For instance, Chollet (2017), Guo et al. (2020), and Cao et al. (2022) have introduced techniques that either decompose or augment convolution layers to reduce FLOPs and improve generalization. These methods pri- marily aim at enhancing efficiency or regularization, leveraging additional layers or decompositions to modify the structure of network without significantly increasing computational complexity. On the other hand, Ding et al. (2019) proposed a different approach by expanding convolutional layers
5

---

Published as a conference paper at ICLR 2025
through addition rather than multiplication, aiming improve robustness against rotational distor- tions in input images. While their method enhances robustness to certain image transformations, does not focus the exploration properties of parameter expansion in non-convex problems.
SGMCMCs for diversity. Building upon the seminal work of Welling & Teh (2011), which in- troduced SGLD
scalable MCMC algorithm based stochastic gradient methods (Robbins & Monro, 1951), ), a range of SGMCMC methods have emerged in the past decade (Ahn et al., 2012; Patterson & Teh, 2013; Chen al., 2014; Ding al., 2014; Ma al., 2015; Li al., 2016). De- spite their theoretical convergence to target posteriors under the Robbins-Monro condition with a decaying step size (Teh et al., 2016; Chen et al., 2015), effectively exploring and exploiting the posterior density of deep neural networks using a single MCMC chain remains challenging due to their multimodal nature. In this context, Zhang et al. (2020) proposed a simple yet effective cyclical step size schedule to enhance the exploration of SGMCMC methods. Intuitively, the larger step size at the beginning of each sampling cycle facilitates better exploration while tolerating simulation enror, whereas the smaller step size towards the end of the cycle ensures more accurate simulation. While the cyclical schedule helps with exploration, it often struggles to fully capture multimodality and typically requires a significant number of update steps to transition between modes (Fort al., 2019; Kim et al., 2024). Thus, improving SGMCMC methods for modern deep neural networks is of research, with recent progress using meta-learning frameworks (Gong al., 2019; Kim et al., 2024). To the best of our knowledge, we are the first to propose the parameter expansion for enhancing the exploration of SGMCMC.
still active
EXPERIMENTS
In this section, we present empirical results demonstrating the effectiveness of the parameter expan- sion strategy proposed in Section 3 for image classification tasks. We aim to show that PX-SGHMC, an SGHMC variant of PXSGMCMC, collects diverse posterior samples and consistently outper- forms baseline methods across various tasks. We compare with the following SGMCMC methods as baselines: SGLD (Welling & Teh, 201 PSGLD (Li et al., 2016), SGHMC (Chen et al., 2014), and SGNHT (Ding et al., 2014). Unless otherwise specified, all methods utilize a cyclical step size schedule. Note that in our method, we set the elements of the semidefinite damping (friction) 7 in Eq. corresponding to P and Q much smaller than those of V, ensuring a stable SGHMC trajectory. For details on implementations and hyperparameters, please refer to Appendices C and D.
We conduct experiments including multimodal distribution (Section 5.1) sampling and Bayesian neural networks on image classification tasks (Sections 5.2 and 5.3). Unless stated otherwise, (1) the reported values are presented as "meanesa," averaged over four trials, and (2) a bold-faced underline highlights the best outcome, while an underline indicates the second- best value. To assess the quality of classifier predictions, we compute classification error (ERR) and negative likeliiid (NLL), while ensemble ambiguity (AMB) quantifies the diversity of ensemble predictions. Refer to Ap- pendix B.1 for the definitions the evaluation metrics.
5.1 ToY RESULTS: MIXTURE OF GAUSSIANS
We begin by comparing SP and EP on the multimodal 2D mixture of 25 Gaussians (MoG), following Zhang al.
Target
SP
EP
(2020). Although this distribution is not a BNN posterior, is sufficient demonstrate the role of the preconditioning induced by EP. The random variables of MoG, x E IR22, are decomposed W; W,W for W1, W2, W3 E R2x2 and HMC sampling also follows the preconditioning on the gra- dients. Specifically, we collect 10,000 samples by running HMC with a single chain, using 10 leapfrog steps and a con- stant step size of 0.05 for both SP and EP. Fig. I illustrates that EP resolves the issue of SP getting trapped in local modes without requiring additional techniques such as tempering or cyclical step sizes (Zhang et al., 2020). As explained in Section 3.2, the implicit step size scaling induced by the preconditioning facili- tates escaping local modes, which suggests strong performance in non-convex problems such as multimodal distributions.
2.79
7:57
12.35
37.11
Figure 1: Toy results. HMC sam- ples with SP and EP. The color rep- resents the negative log probability.

---

Published as a conference paper at ICLR 2025
Table 1: Main results on CIFAR-10 and associated distribution shifts. Evaluation results on C10 (CIFAR-10), C10.1 (CIFAR-10.I), C10.2 (CIFAR-10.2) STL, and C10-C CIFAR-IO-CY, Metries for C10-C are computed using total of 950,000 examples, encompassing 5 intensity levels and 19 corruption types. Please refer to Appendix B.2 for more detailed results regarding C10-C.
Distribution shifts
Metric
Method
C10
C10.1
C10.2
STL
C10-C
SGLD
0.152a0.003
0.258
0.293
0.326x00004
0.309#a.007
PSGLD
0.158
0.277=0.044
0.305a0.003
0.33320.003
0.316=0.006
ERR (L)
SGHMC
0.133:0.002
0.234.0.003
0.277.0.006
0.30620.003
0.292.0.002
SGNHT
2.
0.23620.005
0.273x0.004
0.30720.005
0.29422.00s
PX-SGHMC (ours)
0.121
0.218.0.005
0.257
0.28720.002
0.287.00.08
SGLD
0.47720.009
0.77220.010
0.891 10.005
0,92810.011
0.91320.022
PSGLD
2.501
0,81220.011
0. 2..2820.007
0,95820.006
0.93820.017
NLL (L)
SGHMC
0.422+0.005
0.698-0.006
0.834+000007
0,868+0.01
0.871+0.005
SGNHT
0.425=0.004
0.705
0.833+0.008
0.873+00006
0.877=0.021
PX-SGHMC (ours)
0.388-0.005
0.661-0amt
0.806+0.009
0.819+0004
0.859-0.022
Table -isto--dist detection. Evaluation results for distinguishing in in-distribution in- puts from CIFAR-10 10 and out-of-distribution inputs from SVHN and LSUN based on predictive entropy. This table summarizes evaluation metries, including AUROC, TNR95 (TNR @ TP)=95%)
and TNR99 (TNR @ TPR==99)) For detailed plots, we refer readers to Appendix B.3.
SVHN
LSUN
Method
AUROC (t)
TNRY5 (1)
TNR99 (1)
AUROC (1)
TNR95 (1)
TNR99 (1)
SGLD
0.784
0.536
0.42420004
0.853an.015
0.520sn.uss
0.276=0089
PSGLD
0.745
0.506#0.012
=.408=
0.855a0.009
0.520sa.02s
0.255=0043
SGHMC
0.790+a.051
0.544940.66
0.427=0000
0.860+6.018
0.554+0.025
0.357
SGNHT
0.776+0.014
0.530++0015
0.436-00014
0.86440.009
0.55660.006
0.375-0015
PX-SGHMC (ours)
0.832+0.014
0.632 +0.00
0.514-0.01
0.884-0.007
0.594.0.057
0.405-0006
5.2
MAIN RESULTS: IMAGE CLASSIFICATION TASKS
5.2.1
CIFAR-10
We present results on CIFAR-10 (Krizhevsky et al., 2009) and extensively study in advanced tasks such as robustness analysis and OOD detection. For our experiments, we employ the R20-FRN- Swish architecture, representing ResNet20 with FRN normalization and Swish nonlinearity and adapted from the HMC checkpoints provided by Izmailov et al. (2021).
Results on CIFAR. Table presents the evaluation on the CIFAR-10 test split in the 'C10' column. It shows that PX-SGHMC significantly outperforms other methods in terms of both ERR and NLL..
Robustness to distribution shifts. One of the key selling points of Bayesian methods is that they produce reliable predictions that account for uncertainty, leading us evaluate robustness to dis- tribution shifts (Recht et al., 2019; Taori et al., 2020; Miller et al., 2021), 1). Specifically, we test on natural distribution shifts using CIFAR-like test datasets, including CIFAR-10.1 (Recht al., 2019), CIFAR- 10.2 (Lu et al., 2020), and STL (Coates et al., 201 1), as well as on image corruptions using CIFAR-101 10-C (Hendrycks & Dietterich, 2019). Table 1 shows that our approach not only outper- forms the baseline methods on the in-distribution data but also exhibits significant robustness to distribution shifts. Please refer to Appendix B.2 for further results on the CIFAR- benchmark.
Out-of-distribution detection. Another important task for evaluating predictive uncertainty is the OOD detection. In real-world scenarios, models are likely to encounter random OOD examples, in which they are required to produce uncertain predictions (Hendrycks & Gimpel, 2017; Liang & Li, 2018). Categorical predictions of the classifiers are expected to be closer to being uniform for OOD inputs than for in in-distribution ones. Therefore, we use predictive entropy (Lakshminarayanan et al., 2017) to distinguish between in-distribution examples from CIFAR-10 and OOD examples from SVHN (Netzer al., 2011) and LSUNN (Yu al., 2015). Table demonstrates method
7

---

Published as a conference paper at ICLR 2025
Table 3: Results with data augmentation. Evaluation results on C10 (CIFAR-1O), C100 (CIFAR- 100), and TIN (Tiny Imneem with data augmentation. In this context, we manually set the poste- rior temperature 0.01 to account for the increased data resulting from augmentation.
ERR (L)
NLL (1)
Method
C10
C100
TIN
C10
C100
TIN
SGLD
0.08020.002
0.32620.006
0.546.0.001
0.24620.001
1.18010018
2.27810010
PSGLD
0.097.0.002
0.411.20.007
0.60 20.055
0.30620.001
1 1.54610.035
2.56220015
SGHMC
0.07 +0.00
0.319=0.002
0.538+00.03
0.223+0.002
1.138+0
2.251 2.251+0002 +0.02
SGNHT
0,07420:001
0.335g0.004
0.53620.002
2.230
1.199
2.240z0.013
PX-SGHMC (ours)
0.069.0.0m
0.290-00004
0.498-0.004
0.217+0.005
1.030+0011
2.089-0008
c= d
c=1,=0
c=1,=1
c=2,d=2
109
10a
10a
10a²
210m Be.1
Condition number
10a
10a
10a
10011
10c
100
10
100
10-3
100
25 50 75100
25 50 75100
25 50 75100
25 50
75100
25 50 75100
Cycle index
Cycle index
Cycle Index
Cycle index
Cycle index
Figure 2: Connection between exploration and singular value dynamics. The first and second plots illustrate exploration through unnormalized and normalized Euclidean distances, while the third to fifth plots depict singular value dynamics, represented by the largest and smallest singular values and condition numbers. For the 21 layers, the singular value plots feature 21 transparent lines for each item, with the maximum (or minimum) value highlighted as the representative.
shows greater predictive uncertainty in handling OOD examples, as indicated by metrics associated with the Receiver Operating Characteristic (ROC) curve.
5.2.2 RESULTS WITH DATA AUGMENTATIONS
Although our evaluations in Section 5.2.1 were conducted without data augmentation, practical set- tings typically involve it. Therefore, we present additional comparative results on various image classification datasets, including CIFAR-10, CIFAR-100, and TinyImageNet, using data augmen- tation that includes random cropping and horizontal flipping. Since these augmentations can lead to inaccurate likelihood estimation (Wenzel et al., 2020; Noci et al., 2021; Nabarro et al., 2022), we introduce the concept of a cold posterior in these experiments, i.c., setting T < 1 in Eq. 5 and sampling from the tempered posterior ((
Table 3 presents the results obtained by setting T = 0.01 in Eq. 7. It demonstrates that PX- PX-SGHMC still outperforms the other methods in both ERR and NI.., indicating that the enhanced diversity also applicable to practical scenarios involving data augmentations and posterior tempering across various datasets. Additionally, Appendix B.4 provides ablation results related to the cold posterior effect in the absence of data augmentation, showing that method consistently outperforms the SGHMC baselines across varying temperatures.
5.2.3 EMPIRICAL ANALYSIS
Connection between exploration and singular value dynamics. Theorem 3.2 suggests that the update size of SGLD dynamics at each time step is constrained by the maximum singular value, implying that small singular values may limit exploration. In this regard, a key characteristic of EP in deep linear neural networks is the learning dynamics of singular values; Arora et al. (2019b) showed that the evolution rates of singular values are proportional to their magnitudes raised to the power of 2 - 2/c, where represents the depth of the factorization. In other words, as the depth of matrices increases, larger singular values tend to grow, while smaller singular values shrinks close to zero. Although our setting does not fully align with the theoretical assumptions in Arora et al. (2019b), as we do not consider DLNNs, we empirically demonstrate the connection between exploration and singular value dynamics by varying the number of expanded matrices c+++11
8

---

Published as a conference paper at ICLR 2025
0
(=i,d=i
2,
Expansion
ERR (L)
NLL (L)
AMB (t)
Train ever
0.12
Voldd ever
0.20
c 0, d 0
0.13520005
0.44420007
0. 0.19620004
c 1,
0.127
0.404z0004
0.214
c 1,
06
0.369=000s
0.253=0000
25
Cycle
50
75
1.00
0
Cycle
50
75
100
c 2, d 2
0.114.0001
0.373
0.278
Figure 3: Trace plots for EP. It depicts training and validation errors along with trajectory.
Table 4: Ablation results on EP. Metrics are
computed using the validation split.
Fig. 2 depicts our experimental findings: (i) The first and second plots quantify exploration by com- puting the Euclidean distance between consecutive posterior samples, defined as d( mm+1) = [[0m+1 - Pm/l12- where O. denotes the sample at cycle index m. To exclude the effect of seale invariance in neural network parameters due to normalization layers such as FRN, we also com- pute their normalized version: d(Om. mm+ = ||0m+1 0mll2 /|0mll2 Both unnormalized and normalized distances between consecutive samples increase the number of expanded matrices e == c + 1 grows. (ii) The third and fourth plots illustrate the dynamics singular values by plotting the maximum and minimum singular values of the convolution layer kernels. For each con- volution layer, we compute the singular values of the kernel tensor following Sedghi et al. (2019). In line with the theoretical argument presented in Arora al. (2019b), although our setup does involve DLNNs, we observe that as the number of expanded matrices increases, the largest singular value rises while the smallest singular value decreases.
EP converges faster than SP. Another important property of EP in DLNNs is its accelerated con- vergence toward optima or modes (Arora al., 2019a), which has also been observed in deep convo- lutional networks with nonlinearities (Guo et al., 2020). Building on this, we empirically investigate the local mode convergence of EP within the SGMCMC framework, where faster convergence is particularly crucial for BMA performance due to the slower local mode convergence caused by the injected Gaussian noise in SGMCMC methods (Zhang et al., 2020) compared SGD in DE. Fig. presents trace plots of training and validation errors, showing that both tend to converge more rapidly as the number of expanded matrices increases.
Based on the empirical findings, hypothesize that EP induces large maximum singular value, as shown in ii), which enlarges the upper bound in Theorem 3.2 and breaks the exploration limit, as demonstrated in i). Table 4 additionally presents the validation metrics for each setup and clearly demonstrates that the proposed EP indeed achieves better functional diversity, as indicated by the in- creased AMB. To sum up, PX-SGMCMC effectively enhances both the exploration and exploitation of a single SGHMC chain, resulting in improved BMA measured by ERR and NLL.
5.3 COMPARATIVE ANALYSIS USING HMC CHECKPOINTS
While Section 5.2 presents extensive experimental results by running SGMCMC algorithms from random initializations, we also conduct a comparison using HMC checkpoints provided by Izmailov et al. (2021) as an initialization of parameters in SGMCMC. Specifically, we run both SGHMC and PX-SGHMC starting from the burn-in checkpoint of HMC, employing hyperparameters aligned with those in Izmailov et al. (2021). Further experimental setups can be found in Appendix D.
5.3.1 DIVERSITY ANALYSIS
The diversity of the parameters does not necessarily imply that the diversity of the corresponding functions they represent; for instance, certain weight permutations and sign flips can leave the func- tion invariant (Hecht-Nielsen, 1990; Chen et al., 1993). To effectively approximate the BMA integral in Eq. 3, functional diversity is essential (Wilson & Izmailov, 2020). Therefore, quantify the di- versity of posterior samples using their predictions by computing the ensemble ambiguity (Wood et al., 2023) as well as the variance of predictions (Ortega et al., 2022).
Table clearly shows that PX-SGHMC exhibits (a) superior exploration in the parameter space compared to vanilla SGHMC, evidenced by the average distances between consecutive samples, and (b) higher diversity in predictions, indicated by ensemble ambiguity and variance of predictions comparable to the gold standard HMC. Consequently, similar to HMC, (c) the BMA performance

---

Published as a conference paper at ICLR 2025
Table Diversity analysis using HMC checkpoints. We measure (a) parameter diversity using unnormalized and normalized Euclidean distances (d and d), (b) prediction diversity using ensemble
ambiguity (AMB) and variance (VAR), and (c) individual (IND) and ensemble (ENS) negative log-
likelihoods for 10 posterior samples from each method.
(a)
(b)
(c)
Method
d( (0. 0m-1)
d 10 0m+1)
AMB
VAR
IND
ENS
HMC
322. 840 333
1.377
0.347
0.107
0.800
0.381
SGHMC
60.65 258
0.258
0.162
0.063
0.690
0.464
PX-SGHMC (Ours)
1290,+1.372
1.372=0.150
0.339
0.105
0.739
0.353
HMC
SGHMC
PX-SGHMC
ERR@TTain
0.9
SGHMC
0.6
Cycle index
0.2
0.6
0.8 1.0
(a) Linear connectivity
(b) Two-dimensional subspace
Figure 4: Loss landscape analysis using HMC checkpoints. We visualize (a) linear connectivity between consecutive posterior samples and (b) a two-dimensional subspace spanned by the Oth (diamond), 1st (circle), and 2nd (pentagon) posterior samples. Both plots depict classification error on 1,000 training examples. Note that the 8th HMC sample was rejected and reverted to 7th.
of PX-SGHMC surpasses that of SGHMC, although individual posterior samples exhibit worse performance in terms of negative log-likelihoods.
5.3.2 Loss LANDSCAPE ANALYSIS
In this section, we investigate how effectively PX-SGHMC explores the posterior distribution over parameters compared to both HMC and SGHMC, from the perspective of loss surface geometry (Li et al., 2018). Fig. 4a visualizes the loss barrier between consecutive posterior samples, illustrat- ing how often each method jumps over these barriers. While PX-SGHMC does not jump high as HMC, the larger barriers it crosses compared to SGHMC indicate significantly better explo- ration, consistent with the discussion in Section 5.3. Fig. 4b visualizes two-dimensional subspace spanning the right-most initial position (Oth) and two subsequent posterior samples (Ist and 2nd), demonstrating that the diversity of PX-SGHMC approaches that of the gold standard HMC when sampling from the multi-moda BNN posterior.
6 CONCLUSION
We have presented PX-SGMCMC, simple yet effective parameter expansion technique tailored for SGMCMC, which decomposes each weight matrix in deep neural networks into the product of new strategy provides form of preconditioning on the gradient updates, enhancing the exploration of the posterior energy surface. The extensive experimental results strongly support our claims regarding the improved exploration linked to the singular value dynamics of the weight matrices explained in our theoretical analysis. As a result, the posterior samples obtained through PX-SGMCMC demon- strate increased diversity in both parameter and function spaces, comparable the gold standard HMC, leading to improved predictive uncertainty and enhanced robustness to OOD data.
matrices. Our theoretical analysis shows that the proposed parameter expansion
Limitations. While EP does not increase inference costs, it does require additional training sources in terms memory and computation. In future work, we aim to optimize the reparameteri- zation design minimize these computational overheads while further enhancing diversity.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
This paper does not raise any ethical concerns, as it presents parameter expansion strategy based on the SGMCMC algorithm, which is free from ethical issues.
REPRODUCIBILITY STATEMENT
For the theoretical results in Section 3, we direct readers to Appendix A. All experimental details necessary for reproducibility can be found in Appendix D.
ACKNOWLEDGMENT
This work was partly supported by Institute of Information & communications Technology Planning & (II))) grant funded by the Korea govemmenn(MSIT) (No.RS-2019. 905-2009--1190007 Artifi- cial Intelligence Graduate School Program(KAIST); Lab), and the National Research Foundation of Korea(NRF) grant funded by the Korea govern- ment(MSIT) (NRF-2022R1 A5A708390812 12; No. RS-2023-00229680), We also thank Byoungwoo Park and Hyungi Lee for their thoughtful discussion.
Global AI Frontier
REFERENCES
Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring. In Proceedings of the 29th International Coference on International Con- ference on Machine Learning, 2012.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization.. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssaan Stockholm, Sweden, July 10-15, 2018, 2018. 3, 4, 28
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis gradient descent for deep linear neural networks. In 7th International Conference Learning Represen- tations, ICIR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019a. 3, 3
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems, 2019b. 3, 8,9
Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry Vetrov. Pitfalls of in-domain uncertainty estimation and ensembling in deep learning. In International Conference on Learning Representations, 2020.
Jimmy Lei Ba. Layer normalization. arXiv preprint arXiv:1607.064400, 2016. 5
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ google/ ax. 24
Jinming Cao, Yangyan Li, Mingchao Sun, Ying Chen, Dani Lischinski, Daniel Cohen-Or, Baoquan Chen, and Changhe Tu. Do-conv: Depthwise verrpaaaaeterrize convolutional layer. IEEE Trans. Image Process. 1113726--736 2022. doi: 10. 10.1109971122222.31155332 URL https: / //doi.org/10. 1109/TIP 2022. 3175432.
An Mei Chen, Haw-minn Lu, and Robert Hecht-Nielsen. On the geometry of feedforward neural network error surfaces. Neural computation, 5(6)91109927, 1993.
Changyou Chen, Nan Ding, and Lawrence Carin. On the convergence of stochastic gradient meme algorithms with high-order integrators. In Advances in neural information processing systems, 2015. 6

---

Published as a conference paper at ICLR 2025
Tianqi Chen, Emily B. Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In Proceedings of the 31th International Conference on Machine Learning. 2014. 1,3, 63
François Chollet. Xception: Deep learning with depthwise separable convolutions. In 2017 IEEE Conference on Computer Vision and Puttern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pp. 1800-1807. IEEE Computer Society, 2017. doi: 10.1 10/1C/P..22000 URL
https: //oii. org/oogggoooooo /00. 1109 CVVPR (C1...2017.1.. 195.
Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the 14th international conference on artificial intelligence and statistics, 2011. 7, 25
Alp Kucukelbir David M. Blei and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859 877, 2017. doi: 0198000162145992117.23838
1
Wei Deng, Qi Feng, Liyao Gao, Faming Liang, and Guang Lin. Non-convex learning via replica exchange stochastic gradient MCMC. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, 2020. 2
Nan Ding, Youhan Fang, Ryan Babbush, Changyou Chen, Robert D) Skeel, and Hartmut Neven. Bayesian sampling using stochastic gradient thermostats. Advances in neural information processing systems, 2014. 6, 23
Xiaohan Ding, Yuchen Guo, Guiguang Ding, and Jungong Han. Acnet: Strengthening the kernel skeletons for powerful CNN via asymmetric convolution blocks. In 2019 IEEÉ/CVF International Conference on Computer Vision, ICCV 2019, Seoul, Korea (South), October 27 -Nvemberr 2, 2019, pp. 1911-1920. IEEE, 2019. doi: 10.1 10.109//CCV.2019.00020.0 org/ 10. 11011119//0 ICt///C/..... 2019 00200.
URL https://doi.
Simon Duane, Brian Pendleton, and Duncan Roweth. Hybrid monte carlo. Physics letters B, 195 (2):216-222, 1987.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018. 25
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape per-
spective. arXiv preprint arxiv:1912.0277700 2019. 6
Victor Gallego and David Rios Insua Stochastic gradient meme with repulsive forces, 2020. URL
httts::/arxiv..t wrg/ ablgg///////// http//aass11111 /11// 1ttp//10111.....1.111111
00071.
Adrià Garriga-Alonso and Vincent Fortuin. Exact langevin dynamics with stochastic gradients.
CoRR, abs/2102.01691, 2021. URL https: http///aavvvv...///e/eeeee..av.a org/ abs 22102 01 69191.4 20
Ethan Goan and Clinton Fookes. Bayesian neural networks: An introduction and survey. CoRR, abs/2006. 1102006.12000 12024, 2020. URL https: : http//aavvvvv//// wwwpa/.wwww./////o///oaaaa 12024.
Wenbo Gong, Yingzhen Li, and José Miguel Hernández-Lobato. Meta-learning for stochastic gra- dient MCMC. In International Conference on Learning Representations, 2019. 2,
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in neural information processing systems, 2018.
Shuxuan Guo, Jose M Alvarez, and Mathieu Salzmann. Expandnets: Linear over-parameterization to train compact convolutional networks. In Advances in Neural Information Processing Systems, 2020. 5, 9
Haiyun He, Christina Lee Yu, and Ziv Goldfeld. Informmation-theorrtiic generalization bounds for deep neural networks, 2024. URL https: /ttp:s////a/iiie//eeeeeeeeeeserr www.wiwwwwww.////////a/aaaaa/a/e 0tt.///2444031186660 03176.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, 2015. 25
12

---

Published as a conference paper at ICLR 2025
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016. 2, 25
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Ad- vanced Neural Computers, pp. 129-135. Elsevier, 1990. 9
Dan Hendryeks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. International Conference on Learning Representations, 2019. 7, 25
Dan Hendrycks and Kevin Gimpel.
Gaussian error linear units (gelus). arXiv preprint
arXiv:1606.08415, 2016. 25
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distributioo examples in neural networks. In International Conference on Learning Representations, 2017. 7
Alan M Horowitz. A generalized guided monte carlo algorithm. Physics Letters B, 268(2):247-252
1991. 20
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on Machine Learning, 2015. 5, 25
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In Proceedings of the 38th International Conference Machine Learning, 2021. 5, 9, 20, 22, 24, 25, 26
Seunghyun Kim, Seohyeon Jung, Seonghyeon Kim, Juho Lee. Learning explore for stochastic gradient MCMC. In Proceedings of the 41st International Conference on Machine Learning, 2024. 2, 22, 24
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009. 7, 24
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 2017. 1,7
Thomas Laurent and James von Brecht. Deep linear networks with arbitrary loss: All local minima are global. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmssmssaan Stockholm, Sweden, July 10-15, 2018, 2018.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned stochastic gradient langevin dynamics for deep neural networks. Proceedings of the AAAI conference on artificial intelligence, 2016. 6, 23
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances in neural information processing systems, 2018. 10
Shiyu Liang and Yixuan Li. Enhancing the reliability of out-of distribbiinn image detection in neural networks. In International Conference on Learning Representations, 2018. 7
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. In Advances in Newral Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, 2016.
Shangyun Lu, Bradley Nott, Aaron Olson, Alberto Todeschini, Hossein Vahabi, Yair Carmon, and Ludwig Schmidt. Harder or different? closer look distribution shift in dataset reproduction. In ICML Workshop on Uncertainty and Robustness in Deep Learning, 2020. 7, 25
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme. In
Advances in neural information processing systems, 2015. 1, 2, 6, 23, 24
13

---

Published as a conference paper at ICLR 2025
John P Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishall Shankar, Percy Liang, Yair Carmon, and Ludwig Schmidt. Accuracy on the line: on the strong correla- tion between out-of-distribution and in-distribution generalization. In Proceedings of the 38th International Conference on Machine Learning, 2021.
Seth Nabarro, Stoil Ganev, Adrià Garriga-Alnsso, Vincent Fortuin, Mark van der Wilk, and Lau- rence Aitchison. Data augmentation in bayesian neural networks and the cold posterior ef- fect. In James Cussens and Kun Zhang (eds.), Uncertainty in Artificial Intelligence, Proceed- ings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence, UAI 2022, 1-5 August 2022, Eindhoven, The Netherlands, volume 180 of Proceedings of Machine Learning Re- search, pp. 1434-1444. PMLR, 2022. URL https:/ httpp/ppopcce...../e/ mlr press/v18C/ nabarro22a html.
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated proba- bilities using bayesian binning. In Proceedings of the AAAI conference on artificial intelligence, 2015. 18
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Andrew Y Ng, et al. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011.
Lorenzo Noci, Kevin Roth, Gregor Bachmann, Sebastian Nowozin, and Thomas Hofmann. Dis- entangling the roles of curation, data- autaaammgmettioonn and the prior in the cold posterior effect. In Advances in neural information processing systems, 2021. 8
Luis A Ortega, Rafael Cabañas, and Andres Masegosa. Diversity and generalization in neural net- work ensembles. In International Conference on Artificial Intelligence and Statistics, 2022.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Advances in neural information process- ing systems, 2019.
Sam Patterson and Yee Whye Teh. Stochastic gradient riemannian langevin dynamics on the proba- bility simplex. In Advances in neural information processing systems, 2013. 6
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.059411 2017. 25
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classi- fiers generalize to imagenet? In Proceedings of the 36th International Conference on Machine Learning. 2019. 2,
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951. 6
Frank Rosenblatt. The perceptron: a probabilistic model for information storage and organization in the brain. Psychological review, 65(6))386, 1958. 2
Hanie Sedghi, Vineet Gupta, and Philip M Long. The singular values of convolutional layers. In International Conference Learning Representations, 2019. 9
Saurabh Singh and Shankar Krishnan. Filter response normalization layer: Eliminating batch de- pendence in the training of deep neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2020. 5, 25
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness to natural distribution shifts in image classification. In Advances in Neural Information Processing Systems, 2020. 7
Yee Whye Teh, Alexandre Thiéry, and Sebastian J Vollmer. Consistency and fluctuations for stochas-
tic gradient langevin dynamics. Journal of Machine Learning Research, 17(7), 2016. 6
14

---

Published as a conference paper at ICLR 2025
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning. 2011. 1,3, 6, 23
Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good the bayes posterior in deep neural networks really? In Proceedings of the 37th International Conference on Machine Learning. 2020. 20. 25
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Proceedings of the 37th International Conference Machine Learning, 2020.
Danny Wood, Tingting Mu, Andrew Webb, Henry WJ Reeve, Mikel Lujan, and Gavin Brown. A unified theory of diversity in ensemble learning. Journal of Machine Learning Research, 24 (359):1-49, 2023. 9, 18
Blake Woodworth, Suriya Gunasekar, Jason D Lec, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pp. 3635-33330 1673. PMLR, 2020. 3
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction large-scale image dataset using deep learning with humans in the loop. arXiv preprint 1506.
2015. 7
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small nonlinearities in activation functions create bad local minima in neural networks. In International Conference on Learning Representations, 2019.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient MCMC for bayesian deep learning. In International Conference on Learning Representations, 2020. 2, 6, 9, 22
15

---

Published as a conference paper at ICLR 2025
A PROOFS
A.1 PROOF OF LEMMA 3.1
Proof. The derivative of FF with respect to Wj for every j I.... can be decomposed
OF (W 1y.... W.) W 1:j-1
OF( W
j+lie
(18)
ow,
OW 1:c
Substituting this in Eq. 12, we get
dW (L)
aF(Wiss (e)
W 1-j-1
W
(19)
dt
aw l:c
j+1ie
Therefore, assuming that W.
Wet1:1 I,
dWi:e(t)
dw, W (t)
W 11)--(tt)
(20)
dt
dt
j=1
OF(WI:e 1se (t))
W1s-i(t)W1s-100
W (t) Wt+1.(t)..
(21)
OW Ise
j=1
By taking the vectorization on both sides,
(WW1.s (1c(()
vec
(W (t) (t+l:e(t) & (1s-1((Wigg11()) vec @FF( 1.:e (t))
Σ
(22)
aW 1se
F=1
Px(O VF(((I)).
(23)
A.2 PROOF OF THEOREM 3.2
Proof. For
W, = W Wft1e 3 Wiy-1WT WT 1:j-1
(24)
in Lemma 3.1, the precondition can be described
Px(0) = Σ W3
WI = Wetie =
(25)
j=1
Since Wj is positive semi-definite and symmetric, the singular values of W2 is the same as the absolute eigenvalues of W,. When Way = Deg by by the singular value decomposition,
W ltt W f+1le Wi W 1:j-1
(26)
= (Vi+1ee DJ Dj+1ie 1.e) @ DTg D Dish----s-
(27)
(VJ+1e @ 11 ( -ie @ D 11) (Vi+1.e U1g-1)
(28)
= OAOT
(29)
Therefore, the eigenvalues of W3 are
or(W (WIiee) Jo (W1g-1)2 for =,...,, n, and r' 1, 1,... n,
(30)
where o, is the r-th singular value. The min-max theorem for singular values yields
aW,) (o,(WW ( 11 < m (W;) 2
(31)
(
16

---

Published as a conference paper at ICLR 2025
Using this value, we derive the upper bound from the fact that the operator 12-norm of matrix I. such that Wite =
the same as the maximum singular value. For X = (vec (w 1:c 1=1
Wi W2 We and (t)((t)) (U(((t)) + such that E[|ls] [] the distance between the two adjacent time steps is bounded as
X (1 1) x
(() (t) + B,
(32)
W2Vū(x (t)) + ()
(33)
|m3ll2 v0(x + (,
(34)
1=1
Se EII () (| + (1=000//) + B.5
(35)
Note that & is still the zero-mean Gaussian because we set the noise corresponding W....., - W Wit....., We except for W, zero. Once we take the all of layers and expec- tation on both sides,
E [
+1) 1) X(III]
7
E
O (W (IVU((()))l + +1 + 1 |B
(36)
1=
j=1
S Lel m (e-1) (Ih L8) + eLC.
(37)
B SUPPLEMENTARY RESULTS
B.1 EVALUATION METRICS
Let Zm,i E R represent the categorical logits predicted by the m posterior sample 0, for the 5th data point. The final ensemble prediction, which approximates the BMA integral of the predictive distribution, for the 5th data point is given by:
M Σ
P
o (zm,i))
(38)
M
m=1
where denotes the softmax function, mapping categorical logits to probabilities. Using P, and Em,, well as the ground truth label y; for the 7th data point, calculate the following evaluation metrics for a given set of N data points.
Classification error (ERR). The classification error, often referred to as 0-1 loss, is the primary metric used to evaluate the performance of a classification model:
ERR
Σ
Ys + arg max p.
(k)
(39)
N
where [-] denotes the Iverson bracket.
Negative log-likelihood (NLL). The negative log- log-likelihood of a categorical distribution, commonly known as cross-entropy loss, serves as the key metric for assessing classification model performance in Bayesian literature:
NLL
log p;
(1.)
(40)
N
17

---

Published as a conference paper at ICLR 2025
Table 6: Supplementary results for CIFAR-IO-C. It summarizes the classification error (ERR), negative log-likelihooke (NLL), and expected calibration error (ECE) averaged over 19 corruption types for each intensity level. For a comprehensive overview of the results, we direct readers to Fig. 5, which illustrates the box-and-whisker plots.
Intensity level
Metric
Method
AVG
5
SGLD
0.3
0.20660:079
0.253++0091
0.294
0.341
0.410++.158
PSGLD
0.317
0.21610.076
0.26620.081
0.309
0.360a0.130
0.433
ERR ())
SGHMC
0.29440.140
0..942
0.242a0.091
0.284
0.335
0.4114
SGNHT
0.296s0.136
0.1 195.007777
a0.044
0.28620.066
0.339
0.421
PX-SGHMC (ours)
0.275+0.126
0.180
0.224
0.264
0.315
0.39440.134
SGLD
0.894+0.307
0.623+0.210
0. 0.74880.235
0,863e0310
1.006=0410
1.229e0475
PSGLD
0.94550.383
0.659+0.203
0.792++.217
0.912+0278
1.066-0.374
1 1.29%=0.447
(L)
SGHMC
0.877+04400
0.585
0.715
0.837+0317
0.994-0411
1 1.253+0.500
SGNHT
0.878+0.303
0.585
0.842 794
a.995
1.252+0040
PX-SGHMC (ours)
0.826e0.571
0.550
0.673+0.220
0. 2.786
0.934-0247
1.199=0.400
SGLD
0.082a0.061
0.070
0.06620.024
2.074
0.094_0.071
0. 107...0099
PSGLD
0.07440.047
0.074
0.060
0.059
0.07720.056
0. 10000.077
ECE (L)
SGHMC
0.07620.022
0.064
0.058an.024
0.064
100022
0.1 120.0%
SGNHT
0.072
0.064
0.057
0.060
0.073:0002
0. 105.0.004
PX-SGHMC (ours)
0.066.0.033
0.067a0.11
0.059a0.021
0.059.0019
0.063.00031
0.084:00019
Ensemble ambiguity (AMB). The generalized ambiguity decomposition for the cross-entropy loss given by (Wood et al., 2023):
M
AMB-
log σ (zm,,) (0)
log
Zm,
(41)
M
N
N
M
average loss
ensemble loss
Notably, logit ensembling in the ensemble loss term is essentially equivalent to computing a nor- malized geometric mean of the categorical probabilities. See Wood et al. (2023) for more details.
Expected calibration error (ECE). The expected calibration error with binning is a widely used metric for assessing the calibration of categorical predictions (Naeini et al., 2015):
ECE = |Byl |ace( B₁) con[(B,)]
(42)
N
where Bj represents the jth bin that includes [By] data points with prediction confidence max; p!a)) falling within the interval (G - 1)/J.j/J], Here, acc(B,) indicates the classification accuracy, while (B) refers to the average prediction confidence within the th bin.
B.2 ROBUSTNESS TO COMMON CORRUPTION
Table 6 presents the classification error and uncertainty metrics, including negative log-likelihood and expected calibration еттоr (Nacini et al., 2015), for each level of corruption intensity. Our PX- SGHMC consistently outperforms all baseline methods across all metries, with the number of bins for computing expected calibration етror set to 15. Notably, PX-SGHMC shows lower calibration error with increasing intensity levels, demonstrating enhanced robustness more severely corrupted inputs. Fig. 5 further presents box-and plots illustrating metrics across 19 corruption types for five intensity levels. Overall, PX-SGHMC exhibits better calibration than the other methods.
B.3 OUT-OF-DISTRIBUTION DETECTION
To obtain the ROC curve and associated metrics (i.e., AUROC and TNR at TPR of 95% and 99%, as shown in Table 2), we used 1,000 in-distribution (ID) examples as positives and 1,000 out-of- distribution (OOD) examples as negatives. We manually balanced the number of examples, because
18

---

Published as a conference paper at ICLR 2025
SGLD
PSGLD
SGHMC
SGNHT
PX-SGHMC
0.60
2.00
0.30
1.50
0.40
od
ECE
1.00
0.10
0.50
Intensity level
Intensity level
Intensity level
Figure Supplementary box-and-whiske plots for CIFAR-IO-C. It illustrates the classification error (ERR), negative log-likelihood (NLL), and expected calibration error (ECE) across 19 corrup- tion types for five intensity levels.
SGHMC
PX-SGHMC
ROC curve
0.5
0.5
1.0
Rel, frequency
0.4
0620
1,001
Rel, Rel. frequency
0.4
0.575
1,885
0.3
0.3
TPR
0.5
0.2
0.2
SGHMC
0.1
0.1
PX-SGHMC
0.0
0.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
Predictive entropy
Predictive entropy
FPR
(a) CIFAR-I0 ID and SVHIN OOD.
SGHMC
PX-SGHMC
ROC curve
0.5
Rel. frequency
0.620
1.275
Rel. R frequency
0.4
0.575
11361
0.3
TPR
0.5
0.2
SGHMC
0.1
0.1
ID
OOD
ID
00D
PX-SGHMC
0.0
0.0
0.0
1.0
2.0
0.5
1.0
1.5
2.0
0.0
1.0
Predictive entropy
Predictive entropy
FPR
(b) CIFAR-10 as ID and LSUN as OOD.
Figure 6: Supplementary plots for -ittri--iittitutionn detection. Histograms of predictive tropy for ID and OOD datasets, along with the receiver operating characteristic (ROC) curve mea- suring the separability between ID and OOD.
the ROC curve becomes less reliable when there is an imbalance between positive and negative examples. To the perceptual differences between ID (CIFAR-10) and OOD (SVHN and LSUN) images, please refer to Fig. 10.
In Fig. 6, histograms show how our PX-SGHMC more effectively assigns low predictive entropy to ID inputs and high entropy to OOD inputs compared to the baseline (with SGHMC as a represen- tative), while ROC curves assess the separability between the ID and OOD histograms. It clearly shows that PX-SGHMC more robust to OOD inputs, offering more reliable predictions when encountering OOD inputs in real-world scenarios.
19

---

Published as a conference paper at ICLR 2025
B.4 ADDITIONAL RESULTS WITH COLD POSTERIOR
To obtain valid posterior samples from D)) the temperature should be one in Langevin dynamics and its practical discretized implementations (i.e., T = in Eqs. 5 and 7). However, many works in the Bayesian deep learning literature have, in practice, considered using T" < 1, which is called cold posterior (Wenzel al., 2020). Therefore, we further present comparative results between SGHMC and PX-SGHMC using the cold posterior.
Fig. 7 presents the evaluation results on CIFAR-10, 10, CIFAR- 10.1, CIFAR- 10.2, STL, and CIFAR- 10-C, as in Table 1. is clear that our PX-SGHMC consistently outperforms the SGHMC baseline across all datasets and temperature values considered. These results suggest that our EP strategy functions orthogonally to the modification of the target posterior through posterior tempering.
B.5 ADDITIONAL RESULTS ON ACCEPTANCE PROBABILITY OF GGMC
While SGHMC omits the Metropolis- Hastings correction, Garriga-Alonso & Fortuin (2021) re- cently argued that SGHMC effectively has acceptance probability of zero, the backward tra- jectory is not realizable in practice due to discretization. Expanding on this, they revisited Gradient- Guided Monte Carlo (GGMC; Horowitz, 1991), method that generalizes HMC and SGLD while ensuring a positive acceptance probability. Building on their insights, we further investigate how the proposed EP influences the acceptance probability within the GGMC framework.
Fig. 8 illustrates the following for both EP and SP:
As the peak step size 6o increases, GGMC produces more diverse samples, as indicated by the unnormalized Euclidean distances in (a). This diversity ultimately contributes to improved ensemble predictions, as shown in (c).
However, as the peak step size increases, the simulation error also grows. step size of co = 3 10 which yields good performance in practical applications, results in relatively low acceptance probability of around 25%.
Notably, the proposed EP enhances both exploration and simulation. The implicit step size scaling introduced by EP facilitates improved exploration without compromising the acceptance probabil- ity due to discretization indeed, it may even enhance it. This indicates that the superior exploration capability achieved by EP cannot be replicated solely by increasing the step size.
B.6 COST ANALYSIS AND LOW-RANK VARIANT
We have compiled system logs comparing PX-SGHMC with SGHMC in Table 7. Notably, the logs show that PX-SGHMC exhibits no significant differences both in sampling speed and memory consumption compared to SGHMC in practice.
Moreover, we further implemented a low-rank variant our approach which reduces memory over- head. Specifically, we used low-rank plus diagonal approach to construct the expanded matrices, i.e, we compose new expanded matrix
D + L, L2 for diagonal matrix D E RRxx and the low-rank matrices L1 L2 E Rr with (pp + 2pr). In Table 8, we set to p/8 8 and p/4 4 for the c = 1 setup, resulting in 303,610 and 330,874 parameters during the sampling procedure, respectively. Even with the expanded matrices of the low-rank plus diagonal form, PX-SGHMC continues to outperform SGHMC, highlighting a clear direction for effectively addressing the increased parameter count of our method. Therefore, the design of expanded parameters is left to users with limited memory resources.
p. which reduces the memory from O(p²) to
B.7 ABLATION RESULTS ON BURN-IN PERIOD
In our main experiments, none of the SGMCMC methods utilize a separate burn-in phase (unlike Izmailov et al. (2021)); in other words, even samples from the first cycle are included in the BMA computation. To further analyze convergence, we conducted an additional ablation study on burn-in, following Izmailov et al. (2021), by examining the performance of BMA estimates and individual samples as a function of burn-in length. Fig. 9, the x-axis represents the number burn-in samples, i.e., the length of the burn-in period, while the y-axis shows the individual performance
20

---

Published as a conference paper at ICLR 2025
0.18
0.60
ERRiy
0.16
M
SGHMC PX-SGHMC
NL LLLVV
0.50
T
SGHMC PX-SGHMC
AMB-VVaild
0.14
SGHMC
0.40
0.20
10
10
10
-1
10
10
10
10-2
10
Posteriar temperature
temperature
Posterior temperature
(a) CIFAR-10
0.90
SGHMC PX-SGHMC
I-
SGHMC PX-SGHMC
ERRi@Valld
NI Valid
AMBiVVaild
0.80
0.60
0.25
0.70
0.40
0.23
PX-SGHMC
10
10-3
10a
10-1
10-2
10
Posterior temperature
Posterior temperature
Posterior temperature
(b) CIFAR-IO.1
0.80
0.33
SGHMC PX SGHMC
SGHMC
ERRi@Valid
NL Validl
AMB-availd
PX SGHMC
0.30
0.60
0.90
0.28
PX-SGHMC
0.80
10a
10
10
10-2
10-3
10a
10-2
10
Posterior temperature
Posterior temperature
Posterior temperature
(c) CIFAR-10.2
SGHMC PX-SGHMC
1.10
SGHMC
0.80
ERRidValid
47
PX-SGHMC
AMBaValid
1.00
0.60
0.30
0.90
0.40
-B
SGHMC
PX-SGHMC
0.28
0.80
10a
10
10
10
10
10
10-2
10-3
10a
10
10
10
Posterior temperature
Posterior temperature
Posterior temperature
STI.
0.33
SGHMC PX-SGHMC
1.00
SGHMC PX-SGHMC
ERRi@Valid
AMBaValld
0.90
0.60
0.28
-
SGHMC
0.80
0.25
PX-SGHMC
109
-1
10
10
10a
10-2
10-a
10a
10-1
10-2
10
Posterior temperature
Posterior temperature
Posterior temperature
(e) CIFAR-10-C
Figure 7: Additional results with varying posterior temperature. It depicts evaluation results for SGHMC and PX-SGHMC with cold posterior.
of the 50 samples (in average) obtained after the burn-in period (denoted "IND (1)") and the ensemble performance (denoted as "BMA (50)"). For reference, we also plotted the performance using 100 samples without burn-in period, previously reported in the main text (denoted "BMA (100)"). clearly demonstrates the the enhanced training dynamics introduced by our EP enable higher BMA performance with shorter burn-in period. In other words, when collecting the same
21

---

Published as a conference paper at ICLR 2025
EP (éo 3.0e-07)
EP (Eo = 3.0e-06)
EP (da 3.0e-05)
EP (co 3.0e-04)
SP (4o 3.0e-07)
SP (Eo = 3. 3.0e-06)
SP (da = 3.0ee-05
SP (co 3.0e-04)
Exploration
Simulation
Prediction
1.35
1.00
mlam. m..00+177
0.75
0.75
0.90
0.50
0.50
0.45
0.25
0.25
0.00
0.00
0.00
40
70
100
10
40
100
10
100
Cycle index
Cycle index
Cycle index
Figure 8: Additional results using GGMC. Trace plots illustrating (a) Exploration (higher is better), which measures the normalized Euclidean distance from the previous sample; (b) Simu- lation (higher is better), which represents the acceptance ratio of the sample; and (c) Prediction (higher is better), which assesses the final performance of the ensemble prediction.
Table 7: Cost analysis of SGHMC and PX-SGHMC. It summarizes the costs involved in the sampling procedure for SGHMC and PX-SGHMC with c = = and c = = 2. "Space (in theory)" refers to the number of parameters during the sampling process, while "Space (in practice)" represents the actual GPU memory allocated in our experimental setup using single RTX A6000. "Time (in practice)" denotes the wall-clock time for each cycle, consisting of 5,000 steps.
Method
Space (in theory)
Space (in practice)
Time (in practice)
SGHMC PX-SGHMC (c = d = 1)
274.042
1818 MB
61 sec/cycle
383,098
1838 MB
64 sec/cycle
PX-SGHMC (c 2)
492,154
1852 MB
73 sec/cycle
50 samples, SGHMC requires a much longer burn-in period achieve performance comparable to that of PX-SGHMC.
B.8 COMPARATIVE RESULTS WITH META-LEARNING APPROACH
We further provide comparative results with the Learning to Explore (L2E; Kim al., 2024) method. Since they also adopted the experimental setup of Izmailov et al. (2021) using the R20-FRN-SMe wwish architecture, the results are largely comparable when key experimental configurations such as data augmentation, the number of steps per cycle, and the number of posterior samples are aligned.
Using the official implementation of Kim et al. (2024)), we ran the L2E method using the same setup as in Table 3 of main text, i.e., with data augmentation, cold posterior with T 0.01, 5000 steps per cycle, and 100 posterior samples. Table summarizes the results. Notably, our PX-SGHMC, which applies vanilla SGHMC sampler the expanded parameterization, yield competitive results compared to L2E, despite the latter relying on a more resource-intensive meta-learned sampler.
B.9 ABLATION RESULTS ON STEP SIZE SCHEDULE
In our main experiments, the cyclical step size schedule is used for all SGMCMC methods. Since the cyclical step size schedule itself is designed enhance exploration SGMCMC and improve sam- ple diversity (Zhang et al., 2020), we conducted ablation experiments on SP/EP parameterizations and constant/cyclical step size schedules to more clearly isolate the contribution EP.
Table 10 summarizes the performance (ERR, NLL; lower is better) and functional diversity (AMB; higher preferred) of SGHMC and PX-SGHMC under both constant and cyclical step size sched- ules. Based on "SGHMC w/ constant schedule," the results clearly show that our expanded param- eterization (EP) contributes more significantly to performance improvements than the adoption of the cyclical schedule (CS).
22

---

Published as a conference paper at ICLR 2025
Table Low-rank variant of PX-SGHMC. It summarizes the evaluation results for low-rank variants of PX-SGHMC with c = d = 1. "# Params (sampling)" indicates the number of parameters during the sampling process, while "# Params (inferece)" refers to the number of parameters after merging expanded matrices into the base matrix.
Memory costs
Evaluation metrics
Method
# Params (sampling)
# Params (inference)
ERR (4)
NLL (1)
AMB (f)
SGHMC
274,042 (x1.00)
274,042 (x1.00)
0.131
0.421
0.183
PX-SGIIMC (r p/8)
303,610 (x1.1))
274,042 (x1.00)
0.126
0.401
0.210
PX-SGIMMC (r p/
330,874 (x1.21)
274,042 (x1.00)
0.122
0.385
0.215
PX-SGHMC (r p)
383,098 (x1.40)
274,042 (x1.00)
0.123
0.396
0.242
SGHMC BMA (100)
SGHMC (1)
SGHMC BMA (50)
PX-SGHMC BMA (100)
PX-SGHMC IND (1)
PX-SGHMC BMA (50)
0.20
0.60
ERR
0.18
NLI
0.50
0.15
0.12
0.40
10
20
30
40
50
0
10
20
30
40
50
# burn-in samples
# burn-in samples
Figure 9: Convergence of SGHMC and PX-SGHMC. Performance comparison of individual pos- terior samples (IND) and Bayesian model averaging (BMA) ensembles with 50 samples from each SGHMC and PX-SGHMC chain, evaluated as function of burn-in length. "BMA (100)" represents a burn-in length of 0 with a BMA ensemble of 100 samples, as used in the main experiments of this paper. "IND (1)" and "BMA (50)" refer to average individual sample performance and ensemble performance of 50 samples, respectively.
C ALGORITHMS
In this section, we outline the practical implementation of the SGMCMC algorithms used in our experiments: Stochastic Gradient Langevin Dynamics (SGLD; Welling & Teh, 2011), Stochastic Gradient Hamiltonian Monte Carlo (SGHMC; Chen et al., 2014), Stochastic Gradient Nosé-Hoover Thermostat (SGNHT; Ding al., 2014), and preconditioned SGLD (pSGLD; Li et al., 2016). Addi- tionally, we experimented with Stochastic Gradient Riemann Hamiltonian Monte Carlo (SGRHMC; Ma et al., 2015) using diagonal empirical Fisher and RMSProp estimates for the preconditioner. However, within the hyperparameter range explored, it demonstrated significantly lower perfor- mance than SGLD, leading us to exclude it from further experiments.
First, we present the hyperparameters that are consistent across all algorithms:
M: the number of sampling cycles, representing the total number posterior samples.
T: the number of updates per cycle, resulting in x T total updates.
4e the step size at time step t. It can follow 'Cyelical' with peak learning rate of defined as a 60 cos T
mod(t,T)) +1] or remain 'Constant', i.e., Vt C = co-
x2: the variance of the zero-mean Gaussian prior over the neural network parameters, where p(0) = (0;
Bm the mini-batch at the hh time step of the mh cycle, with size of [B].
23

---

Published as a conference paper at ICLR 2025
Table 9: Comparative results with data augmentation. Evaluation results on C10 (CIFAR-IO), C100 (CIFAR-100) and TIN (TinylmageNett with data augmentation. In this context, we manually set the posterior temperature 0.01 to account for the increased data resulting from augmentation.
ERR (J)
NLL (1)
Method
C10
C100
TIN
C10
C100
TIN
SGHMC PX-SGHME (ours) 1.2H (Kim et al., 2024)
0.07 L40,001
1.3
0.538
0.223
1.13820.008
2.25
0.06940.001
0.290
0.498
0.217
1.030a0.011
2.08920.008
0.071 +3.011
0.268+0.002
0.488-0
0.232++0.00
0.980-0.006
2.062x0.011
Table 10: Additional results with a constant step size. Classification emror (ERR), negative log- likelihood (NLL), and ensemble ambiguity (AMB) on the test split of CIFAR-10, comparing the use of the cyclical schedule (CS) and our proposed expanded parameterization (EP).
Components
Evaluation metrics
Label
CS
EP
ERR (L)
NLL (4)
AMB (t)
SGHMC w/ constant schedule
2.
0.441
0.20920.001
SGHMC w/ cyclical schedule
0. 0..00000002
0.422+0.005
0.186=0.004
PX-SGHMC w/ constant schedule
0.115
0.379
0.232+0.003
PX-SGHMC w/ cyclical schedule
0.121
0.388
0.242:0.002
Next, briefly summarize the additional components introduced in each method. For more in- depth exploration of SGMCMC methods, we refer readers to Ma et al. (2015) and references therein, which offer concise summary from the perspective of stochastic differential equations.
PSGLD introduces adaptive preconditioners from optimization, e.g., RMSProp.
SGHMC introduces the friction matrix to mitigate the noise from mini-batch gradients. While the friction term is originally matrix. it is often implemented as scalar value in practice (C = CT). Also, the gradient noise estimate is set to zero (B = 0), and the mass matrix is defined as the identity matrix (M in practical implementations.
SGNHT introduces an auxiliary thermostat variable & to maintain thermal equilibrium of the system. In practical implementations, it can be interpreted making the friction term used for momentum decay in SGHMC learnable. Intuitively, when the mean kinetic energy exceeds 1/2, & increases, leading to greater friction on the momentum.
Algorithms 1, and 4 summarize our practical implementations of SGLD, PSGLD, SGHMC, and SGNHT, while Appendix provides a detailed hyperparameter setup for each method used in our experiments.
D EXPERIMENTAL DETAILS
D.1 SOFTWARE AND HARDWARE
We built our experimental code using JAX (Bradbury et al., 2018), which is licensed under Apache- 2.0.² All experiments conducted machines equipped with an RTX 2080, RTX 3090, or RTX A6000. The code is available at https://gituub.com/ss- www.oo/////iiiigggg/ggmmmm
D.2 IMAGE CLASSIFICATION ON CIFAR
Dataset. CIFAR-10 10 (Krizhevsky et al., 2009) is dataset comprising 32 x 32 x 3 images classified into 10 categories. We utilized 40,960 training examples and 9,040 validation examples based on the HMC settings from Izmailov al. (2021), with the final evaluation conducted 10,000 test
24

---

Published as a conference paper at ICLR 2025
Algorithm 1:
Practical implementation of SGLD
Require: the hyperparameters mentioned above.
Ensure: a set of posterior samples θ <- ()-
Initialize position 0,,7 from scratch or pre-set values.
for m = 1,2,..., MM do
Initialize m, m-TT for t 1,2..... .T do
2 No, I).
Im.t Veú (0; (...))m0000m...
0 m,, m 1 - mmm,t + mm.
end
Collect + + U (0m.T}.
end return
Algorithm 2: Practical implementation pSGLD
Require: smoothing factor B and the hyperparameters mentioned above.
Ensure: a set of posterior samples θ t ().
Initialize position o,,r from scratch or pre-set values.
for 1,2,..., MM do
Initialize 0, & mm1IT
for 1,2, ................................ do
2m,1 I).
Im.t VoŻ Bm,,BB,) 10m..0.00m..1111
Pm,t 3Vm,t-1 + B)g,,t
O Om,t 0 6.gm, @ (mmit + + 2m @ (w
01/2
+ 8)1//2
end
Collect θ + U U T..
end return
+
examples. For a comprehensive evaluation, we also employed additional datasets, including STL- 10 (Coates et al., 201 1). CIFAR-10. (Recht et al., 2019), CIFAR-10.2 (Lu et al., 2020), and CIFAR- 10-C (Hendrycks & Dietterich, 2019). We refer readers to the corresponding papers for more details about each dataset. In our main experiments detailed in Section 5.2.1, we did not apply any data augmentations.
Network. We conducted our experiments using R20-FRN-Swish, HMC checkpoints provided by (Izmailov et al., 2021) were publicly available. The model is a modified version of the 20-layer residual network with projection shortcuts (He al., 2016), incorporating Filter Response Normal- ization (FRN; Singh & Krishnan, 2020) the normalization layer and Swish (Hendrycks & Gim- pel, 2016; Elfwing al., 2018: Ramachandran et al., 2017) the activation function. Substituting Batch Normalization (BN; Toffe & Szegedy, 2015) with FRN removes the mini-batch dependencies between training examples, while using Swish results smoother posterior surface, facilitating clearer Bayesian interpretation (Wenzel et al., 2020; Izmailov al., 2021).
Running from seratch. In the first setting of the CIFAR experiments, SGMCMC is executed from random initialization to collect posterior samples in a 'from scratch' manner. This represents the most basic setup, requiring SGMCMC methods to quickly reach low posterior energy regions while gathering functionally diverse posterior samples. Starting from the He normal initialization (He et al., 2015), SGMCMC methods were allocated 5,000 steps per sampling cycle (approximately epochs) generate total of 100 samples. Table 11 provides detailed hyperparameters.
25

---

Collect + U (0 end Published as a conference paper at ICLR 2025
Algorithm 3: Practical implementation of SGHMC
Require: constant friction value ? and the hyperparameters mentioned above.
Ensure: a set of posterior samples + ()-
Initialize position 0,,r from scratch or pre-set values.
Initialize T0, f- 0.
for m 1, ..2.....4 .M do
Initialize m0, Tm.u) .) + (0m-1,r,0m O mmr)
for 1,2...... 2......... T, do
2m,t I).
Im.t Voū 10.00.0mmmm..
Mm, (1 -m))mm.-111 + cmm, + 22
0 m,t 0 m mmt.
end
return
Algorithm 4: Practical implementation of SGNHT
Require: initial thermostat value & and the hyperparameters mentioned above. Ensure: a of posterior samples + (.
Initialize position 6, from scratch or pre-set values. Initialize To,7 &- and forr + 5. .M do
for m Initialize m6,, Fm.0, Sm,o)) f- (fm-1,7, 0m-1,77 0m-111,5m-117) for t .,,,..., do
Zm,t 1).
Im.t Voū 1000))0==0m.-111
Pm, = (1 - 1m, + e.gm,t + 22
θ 0m, 0 1 CmTm.
&r f--1 tel 11-1 1), where n is the dimension of r:-1-
end
Collect θ + U (.
end return
+
Running from HMC burn-in. The second setting of the CIFAR experiments involves running SGMCMC from HMC burn-in initialization to analyze the dynamics of SGMCMC methods in com- parison with the gold-standard HMC. Specifically, adopted the 50th HMC checkpoint provided by Izmailov et al. (2021), they designated 50 as the burn-in iteration. To minimize mini-batch noise as much as possible within our computational constraints, a large mini-batch size 4,096 was employed, aligning with HMC's use of full data to compute gradients. Consequently, using the 50th HMC sample as the initial position, the both SGHMC and PX-SGHMC methods were allocated 70,248 steps per sampling cycle (approximately 7025 epochs), matching the 70,248 leapfrog steps of HMC, to generate a total 10 samples. We also use the constant step size of e = 10-b and prior variance of 2² 0.2, in line with HMC.
D.3 IMAGE CLASSIFICATION WITH DATA AUGMENTATION
Dataset. For CIFAR-100, used 40,960 training examples and 9,040 validation examples, sistent with CIFAR-10. For TinylmageNet, we employed 81,920 training examples and 18,080 val- idation examples, resizing images from 64 x 64 x 32 32 x 3. In Section 5.2.2, we employed
26

---

Published as a conference paper at ICLR 2025
Table I1: Hyperparameters for CIFAR. summarizes the hyperparameters for each method used in our main evaluation results on the CIFAR experiments (i.e.. Tables and 2). If a hyperaparameter was manually set without tuning, it indicated with dash in the 'Search Space' column.
Method
Hyperparameter
Value
Search Space
Notation
SGLD
Peak Step Size
1 x 10-5
{3 x 10 1 10
k16
Jk-4
60 σ²
Prior Variance
0.2
{0.5, 0.2,0 0.05, 0.02, 0.01)
PSGLD
Peak Step Size
3 x 10
x 10 .1= 10
60
Prior Variance
0.2
[0.5, 0.2, 0.1, 0.05, 0.02. 0.01)
x²
Smoothing Factor
0.99
{0.9, 0.99, 0.999}
8
SGHMC
Friction
100
(1, 100, 1000) (3 x 10 1 K, {0.5, 0.2,0.1,,00000 0.02, 0.01)
Peak Step Size
3 10-4
10 10-
60
Prior Variance
0.05
x²
PX-SGHMC
Friction for V
100
{1, 100, 1000)
Friction for
Peak Step Size
1 x 10-4
(3 10- 10 K=
60
Prior Variance
0.02
(0.5, 0.1,0.0. 0.02, 0.01}
σ²
Shared
Batch Size
256 Cyclical 5000 x 100
[B]
Step Size Schedule
Total Updates
T
Total Samples
100
M
standard data augmentation policy that includes random cropping of 32 pixels with padding of pixels and random horizontal flipping.
Network. We employ R20-FRN-Swi Swish, consistent with the CIFAR-10 experiments.
Running from scratch. We utilize the same hyperparameters as the CIFAR- 10 experiments.
D.4 DATASET DETAILS
Fig. 10 visualizes example images from datasets considered:
CIFAR-10 (unknown license); https:/wwww.cs.orrrtoocll//ttttt http://tt.....kkiiiiicccaaa.amm. kriz/cifar.html
CIFAR-10.1 under the MIT license; http://ghhhu..cmmmmooodyyyycc//a
CIFAR-10.2 (unknown license); http://gthhh..ommmmnndeyyaaaeeen
STL (unknown license);
SVHN (unknown license); https//giinbbccmm/cccooorseereee
LSUN (unknown license); http://gthbbocom//scccootttsssrr
CIFAR- (unknown license); https//wwww.csslooomoo.d/////// http://wwoooooooooooni/i//iaaaaa
TinyImageNet (unknown license); https://www..aagglccco/////ttttt
E CONCEPTUAL ILLUSTRATION FOR EFFECT OF PARAMETER EXPANSION
The main motivation for our method is the well-known effects deep linear neural networks, which can be interpreted
implicit acceleration of training induced by gradient updates in such net- works (Arora et al., 2018). We conceptually illustrated this in Fig. 11.
At a high level, the preconditioning matrix can be understood evoke an effect akin to adaptive step size scaling, which varies across different components of the parameters. This contrasts with simply increasing the step size, which scales up updates along all axes of parameters by the same factor. Our parameter expansion scales the update along each eigenvector of the preconditioning matrix proportionally to its eigenvalue. This not only amplifies the gradient but also changes its direction.
27

---

Published as a conference paper at ICLR 2025
(a) CIFAR-II
(b) CIFAR-IO.1
(c) CIFAR-10.2
5
2
A
33
2
0
19
19
6
(d) STL
(e) SVHN
LSUN
(g) CIFAR-100
(h) TinylmageNet
Figure 10: Sample images from datasets. It shows the first 16 test images from CIFAR-like datasets, including (a) CIFAR-10, 10, (b) CIFAR-10.1. (c) C1..R-10.2, and (d) STI, as well as out-of- distribution datasets such (e) SVHN and LSUN. Additionally, (g) CIFAR-I00 and (h) Tiny- ImageNet datasets are utilized in experiments involving data augmentation.
The eigenvalues the preconditioner can be mathematically described by the singular values merged weight matrix under certain assumptions, including the initialization of the weight matrices as described in Claim of Arora et al. (2018).
28

---

Published as a conference paper at ICLR 2025
(W1: I: W2 W2)-spcc
W
V.. U/
T/w-U77 47
We
W-space
W-space
VwU
Vw(
W.W. U + VW, UW,
PV W. V..
(a) SP (i.e., SGHMC)
(b) EP (i.e., PX-SGHMC)
Figure 11: Conceptual illustration comparing SP and EP. A simplified illustration of the up- date rules for standard parameterization (SP) and expanded parameterization (EP). The gradient computed in the expanded (W1, W2)-space acts as preconditioned gradient PVwU, where P represents a preconditioner defined in Lemma 3.1, in the original W-space. Compared to the orig- inal gradient U, the preconditioned gradient is expected to enable more efficient updates by better aligning with the geometry of the parameter space.
29


Reviews:
[{"review_id": "DgAh4GPaC8", "reviewer": "Reviewer_GmGR", "summary": "This paper proposes a new SGMCMC method using the idea of parameter expansion to improve the diversity of samples. Specifically, the weight and bias terms in neural networks are reparameterized using a sequence of matrix multiplication. The upper bound on the differences between two consecutive time steps is given. The experiments consider distribution shifts and of-distribution detection on CIFAR datasets.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper studies an important problem in SGMCMC, which is to improve sample diversity.\n- The method is simple and the algorithm is easy to implement in practice. \n- The empirical results show improvement over existing SGMCMC methods.", "weaknesses": "- The main concern I have is the motivation of the proposed method. It is unclear why parameter expansion can solve the sample diversity issue. The paper did not explain the intuition and only explain it through Theorem 3.2. However, the theorem did not provide a convincing motivation either.\n- In Theorem 3.2, c and d have similar effects as the step size. I believe the values of c and d cannot be too large to achieve good convergence. If so, how is it different from using larger step sizes? Also, does Theorem 3.2 include the standard SGMCMC as a special case? In experiments, the authors use c=0,d=0 to refer to the standard SGMCMC, but the formulation in Eq.9 suggests c,d>=1.\n- Another main concern is that the required additional memory and computation will increase significantly. For an m by n W, we now will have c m by m and d n by n matrices, which at least increases O(c+d) costs. \n- In experiments, the proposed method did not compare with SGMCMC methods that aim to improve sample diversity, such as cyclical SGMCMC. \n- Fig 1 seems to show that EP can find modes but cannot converge to the target distribution, as the bottom left region seems to have more density. \n- Is the proposed parameter expansion implemented upon cyclical SGMCMC? This is because results in Fig 3 and 4 suggest so. Also, the authors use the cycle index in several figures. However, it is not explicitly mentioned anywhere in the paper whether PX-SGMCMC is built upon cSGMCMC.\n- The empirical performance does not match typical benchmarks. E.g. err 29% on C100 where the neural network with the same architecture typically can achieve <20%.", "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": "", "extra_scores": {}}, {"review_id": "g3xrl1bIxs", "reviewer": "Reviewer_Awai", "summary": "In this paper, the authors address the issue of limited sample diversity when using Stochastic Gradient Markov Chain Monte Carlo (SGMCMC) for posterior sampling in Bayseian Neural Networks (BNNs). They propose a novel expanded parametrization (EP) of neural networks that relies on a decomposition of the network weights for producing diverse sample weights. They provide theoretical and empirical validation of the proposed approach.", "questions": "* Have the authors tried comparing the proposed method to one of the meta-learning approaches proposed for increasing sample diversity (Gong et al., 2019; Kim et al., 2024)?\n* Do the authors have any ideas on modifying the EP design to address computational overhead? For instance, I’m curious to know if introducing low-rank constraints in these matrices significantly impact the results?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is well structured and the writing is easy to follow. The approach is fairly practical, easy to adopt for general neural network architectures and circumvents the need for high compute resources associated with running multiple chains. Moreover, the theoretical analysis on the gradient flow of EP and the corresponding bound on distance between consecutive samples provides interesting insights on the relation between the induced network “depth”, maximum singular value of the matrices and exploration during sampling. The empirical evaluations clearly demonstrate the strength of the proposed approach for robustness to distribution shifts, ood performance and further validate the theoretical bound.", "weaknesses": "* One of the weaknesses, as noted in the paper, is the additional overhead introduced by the EP. Since the sample diversity depends on the depth of the EP, this could potentially be prohibitive when dealing with large networks. \n* The authors note that one of the reasons why less principled deep ensembles perform better than SGMCMC is better sample diversity. It would be helpful to see how the proposed approach compares to deep ensembles.", "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": "", "extra_scores": {}}, {"review_id": "FR94xbF1ki", "reviewer": "Reviewer_53Wx", "summary": "The authors propose an approach to enhance sample diversity in SGMCMC for Bayesian neural networks.  Their approach is based on decomposing the weight matrices of each BNN layer into a product matrices.  They provide theoretical justification for this approach by showing that the euclidean distance between consecutive samples depends on the number of matrices used in the product construction.  They validate their approach on toy data and several classification tasks.  their results show empirical correspondence to the derived theorem, and show their method helps to provide robustness to distribution shifts, can enhance OOD detection and increase sample diversity.", "questions": "I am confused about the variables introduced in Eq.11 — in particular, it could be helpful to see how they relate to $r$ for the langevin dynamics of the classic unaltered network.\n\nwhile fig.2 shows that the maximum singular value increases while the minimum singular value decreases, how does the condition number change?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "the paper is clearly written and easy to read.  it is impressive that the authors provide theoretical justification for their proposed parameter expanded SGMCMC by showing that the expected euclidean distance between parameters is bounded above by a quantity depending on the depth of the weight matrix parameterization.  \n\nthe authors provide a few variants of EP for different architectures showing that it can be extended beyond linear layers in theory and in practice (as justified by their latter experiments).  they provide extensive experimental/empirical results that show their proposed parameterization can enhance sampling quality in BNNs.  they examine several properties like robustness to distribution shift, out of distribution detection, and show their method also provides improvements even when data augmentation is used.  the number of experiments and their coverage of standard questions readers/researchers might be interested in is both comprehensive and extensive.", "weaknesses": "in section 3, it might be useful to hold the reader’s hand more and expand on lines 168-174 in a way that makes statements about the preconditioning aspect more precise and perhaps easier to visually parse.  for example i am not sure what was meant by the preconditioning producing `extraordinary directions of gradient steps.’\n\n\nthe authors remark about increased parameter counts during training, but it might be helpful to see wallclock times as well as burn in times.  maybe it is possible that drawing single samples is slower, but the enhanced training dynamics lead to a smaller burn in period, which could be another plus for the parameter expanded neural network layers if that difference were sufficiently large.", "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": "", "extra_scores": {}}, {"review_id": "7ctUHrCFa9", "reviewer": "Reviewer_Sd8N", "summary": "In this paper, the authors studied sampling strategy used in Bayesian Neural Networks (BNNs), and proposed a new approach to enhance sampling diversity in stochastic gradient Markov chain Monte Carlo (SGMCMC). which is one of the most commonly used sampling methods in BNNs. Their proposed approach , known as PX-SGMCMC, can better explore the target parameter space by decomposing the weight matrices int o a produce of matrices. While doing this decomposition, their approach does not increase sampling cost compared against the original SGMCMC. They evaluated PX-SGMCMC on image data, and show that  PX-SGMCMC can indeed achieve more diverse samples, leading to outperformance in uncertainty characterization and robustness.", "questions": "1. How to interpret the different colors shown as the distribution modes? It shows a color scale but I don't know what the values indicate.\n2. What do the authors think are the possible ways of reducing the computational costs for PX-SGMCMC? The answer does not have to be concrete, it just will be interesting to know conceptually how to optimize the cost in this case.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In general, this paper is well-structured and the writing is in good quality. It's great that the authors provided detailed mah derivations for their approach, and to my best knowledge, the proof is sound. \nIn term of originality, because sample diversity can significantly influence the methods such as SGMCMC, I believe the idea from this paper is well motivated. And the authors provided comprehensive evaluations on multiple image datasets and it is impressive that the PX-SGMCMC can achieve better sample diversity.", "weaknesses": "The main concern is how this work is distinguished from the prior work on SGMCMC. Because there are many prior methods (such as the cited work in the Related Work section), I would hope to better understand their unique contributions to this topic.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)