Paper: Entropy-MCMC: Sampling from Flat Basins with Ease (ICLR-2024)

Reviewer comment (Theory):
The major weakness of this work is the lack of a rigorous/theoretical justification why the proposed sampling scheme encourages the sampler to explore flat basins. Towards the end of section 4.1, the authors mention that "$\theta_a$ provides additional paths for $\theta$ to traverse, allowing $\theta$ to reach flat modes more efficiently." However, this argument is not rigorously reasoned in the paper. The only support for this claim seems to be the synthetic examples in section 6.1, which is purely empirical and can have some issues (see Questions section). This claim needs careful theoretical reasoning (even though it can be challenging due to the need to understand the mixing speed of the MCMC chain in non-log-concave settings). This is crucial since it is a key motivation for the methodology. A possible direction to look into is the generalization bound of SGLD, as mentioned in reference [1].

Author response:
Q1: The lack of a rigorous/theoretical justification why the proposed sampling scheme encourages the sampler to explore flat basins. In my opinion, this claim needs careful theoretical reasoning. This is crucial since it is a key motivation for the methodology. 

A1: We want to clarify that the proposed sampling method does have theoretical justification for the convergence to flat basins. As shown in Lemma 1, the stationary distribution of $\theta_a$ in EMCMC is guaranteed to be the flat target distribution. Consequently, the samples of $\theta_a$ are guaranteed to be from flat basins asymptotically. 

Although there is no theoretical guarantee that $\theta$ will visit flat basins before sharp modes, $\theta$ is guaranteed to eventually visit flat basins since its stationary distribution is the original target distribution. 

It is important to clarify that we did not claim that $\theta$ is guaranteed to visit flat basins before sharp mode. We claim that $\theta_a$ will _encourage_  $\theta$ to visit flat basins through the interaction term $\frac{1}{2\eta}\\|\theta-\theta_a\\|^2$. This is also verified empirically. See A2 for details. 

As acknowledged by the reviewer, proving non-asymptotic behaviors of MCMC methods can be challenging. In fact, many existing papers on improving the non-asymptotic convergence of MCMC do not have rigorous guarantees. For example, papers on improving sampling methods on multi-modal distributions usually only have an asymptotic guarantee and do not have the guarantee to explore multiple modes in a finite time (non-asymptotic convergence), e.g. [4,6,7,8]. 

While we agree that theoretically proving non-asymptotic behaviors of our method, i.e. $\theta$ converges to flat basins before sharp modes, is a very interesting question, we emphasize that it is not crucial for the proposed methodology. We have provided theoretical guarantees on the convergence to flat basins (Lemma 1) and its fast convergence rate (Theorems 1,2&3). Moreover, besides theoretical results, we provided extensive experimental results to demonstrate the superiority of our method, verifying that it can sample from flat basins and achieve outstanding performance on diverse tasks.

We will clarify this point further in the final version. We hope the significance, novelty, and potential impact of this work can be considered by the reviewer in the final assessment.

Q2: The only support for this claim seems to be the synthetic examples in section 6.1

A2: Besides the results in section 6.1, Fig. 5(c) also demonstrates that $\theta$ and $\theta_a$ both converge to flat modes under practical budgets. In particular, at the same iteration, $\theta$ and $\theta_a$ stay in the same flat mode while maintaining diversity.

Reviewer comment (Theory):
It would be better to elaborate why the result is independent of the choice of $\eta$. To my understanding, if $\eta \to 0$, the smoothed target will be identical to the exact target, so there is no effect on "biasing towards the flat region," and with excessively large $\eta$, the convergence of the MCMC will be very slow. It will also be helpful to visualize the slice of the target density and smoothed target density (with various choices of $\eta$) on the slice of (x, x).

Author response:
Q6: It would be better to elaborate why the result is independent from the choice of $\eta$.

A6: We did not claim that the results are independent of $\eta$. What we stated in Section 6.1 is “independent of initialization with appropriate $\eta$”. What we mean is the following: if choosing an appropriate $\eta$, EMCMC will find the flat mode no matter how it is initialized. This is because the stationary distribution of $\theta_a$ is the flattened distribution, which removes the sharp modes. Through the interaction term, $\theta_a$ will encourage $\theta$ to flat basins. We will clarify this in the final version.

Q7: It will also be helpful to visualize the slice of the target density and smoothed target density (with various choices of $\eta$).

A7: Please check the following link for this new visualization. We will add it to the final version.

https://anonymous.4open.science/r/EMCMC-FF35/dis_more_eta.pdf

