Thanks for constructive comments that can improve the quality of our paper. We would like to answer your comments.

Underdamped Langevin/ Hamiltonian Monte Carlo is a good method for accelerating the overdamped alternative in terms of mixing rates. However, I believe it is far from explorative enough compared to other baseline methods, such as the replica-exchange based approaches [1,2]. I am not fully convinced if optimizing the kinetic energy is sufficient enough to solve the exploration problem. Your model input is (θ,r), I am not even sure it is really learning anything useful.

Firstly, we develop our method to enhance the exploration of single-chain MCMC from the meta-training stage. Considering the scale of modern machine learning, improving the exploration of single-chain methods is advantageous due to the significant memory costs associated with parallel chain-based approaches. Therefore, directly comparing the exploration of replica-exchange based methods that use multiple independent chains with our method seems unfair for our method. However, following your suggestion, we are now implementing the method proposed in [3]. We will upload the comparison with [3] as soon as experiments are completed.

We have added the advantages of parameterizing kinetic energy term in overall response and Appendix A.1. Also, our learned sampler takes not only $(\theta,r)$ as inputs but also takes the gradient and moving average of the gradient which encode the posterior information as inputs. We explained details for inputs of neural networks in section 3.1 and Appendix E.1.

Empirically, your baselines such as deep ensemble and cyclical SGMCMC are pretty weak in terms of multi-modal simulations (although their optimization performance is acceptable).

DE is known as one of the most effective ensemble methods, including Bayesian methods, in capturing both weight diversity and functional diversity [4]. While it is difficult to say DE is good at simulating BNNs posterior as it is a non-Bayesian method, it is a powerful approach at least in capturing multi-modality in the function and weight space. Also, CSGMCMC are currently widely used by practitioners as a baseline for efficient simulation of BNNs posterior with a single chain which is deeply related to our method. But, comparing with other methods will also be interesting, we will upload the comparison with [3] as soon as experiments are completed.

[1] Non-convex Learning via Replica Exchange Stochastic Gradient MCMC. ICML'20

[2] Non-reversible Parallel Tempering for Deep Posterior Approximation. AAAI'23.

[3] Interacting Contour Stochastic Gradient Langevin Dynamics. ICLR'22

[4] Ovadia, Yaniv, et al. "Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift." Advances in neural information processing systems 32 (2019).

Deep Ensemble is a non-Bayesian method

I may respectfully disagree with this (although some debate exists) and I believe deep ensemble is also a Bayesian method. First, SGD is an approximate Bayesian method to capture one major mode [1]; second, deep ensemble uses different initialization to capture different modes, which is why I believe that the exploration solely driven by different initialization may not be sufficient to tackle explorations. Nevertheless, the straightforward advantage is that it is easy to implement, which is why it is so popular. But fundamentally it is not a principled method.

[1] Stochastic Gradient Descent as Approximate Bayesian Inference. JMLR'17.

a single chain that is deeply related to our method

I understand why the authors didn't compare the multiple-chain baselines. However, in my experience, they are really crucial to address the exploration problem. In my limited experience in motivating explorations, how to appropriately set up different learning rates and temperatures is crucial (in theory and also well supported in practice) and more important than the rest of the factors, such as momentum, etc. So these references on simulated tempering (one chain version) and parallel tempering (multi-chain version) cannot be missing in any papers that aim to facilitate exploration.

Popularity is not the unique criterion as a baseline.

Many popular Github repos also work quite well, but they are not published as papers. Popularity is not the first principle to choose baselines in academic papers.

Summary

I will slightly increase my rating from 3 to 5 to thank the authors including one more baseline.

