Paper: Learning to Explore for Stochastic Gradient MCMC (ICLR-2024)

Reviewer comment (Experiments):
Underdamped Langevin/ Hamiltonian Monte Carlo is a good method for accelerating the overdamped alternative in terms of mixing rates. However, it is far from explorative enough compared to other baseline methods, such as the replica-exchange based approaches [1,2]. It is not fully convincing that optimizing the kinetic energy is sufficient to solve the exploration problem. The model input is $(\theta, r)$, and it is uncertain whether it is really learning anything useful.

Author response:
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.

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

Author response:
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.

Reviewer comment (Experiments):
There is interest in seeing how the algorithm compares to [2] Non-reversible Parallel Tempering for Deep Posterior Approximation (AAAI'23) or [3] Interacting Contour Stochastic Gradient Langevin Dynamics (ICLR'22).

Author response:
We comare our method with ICSGLD[1] in CIFAR-10 and CIFAR-100. For both methods, we do not use data augmentation and fix temperature hyperparameter to be 1. We collect 10 samples for both methods. We fix other hyperparameters same as in [1]. As shown in tables below , We confirm that L2E shows better predictive accuracy and predictive diversity than ICSGLD on CIFAR-10 and CIFAR-100. This results implies the strong exploration property of L2E shows even compared to parallel-chain based SGMCMC method.

- CIFAR10 / ResNet-20
|Method | ACC | NLL | ECE | KLD|
|------ | :-----------: | :-----------: | :-----------: | :-----------:|
|ICSGLD | 0.883 ± 0.001 | 0.349 ± 0.002 | 0.079 ± 0.005 | 0.628 ± 0.087|
|ours(L2E) | 0.892 ± 0.003 | 0.334 ± 0.008 | 0.041 ± 0.002 | 0.639 ± 0.013|

- CIFAR100 / ResNet-20
|Method | ACC | NLL | ECE | KLD|
|------ | :-----------: | :-----------: | :-----------: | :-----------:|
|ICSGLD | 0.655 ± 0.002 | 1.251 ± 0.005 | 0.087 ± 0.009 | 0.874 ± 0.011|
|ours(L2E) | 0.668 ± 0.013 | 1.211 ± 0.003 | 0.097 ± 0.002 | 1.412 ± 0.007|


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

Reviewer comment (Evaluation):
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 in 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.

Author response:
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.

Reviewer comment (Experiments):
The deliberate choice to not use data augmentations is a little concerning. I understand that it does not neatly fit into the Bayesian perspective, but when the same datasets perform significantly better with data augmentation (including with SGMCMC), it makes the comparisons incomplete. It would not be too hard to add data augmentations as is into the same training setup, and make the comparison on these datasets more fair to modern deep learning that achieves both better accuracy and better calibration properties.

Author response:
Since including data augmentation violates i.i.d assumption of the dataset which is commonly assumed for Bayesian methods, this can lead to model misspecification problem. Therefore, prior work such as [1] argued the incompatibility between Bayesian methods and data augmentation. However, in modern machine learning, data augmentation is an indispensable technique. Thus, in response to your request, we conduct experiments with data augmentation on CIFAR-10 and CIFAR-100. Please refer to Appendix C.2 for full details and results of the experiments. Since using data augmentation introduces cold posterior effect [2] for Bayesian methods such as CSGMCMC and L2E, we additionally tune the temperature hyperparameter for these methods.

According to Tables, in terms of predictive accuracy, DE outperforms other baselines on CIFAR-10 and L2E beats other baselines on CIFAR-100. With data augmentation, it is not very surprising that DE outperforms other Bayesian methods like L2E and CSGMCMC in terms of predictive accuracy since they suffer from model misspecification and temperature tuning can partially handle this problem [3]. This can explain why L2E shows the worst NLL in CIFAR-10 experiment since likelihood does not fit with the observation. When it comes to predictive diversity, L2E significantly outperforms baselines on both experiments. Although applying data augmentation introduces significant variations to the posterior landscape, we confirm that L2E still maintains the exploration property. To sum up, we argue that L2E is still a practical method even with data augmentation since it shows competitive predictive performance and efficiently explores the posterior landscape. 

[1] Izmailov, Pavel, et al. "What are Bayesian neural network posteriors really like?." International conference on machine learning. PMLR, 2021. 

[2]  Wenzel, Florian, et al. "How good is the bayes posterior in deep neural networks really?." 

