Microcanonical Langevin Ensembles: Advancing the Sampling of Bayesian Neural Networks

Dear Reviewer 83Bk,

Thank you for your positive feedback. We appreciate your thoughtful review and are pleased that you found the paper engaging and well-explained. Below, we address your questions and suggestions point by point.

Question on Time Complexity

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

The reviewer is correct that the sampling time needs to be added to the optimization time of DE, effectively doubling the DE time. We will make this more clear in an updated version of the manuscript. Note, however, that this is still the property we intended to highlight. Maintaining good sample quality while achieving sampling times that match optimization times is unprecedented for sampling-based inference in BNNs. Being able to sample at the same speed required by Adam-type optimization makes sampling strategies a viable alternative to optimization routines in deep learning applications and thereby creates an entirely new training paradigm of neural networks.

Comments on Weaknesses

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

We agree that examining the role of each adjustment is valuable:

• a) DE warmstart: While the reviewer is correct that the DE startup notably improves the method, our reported speedups are a comparison between two methods that both use DE initialization. Hence, it is not correct to say that “the main speed-up is due to the warm starting using the DE” as both methods use warm starts via DEs.
• b) MCLMC without adjustment vs. with adjustment: We did not provide such a study in the original submission as MCLMC without adjustment for BNNs will either be inefficient or in most cases not work at all. Thus, adjustments, such as the proposed tuning initialization, are integral to the numerical and overall feasibility of MILE in the BNN setting and cannot be entirely ablated. We however agree that it would be essential to also show this in the paper. To this end, we have conducted a small benchmark study, sampling 100 chains of MILE and 100 naive MCLMC without adjustments for different datasets and dimensions of the parameter vector/model. We compare how many runs lead to numerical issues rendering all obtained samples useless (NaN).

• c) Tuning parameters: as discussed in 4.3, tuning parameters do not notably affect the procedure once the adjustments from b) are incorporated

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

We thank the reviewer for pointing this out. While the focus of our work is on empirical efficiency rather than guaranteed convergence, we acknowledge the importance of discussing MILE’s convergence properties. To check convergence, we examined local mixing using split-chain Rhat (see, e.g., Appendix A.3), which indicates improved local convergence over competitor methods. Apart from expanding Appendix A.3 to better convey our convergence diagnostics and rationale, we will further add a discussion on the theoretical convergence of our method. While the reviewer is correct that we do not provide any theoretical guarantees at this point, a theoretical argument would “only” require the analysis of the effect of initialization as, e.g., done in Wild et al. (2024) for deep ensembles.

We thank the reviewer again for the great suggestions and comments. We think these helped us to further improve our manuscript. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References:

Fan et al., 2024: Path-Guided Particle-based Sampling, ICML 2024

Tian et al., 2024: Liouville Flow Importance Sampler, ICML 2024

Wild et al., 2024: A rigorous link between deep ensembles and (variational) Bayesian methods, NeurIPS 2024

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

Including “Path-Guided Particle-based Sampling” (PGPS) as a competitor method is an excellent suggestion. We will discuss PGPS in the related work section together with another related paper from ICML on Liouville Flow Importance Sampler (Tian et al., 2024). We also conducted an empirical comparison, using the experimental setup from Section 5.2.1 of their paper, which we will add to our manuscript. We performed BNN inference on seven UCI classification datasets, evaluating average negative log-likelihood (NLL) and accuracy. Notably, MILE completed all experiments and replications in under 5 minutes on a small PC CPU, with many tasks taking less than 1 minute. The results, presented in the table below, demonstrate that MILE achieves comparable or better accuracy while consistently outperforming PGPS in NLL across most cases, thus, again highlighting its competitive performance.

I appreciate the author's thorough comparison with the PGPS method and clarifying my other questions. Indeed MILE shows competitive performance vs PGPS (in most cases MILE wins). The very small runtime for MILE is also very interesting and supports the main contribution of the paper. Can the authors specify the hyperparameters used for PGPS? Also, please comment on the efficiency of MILE vs PGPS.

Since this was one of my main concerns for the paper, I will increase the score.

Dear Reviewer k13v,

Thank you for your positive feedback. Below, we respond to your comments and questions one by one.

Questions

What is d in (d−1)−1 in Eq. 3?

The d in Eq. 3 refers to the dimensionality of the parameter vector θ, as noted at the start of Section 2. We will better clarify this in the text.

Comments on Weaknesses

The authors claim that their method is “tuning-free”, but in 3.2, several parameters are still tuned.