Reviewer comment (Evaluation):
For the image classification experiment, how do you compare the results obtained from sampling-based and optimization-based algorithms? Specifically, do you obtain certain point estimates of NN's weights from MCMC samples, and compute metrics on the test dataset using the NN with the estimated weights?

Author response:
Q8: For the image classification experiments, how do you compare the results obtained from sampling-based and optimization-based algorithms?

A8: Each MCMC sample $\theta_j$ is a neural network model. If the algorithm collects $M$ samples in total, there will be $M$ neural network weights saved in the memory. In the classification experiment, each model $\theta_j$ will output a probability $p(y|x,\theta_j)$, and the final output of sampling-based methods would be:
$$
p(y|x,D)\approx\sum_{j=1}^M p(y|x,\theta_j).
$$
This is called Bayesian model averaging, which is widely used in Bayesian deep learning [4,5]. In the classification experiment, accuracy, NLL, and ECE of sampling-based methods are obtained via Bayesian model averaging, following previous works [4,5].

Reviewer comment (Experiments):
Is there a systematic way of choosing the variance parameter $\eta$?

Author response:
Q9: Is there a systematic way of choosing the variance parameter $\eta$?

A9: $\eta$ is chosen by hyperparameter tuning via cross-validation. We found that the performance is not very sensitive to $\eta$ and we only need to determine the appropriate scale (e.g., considering $\eta \in \\{10^{-1},10^{-2},10^{-3},10^{4}\\}$). For example, the accuracy in the scale of $10^{-4}$ is shown in the following link, where no significant difference is found. 

https://anonymous.4open.science/r/EMCMC-FF35/eta_CIFAR10_grid.pdf

We agree that a more systematic way of choosing $\eta$ is an interesting question and leave it for future work.

Reviewer comment (Novelty):
There is a lack of novelty with respect to Baldassi et al. PNAS '16, only partly justified by the focus on the Bayesian setting.

Author response:
We want to emphasize that our method, EMCMC, substantially differs from Baldassi et al.'s approach and is not a special case of the Robust Ensemble (RE) with y=1 replica. There are several key differences between EMCMC and RE in Baldassi et al.: 1) In RE, all model replicas play identical roles and possess the same stationary distribution.  In EMCMC, $\theta$ and $\theta_a$ have different stationary distributions. 2) RE approximates the integral in local entropy by a finite set of $y$ replicas of the model, with exactness achieved only in the limit of  $y\rightarrow\infty$. In contrast, EMCMC does not make any approximation by using a joint distribution in a doubled parameter space. The stationary distributions of $\theta$ and $\theta_a$ in EMCMC are guaranteed to be the original target distribution and flat target distribution respectively. 3) RE is an optimization method, aiming to find a flat optimum as the point estimation. In contrast, EMCMC is a sampling method, aiming to sample from the flat basins within the posterior.

In summary, EMCMC and RE are developed from distinct ideas and they use local entropy in significantly different ways. To the best of our knowledge, the proposed auxiliary guiding variable and the joint distribution's form are novel and non-trivial solutions to flatness-aware learning. We will include these comparisons in the final version.

Reviewer comment (Experiments):
There is some hyper-parameter tuning, e.g. eta and T, carried on the test set instead of a validation set. This is not great practice.

Author response:
We want to clarify that the chosen values for $\eta$ and $T$ are tuned via cross-validation, following the prior work [3]. The results presented in Table 4 and Fig. 6 serve the purpose of ablation studies, showing the influence of these hyper-parameters on performance (i.e. test accuracy). We will clarify this in the final version.

Reviewer comment (Theory):
Can the author comment on the possibility of enhancing the effect of attraction toward flat minima. E.g. in the framework of Baldassi et al. PNAS '16 this would mean increasing the number of replicas used. Can more replicas be used by also having one of the marginal equal to the origin measure?

Author response:
As discussed in A1, EMCMC differs from the methods that approximate the integral in local entropy using a finite set of replicas. EMCMC avoids the integral approximation by using a joint distribution. The stationary distributions of $\theta$ and $\theta_a$ in EMCMC are guaranteed to be the original target distribution and flat target distribution respectively. Therefore, EMCMC does not need to use more “replicas” for a more precise approximation of the integral.

Reviewer comment (Experiments):
In some of the experiments (or maybe all) experiments the author collect samples of both theta and theta_a. Since the spirit of the paper is to perform Bayesian sampling, shouldn't they collect only theta samples?