[3] Kapoor, Sanyam, et al. "On uncertainty, tempering, and data augmentation in bayesian classification." Advances in Neural Information Processing Systems 35 (2022): 18211-18225

- CIFAR10 / ResNet-20
|Method | ACC | NLL | ECE | KLD|
|------ | :-----------: | :-----------: | :-----------: | :-----------:|
|DE | 0.927 ± 0.001 | 0.218 ± 0.002 | 0.021 ± 0.001 | 0.238 ± 0.001|
|CSGMCMC | 0.923 ± 0.000 | 0.239 ± 0.003 | 0.014 ± 0.000 | 0.125 ± 0.001|
|ours(L2E) | 0.924 ± 0.003 | 0.249 ± 0.001 | 0.042 ± 0.002 | 0.332 ± 0.003|

- CIFAR100 / ResNet-20
|Method | ACC | NLL | ECE | KLD|
|------ | :-----------: | :-----------: | :-----------: | :-----------:|
|DE | 0.707 ± 0.001 | 1.053 ± 0.002 | 0.096 ± 0.005 | 0.515 ± 0.002|
|CSGMCMC | 0.697 ± 0.001 | 1.066 ± 0.007 | 0.064 ± 0.001 | 0.373 ± 0.005|
|ours(L2E) | 0.710 ± 0.001 | 1.053 ± 0.005 | 0.053 ± 0.001 | 0.911 ± 0.007|

Reviewer comment (Experiments):
The comparisons to DE and cSGMCMC looks to me a little unfair. L2E sees more variety of data by design, and one could then posit that the better OOD detection is not that surprising and simply a matter of having seen broader distributions at meta-train time.

Author response:
DE is currently a widely referenced method for OOD detection [1, 2] and it is not that strange to compare our approach with DE and CSGMCMC even though L2E was trained on additional datasets during the meta-training. We believe this comparison is realistic and fair because we did not explicitly meta-train L2E for OOD detection during meta-training. Furthermore, we conduct additional ablation experiments to assess whether seeing more datasets during meta-training affects the performance of OOD detection. In the following Table, we confirm that seeing more datasets during the meta-training stage do not necessarily lead to better OOD detection performance. Small L2E which is meta-trained only on Fashion-MNIST shows almost similar OOD detection performance with L2E or even better performance for some datasets. We think that meta-training on external datasets can rather lead to overconfidence of model on meta-training datasets which can potentially decrease OOD detection performance.  Also, we have added 1-D synthetic regression task in Appendix C.1. and L2E also shows competitive performance at capturing epistemic uncertainty in regression task. Since 1-D synthetic dataset is completely different from the meta-training datasets, it is difficult to say that L2E shows good OOD detection performance due to the external sources of the dataset. 

|In-dist|OOD|Small L2E|L2E
|:---:|:---:|:---:|:---:|
|CIFAR10|CIFAR100|**0.855±0.001**|0.851±0.001|
|   |SVHN|0.929±0.001|**0.931±0.003**|
|   |Tiny-ImageNet|**0.855±0.001**|0.853±0.001|
|CIFAR100|CIFAR10|0.789±0.003|**0.792±0.002**|
|   |SVHN|0.912±0.004|**0.917±0.003**|
|   |Tiny-ImageNet|0.777±0.001|**0.782±0.003**|


[1] Lakshminarayanan, Balaji, Alexander Pritzel, and Charles Blundell. "Simple and scalable predictive uncertainty estimation using deep ensembles." Advances in neural information processing systems 30 (2017).

[2] 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).

Reviewer comment (Experiments):
It would be really helpful if we could have a toy illustration to build intuitions. For instance, meta-learning datasets sampled from 1-D sines and cosines. It would also reveal how the method behaves with "in-between" uncertainty, i.e. parts of the input space not in the training dataset. It would also help to see how the sampler behaves beyond the training dataset, when say you change to say a different function family of polynomials.

Author response:
To visually check whether L2E can capture the epistemic uncertainty, we have added 1-D synthetic regression experiment in Appendix C.1. We use DE and HMC as baseline methods. Please refer to Appendix C.1 for experimental setup and Figure 7 for the regression results. We find that even without meta-training on synthetic 1-D dataset generated from particular polynomials, our method captures ‘in-between’ uncertainty well. As shown in Figure 7, L2E produces diverse predictions for out-of-distribution data points. HMC shows the best performance in capturing ‘in-between’ uncertainty which aligns with results from other works such as [1], but we believe that L2E still has sufficient capabilities to capture ‘in-between’ uncertainty in various tasks.

