Parameter Expanded Stochastic Gradient Markov Chain Monte Carlo

Authors: Hyunsu Kim, Giung Nam, Chulhee Yun, Hongseok Yang, Juho Lee
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: SGMCMC, Bayesian Neural Network, Parameter Expansion

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730608494655

thread_type:
official_review

<thread_object>
creator: reviewer_sd8n

<content>
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.

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. 
In 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.

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.
2. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730643189763

thread_type:
official_review

<thread_object>
creator: reviewer_53wx

<content>
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.

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.  

the 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.’


the 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.

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.

while fig.2 shows that the maximum singular value increases while the minimum singular value decreases, how does the condition number change?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648035909

thread_type:
official_review

<thread_object>
creator: reviewer_awai

<content>
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.

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. 
* 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.

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)?
* 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731123760502

thread_type:
official_review

<thread_object>
creator: reviewer_gmgr

<content>
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.

strengths:
- The paper studies an important problem in SGMCMC, which is to improve sample diversity.
- The method is simple and the algorithm is easy to implement in practice. 
- 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.
- 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.
- 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. 
- In experiments, the proposed method did not compare with SGMCMC methods that aim to improve sample diversity, such as cyclical SGMCMC. 
- 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. 
- 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.
- 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%.

questions:
NA
</content>
</thread_object>

</thread>
<thread>
creation date:
1732000061618

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all the reviewers for their time and effort in evaluating our paper. We are pleased to see a shared recognition of its high quality. In particular, the reviewers noted that our paper is studying an important problem (GmGR), well-structured (Awai, Sd8N), clearly written (53Wx), easy to follow (Awai), easy to read (53Wx), and in good quality (Sd8N). We are also delighted that our idea and methodology received favorable feedback, being described as practical (GmGR, Awai), interesting (Awai), well motivated (Sd8N), impressive (53Wx, Sd8N), and sound (Sd8N). Additionally, we are pleased that the reviewers appreciated the extensiveness (53Wx) and comprehensiveness (53Wx, Sd8N) of our experiments, along with the theoretical aspects of our work (Awai, 53Wx, Sd8N).

We have addressed all of the reviewers' comments and, in accordance with ICLR policy, we have uploaded the revised PDF with changes from the initial version highlighted in blue for clarity. The key updates are as follows:

- (Table 3) We noticed that the pSGLD baseline results were accidentally omitted from Table 3, and we have now included them.
- (Figure 2) We have added the condition number plots to address the concerns raised by Reviewer 53Wx. These plots further clarify that EP enhances exploration without compromising convergence.
- (Appendix B.6 and B.7) All reviewers have raised concerns regarding computational overhead. In response, we have provided practical details on GPU memory allocation and wall-clock time during our experiments. Additionally, following constructive feedback from Reviewer Awai, we proposed representing the expanded matrices in a low-rank plus diagonal form to effectively reduce memory overhead. Furthermore, based on insightful comments from Reviewer 53Wx, we emphasized that introducing EP is more efficient than adding a burn-in period when considering actual wall-clock time.
- (Appendix B.8) Following Reviewer Awai's comment, we have conducted an additional comparison with the recent meta-learning approach by Kim et al. (2024) for improving SGHMC, using their official implementation to ensure reliable experimentation. By simply modifying the parameterization, we achieve competitive performance with their highly costly meta-learned sampler.
- (Appendix B.9) To clearly isolate the contribution of EP from the cyclical step size schedule, as pointed out by Reviewer GmGR, we present additional experimental results using a constant step size schedule. The additional results demonstrate that EP improves SGMCMC in a way that is orthogonal to the cyclical step size schedule.
- (Appendix E) We have attached a conceptual illustration as Figure 11 that provides an intuition why the parameter expansion induces the implicit step size scaling by comparing with the standard SGMCMC gradient update.

We will continue to make further revisions, taking into account any additional concerns raised by the reviewers. Once again, we sincerely appreciate the time and effort the reviewers have dedicated to our work!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732508172635

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

Thank you again for your valuable feedback and thoughtful questions. We have provided responses addressing your concerns and shared an updated version of our manuscript.

As we approach the final days of the discussion period (with approximately two days remaining), we kindly encourage your participation in the discussion. Your additional comments and responses would greatly assist us in refining our work further.

We sincerely appreciate your time and effort in helping us improve our submission.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732532441080

thread_type:
official_comment

<thread_object>
creator: area_chair_ha8d

<content>
title:
Discussions between reviewers and authors

comment:
Time for discussions as author feedback is in. I encourage all the reviewers to reply. You should treat the paper that you're reviewing in the same way as you'd like your submission to be treated :)
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes PX-SGMCMC, where the main idea is to reparameterise the target random variable (in this work the neural network weights) as a product of several matrices and then conducting SGMCMC sampling in that augmented space ("parameter expansion"). Theoretical justification of the approach is provided. Experimentally the proposed sampler performed better in terms of improving sample diversity compared with a number of pre-conditioned SGMCMC methods, without incurring too much higher computational cost.

Reviewers overall commended on the paper's novelty and solid presentations, although they initially also raised questions regarding: the motivation of the approach, how the proposed approach compare with other pre-conditioning techniques (e.g., meta-learned ones), and the memory cost. During rebuttal the authors have improved paper with more explanations, and added experimental results regarding the requested comparisons and memory reduction.

After a brief read I think the paper is worth highlighting due to its novelty and solid development (in both theory and also extensive experiments). It should be of great interest for the Bayesian machine learning community.