Microcanonical Langevin Ensembles: Advancing the Sampling of Bayesian Neural Networks

Authors: Emanuel Sommer, Jakob Robnik, Giorgi Nozadze, Uros Seljak, David Rügamer
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Sampling, Bayesian Neural Networks

Abstract:
Despite recent advances, sampling-based inference for Bayesian Neural Networks (BNNs) remains a significant challenge in probabilistic deep learning. While sampling-based approaches do not require a variational distribution assumption, current state-of-the-art samplers still struggle to navigate the complex and highly multimodal posteriors of BNNs. As a consequence, sampling still requires considerably longer inference times than non-Bayesian methods even for small neural networks, despite recent advances in making software implementations more efficient. Besides the difficulty of finding high-probability regions, the time until samplers provide sufficient exploration of these areas remains unpredictable. To tackle these challenges, we introduce an ensembling approach that leverages strategies from optimization and a recently proposed sampler called Microcanonical Langevin Monte Carlo (MCLMC) for efficient, robust and predictable sampling performance. Compared to approaches based on the state-of-the-art No-U-Turn Sampler, our approach delivers substantial speedups up to an order of magnitude, while maintaining or improving predictive performance and uncertainty quantification across diverse tasks and data modalities. The suggested Microcanonical Langevin Ensembles and modifications to MCLMC additionally enhance the method's predictability in resource requirements, facilitating easier parallelization. All in all, the proposed method offers a promising direction for practical, scalable inference for BNNs.

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

Reviews and discussion:
<thread>
creation date:
1730272118401

thread_type:
official_review

<thread_object>
creator: reviewer_83bk

<content>
summary:
The paper proposes a Microcanonical Langevin Ensemble (MILE) approach which adapts the MCLMC to high-dimensional posteriors common in modern Bayesian deep neural networks (BNN). The paper integrates optimization strategies from deep ensembles to carefully adjust components of MCLMC to make it scalable. Through extensive experiments, the paper shows the superiority of MILE in prediction and UQ on popular benchmarks.

strengths:
I enjoyed reading the paper. Extensive experiments have been conducted where MILE has shown superior performance compared to its competitors. The main strengths of this paper are –
1.	This paper is a significant step towards scalable full Bayesian inference with deep neural networks.
2.	The paper explains the key contributions clearly.
3.	The paper is well written. The authors presented their approach clearly with an intuitive explanation of the key components. 
4.	The authors presented extensive experiments to support their key claim which is the scalability of their approach for sampling high-dimensional Bayesian posteriors.

weaknesses:
The paper can significantly improve if the authors can discuss/expand on the following points –
(1)	The paper discusses careful tunings of the components of MCLMC.  An ablation study is needed to understand which adjustments are more important than others. My guess is the main speed-up is due to the warm starting using the DE. However, this demands a proper study.
(2)	The work lacks theoretical justification. However, this is not a critical point for the paper as the authors have presented extensive experiments in support of their key claims. However, a discussion related to the convergence of the approach can improve the paper.   
(3)	The numerical section seems to be lacking more competitive methods. “Path-Guided Particle-based Sampling” (ICML 24) can be another competitor approach that is proven to draw efficient samples from multi-modal Bayesian posteriors.

questions:
In the “Performance results” paragraph the last sentence claims that “It is noteworthy that this is a big step for sampling-based inference, yielding a time complexity comparable to DE, while providing better and more principled uncertainty measures.” However, the time measured is on top of the DE fit as claimed in the caption of Table 1. Then the total time for MILE should be twice of DE. Please clarify.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730523407703

thread_type:
official_review

<thread_object>
creator: reviewer_k13v

<content>
summary:
This paper adapts the MCLMC algorithm to sample from BNN posteriors. The authors propose a series of changes to the 3-stage tuning scheme in the original paper of MCLMC algorithm and name the adapted algorithm MILE. The authors show that by doing these, MILE demonstrates superior predictive performance, improved uncertainty quantification and improved runtime.

strengths:
1. The paper is clear and detailed in terms of related work and the method.
2. The empirical results look comprehensive and solid.

weaknesses:
1. The authors claim that their method is “tuning-free”, but in 3.2, several parameters are still tuned. Maybe I misunderstood what the authors mean by “tuning-free”

2. I feel that the author could elaborate more on which of the benefits of MILE mentioned in section 4 is inherited from the MCLMC algorithm, and which of them result from the authors’ adaptation. 