[1] D'Angelo, Francesco, Vincent Fortuin, and Florian Wenzel. "On stein variational neural network ensembles." arXiv preprint arXiv:2106.10760 (2021).

Reviewer comment (Novelty):
The choice of citation in the last paragraph on Page 1 "Zhang et. al., 2020" alone is odd. I would recommend citing all the other works.

Author response:
Thanks for pointing this out. We have added other works in the last paragraph on Page 1.

Reviewer comment (Novelty):
The authors have made the choice of not having an explicit related work section. I think that is alright, but other work must be contextualized somewhere in the main text. Part of it is done in

Author response:
Thanks for pointing this out. We have added related work section in the Appendix D.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Novelty):
More importantly, the proposed method is essentially very similar to Gong et al, which uses the same formulation for the SGMCMC class, but with 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.

Author response:
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.

Reviewer comment (Novelty):
Gong et al has proposed several meta objectives. Again, the authors did not discuss their meta objective with existing ones.

Author response:
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.

Reviewer comment (Evaluation):
Although the experimental results of the proposed method are promising, it is not clear why the proposed meta-learning method can lead to these empirical improvements. What kind of learning updates did the meta-learning algorithm learn in the end? Why does the method show a better exploration-exploitation balance? How does it achieve that?

Author response:
We have added additional analysis of the behavior of L2E in Appendix B. Firstly, L2E shows good exploration-exploitation balance since L2E makes considerable amount of updates even in the local minima without increasing loss. This is different from other gradient-based MCMC methods where the update amount tends to decrease at minima. Considering that L2E collects samples from different modalities, we can say that our parameterized kinetic energy gradients learned to move around high density regions among different modalities.

Reviewer comment (Experiments):
The meta-learning tasks are classification on MNIST, Fashion-MNIST, EMNIST and MedMNIST and the downstream tasks are classification on Fashion-MNIST, CIFAR and Tiny ImageNet. Since Fashion-MNIST has appeared in meta-learning tasks, is it reasonable to use it as a test task? Are CIFAR and Tiny ImageNet similar to MNIST to be considered as test tasks?

Author response:
We added the Fashion-MNIST experiment to assess the performance of L2E on the in-distribution dataset included in meta-training tasks. However, since model architecture used for meta-training is different from that used in Fashion-MNIST experiment, we can still evaluate the generalization performance of L2E with respect to model architecture. Importantly, datasets such as CIFAR and Tiny-ImageNet are significantly different from the datasets in meta-training tasks, enhancing their value as test tasks. Evaluating L2E on completely different datasets and architectures can effectively showcase the transferability and generalization capabilities of L2E.

Reviewer comment (Evaluation):
Can you explain T-SNE visualization of learning trajectories? Why does L2E’s trajectory look very different from DE and CSGMCMC?

Author response:
T-SNE visualization of learning trajectories is the method demonstrated by [1]. The authors of this paper aim to visualize the similarity in function space for each trajectory of DE. Therefore, we used T-SNE to visualize how L2E traverses the function space and collects samples. Figure 3 shows that the functions explored by L2E are far from the initial function, while simultaneously covering a wide range in function space within a single trajectory. The diversity of collected samples is supported by the results from KLD values in Table 1, cosine similarity heatmap in Figure 5 and linear interpolation plot in Figure 6. 

The difference in T-SNE plot between DE, CSGMCMC, and L2E is expected. DE starts from different random initializations for each trajectories so that there are several truncated trajectories of functions in T-SNE plot of DE. CSGMCMC makes artificial spikes of loss in the local minima through the cyclic learning rate schedule so the trajectories deviate from the highlighted region in Figure 3 with a large learning rate. In contrast, L2E does not initialize the chain or makes artificial spikes of loss while traversing through various regions of posterior distribution. In Figure 20, we show that L2E consistently keeps the loss and NLL as nearly constant while obtaining diverse samples from different modalities.

[1] Fort, Stanislav, Huiyi Hu, and Balaji Lakshminarayanan. "Deep ensembles: A loss landscape perspective." arXiv preprint arXiv:1912.02757 (2019).

Reviewer comment (Experiments):
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.

Author response:
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?

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.