I would suggest the authors include more relevant experiments to justify your work in your next revision. For example, instead of running so many large-scale experiments, I am actually more interested to see how your algorithm performs on the 25-mode 2D mixture distribution with unequal weights (figure 3 in [1]). If your algorithm can achieve similar performance without introducing too much cost, I will be more convinced if the algorithm is good or not. (I don't expect your algorithm to outperform replica-exchange.)

[1] INTERACTING CONTOUR STOCHASTIC GRADIENT LANGEVIN DYNAMICS

Thanks for constructive comments that can improve the quality of our paper. We would like to answer your comments.

I think the first and the biggest weakness of the work lies in the way it is framed. It seems that you are trying to compete with deep ensembles, and other SGMCMC variants. In fact, the message I get from the experimental results is that L2E is that it performs pretty much worse is most cases and at worse computational cost. But one could argue, the promise of the method is rather in the generalizability of the meta-learning procedure to unseen datasets. In a sense, it is pretraining for SGMCMC, which to me is the most interesting. Unfortunately, this is not how the authors position this paper.

Would you elaborate why L2E performs pretty much worse than other methods? According to Table 1, we outperform other baselines in general. Only DE can match the predictive accuracy in some experiments and also L2E shows the best performance in OOD detection. Except for distribution-shift experiment where Bayesian methods can often fail [1], L2E shows competitive experiments with DE and outperforms CSGMCMC. From a computational cost perspective, we only meta-train once and L2E only takes about 1.2~1.5 times more wall clock time for one step update than CSGMCMC according to Table 18. Also, L2E is potentially much more efficient than other baselines since it is robust to the choice of thinning interval. In Figure 19, we plot the performance of the sampler with varying cycle length (thinning interval). Although L2E uses more computational cost for one step update, L2E can achieve comparable accuracy with other baselines with smaller total computational cost. We believe L2E excels in both aspects (performance and generalization), as it achieves better performance with reasonable computation costs even on unseen datasets compared to other methods.

[1] Izmailov, Pavel, et al. "Dangers of Bayesian model averaging under covariate shift." Advances in Neural Information Processing Systems 34 (2021): 3309-3322. .

we confirm that seeing more datasets during the meta-training stage do not necessarily lead to better OOD detection performance.

Thank you for conducting this experiment.

I have a different reading of the results here. On CIFAR-100, a broader dataset, large version seems to win always. Of course, it would be hard to extrapolate, but it would appear that a more diverse dataset leads to better sampler and better uncertainty estimation. Wouldn't that be the ideal scenario for L2E anyways - it improves its uncertainty estimation with more diverse meta-training?

In any case, what I was hinting at perhaps not a criticism of the approach, but more of trying to understand where the benefits come from - parametrizing the kinetic energy term or the choice of diverse datasets. In the current work, both these choices haven't quite been studied independently.

The motivation of the proposed method is weak. Since the main difference compared with Gong et al is the parameterization of the kinetic energy term, it is important to clearly motivate this choice. The authors did not mention at all in the paper why they choose to learn the kinetic energy term rather than the curl matrix and the diffusion matrix. The advantages of doing so are not discussed.

Firstly, we would like to say that we noted the advantages of the design of kinetic energy over the design of diffusion and curl matrices in Appendix F along with ablation study. Also, we have added the motivation to learn the kinetic energy term in Appendix A.1. Parameterization of diffusion matrix $D_f$ and curl matrix $Q_f$ suggested in Gong et al., has two limitations to achieve our goal, efficient exploration of multi-modal large scale BNNs posterior distribution. First, learning $D_f$ and $Q_f$ to be dependent on $z$ harms scalability of algorithm since it introduces a new correction term $\Gamma_i(z)$ which forces to compute the gradient of $D_f$ and $Q_f$ with respect to $z$. This makes significant computational burden as the dimension of $z$ increases.

Additionally, since $Q_f$ which governs the acceleration of $\theta$ is limited to operating as multipliers for the energy gradient and momentum, learning $D_f$ and $Q_f$ is not the best parameterization for learning sampler to effectively explore multi-modal distribution. In low energy regions where the norm of gradient and momentum are extremely small, it is difficult to make reasonable amount of update of $\theta$ for exploration by multiplying $Q_f$ to momentum and gradient. By contrast, In the update rule of L2E, two parameterized gradients $\alpha_\phi, \beta_\phi$ are added to the energy gradient and $\beta_\phi$ directly updates the $\theta$. This is more suitable for controlling the magnitude and direction of update to enhance exploration property in low energy regions.

Gong et al has proposed several meta objectives. Again, the authors did not discuss their meta objective with existing ones.

To the best of our knowledge, Gong et al did not propose meta-objectives other than the meta-objective they actually used. However, we have added a discussion on the meta-objective proposed by Gong et al and other related papers in Appendix A.3. Meta-objective of Gong et al aims to reduce the KL divergence between the target density $\pi$ and the marginal distribution $q(\theta|D)$ at time $t$. However, $q(\theta|D)$ is intractable since it is an unknown density function, computing the gradient of $q(\theta|D)$ requires a gradient estimator. They used stein-gradient estimator which requires multiple independent chains, but this undermines the scalability of the algorithm. Also, this objective does not lead the learned sampler to explore multi-modal distribution. Figure 3 in Gong et al showed that learned sampler fastly converges to low energy region, but learned curl matrix restricts the amount of update in low energy region to prevent divergence which harms the exploration in low energy region.

In experiments, the authors did not compare with Gong et al, which is a very related method.

In Appendix A.4 of revised version, we follow the experimental setup of Gong et al, and compare our method with Meta-SGMCMC (Gong et al). For experimental details and results, please refer to Appendix A.4. To summarize the experimental results, we have confirmed that L2E outperforms Meta-SGMCMC in various generalization experiments in terms of ACC and NLL. In CIFAR-10 experiment, despite L2E was not trained on the CIFAR-10 during the learning process, L2E significantly outperforms with a wide margin without being trained on CIFAR-10, indicating that our approach better generalizes to unseen datasets compared to Meta-SGMCMC.

Thanks for constructive comments that can improve the quality of our paper. We would like to answer your comments. Also, we revise our paper with comprehensive comparison with Gong et al., in Appendix A. Please refer to Appendix A for more detailed information.

The paper did not mention at all at the beginning that, there already exist meta-learning methods for SGMCMC, such as Gong et al. Only in Section 3.1, the authors first briefly mention that paper. The presentation is misleading and may give the impression that this paper is the first to study meta-learning for SGMCMC.

Sorry for not faithfully discussing Gong et al., in our main paper, we did not intend to be misleading in acknowledging its contribution. We could not put enough discussion in the main paper due to the page limit. In the revised version, we added the discussion about Gong et al., and devoted an additional section in the appendix to compare ours to Gong et al., in various aspects.

More importantly, the proposed method is essentially very similar to Gong et al, which uses the same formulation for the SGMCMC class, but will slightly different parameterization. Gong et al learns the curl matrix Q and the diffusion matrix D whereas the proposed method learns the kinetic energy term. Given the similarity, it is important to clearly state the difference compared with Gong et al.

We have added Table 6 and Appendix A to state the difference between Gong et al. Please refer to Appendix A and Table 6 for more detailed explanation.

For what you have pointed out, especially the statement that ours is very similar to Gong et al., because they share the same formulation for the SGMCMC class and differ only in parameterization. We respectfully disagree with this. The SGMCMC class is a complete class, so any convergent SGMCMC algorithm must belong to this class, so whatever SGMCMC algorithm we try to learn we always have to be in the class both ours and Gong et al., are following. Moreover, as we have detailed in Appendix A, the fact that we are learning kinetic energy is a significant innovation, bringing noticeable gains both in terms of the exploration behaviour of the sampler and efficient training.

Moreover, we have different goals with Gong et al. Their purpose was to learn the sampler that fastly converges to high density region with low bias. However, our approach is specifically designed with more concrete purpose of effectively simulating multi-modal BNNs posterior distribution and also generalizing to unseen problems for practicality. This goal was not achieved by Gong et al because Figure 3 in Gong et al showed that learned sampler fastly converges to low energy region, but learned friction term $D_f$ restricts the amount of update in low energy region to prevent divergence which harms the exploration of the sampler. Also, we aim to learn more generalizable and scalable sampler so we use different parameterization and meta-learning strategy with Gong et al. Therefore, although we share some similar features with Gong et al., it is difficult to say our method is essentially very similar to Gong et al since they could not achieve our goals.

Thanks for your positive view supporting our method that enhances exploration of SGMCMC through learning kinetic energy terms. Our answer for your concern is as follows:

While the experimental results are useful, I think the paper also needs an ablation study. We need to understand the impact of the parameterized gradients v/s transferability of the posterior information across the tasks.

We understand your question as asking “What is the common knowledge that parameterized gradients learn for the exploration of posterior distribution of different tasks”. If this is your question, we have added an analysis of the parameterized gradient during the evaluation stage in Appendix B for answering this. We track the norm of parameterized gradient, $||\beta||^2 = \frac{||\Delta\theta||}{\epsilon}^2$ we find that its magnitude is larger in local minima than in the early stages of training. This tendency is different from other gradient-based MCMC methods where the update amount decreases at minima. Additionally, we notice that L2E actively updates $\theta$ at minima while maintaining loss as nearly constant. This trend is consistently observed in both CIFAR-10 and CIFAR-100, implying that L2E learns some common knowledge of posterior information across tasks for efficient exploration in low loss regions.

If this is not an appropriate answer to your question, could you please explain your question again in more detail?

Understanding the compute requirement at train time is equally important. We don't know how much training time is required per step and in total and how it compares with other baselines that the authors have compared with. It'll also be useful to know the additional # of parameters added through the gradient parameterization.

For meta-training, we use approximately 10 hours for 2000 updates of meta-parameters on a single NVIDIA RTX A6000 GPU. It is important to note that we meta-train only once for all experiments in the paper since we aimed to meta-train the sampler which is transferable to various tasks. Therefore, it is unfair for our method to compare the wall clock time including meta-training time with other baselines.

Our meta-learner consists of 2 layer MLP with 32 hidden nodes. This neural network takes a 9-dimensional input and output 2-dimensional vector so the additional number of parameters is 1442. This neural network is applied to each dimension of the inner-parameter independently. L2E takes about 1.2~1.5 times more wall clock time than CSGMCMC for one step update according to Table 18.

I think there will be some tradeoff between the number of samples we generate in the inner loop to compute L and final convergence in terms of how quickly we converge and quality of the posterior samples. Running an ablation on the number of samples we generate in the inner loop would also be useful for practitioners who would be interested in using this approach.

We have added ablation experiments in Appendix C.3. for the number of samples in the inner loop. As you expected, we confirm that generating a smaller number of samples during inner-loop harms the convergence of the sampler in evaluation tasks. When the inner loop is too short, the sampler may not sufficiently learn information about the loss surface in low-loss regions during meta-training. Then, the sampler will fail to learn desirable properties such as exploration in minima or fast convergence.

I do like this paper overall and I think it has shown some interesting results. I have listed my comments that I'd like the authors to address in the weaknesses section. Also, it'll be useful if the author can provide some insights on how to speed up their algorithm and make it more efficient. For e.g., would warm starting with an SGD estimate in the meta-learning framework help?

We can speed up L2E by reducing the thinning interval. As shown in Figure 19, while CSGMCMC shows performance drop as the thinning interval decreases, L2E can maintain a similar level of accuracy and diversity with shorter thinning interval. Specifically, L2E maintains performance even when reducing the thinning interval from 50 epochs to 10 epochs for both CIFAR-10 and CIFAR-100. L2E outperforms CSGMCMC's 5100 epochs performance with just 1100 epochs.

Furthermore, L2E may benefit from using burn-in epochs with SGD, as it allows for warm starting with SGD and immediate sample collection using L2E to traverse high density regions between different modes. Additionally, since our approach shares some similarities with other SGMCMC in parameterization, we can apply techniques such as LR warm-up or LR scheduling. In other words, existing techniques for accelerating training can be applied to L2E so L2E has great potential in terms of efficiency and scalability. In conclusion, based on these insights and results, L2E has numerous ways to achieve speed up and potentially efficiency.