Minor:
1. The acronym ESS first appear in 3.2.3 without any explanation

questions:
1. What is d in $(d-1)^{-1}$ in Eq. 3?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722622952

thread_type:
official_review

<thread_object>
creator: reviewer_qcdd

<content>
summary:
The authors propose an adaptation of the Microcanonical Langevin Monte Carlo method for Bayesian Neural Networks.

strengths:
Clarity:
The paper is clearly written and well-organized.

Quality:
The evaluation is conducted meticulously, with numerous ablations considered. The results effectively justify the proposed method.

weaknesses:
Novelty:
The paper appears to be an incremental modification of the Microcanonical Langevin Monte Carlo method.

Significance:
Proposing efficient sampling methods for Bayesian Neural Networks is an important problem. However, the main results of the paper are achieved through careful parameter settings, particularly in establishing ensemble methods, tuning size, step size, energy variance scheduler, sample size, and so on. Each of these parameters is known to be crucial for obtaining better results. This makes the paper seem somewhat ad-hoc to me, lacking sufficient significance for developing improved BNN samplers.

questions:
The authors proposed a set of techniques to adapt the sampler for Bayesian Neural Network solvers. Is my understanding correct that to adapt Microcanonical Langevin Monte Carlo to BNNs, all we need to do is configure the parameters of the MCLMC and treat it as an ensemble method to reduce initialization error?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730726944435

thread_type:
official_review

<thread_object>
creator: reviewer_e8p6

<content>
summary:
This interesting paper considers the application of microcanonical Langevin Monte Carlo, a variant of HMC, to sample from Bayesian neural network posteriors. The main idea of microcanonical LMC is that the velocity's norm is fixed and not changing unlike in HMC, which is supposedly helping with more stable exploration and allows for larger step sizes in the presence of steep landscapes. The original algorithm was proposed in the 2023 paper "Fluctuation without dissipation: Microcanonical Langevin Monte Carlo". The main contribution here is to consider an ensembled variant of that algorithm, and extensive numerical experiments on various Bayesian neural networks to show the practical performance of this method.

strengths:
Microcanonical Langevin Monte Carlo is an interesting idea, and the numerical performance shows good improvements over NUTS, with impressive results in terms of predictive log posterior values.

weaknesses:
The only novelty in terms of algorithmic development is the use of ensembling, Microcanical LMC was already proposed in the previous paper "Fluctuation without dissipation: Microcanonical Langevin Monte Carlo".

The method is not using minibatches, but instead, each step needs a full gradient, meaning that the cost of Bayesian neural networks is at least 2-3 orders of magnitude higher than using deterministic neural networks. Hence the practicality of the algorithm at this point is questionable.

There was no comparison done with the earlier method "Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting" by Cobb et al., who managed to get HMC working with no bias using minibatches at each step, and an accept/reject step. It would be interesting to explore whether such a variant could be extended to Microcanonical LMC.