"Tuning-Free" Terminology: While our method does involve a tuning phase, it is designed to be automated so that no manual tuning by the practitioner is required. This allows MILE to function as a reliable, off-the-shelf learning approach. We will try to highlight this again in the manuscript to avoid misunderstandings. Thank you for bringing this up.

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.

We adapted MCLMC on multiple fronts.

• Multimodality: Without the DE initialization, MCLMC will not be able to explore the posterior of BNNs sufficiently fast.
• Robustness, numerical stability in high dimensions: Without our adjustments, MCLMC will not only suffer from insufficient exploration but fail to produce meaningful samples. To emphasize this, we ran an ablation study requested by Reviewer 83Bk, showing how often MCLMC will fail in an application to BNNs without our adjustment. For this, we sample 100 chains of MILE and 100 naive MCLMC without adjustments for different datasets and dimensions of the parameter vector/model. We compare how many runs lead to numerical issues rendering all obtained samples useless (NaN).

• Enhanced Scalability: We improve scalability by addressing critical bottlenecks, such as the FFT component within the tuning phase, further strengthening MILE's suitability for practical BNN applications.
• MILE’s speed and performance: The speed at which MILE achieves state-of-the-art performance on various BNN inference tasks is unprecedented for full sampling-based inference and can be considered a novelty by itself. This not only becomes clear from the comparison with BDE in the paper but also in the newly added comparison with Split HMC provided to Reviewer e8p6.

Minor comments

ESS acronym not defined: Thank you for pointing this out. We will define ESS (Effective Sample Size) when it first appears in Section 3.2.3.

---

Thank you for carefully reading our manuscript and helping us to refine our explanations and terminology in the manuscript. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

Dear Reviewer e8p6,

Thank you for your thorough and thoughtful review. Below, we address each of your comments and questions in detail.

Algorithm Description and Dynamics

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.

1. Algorithm Implementation: We follow the dynamics proposed by [1] and adapt the MCLMC implementation from [2]. For each sample, our implementation requires two gradient evaluations as we use the minimal norm integrator.
2. Invariant Distribution: Our method is not a fully exact scheme. This is an intentional design choice, leveraging the resulting efficiency gains. We will further clarify this point in the manuscript.
3. Accept/Reject Step: We adopt an unadjusted sampling approach without MH correction, prioritizing computational efficiency over exact discretization error correction. This is discussed in the background section of the paper. We argue that the discretization error targeted by the MH correction is comparatively small compared to the initialization and Monte Carlo error. Hence, there is little benefit of using a correction in light of its computational cost. We will revise the manuscript to better highlight the rationale behind our approach.

Comparison with Cobb et al. (2021) [3]

How does the method compares with "Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting"?

We thank the reviewer for the additional reference. We will include the below discussion in a revised version of the manuscript:

1. Limitations of Symmetric Split HMC: Symmetric Split HMC advances HMC, but inherits the same hyperparameter sensitivity (e.g., depends on trajectory length and step size). As discussed in detail in [4], these hyperparameters can limit the application in Bayesian neural network inference. In [3], the authors use Bayesian Optimization to derive hyperparameters which introduces further complexity and a significant computational burden. Unlike MILE, Symmetric Split HMC employs an MH correction step, which further increases the computational costs. Another downside of [3] is that with an increased number of batches, the computational requirements increase drastically. Both approaches have merit, but their main goal and contribution differ considerably. Symmetric Split HMC focuses on memory scalability, while MILE optimizes speed and performance.
2. Empirical Comparison: We have conducted a benchmark with Symmetric Split HMC (using the official hamiltorch implementation) on our CNN (v2) architecture on Fashion-MNIST using the optimized parameters for the same task from Section 5.3 of [3]. All models are run on the same hardware to ensure runtime comparability. The results from this benchmark are displayed in the table below. We found that for a fixed amount of posterior samples, Symmetric Split HMC performs better for smaller batch sizes. However, as noted earlier and confirmed empirically, runtime increases dramatically for smaller batch sizes. We chose a batch size of 64 and ran Symmetric Split HMC for 200 samples. This required 15.5 hours, after which we had to stop the routine. The intended goal of sampling 1000 samples (as in our paper) would take more than 3 days. For larger batch sizes, we were able to generate 1000 posterior samples for Symmetric Split HMC, but without a notable gain in performance. Regardless of the specification, it becomes clear that MILE achieves considerably better performance in a fraction of the required time for Symmetric Split HMC, without even considering the cost of the Bayesian Optimization necessary for the practical application of Symmetric Split HMC.

Additional Comments