Author response:
We have discussed it in Section 4.2 and Appendix C.3. Specifically, we collect both $\theta$ and $\theta_a$ in order to obtain more high-quality and diverse samples in a given time budget. The distribution of $\theta$ is the Bayesian posterior only given infinite time and computation budgets (since MCMC converges _asymptotically_ to the target distribution). Therefore, under a practical budget, we mix $\theta$ and $\theta_a$ in the sample set since both usually are from low-energy and flat regions. 

As shown in Appendix C.3, collecting both $\theta$ and $\theta_a$ achieves the best performance compared to using $\theta$ or $\theta_a$ only. We will add the above clarification in the final version.

Reviewer comment (Presentation):
Maybe more remands to the appendices, e.g. to appendix B for proofs, are needed in the main text.

Author response:
Thanks for the suggestion! We will add more remands in the final version.

Reviewer comment (Experiments):
Related to the previous question, did the author consider some adaptive tuning scheme for eta, such as the one proposed in the "Focusing" section of Pittorino et al ICLR '21?

Author response:
We did not consider adaptive tuning techniques for $\eta$ in our method. Our experiments have demonstrated that a constant value of $\eta$ already yields improvements over the baseline methods. We agree that using an adaptive tuning scheme for $\eta$ is a very interesting direction and leave it for future work.

Reviewer comment (Writing):
Typo in eq. 16, there should be a +.

Author response:
Thanks for pointing this out! We will correct it in the final version.

Reviewer comment (Reproducibility):
I wonder if after the burn-in the samples are collected at each iteration.

Author response:
We collect samples at the end of each epoch after burn-in, similar to previous works [1]. This choice is mainly due to memory constraints for deep neural networks.

Reviewer comment (Theory):
The theoretical results focus on a very restricted case of strong convexity. Although analysis of this situation provides interesting context for the relative abilities of the proposed method and existing methods, nothing can be firmly concluded in realistic settings.

Author response:
Strong convexity is a common assumption in analyzing the convergence of stochastic gradient MCMC, e.g. in [5,6,7]. Our theoretical results follow this line of works and provide convergence guarantees for the proposed method. Under the strongly convex setting, there are realistic tasks (e.g., logistic regression in Section 6.2) and we can firmly conclude that our method converges faster than baselines on these tasks. We agree that the convergence analysis under non-convex settings is a very interesting research question and we leave it to future work.

Reviewer comment (Novelty):
The degree of novelty is not especially high, as the method is a straightforward change to the approach from Chaudhary et. al 2019 and the experimental settings are fairly commonplace. On the other hand, the simplicity of the proposed method is part of its appeal.

Author response:
We would like to emphasize that our method is not a straightforward change from Chaudhary et. al 2019. In fact,  it is highly non-trivial to develop an efficient flatness-aware method, as biasing toward the flat basins often introduces substantial computational overhead [10, 4]. 

A key issue with Chaudhary et al. 2019 approach is its usage of nested Markov chains with Monte Carlo approximation, which significantly reduces convergence speed and estimation accuracy. In contrast, our method, by using an auxiliary variable, guarantees to converge to flat basins with fast speed and minimal overhead. To the best of our knowledge, the proposed auxiliary guiding variable and the joint distribution's form are novel and non-trivial solutions to flatness-aware learning. We believe that the simplicity of our method should not diminish its novelty. We appreciate your acknowledgment of the benefits from the simplicity of our approach.

Reviewer comment (Theory):
Many works find that flat minima lead to better generalization, but some works such as Dinh et al. 2017 (cited in the paper) claim that flatness is not necessarily for good generalization. Can the authors elaborate on the reasons for believing that flatness leads to better generalization and discuss whether generalization can be (or cannot be) achieved without flatness?

Author response:
In the second paragraph of the introduction, we mentioned the reason for flatness leads to better generalization from a Bayesian perspective. Specifically, to make prediction for testing data point $x^*$ in Bayesian learning, we need to obtain the following predictive distribution:
$$
p(y^*|x^*,D)=\int p(y^*|x^*,\theta)p(\theta|D)d\theta.
$$
Samples of $\theta$ from the flat basin within the posterior occupy a substantial volume and significantly contribute to this integral. Therefore, obtaining $\theta$ from the flat basin can lead to a more precise estimation of the posterior and the predictive distribution. This will lead to better generalization according to PAC-Bayes bounds [2, 3]. 

However, comprehensively justifying the relationship between flatness and generalization still remains an open question and beyond the scope of this work.