A major problem is that the implemented algorithm is not clearly described in the paper, and it's not clear whether a Metropolis/Hastings step is used or not, but the authors did not claim they used one so I presume it's not used. The previous paper "Fluctuation without dissipation: Microcanonical Langevin Monte Carlo", (https://arxiv.org/pdf/2303.18221) claims that their Euler-Mayurama discretization (15) exactly preserves the invariant distribution, so there is no need for accept/reject step. The precise form of the implementation of (15) is not stated there either. I am very skeptical that an explicit discretization only using a single gradient evaluation per iteration can exactly preserve the invariant distribution, hence I am on the opinion that this is not a fully explicit scheme, unless the authors convince me otherwise in the rebuttal.

questions:
State the implemented algorithm precisely, including how many gradient evaluations are used per step, and state whether it preserves the invariant distribution. State whether accept/reject step is needed or not.

How does the method compares with "Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting"?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732296548844

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Discussion and Key Updates – By the Authors

comment:
**Dear Reviewers,**

As we approach the final quarter of the discussion phase, we wish to summarize the key points discussed and extend our gratitude to reviewer 83Bk for acknowledging our responses and increasing the score. Furthermore, we have submitted an improved manuscript based on your constructive comments and also included the additional code for the new experiments in the anonymous repository.

**Key Updates**

- **Additional comparisons**: Benchmarked MILE against symmetric Split HMC and PGPS, showcasing MILE’s superior speed and performance (Section 1, Appendix A.1.3).
- **Algorithm and Convergence Clarifications**: Detailed our algorithmic choices, in particular the one of unadjusted sampling and the resulting efficiency advantages (Sections 2, 3, Appendix A.3).
- **Highlight MILE’s Novelty**: We highlight MILE as a reliable off-the-shelf method for BNN inference, with critical adaptations for multimodality and robustness, as shown in a new benchmark (Section 2, Appendix A.1.2) where naive MCLMC fails up to 86% of runs while MILE achieves a failure rate of 0%. These novelties, alongside its unprecedented speed, and state-of-the-art performance, underline MILE’s practical appeal. In order to better convey this, we reformulated our contributions in Section 1, parts of Section 4 and the Discussion in Section 5.

We kindly invite the remaining reviewers to share their thoughts on whether our responses have adequately addressed their questions and concerns, as well as to let us know if there is anything else we should clarify or improve. Once again, we sincerely thank all reviewers for their valuable feedback, which has helped us enhance the paper by showcasing the strengths of our framework through additional benchmarks and more effectively highlighting our contributions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733128697240

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final Reminder & Summary

comment:
Dear Reviewers,

Thank you again for your feedback and engagement throughout this process. As the **discussion period closes in about 24 hours**, we would like to kindly remind you to share any final thoughts on our responses and the additional experiments provided.

To summarize:

- We addressed all concerns raised, providing clarifications on theoretical aspects, highlighting our empirical contributions, and demonstrating the robustness of MILE compared to existing methods.
- Misunderstandings about related work and our methodology were clarified, and we pointed to our comprehensive, anonymized code repository that supports reproducibility.
- Further experiments included during the rebuttal phase underline the practical significance of our approach, offering fast and accurate off-the-shelf sampling for BNNs, which is a significant contribution to the field.

We are confident that these responses fully address the concerns raised and hope they support a fair and constructive evaluation of the work.

Thank you for your time and consideration.

Sincerely,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733243427826

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Closing Summary by the Authors

comment:
Dear Reviewers,

As the discussion phase concludes, we would like to summarize our method in one sentence to again highlight its significance and advantages:

***MILE** is a robust off-the-shelf method that generates high-quality posterior samples of BNNs in a fraction of the runtime of competing methods, delivering superior performance, in particular over sequential state-of-the-art samplers like NUTS and MCLMC.*

We have visualized this behavior in this **GIF**: [https://anonymous.4open.science/r/MILE-1CC1/README.md](https://anonymous.4open.science/r/MILE-1CC1/README.md), which is now included in the README of our repository. 

We would further like to highlight that the same repository contains a fully reproducible setup of our experiments, including the benchmarks with new comparison methods as requested by Reviewer e8p6 and 83Bk.

We hope that we have eliminated all the remaining doubts and again thank all reviewers for their engagement and time investment. 

Sincerely,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes MILE, an adaptation of the Microcanonical Langevin Monte Carlo (MCLMC) algorithm for sampling from Bayesian Neural Network (BNN) posteriors. The key claims are:

MILE achieves speedups compared to NUTS while maintaining or improving predictive performance
The method provides more reliable and predictable sampling through ensembling and careful tuning of MCLMC
MILE enables practical, scalable sampling-based inference for BNNs with runtime complexity comparable to deterministic optimization

Strengths

This is an BNN sampling is an importnat problem (IMO)
Extensive experimental validation on multiple datasets demonstrates improvements over existing methods
The authors provide thorough ablation studies showing the benefit of their adaptations to MCLMC
The work is well-presented with clear explanations of technical concepts

Weaknesses

Limited theoretical analysis of convergence properties
Some reviewers felt the novelty was incremental since it builds on MCLMC
The precise algorithmic details of the discretization scheme were  unclear
While claiming to be "tuning-free", the method still requires some parameter tuning, though automated
Full batch gradient computations are required

I recommend acceptance of this paper because:

It makes a meaningfully practical contribution to (almost) black-box BNN sampling
The empirical results are thorough, showing clear improvements over existing methods
The authors were responsive during rebuttal and provided additional experiments and clarifications
While building on existing work, the adaptations are non-trivial and necessary for making MCLMC work effectively for BNNs
I think the full-dataset gradient computation could be addressed in future work using existing techniques.