only novelty in terms of algorithmic development is the use of ensembling

Novelty Beyond Ensembling: While ensembling is a key component, our method also incorporates several essential modifications of MCLMC and algorithmic adjustments that enable the sampler to be applied to BNNs in a numerically stable, more scalable, and efficient way. To emphasize this, we ran an ablation study requested by Reviewer 83Bk, showing how often MCLMC will fail in an application to BNNs without our adjustment. For this, we sample 100 chains of MILE and 100 naive MCLMC without adjustments for different datasets and dimensions of the parameter vector/model. We compare how many runs lead to numerical issues rendering all obtained samples useless (NaN).

Furthermore, as also shown in the comparison with Symmetric Split HMC, the speed at which MILE achieves state-of-the-art performance on various BNN inference tasks is unprecedented for full sampling-based inference. This is achieved by optimizing each step in our proposed pipeline and can be considered a novelty by itself. We will expand on these aspects to highlight the contribution of our work more clearly.

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.

Limitation on Minibatching: Exploring minibatch variants is indeed a promising future direction, which we made explicit in the paper’s discussion section. Our current focus, however, is on applications with manageable dataset sizes and advancing the scalability of full-batch samplers. Stochastic samplers come with various challenges that are worth studying on their own and would potentially dilute the performance gains of MILEs. In addition, various recent stochastic sampling approaches still struggle with being robust enough to be practically applicable without an extensive overhead in hyperparameter tuning [5-6].

---

Summary

We will add a discussion of Cobb et al. (2021) [3] to the related works section and add the benchmark study to the Appendix. In addition, we will incorporate the above points to make our contributions more clear. We believe that this will further strengthen the paper. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify. Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.

References

[1] Robnik, J., & Seljak, U. (2024). Fluctuation without dissipation: Microcanonical Langevin Monte Carlo. In Symposium on Advances in Approximate Bayesian Inference (pp. 111–126). PMLR.

[2] Cabezas, A., Corenflos, A., Lao, J., & Louf, R. (2024). BlackJAX: Composable Bayesian inference in JAX. arXiv.

[3] Cobb, A. D., & Jalaian, B. (2021, December). Scaling Hamiltonian Monte Carlo inference for Bayesian neural networks with symmetric splitting. In Uncertainty in Artificial Intelligence (pp. 675-685). PMLR.

[4] Sommer, E., Wimmer, L., Papamarkou, T., Bothmann, L., Bischl, B., & Rügamer, D. (2024). Connecting the Dots: Is Mode-Connectedness the Key to Feasible Sample-Based Inference in Bayesian Neural Networks? Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:45988-46018.

[5] Kim, S., Jung, S., Kim, S. ; Lee, J.. (2024). Learning to Explore for Stochastic Gradient MCMC. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:24015-24039

[6] Yi-An Ma, Tianqi Chen, and Emily B. Fox. 2015. A complete recipe for stochastic gradient MCMC. In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2 (NIPS'15). MIT Press, Cambridge, MA, USA, 2917–2925.

Dear Reviewer e8p6,

Thank you for taking the time to read our response and providing further feedback. Please find below our response.

It could be that you did not notice this, but that paper is full of mistakes, such as the claim of Theorem 3 that the Euler-Mayurama discretization exactly preserves the invariant distribution is completely wrong.

There might have been a misunderstanding. The theorem is correct, but maybe it was not clear what its claim is. The claim is that if you do Euler-Mayurama discretization of the stochastic-deterministic part of dynamics AND you were able to exactly solve both separate parts, then you would have no bias. In practice of course you do not know how to exactly solve the deterministic part, so you have to do further splitting to momentum and position updates which introduces bias. So the claim of the theorem is not that there is no bias, but that the stochastic part of the dynamics does not contribute to the bias. If you think that there are additional mistakes in that paper, please let us know. We will do our best to address them.

The code for the paper of Robnik and Seljak attempts to solve the SDE accurately by using an adaptive step size high order Runge-Kutte method,

There might be another misunderstanding as this statement is also not correct as the deterministic part is solved by a splitting method, such as leapfrog or Omelyan integrator. It is not of the Runge-Kutte type and is not a higher order method (it is second order, as most integrators used in MCMC). While these explanations are stated in great detail in the “Microcanonical Hamiltonian Monte Carlo” paper by Robnik et al., we are happy to reiterate those in our paper in an additional section in the appendix.

I strongly feel that the precise statement of the discretization of the SDE should be included in the paper, otherwise it is absolutely unclear what is implemented in the code. You cannot expect readers to look at thousands of lines of code to figure out what is the key algorithm of the paper actually doing.

We absolutely agree with the reviewer and are thankful for this suggestion. We will make this more clear in the updated version of the manuscript.

It is commendable that you attempted to compare the method with Symmetric Split HMC, but I am skeptical about the fact that the differences between the performance in terms of accuracy are so significant, as the Symmetric Split HMC has no bias (due to the Metropolis/Hastings step).

did not do any experiments to evaluate the bias of the method.

As we explain in the text, having no bias at the cost of the Metropolis test is actually harmful for the performance of the algorithm, because it increases the variance. If the bias is small, but controlled, as it is in MCLMC by the energy error, the resulting algorithm achieves better performance. Our comparisons with BDEs (HMC-based sampling with correction step) confirm this hypothesis, showing that MCLMC indeed performs on par with (or sometimes even slightly better than) HMC. Experiments explicitly evaluating the bias of MCLMC are given in the respective paper “Microcanonical Hamiltonian Monte Carlo” by Robnik et al., see for example Figure 4.

With better tuning of the step size and the number of steps, as well as the batch size, you should be able to obtain similar accuracy

We think that a useful sampling algorithm should work out of the box, with no need to additionally tune its hyperparameters on a problem-to-problem basis. This is exactly what we achieve with MILE and what we aimed to accomplish with our paper. As stated in our revision, Symmetric Split HMC can be a very useful approach and we agree that there might be hyperparameters for which a better performance than the one reported in our previous response can be obtained. In our above experiments, we used the hyperparameters provided by the authors as these already represent meaningful and tuned hyperparameters. However, as stated in our rebuttal, we had to limit sampling times to a certain budget to be able to respond within the rebuttal phase. We are more than happy to run Symmetric Split HMC including hyperparameter optimization and without budget constraint after the rebuttal phase to demonstrate the reviewer’s statement.

---

---

Thank you again for your feedback. Please let us know if our response addresses your concerns, and if there is anything else we need to clarify.

Dear Reviewer qCdd,

Thank you for your feedback. Below, we respond to your comments on the novelty and significance of our approach.

Weaknesses and Questions

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

We thank the reviewer for this thoughtful comment but politely disagree. We adapted MCLMC on multiple fronts.

• Multimodality: Without the DE initialization, MCLMC will not be able to explore the posterior of BNNs sufficiently fast.
• Robustness, numerical stability in high dimensions: Without our adjustments, MCLMC will not only suffer from insufficient exploration but fail to produce meaningful samples. To emphasize this, we ran an ablation study requested by Reviewer 83Bk, showing how often MCLMC will fail in an application to BNNs without our adjustment. For this, we sample 100 chains of MILE and 100 naive MCLMC without adjustments for different datasets and dimensions of the parameter vector/model. We compare how many runs lead to numerical issues rendering all obtained samples useless (NaN).

• Enhanced Scalability: We improve scalability by addressing critical bottlenecks, such as the FFT component within the tuning phase, further strengthening MILE's suitability for practical BNN applications.
• MILE’s speed and performance: The speed at which MILE achieves state-of-the-art performance on various BNN inference tasks is unprecedented for full sampling-based inference and can be considered a novelty by itself. This not only becomes clear from the comparison with BDE in the paper but also in the newly added comparison with Split HMC provided to Reviewer e8p6.

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 [...] This makes the paper seem somewhat ad-hoc to me

MILE is designed as a reliable off-the-shelf option, similar to NUTS and Adam, offering high performance with minimal tuning required for reliable, scalable BNN inference. While we use the approach of Sommer et al. (2024) to induce enough exploration in the sampling procedure, other warm start options are potentially possible. If these options are equally or less computationally costly than DE starting values, MILE enables good sample quality for highly multimodal posteriors while achieving sampling times that match optimization times. This is unprecedented for sampling-based inference in BNNs. Being able to sample at the same speed required by Adam-type optimization makes sampling strategies a viable alternative to optimization routines in deep learning applications and thereby creates an entirely new training paradigm of neural networks. So to answer the reviewer’s question: the ingredients are MCLMC and ensembles similar to Adam being a combination of momentum and RMSprop. However, it takes considerable effort to adapt and combine these methods in order to perform on a scale that yields overhead the community is willing to accept when optimizing (sampling) neural networks. We strongly believe that the two components to achieve this are MCLMC and an ensemble of diverse starting values.

---

Thank you again for your constructive feedback. Please let us know if our response addresses your questions and concerns, and if there is anything else we need to clarify.