Entropy-MCMC: Sampling from Flat Basins with Ease
Authors: Bolian Li
Ruqi Zhang
Venue: ICLR 2024 poster (2024)
Keywords: MCMC
Bayesian Deep Learning
Flatness-aware Learning
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We propose a practical MCMC algorithm to sample from the flat basins of deep neural network posteriors.

Abstract: Bayesian deep learning counts on the quality of posterior distribution estimation. However, the posterior of deep neural networks is highly multi-modal in nature, with local modes exhibiting varying generalization performance. Given a practical budget, targeting at the original posterior can lead to suboptimal performance, as some samples may become trapped in "bad" modes and suffer from overfitting. Leveraging the observation that "good" modes with low generalization error often reside in flat basins of the energy landscape, we propose to bias sampling on the posterior toward these flat regions. Specifically, we introduce an auxiliary guiding variable, the stationary distribution of which resembles a smoothed posterior free from sharp modes, to lead the MCMC sampler to flat basins. By integrating this guiding variable with the model parameter, we create a simple joint distribution that enables efficient sampling with minimal computational overhead. We prove the convergence of our method and further show that it converges faster than several existing flatness-aware methods in the strongly convex setting. Empirical results demonstrate that our method can successfully sample from flat basins of the posterior, and outperforms all compared baselines on multiple benchmarks including classification, calibration, and out-of-distribution detection.

Published as a conference paper at ICLR 2024
ENTROPY-MCMC. SAMPLING FROM FLAT BASINS WITH EASE
Bolian Li, Ruqi Zhang
Department of Computer Science, Purdue University, USA (1 (1i4468, ruq. @ppuduue
ABSTRACT
Bayesian deep learning counts on the quality of posterior distribution estimation. However, the posterior of deep neural networks is highly multi-modal in nature, with local modes exhibiting varying generalization performance. Given practi- cal budget, targeting at the original posterior can lead to suboptimal performance, as some samples may become trapped in "bad" modes and suffer from overfitting. Leveraging the observation that "good" modes with low generalization error often reside in flat basins of the energy landscape, propose to bias sampling the posterior toward these flat regions. Specifically, we introduce an auxiliary guid- ing variable, the stationary distribution of which resembles smoothed posterior free from sharp modes, to lead the MCMC sampler to flat basins. By integrating this guiding variable with the model parameter, we create simple joint distri- bution that enables efficient sampling with minimal computational overhead. We prove the convergence of our method and further show that converges faster than several existing flatness-aware methods in the strongly convex setting. Empirical results demonstrate that our method can successfully sample from flat basins of the posterior, and outperforms all compared baselines on multiple benchmarks including classification, calibration, and out-of-distribution detection.
INTRODUCTION
The effectiveness of Bayesian neural networks relies heavily on the quality of posterior distribution estimation. However, achieving an accurate estimation of the full posterior is extremely difficult due its high-dimensional and highly multi-modal nature (Zhang et al. 2020b Izmailov et al. 2021) Moreover, the numerous modes in the energy landscape typically exhibit varying general izzatin performance. Flat modes often show superior accuracy and robustness, whereas sharp modes tend to have high generalization errors (Hochreiter & Schmidhuber] 1997; Keskar et al. [001]] Bahri et al., 2022 This connection between the geometry of energy landscape and generalization has spurred many works in optimization, ranging from theoretical understanding (Neyshabur et al., 2017 Dinh
ct al. 2017: Dziugaite & Roy| 2018 Jiang et al.] 2019a) to new optimization algorithms (Mobahi
2016
Izmailov et al.) 2018 Chaudhari et al. 2019 Foret et al. 2020).
However, most of the existing Bayesian methods are not aware of the flatness in the energy landscape during posterior inference (Welling & Teh 2011 Chen et al.) 2014 Ma al. 2015 Zhang al.) 2020b). Their inference strategies are usually energy-oriented and cannot distinguish between flat and sharp modes that have the same energy values. This limitation can significantly undermine their generalization performance, particularly in practical situations where capturing the full posterior is too costly. light of this, we contend that prioritizing the capture of flat modes is essential when conducting posterior inference for Bayesian neural networks. This is advantageous for improved generalization as justified by previous works (Hochreiter & Schmidhuber] 1997 Keskar al. 2017 Bahri et al., 2022). It can further be rationalized from Bayesian marginalization perspective: within the flat basin, each model configuration occupies a substantial volume and contributes signif- icantly to more precise estimation of the predictive distribution (Bishop) 2006). Moreover, existing flatness- aware methods often rely on a single solution to represent the entire flat basin (Chaudhari et al. 2019 Foret et al. 2020), ignoring the fact that the flat basin contains many high- perffrmnfn models. Therefore, Bayesian marginalization can potentially offer significant improvements over flatness-aware optimization by sampling from the flat basins (Wilson! 20200 Huang al. 2020)

---

Published as a conference paper at ICLR 2024
0
p())
Sharp mode
p(
Flat mode
10 0.)
(e)
faine gradient direction
Vf(e) original gradien dinesiio
(a) Training dynamics at each step
(b) Original and smoothed posterior distributions
Figure I: Illustration of Entropy-MCMC. (a) shows how the guiding variable O, pulls 0 toward flat basins; (b) shows two posterior distributions, where p(O, is smoothed distribution transformed from p ()) and only keeps flat modes. Entropy-MCMC prioritizes flat modes by leveraging the guiding variable O2 from the smoothed posterior as a form of regularization.
Prioritizing flat basins during posterior inference poses an additional challenge to Bayesian infer- ence. Even for single point estimation, explicitly biasing toward the flat basins will introduce sub- stantial computational overhead, inducing nested loops (Chaudhari et al., 2019 Dziugaite & Roy 2018), doubled gradients calculation (Foret et al. 2020; Möllenhoff & Khan 2022) or min-max problems (Foret et al., 2020). The efficiency problem needs to be addressed before any flatness- aware Bayesian method becomes practical for deep neural networks.
In this paper, we propose efficient sampling algorithm to explicitly prioritize flat basins in the en- engy landscape of deep neural networks. Specifically, we introduce an auxiliary guiding variable 0, into the Markov chain to pull model parameters toward flat basins at cach updating step (Fig. Lak. O, is sampled from a smoothed posterior distribution which climinates sharp modes based on local entropy (Baldassi al. 2016) (Fig. 1b Ou can also be viewed as being achieved by Gaussian convolution, a common technique in diffusion models (Sohl-Dickstein et al.) [2015] [ong & Ermon 2019). Our method enjoys simple joint distribution of and On, and the computational overhead is similar to Stochastic gradient Langevin dynamics (SGLD) (Welling & Teh 2011). Theoretically, prove that our method is guaranteed converge faster than some common flatness-aware meth- ods (Chaudhari et al., 2019 Dziugaite & Roy 2018) in the strongly convex setting. Empirically, we demonstrate that our method successfully finds flat basins efficiently across multiple tasks. Our main contributions are summarized as follows:
We propose Entropy-MCMC (EMCMC) for sampling from flat basins in the energy land- scape of deep neural networks. EMCMC utilizes an auxiliary guiding variable and simple joint distribution to efficiently steer the model toward flat basins.
We prove the convergence of EMCMC and further show that converges faster than several existing flatness-awan methods in the strongly convex setting.
We provide extensive experimental results to demonstrate the advantages EMCMC in sampling from flat basins. EMCMC outperforms all compared baselines on classification, calibration, and out-t--distribution detection with comparable overhead akin to SGLD. We release the code htttps://githu.. con.
2 RELATED WORKS
Flatness-aware Optimization. The concept of flatness in the energy landscape was first stud- ied by HHochreiter Schmidhuber (994)), and its connection with generalization was then em- pirically discussed by Keskar et al. (201)): Dinh et al. (2017); Jiang et al.] (2019b). To pursue flatness for better generalization, Baldassi et al. (2015) proposed the local entropy to measure the flatness of local modes, Baldassi al. (2016) used "replicated" models to implement local entropy, Entropy-SGD (Chaudhari et al., 2019) introduced a nested chain to approximate the local entropy, optimizer to minimize the worst-case
the current model. bSAM (Möllenhoff & Khan 2022) further improved SAM with a Bayes optimal convex
SAM Foret 2020) developed
lower bound, IPF (Bisla et al.] 2022) introduced low-pass filter to actively search flat basins, and
2

---

Published as a conference paper at ICLR 2024
SWA (Izmailov al. 2018) found that averaging weights along the trajectory of SGD training can also find flatter modes. Our Entropy-MCMr follows the local entropy measurement and collects more than single point to fully exploit the flat basins. For detailed comparisons with prior works considering local entropy. please refer to Appendix B
MCMC on Deep Neural Networks. Markov chain Monte Carlo is a class of general and prac- tical sampling algorithms (Andrieu et al., 2003), which has been applied to infer Bayesian neural network posteriors (Neal 2012). SGMCMC SGMlMC(Weling & Teh 2011] et al. 2015) methods use the mini-batching technique to adapt MCMC to deep neural networks. SGHMC (Chn et al. 2014) ex- ploited the second-order Langevin dynamies calibrate the stochastic estimates of HMC gradients. cSGMCMC (Zhang et al., 2020b) further improves sampling efficiency by leveraging a cyclical step size schedule. Symmetric Split HMC (Cobb & Jalaian 202I) developed a way to apply HMC to deep neural networks without stochastic gradients. Our Entropy-MCMC builds upon the SGMCMC framework and is designed to favor the flat basins the energy landscape during sampling.
PRELIMINARIES
Flatness-aware Optimization. One common flatness-ar aware optimization technique is to use the concept of local entropy, which measures the geometric properties of the energy landscape (Baldassi al. 2016, Chaudhari et al. 2019). The local entropy computed by:
1 10 21
F(O; m)) = log
exp (((0)
0'112 d0'
(1)
where ʃ()) is the loss function computed over the entire dataset and 17 is a scalar. The local entropy of point 0 is determined by its neighbors weighted by their distances, which considers the volume of local modes. Previous optimization methods minimize F(0;n) 7)) to find the flat minimum.
SGMCMC. Given a dataset D, a neural network with parameters E Rd the prior p(0) and the likelihood p())0). we can use Markov chain Monte Carlo (MCMC) sample from the posterior p(0[D) & exp(-U(0)), where the energy function is (0) = - ΣnED log p( () log p()). However, the computational cost for MCMC with large-scale data is too high be practical. SGM- CMC tackles this problem by stochastic gradient VU= based on subset of data E S D. We use Stochastic Gradient Langevin Dynamics (SGLD) (Welling Teh, 201T) in the paper as the backbone MCMC algorithm, which has the following updating rule:
0 0- 0 - vVo + -2
where a is the step size and e is standard Gaussian noise. Our method can also be implemented by other SGMCMC methods. During testing, Bayesian marginalization is performed to make predic- tions based on the sample set collected during sampling S ) and the predictive distribution obtained by p(ylæ, D) = plyla, 0
CoesPly'as
ENTROPY-MCMC
In this section, we present the Entropy-MCMC algorithm. We introduce the guiding variable 0, obtained from the local entropy in Section 4.1 and discuss the sampling strategy in Section 4.2
4.1 FROM LOCAL ENTROPY TO FLAT POSTERIOR
While flat basins in the energy landscape are shown to be of good generalization (Hochreiter & Schmidhuber 1997 Keskar et al., 2017; Bahri et al. 2022), finding such regions is still challenging due to the highly model typically needs extra force to keep the sampler away from sharp modes (Chaudhari et al. 2019 Foret et al. 2020). To bias sampling to flat basins, we look into the local entropy (Eq. which can eliminate the sharp modes in the energy landscape (Chaudhari et al. 2019).
nature of the DNN energy landscape. The updating direction of the
We begin by the original posterior distribution p(0 DD) x (f))) = elog (D)0) + log p0)), which contains both sharp and flat modes. By replacing the original loss function with
3

---

Published as a conference paper at ICLR 2024
local entropy, obtain smoothed posterior distribution in terms of a new variable 0a:
/ exp [ f(0)
1
p p(O,lD) & exp:F(ou:1)
0
0all 00-l
do.
(3)
21
The effect of local entropy on this new posterior is visualized in Fig. lb The new posterior mea- sures both the depth and flatness of the mode in pD) by considering surrounding energy values. Thereby, p0, (D) is expected to primarily capture flat modes in the energy landscape, which can be used as the desired external force to revise the updating directions of the model parameter O. Moreover, the smoothed posterior p D)) can be regarded being obtained through Gaussian con- volution, a common approach in diffusion models (Sohl-Dickstein et al. 2013 Song & Ermon, 2019). We also show the effect of hyper- syper-parameter 17 on the flatness of p(o.. D) in Appendix |A.4
However, the complex integral in Eq. 3 requires marginalization on the model parameter 6, which poses non-trivial challenge. Previous works using local entropy usually adopt an inner Markov chain for approximation (Chaudhari et al. 2019 Dziugaite & Roy 2018), which sacrifices the accu- racy in local entropy computation and induces computationally expensive nested loops training- We tackle this challenge in simple yet principled manner, eliminating the need for nested loops or approximation. This is achieved by coupling ~ p(O and 0, ~ p(0,(D) into a joint posterior distribution which enjoys simple form, as discussed in Lemma 1
Lemma 1. Assume j0" TT E Rd2d and i has the following distribution:
p D) p(0, (D) exp (
1 1 0all
2n
Then the marginal distributions of 0 and 02 are the original posterior p(0 D) and p(0,)D) (Eq. Further, the density p(0 D)) integrates to a finite quantity and thus it is mathematically well- well-defied.
This joint posterior offers three key advantages: i) by coupling 0 and O.. we avoid the intricate inte- gral computation, and thus remove the requirement of expensive nested training loops and mitigate the MC approximation error; ii) the joint posterior turns out to be surprisingly simple, making it easy to sample from both empirically and theoretically (details discussed in Sections and and (5): iii) after coupling, O₂ provides additional paths for 0 to traverse, making 0 reach flat modes efficiently.
4.2 SAMPLING FROM FLAT BASINS
We discuss how to sample from the joint posterior distribution (Eq. 4 in this section. We adopt SGLD (Welling & Teh, 2011). simple stochastic gradient MCMC algorithm that is suitable for deep neural networks, as the backbone of EMCMC sampling. More advanced MCMC algorithms can also be combined with our method. The energy function of the joint parameter variable 0 is (()) f(0) + 1 1|0 1a and thus its gradients is given by:
VOU( VOU(O) = Vo. (i)
]
+ ()) + 0a) ( 0)
].
(5)
For the model parameter 0, the original gradient direction vef(O) revised by 1 (0 0.) to get the flatness-aware gradient direction Vol(() as visualized in Fig. la Importantly, the practical implementation does not require computing Vo. (i)) through back bock-propagation, as we can utilize the analytical expression presented in Eq. 5 Therefore, despite Ø being in 2d dimension, our cost of gradient computation is essentially the same as d-dimensional models (e.g., standard SGLD).
With the form of the gradients in Eq. 5 the training procedure of EMCMC is straightforward us- ing the SGLD updating rule in Eq. 2 The details are summarized in Algorithm 1 At testing stage, the collected samples S are used to approximate the predictive distribution p(y D) s Ee, es x 0.). Our choice of sampling from the joint posterior distribution using SGLD, rather than Gibbs-like approach (Gelfandd 2000). is motivated by SGLD's ability simultaneously up- date both f and 0.. which is more efficient than alternative updating (see Appendix A for a detailed
Athhouhh we refer to p p(0 (D) as joint "posterior" to denote its dependency on the dataset, it is obtained through coupling rather than Bayes' rule. Thus, it does not have explicit prior distribution.
4

---

Published as a conference paper at ICLR 2024
discussion). For the sample set S, we collect both 0 and 0₂ after the burn-in period in order to obtain more high-quality and diverse samples in a finite time budget (see Appendix [.2 for the evidences that 0 and 0, find the same mode and Appendix E.3 for performance justification).
In summary, thanks to EMCMC's simple joint distribution, conducting sampling in EMCMC is straightforward, and its computational cost is comparable that of standard SGLD. Despite its algorithmic simplicity and computational efficiency, EMCMC is guaranteed to bias sampling to flat basins and obtain samples with enhanced generalization and robustness.
Algorithm 1: Entropy-MCMC
The model parameter O, guiding variable 0, E O and dataset D {(x), vi)) Results: Collected samples S C O;
Inputs:
6 O2
00;
/* Initialize */
for each iteration do
E - A mini-matte batihb sampled from D:
Uz - - log (( - log p(0) ()+ + 1
2n 110 02l
< a VoU= + 22. 61:
6,,2 */
O₂ - 02 - a V. 0. Uz + 22 E2
if after burn-in then - U (0, 0,)
/* Collect samples */
end
end
THEORETICAL ANALYSIS
In this section, we provide theoretical analysis the convergence rate of Entropy-MCMM and compare it with previous local-entropy-base methods including Entropy-SGD (Chaudhari et al. 2019) and Entropy-SGIDD (Dziugaite & Roy 2018) (used theoretical tool in the literature rather than a practical algorithm). We leverage the 2- 2-sssersteiinn distance bounds of SGLD, which as- sumes the target distribution be smooth and strongly log-concave (Dalalyan & Karagulyan 2019). While the target distribution in this case is unimodal, it still reveals the superior convergence rate of EMCMC compared with existing Ratness-aware methods. We leave the theoretical analysis on non-log-concave concave distributions for future work. Specifically, we have the following assumptions for the loss function f(.)) and stochastic gradients
Assumption 1. The loss function f() in the original posterior distribution æ = p(0)D) x exp (-) is Msmooohh and m-strongly convex (i.e., mI (0') = MT).
Assumption 2. The variance of stochastic gradients is bounded by E|| V f(0) (z) for some constant o
To establish the convergence analysis for EMCMC, we first observe that the smoothness and convex-
ity of the joint posterior distribution Tiaim (0, (0,0.) p00 (D) in Eq. 4 is the same the original posterior p( (D) which is formally stated in Lemma 2
Lemma 2. If A'Assmpption 7 holds and m 1/n M, then the energy function in the joint posterior distribution Jjoin O 0.) = p(0, (a( is also M-moooth and m-strongly convex.
With the convergence bound of SGLD established Dalalyan & Karagulyan (001)), we derive the convergence bound for EMCMC in Theorem [
Theorem 1. Under Assumptions [ let no be the initial distribution and uk be the distribution obtained by EMCMC after iterations. If S 1/1 S M and the step size a S 222mm + M), the
2- Wasserstein distance between μ. and Tjoi will have the following upper bound:
o (2ad)
W2(KK, Tjooi)) S - am)) Wa(no T) + (2ad) 1/2
(6)
1.65M + aamm
Assumption 2 are only for the convergence analysis. Our method and experiments not restricted to the strong convexity.
5

---

Published as a conference paper at ICLR 2024
Comparing Theorem 1 with the convergence bound of SGLD obtained by Dalalyan & Karagulyan (2019), the only difference is the doubling of the dimension, from d to 2d. Theorem implies that the convergence rate of EMCMC will have at most minor slowdown by constant factor compared to SGLDD while ensuring sampling from flat basins.
In contrast, previous local-entropy- -aased methods often substantially slow down the convergence to bias toward flat basins. For example, consider Entropy-SGD (Chaudhari et al. 2019) which mini-
mizes flattened loss function fna
F(0: 17) log Se exp -
L ll 011 do'
2ny
We discuss the convergence bound of Entropy- Entropy-SGD in Theorem 2 which shows how the presence of the integral (and the nested Markov chain induced by it) slows down the convergence.
Theorem 2. Consider running Entropy-SGD minimize the flattened loss function fia(() under Assumptions 7 and 2 Assume the inner Markov chain runs I. iterations and the 2-Wasserstein distance between the initial and target distributions always bounded by K. Let Sha represent the global minimum value of [aa(() and Es == (. -jar If the step size a S 2/ + M), then we have the following upper bound:
FK S (1
(1-1+
Cmn
K
En
+ 1)M) A(1
1+ nM where A? = amm)* 1.65
2m
[M+1/y)
(ad) 1/2 10
o² oadd
m+1//
Another example is Entropy SGLD (Dziugaite & Roy, 2018), , a theoretical tool established to ana- lyze Entropy-SGD. Its main distinction with Entropy-SoC is the SGLD updating instead of SGD updating in the outer loop. The convergence bound for Matropy-SGLan is established Theorem
Entropy-SGLD
Theorem 3. Consider running Entropy-SGLD to sample from Tga x exp F(0; 1)) under As- Assume the inner Markov chuin runs L. iterations and the 2- 2-Wasserstein dis- tance between initial and target distributions always bounded by k. Let vo be the initial distri- bution and vk be the distribution obtained by Entropy-SGLD after K iterations. If the step size a & (mmm=+ + M). then:
sumptions D and [
1 1 nM
(M (M/m)(ad) A(1+ nM)
Wa(uK, Tyaa) S (1-am))* K. W2( + +166
+ mym
m
where A = (1-aam) K + 1.65
M+1/g m++1/n
(ad)1/2
(add) 12
L.6G(M+1nn)+ m++/n
The complete proof of theorems is in Appendix C Comparing Theorem I and the convergence rates of Entropy-SGD and Entropy-SGLD algorithms are significantly hindered due the presence of the nested Markov chains, which induces large and complicated error term A. Since a and a are typically very small, the third term in Theorem will be much smaller than both the third term in Theorem and the second term in Thcorem 2
observe that
To summarize, the theoretical analysis provides rigorous guarantees on the convergence of Entropy- MCMC and further demonstrates the superior convergence rate of EnCropy-MCMCC compared to previous methods in the strongly convex setting.
6 EXPERIMENTS
We conduct comprehensive experiments to show the superiority of EMCMC. Section 6. and 6.3 demonstrate that EMCMC can successfully sample from flat basins. Section 6.2 verifies the fast convergence of EMCMC. on multiple benchmarks. Following Zhang al. (2020b), we adopt cyclical step size schedule for all sampling methods. For more implementation details, please refer to Appendix
6.55
the outstanding performance of EMCMC
6.1 SYNTHETIC EXAMPLES
To demonstrate EMCMC's capability to sample from flat basins, we construct a two-mode energy landscape [ -1JT. 0.5l) [ 1j I) containing sharp and flat mode. To make the

---

Published as a conference paper at ICLR 2024
(a) SGD
(b) SGLD
(c) EMCMC (0)
(d) EMCMC (0.)
Figure 2: Sampling trajectories on a synthetic energy landscape with sharp (lower left) and flat (top
right) modes. The initial point located at the ridge of two modes. EMCMC successfully biases toward the flat mode whereas SGD and SGLD are trapped in the sharp mode.
99.23
SaD
SGLD
99.00
Entropy SGLD
Training NL1
04
EMCMC
Testing (CC())
08.55 50
6.3
08.22
98
0.2
Entopy-SGi
SGLD
0.1
17.50-
Entropy-SGLa.
07.25
60
Figure 3: Logistic regression on MNIST in terms of training NLL and testing accuracy (repeated 10 times). EMCMC converges faster than others, which is consistent with our theoretical analysis.
case challenging, we set the initial point at (-0.2, -0.2), the ridge of the two modes which has no strong preference for either mode. The settings for this experiment are: 1) 0.5, a 5x 10-3, 1000 iterations, and collecting samples per 10 iterations. Fig. 2 shows that the proposed EMCMC finds the flat basin while SGD and SGLD still prefer the sharp mode due the slightly larger gradients coming from the sharp mode. From Fig. (c)) we see that the samples of 6, are always around the flat mode, showing its ability to eliminate the sharp mode. Although visits the sharp mode in the first few iterations, it subsequently inclines toward the flat mode, illustrating the influence of gradient revision by the guiding variable 0.. choosing an appropriate 1,, EMCMC will find the flat mode matter how it is initialized. This is due to the stationary distribution of O.. which is flattened and removes the sharp mode. Through the interaction term, O, will encourage the flat mode. We also show the results for different initialization in Appendix D.I
6.2 LOGISTIC REGRESSION
To verify the theoretical results on convergence rates in Section 5 we conduct logistic regres- sion MNIST (LeCun 1998) compare EMCMC with Entropy- SGD Chaudhari al. (2019). SGLD (Welling& & Teh 2011) and Entropy- (Dziugaite & Roy 2018). We follow Mac laurin
Adams 2015) and Zhang al. (2020a) to use subset containing 7s and and the resulting pos-
terior strongly log-concave, satisfying the assumptions in Section 5 Fig. 3 shows that FMCMC converges faster than Entropy-SG(L))
demonstrating the advantage of using simple joint distri- bution without the need for nested loops MC approximation, which verifies Theorems I 2
Besides, while EMCMC and SGLD share similar convergence rates, EMCMC achieves better gen- eralization as shown by its higher test accuracy. This suggests that EMCMC is potentially beneficial in unimodal distributions under limited budgets due to finding samples with high volumes.
6.3 FLATNESS ANALYSIS ON DEEP NEURAL NETWORKS
We perform flatness analysis with ResNet18 He et al. 2016) CIFARI00 (Krizhevsky 2009 We use the last sample of SGD, SGLD and EMCMC (averaged result from 0 and O..) respectively, and each experiment is repeated times to report the averaged scores.
Eigenspectrum of Hessian. The Hessian matrix of the model parameter measures the second- order gradients of a local mode on the energy landscape. Smaller eigenvalues of Hessian indicate a flatter local geometry (Chaudhari et al. 2019 Foret et al. 2020). Since computing the exact Hessian
A set of maaalmmximmmm points with zero gradients in all directions.

---

Published as a conference paper at ICLR 2024
10
10
Amax
9.2 10
Amax 3.7 x 10
10
Amax 1.3
As 4.9 x 10
Aa = 3.6 x 10
10
As 1.1 x 10
10
16
0.00
(a) SGD
(b) SGLD
(c) EMCMC
Figure 4: Eigenspectrum of Hessian matrices of ResNet18 CIFARIO0. r-axis: eigenvalue, y- axis: frequency. A nearly all-zero eigenspectrum indicates a local mode that is flat in all directions. EMCMC successfully finds such flat modes with significantly smaller eigenvalues.
SGD
8GD
SGLD
SOLD
EMCNNC
Toeting NLI
0.
DMCMC
nadem
from diccciio)
(a) 0 - Random
(b) + Random
(c) 0 O..
Figure 5: Parameter space interpolation of R 8 CIFARIO0. Exploring the neighborhood of
local modes from 9 to (a)-(b): a random direction in the parameter space, and (c): 6g. (a) and (b) show that EMCMC has the lowest and the most flat NLL and error curves. (c) shows that 0 and 6. converge to the same flat mode while maintaining diversity.
of deep neural networks is extremely costly due to the dimensionality (Luo et al., 2023), we use the diagonal Fisher information matrix (Wasserman] 2004) to approximate its eigenspectrum:
() >a]² (iag(I(O))
(VU EVU)²
where A1,.... Ad are eigenvalues of the Hessian. Fig. shows the eigenspectra of local modes discov- ered by different algorithms. The eigenvalues of EMCMC are much smaller compared with SGD and SGLD, indicating that the local geometry of EMCMC samples is flatter. The eigenspectrum comparison verifies the effectiveness of EMCMC to find and sample from flat basins.
Parameter Space Interpolation. Another way to measure the flatness of local modes directly interpolating their neighborhood in the parameter space (Izmailov et al. 20188.. Local modes located in flat basins are expected have larger widths and better generalization performance (Keskar al.,
2017 Chaudhari et al. 2019). The interpolation begins at 0 and ends at 0, (a random point near or Og 0z). The interpolated point 0g is computed by:
0,- (1 6/110 0e 0 (8/ 0. O..
(10)
where 8 is the Euclidean distance from 0 to Os. Fig. 5a and 5b show the training NLL and test- ing error respectively. The neighborhood of EMCMC maintains consistently lower NLL and er- rors compared with SGD and SGLD, demonstrating that EMCMC samples are from flatter modes. Furthermore, Fig. 5c visualizes the interpolation between 0 and O., revealing that both variables essentially converge the same flat mode while maintaining diversity. This justifies the benefit of collecting both of them as samples to obtain a diverse set of high-performing samples.
6.4 IMAGE CLASSIFICATION
We conduct classification experiments on CIFAR (Krizhevsky) 2009). corrupted CIFAR (Hendrycks)
& Dietterich, 2019b) and ImageNet (Deng et al. 2009), to compare EMCMC with both flatness- aware optimization methods (Entropy-a SGD (Chaudhari and bSAM (Möllenhoff & Khan 2022)) and MCMC methods (SGLD (Welling & Teh 201 and Entropy-SGi.n (Dziugaite Roy (2088). We use ResNet18 and ResNer50 (He et 201 for CI- FAR and ImageNet respectively. All sampling algorithms collect total of 16 samples for Bayesian marginalization, and all entries are repeated 3 times to report the tstd Table shows the re- sults the 3 datasets, in which EMCMC significantly outperforms all baselines. The classification results strongly suggest that by sampling from flat basins, Bayesian neural networks can achieve outstanding performance and EMCMC is an effective and efficient method to do
al., 2019). SAM (Foret et al. 2020)
The results for corrupted CIFAR (Hendrycks Dietterich) 2019a) are shown in Table show the robustness of EMCMC against multiple types of noises. The results averaged over all noise types,
8

---

Published as a conference paper at ICLR 2024
Table 1: Classification results on (a) CIFARIO/100, (b) corrupted CIFAR and (c) ImageNet, mea- sured by NLL and accuracy. EMCMC outperforms all compared baselines.
(a) CIFARIO and CIFARI00
CIFARTO
CIFARTOU
Method
1%)
NLL 1
(%)
NLL
SGD
JL.77
0.205 0.015
76.49 0.27
0.935 1 0.021
Entropy-SGID
95.11 +
0.184 + 0.020
77.45 + 0.03
0.895 + 0.009
SAM
95.25 0.12
0.166 1 0.005
78.41 1: 0.22
0.876 1 0.007
bSAM
95.53 0.09
0.165 + 0.002
78.92 0.25
0.870 1: 0.005
SGLD
9014
0.167 0.0011
78.79 0.35
0.854 E 0.031
Entropy-SGIA SGLD
04.46
0.194 0.020
77.98 = 0.39
0.897 + 0.027
EMCMC
95.69 0.06
0.162 + 0.002
79.16 + 0.07
0.840 + 0.004
(b) Corrupted CIFAR (ACC (%) 1)
(c) ImageNet
Severity
2
3
4
5
Metric
NNL 1
Top- (96)
Top-5 (%)
SGD
88.43
82.43
76.20
67.93
55.81
SGD
0.960
76.046
92. 776
SGLD
88.61
82.46
76.49
69.19
56.98
SGLD
0.921
76.676
93.174
EMCMC
88.87
83.27
77.44
70.31
58.17
EMCMC
0.895
77.096
93.424
Table OOD detection on CIFAR-SVHN. The predictive uncertainty quantified by EMCMC the best among the compared algorithms.
CIFARTO-VHHN
CIFARTOO-SHHI
Method SGD
AUROC (%)
AUPR (%)
AUROC (%)
AUPR (%)
98.30
99.24
71.96
84.08
Entropy-SGD
98.71
99.37
79.15
86.92
SAM
94.23
95.67
74.56
84.61
SGLD
97.66
98.64
72.51
83.35
Entropy- SGLD
90.07
91.80
71.83
82.89
EMCMC
98.15
99.04
81.14
87.18
and the severity level refers to the strength of noise added to the original data. EMCMC consistently outperforms all compared baselines across all severity levels, indicating that samples from flat basins are more robust to noise. The results for individual noise types are shown in Appendix D.3
6.5 UNCERTAINTY AND OOD DETECTION
To illustrate how predictive uncertainty estimation benefits from flat basins, we evaluate EMCMC on out-of-distributio (OOD) detection. We train each model CIFAR and quantify uncertainty using the entropy of predictive distributions (Malinin & Gales 2018). Then we use the uncertainty to detect SVHN samples in a joint testing combined by CIFAR and SVHN (Netzer al. 2011). We evaluate each algorithm with Area under ROC Curve (AUROC) (MeClish) 1989) and Area under Precision-Recal curve (AUPR) (Olson & Delen 2008). All other settings remain the same as the classification experiments. Table the evaluation results, where EMCMC outperforms nearly all baselines, especially when trained CIFARI00. This indicates that predictive uncertainty estimation is more accurate if the samples are from flat basins of the posterior. The confidence calibration experiments are shown Appendix D.4
CONCLUSION AND DISCUSSION
We propose a practical MCMC algorithm sample from flat basins of DNN posterior distributions. Specifically, we introduce a guiding variable based on the local entropy to steer the MCMC sampler toward flat basins. The joint distribution of this variable and the model parameter enjoys a simple form which enables efficient sampling. We prove the fast convergence rate of our method compared with two existing fatness-aware methods. Comprehensive experiments demonstrate the superiority of our method, verifying that it can sample from flat basins and achieve outstanding performance on diverse tasks. Our method is mathematically simple and computationally efficient, allowing for adoption drop-in replacement for standard sampling methods such SGLD.
The results hold promise for both Bayesian methods and deep learning generalization. On the one hand, we demonsirate that explicitly considering flatness in Bayesian deep learning significantly improve generalization, robustness, and uncertainty estimation, especially under practical computa- tional constraints. On the other hand, we highlight the value of marginalizing over flat basins in the energy landscape, as a means to attain further performance improvements compared to single point optimization methods.

---

Published as a conference paper at ICLR 2024
REFERENCES
Ahmad Ajallocian and Sebastian U. Stich. On the convergence of sgd with biased gradients. Jour- nal of Machine Learning Research, 2020. URL https://appt www.sstssscccsooleeee.e org
Corpus ID: 1::358:582341
Christophe Andrieu, Nando De Freitas, Arnaud Doucet, and Michael I Jordan. An introduction to meme for machine learning. Machine learning, 50:5-43, 2003.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness- aware minimization improves language model generalization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp- 7360-7371, 2022.
Carlo Baldassi, Alessandro Ingrosso, Carlo Lucibello, Luca Saglietti, and Riccardo Zecchina. Sub- dominant dense clusters allow for simple learning and high computational performance in neural networks with discrete synapses. Physical review letters,
2015.
Carlo Baldassi, Christian Borgs, Jennifer T Chayes, Alessandro Ingrosso, Carlo Lucibello, Luca Saglietti, and Riccardo Zecchina. Unreasonable effectiveness of learning neural networks: From accessible states and robust ensembles to basic algorithmic schemes. Proceedings the National Academy of Sciences, E13(58):E7655-ET)) E7662, 2016.
Christopher M Bishop. Pattern recognition and machine learning, volume 4. Springer, 2006.
Devansh Bisla, Jing Wang, and Anna Choromanska. Low-pass filtering sgd for recovering flat optima in the deep learning optimization landscape. In International Conference on Artificial Intelligence and Statistics, pp. 8299-8333 8339. PMLR, 2022.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):
124018, 2019.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pp. 1683-1691. PMLR, 2014.
Adam Cobb and Brian Jalaian. Scaling hamiltonian monte carlo inference for bayesian neural networks with symmetric splitting. In Uncertainty in Artificial Intelligence, pp. 685. PMLR, 2021.
Amak Dalalyan and Avetik Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications. 129(1)):5278-3)3 53118, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-258 255. Ieee, 2009.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Learning, pp. 1019-1028. PMLR, 2017.
Gintare Karolina Dziugaite and Daniel Roy. Entropy-sgd optimizes the prior of pac-bayes bound: Generalization properties entropy-sgd and lata-dependent priors. In International Conference Machine Learning, pp. 1377-1386. PMLR, 2018.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations, 2020.
Alan Gelfand. Gibbs sampling. Journal of the American statistical Association, 95(452):1300 1304, 2000.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
10

---

Published as a conference paper at ICLR 2024
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common ruptions and perturbations. In International Conference on Learning Representations, 2018.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. Pruceedings of the International Conference on Learning Represen- lutions, 2019a.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In International Conference on Learning Representations, 2019b.
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and distr--dist examples in neural networks. In International Conference on Learning Representations, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in neural information processing systems, 1994.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9()):1-4,, 1997.
W Ronny Huang, Zeyad Ali Sami Emam, Micah Goldblum, Liam H Fowl, JK Terry. Furong Huang, and Tom Goldstein. Understanding generalization through visualizations. In "I Can't ' Believe Not "exter""NeurlPP" workshop, 2020.
Izmailov, AG Wilson, D Podoprikhin, Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 876-885, 2018.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffmann, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? International conference on machine learn- ing. pp. 4629-4640 PMLR, 2021.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2019a.
Yiding Jiang. Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where find them. In International Conference Learning Representations, 2019b.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang. Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, 2017.
A Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Tront, 2009.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. con/exdb/mmist// 1998.
Yin Tat Lee, Ruoqi Shen, and Kevin Tian. Structured logconcave sampling with a restricted gaussian oracle. In Conference on Learning Theory, pp. 2993-3050, PMLR, 2021.
Mu Li, Tong Zhang, Yuqiang Chen, and Alexander J Smola. Efficient mini-batch training for stochastic optimization. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 661 661-670, 2014.
Xinyu Luo, Christopher Musco, and Cas Widdershoven. Dimensionality reduction for general KDE mode finding. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research. pp. 23067-23082. PMLR, 2023. URL httpp::/ppocceeaaaii .tt.sresss/.//p/// 12202/1u623c.nm html
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme. Advances in neural information processing systems, 28, 2015.

---

Published as a conference paper at ICLR 2024
Dougal Maclaurin and Ryan P Adams. Firefly monte carlo: exact meme with subsets of data. In Proceedings of the 24th International Conference on Artificial Intelligence, pp. 4289-4295, 2015.
Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. Advances in neural information processing systems, 31, 2018.
Donna Katzman McClish. Analyzing portion of the roc curve. Medical decision making, 9(3): 190-195, 1989.
Hossein Mobahi. Training recurrent neural networks by diffusion. arXiv preprint arXiv:160100111.. 2016.
Thomas Möllenhoff and Mohammad Emtiyaz Khan. Sam as an optimal relaxation of bayes. In The Eleventh International Conference Learning Representations, 2022.
Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated proba- bilities using bayesian binning. Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. NIPS Workshop on Deep Learning and Unsupervised Feature Learning. 2011.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization in deep learning, Advances in neural information processing systems, 30, 2017.
David L Olson and Dursun Delen. Advanced data mining techniques. Springer Science & Business Media, 2008
Marcelo Pereyra. Proximal markov chain monte carlo algorithms. Statistics and Computing, 26: 745-760, 2016.
Fabrizio Pittorino, Carlo Lucibello, Christoph Feinauer, Gabriele Perugini, Carlo Baldassi, Elizaveta Demyanenko, and Riccardo Zecchina. Entropic gradient descent algorithms and wide flat minima. In International Conference on Learning Representations, 2020.
Saurabh Singh and Shankar Krishnan. Filter response normalization layer: Eliminating batch de- pendence in the training of deep neural networks. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 11237-11246, 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265. PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Larry Wasserman. All of statistics: a concise course in statistical inference, volume 26. Springer, 2004.
Max Welling and Yee W Tch. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-11), pp. 681-688, 2011.
Andrew Gordon Wilson. The case for bayesian deep learning. arXiv preprint arXiv:2001.10995. 2020.
Ruqi Zhang, A Feder Cooper, and Christopher De Sa. Asymptotically optimal exact minibatch metropolis--hstingg. Advances in Neural Information Processing Systems, 13:19500-19510, 2020a.
Ruqi Zhang, Chumyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. In International Conference on Learning Representations, 2020b.
12

---

Published as a conference paper at ICLR 2024
A ALGORITHM DETAILS
We list some details of the proposed Entropy-MCMC in this section, to help understand our code and reproduction. As discussed in Section 4.2 the updating rule of Entropy-MCMC can be written as:
0 0 vV + V2a E,
(11)
which is full-batch version. We will show how to apply modern deep learning techniques like mini-batching and temperature the updating policy in the following sections.
A.1 MINI-BATCHING
We adopt the standard mini-batching technique in our method, which samples a subset of data points per iteration (Li et al. 2014). We assume the entire dataset to be 'D {( V: ) Then a batch
sampled from D E {2 (1) A D with M N. For the entire dataset, the loss function
is computed by:
N log p(y a;; 0) - logp(())
((0) x
Σ
(12)
==1
and to balance the updating stride per iteration, the loss function for mini-batch is:
N
(=
log p(u [s, 0) log P(O).
(13)
M
Therefore, if we average the mini-bated loss over all data points, we can obtain the following form:
1 logp(u/z,,00 0) N log)
(0)
(14)
M
we regard the averaging process as a modification on the stepsize a (i.e., ā = a/N), we will have the following form for the updating policy:
A - V (0) + v
2
- V
(0)
0
0a
2
c
2m
á
(15)
1
2 aN
V
(z
110 02
1
N
Therefore, the updating rule in Eq
can be equivalently written as:
i 0- 0
(16)
A.2 DATA AUGMENTATION AND TEMPERATURE
We apply data augmentation, which is commonly used in deep neural networks, and compare all methods with data augmentation in the main text. Here, we additionally compare the classification results without data augmentation in Table 3 to demonstrate the effectiveness of EMCMC in this case.
Table 3: Comparison data augmentation of 3 baselines on CIFAR10. EMCMC outperforms previous methods with and without data augmentation.
Augmentation
SGD
SGLD
EMCMC
89.60
89.24
89.87
95.59
95.64
95.79
Besides, in the updating policy, a noise term is introduced to add randomness to the sampling pro- However, in mini-batch training, the effect of noise will be amplified that the stationary
13

---

Published as a conference paper at ICLR 2024
distribution of might be far away from the true posterior distribution Zhang al. 2020b Izmallov et al. 2021 Therefore, we also introduce a system temperature to address this problem.
Formally, the posterior distribution is tempered to be p(0) D) exp(-U(0)/T) with averagely sharpened energy landscape. Similarly, we can regard the temperature effect as a new stepsize or az ax/T, and the updating policy would be:
Ab aa. Vi
(=( + 2n 11 0ul1²)
/T -
VVa
aN
(17)
1 2 0all
2T oer 6 0
-or V
[=(0)
To empirically determine how temperature influences classification performances, we compare different temperature levels in Table 4 We compare different temperature magnitudes for both SGLD (Welling & Tch (201]) and EMCMC. With smaller temperatures, both sampling algorithms improve significantly from 10 10-4, and EMCMC outperforms SGLD under most of temper- atures. These findings justify the statement that temperature is needed for good results under data augmentation (Izmailov al. 2021) and the superiority of EMCMC against other SG-MCMC al- gorithms.
Table 4: Test accuracy (%) comparison on different temperature magnitudes, with data augmenta- tion, on CIFARIO.
Temperature
10"
10-1
10-2
10
10
10-5
10
SGLD
82.89
91.67
94.12
94.95
94.86
94.88
95.17
EMCMC
81.93
91.67
94.47
94.98
95.06
95.06
95.18
A.3 GIBBS-LIKE UPDATING PROCEDURE
Instead of jointly updating, we can also choose to alternatively update and 0g. The conditional distribution for the model O is:
p(0 D) = p(0, Oa(
1 exp (
x Ze.
-f(0)
100 0a113
(18)
p(Dx(D)
21
where Ze. exp F(),: a constant. While for the guiding variable O.. its conditional distribu- tion is:
p(0, D) p(0, 0 (D)
exp
0 0all
(19)
D(0 Zo
2m
where Zo exp (ff is constant. Therefore, with Guassian noise, 0,, is equivalently sampled from ((0,TI). and the variance 17 controls the expected distance between and 6.. obtain samples from the joint distribution, The advantage of doing Gibbs-like updating is that sampling 0, can be done exactly. This Gibbs- like updating also shares similarities with the proximal sampling methods (Pereyra 2016 Lee al. 2021).
sample from 0(0)0,, D) and p(0, D) alternatively.
Empirically, observe that joint updating yields superior performance compared Gibbs-like updating due to the efficiency of updating both and 0, at the same time.
A.4 EFFECT OF THE VARIANCE TERM ON THE SMOOTHED TARGET DISTRIBUTIONS
We visualize flattened target distributions p(O,l under different values of in Fig. 6 As 17 be- comes smaller, the target distribution will be flatter.
14

---

Published as a conference paper at ICLR 2024
p(e)p)
p[0, n= 10
p(o, 10
p(e, n 10
p(o,(0), n 10
Figure 6: Effect of on the flattened target distribution. Smaller will make the original distribution flatter.
B DETAILED COMPARISON WITH RELATED WORKS
B.1 COMPARISON WITH ROBUST ENSEMBLE
Intropy-MCMM. substantially differs from Baldassi et al. (016) approach and not special case of the Robust Ensemble (RE) with y replica. There are several key differences between Entropy- MCMC and RE in Baldassi et al. (016):
In Baldassi et al. (2016), the definition of Robust Ensemble, as outlined in Eq. 4, traces out the reference (0, in our paper) and only has y identical replicas the system, resembling ensemble methods. This can be further verified by Pittorino et al. (020)) As shown in Algorithm 2 in PPittorino et al. (2200) Robust Ensemble with = will essentially be the standard SGD. In contrast, Entropy-MCMC has two variables performing different roles, unlike ensemble methods. 0g cannot be traced out and plays a crucial role in leading flat basins during training. Furthermore, Robust Ensemble typically requires at least four models (i.c., y 2 3) while EMCMC only needs two models. Entropy-MCMC greatly reduces computational costs, especially when using large deep neural networks.
Though Eq. 3 with y in Baldassi et al. (016) is equivalent to the proposed joint dis- tribution, Baldassi et al. (2016) did not consider the marginal distribution of replicas. In contrast, our work highlights that the marginal of the replica is the original target distribu- tion. This distinction arises because our joint distribution is derived by coupling the original and flattened distributions, different motivation from Baldassi al. (2016). Moreover, to the best of our knowledge, Eq. 3 has not been practically applied (as seen in Baldassi et al. (2016) and Pittorino et al. (2020)) primarily serving as an intermediate step deriving Robust Ensemble (which uses identical replicas).
RE is an optimization method that aims to find a flat optimum as the point estimation. In contrast, PatropyyMCMM is sampling method that aims to sample from the flat basins of the posterior distribution.
In summary, Entropy-MCMC 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 's form are novel and non-trivial solutions to flatness aware learning.
15

---

Published as a conference paper at ICLR 2024
B.2 COMPARISON WITH ENTROPY-SGD
Entropy-MCMCC is not a straightforward change from Chaudhari et al. (2019) is highly non-trivial develop an efficient flatness-aware method, as biasing toward the flat basins often introduces substantial computational overhead (Chaudhari et al. 2019 Foret et al. 2020).
key issue with Chaudhari et (2019) approach is its usage of nested Markov chains with Monte Carlo approximation, which significantly reduces convergence speed and estimation accuracy. In contrast, Matropy-MCMC. with an auxiliary guiding variable, guarantees to converge to flat basins with fast speed and minimal overhead.
C PROOF THEOREMS
C.I LEMMA
Proof. Assume = 10" is sampled from the joint posterior distribution:
1 = 0al1
0 D) p(0,0, exp -s(0)
(20)
2nj
Then the marginal distribution for 0 is:
p(0 = / p00,0, 0, D)do,
= (2TT)) z- / exp f((0)
1
1|0
0al do,
2n
(21)
2- ex=eep((((())
exp
0
0z1 d0a
2m
= = (x-xe((-f(0))
where 7 -eex((- ('ep((_((())ee is the normalizing constant, and it is obtained by:
1/// exp -f(0) - 2n -110 0 0a113 d0.d0 = (2mm) $ exp(- = (2mm) (22)
This verifies that the joint posterior distribution p(0, 0, (D) is mathematically well-defined larly, the marginal distribution for 0g is:
Simi-
p(O,
p(0, 0 dD
(23)
exp -(())
0 0al1
d0
2m
exp F(O.;n)
C.2
LEMMA 2
Proof. Note that we have
V2 1
22log log Kjoint"
11
and after a row reduction, we get
(2((0) 0 --1 17
1
*The exact form of the joint posterior 000,0. ( (2mn) 4 (xp(-f(()) 1 0 0.113)
16

---

Published as a conference paper at ICLR 2024
The eigenvalues for this matrix are the eigenvalues of V f(6) and 1/n. By the assumption 1 1/n $ M, we have
V
f(0)
0 1
ml
MI,
which means V2 log Jioint is also -Msmooth and m-strongly convex function.
In most scenarios, m often takes on very small values and M tends to be very large. Therefore, the assumption m Σ 1/n S M mild.
C.3 PROOF OF THEOREM I
Proof. The proof relies on Theorem 4 from Dalalyan & Karagulyan (2019). Lemma 2 has already provided us with smoothness and strong convexity parameters for Tjoint- will now address the bias and variance of stochastic gradient estimation. The stochastic gradient given by [V] + ;, (6' 0). - ( 0)1² As ff())) is unbiased and has a variance of o², the stochastic gradient our method is also unbiased and has variance o². Combining the above results, we are ready to apply Theorem (Dalalyan & Karagulyan, 2019) and obtain
d²(2d))
W2(KK, Wjjin) S (1 am) W2(uo, T) + 1.655M/m))22 1.65M + +aym
C.4 THEOREM
Proof. Let T T'0' x exp(- 2n 110 -013). is easy to see that m 1mmm++1mm +1/ 3 (=(-log <
M + 1/n. Based on Theorem 4 in Dalalyan & Karagulyan (2019), the 2- 2-Wasserstein distance for the inner Markov chain is
JM + 1/n)
(ad))
We(Se, & (1 am)' + 1.65
(ad) 1/2
m 1/m
1.65 + 1/1) + t m m 1/n
M + 1/n
1.65(M + (² 1/2
< (1 am)) + 1.65
(ad) 1/2
+
+ 1/1
+ 1/11
Now we consider the convergence of the outer Markov chain. We denote Tatt (0) x F(P: From Chaudhari al. (2019). know that
inf
ml 1 -2 2 Thai I sup
MI.
I + 2 f(O)
1 f(0)
Since ml 1 52)
MI, it follows
1
inf
1+
2 1+ nM
sup
I+ 1 ((0)
$ 1 mm
Therefore,
m
1 1 2
log Tiat
M 1 + 17m
/.
1 nM
The update rule of the outer SGLD
= - a/1((0 (0') /20g.
17

---

Published as a conference paper at ICLR 2024
The gradient estimation can be written as 0 - Ex [8]] + (E [0'] Egg, (8')) which can be regarded as the true gradient 0 E [6'] plus some noise (E,r [6'] Eg,, [01)). The bias of the noise can be bounded as follows
|1 10 -8 (d. 0%c
E [0] Eçg 10110111
/ 110- - ltall1 d
Since the inequality holds for any we can take the infimum over all possible distributions to conclude
E [0'] E,, 10'11 W2(CL,
Furthermore, we note that the variance of the noise zero. Therefore, by applying Theorem 4 in Dalalyan & Karagulyan (2019) we get
1 ) (M/m)(ad) A(1 + mMM
W2 (2K Raa) S (1 -am) W2(wo, Taat) 1.65
(++ (mm
m
C.5 THEOREM
Proof. Compared to Entropy-SGLDD, the only difference with Entropy-SGD is doing SGD updates instead of SGLD updates the outer loop. Therefore, the analysis for the inner Markov chain
remains the same in Theorem 3 To analyze the error of SGD in the outer loop, follow the
results in Ajalloeian Stich (2020). Since the strongly convex parameter for fua 149
m by
Section 4.2 and Assumption in Ajallocian & Stich (2020), we know that
E.
Ett1
v fat (n)
A
2
m
Fx.
F1+1
E, S
2 A
+ M
a
am
E141 S (1
E. + aaA
1 + nM
By unrolling the recursion, obtain
EK I (1
(1 -1-1 1+ nM
am
K
A(1 + MM)
Eg
2m
D ADDITIONAL EXPERIMENTAL RESULTS
We list the additional experimental results in this section, to demonstrate the superiority of our method and show some interesting findings.
D. ADDITIONAL SYNTHETIC EXAMPLES
To demonstrate that EMCMC can bias toward the flat mode under random initialization, we conduct additional synthetic experiments under two different initialization settings. Specifically, we set the initial point to be (-0.4, -0.4) to prefer the sharp mode (Fig. 7 and (0.0, 0.0) to prefer the flat mode (Fig. (8). EMCMC can find the flat mode under all initialization settings, while SGD and SGLD are heavily affected by the choices of initialization.
Besides, we also conduct experiments of running EMCMC for a sufficiently long time to see whether it can converge to the true target distribution (both modes). The results shown in Fig, 9 successfully finds both modes and f., still samples from the flat mode, both converging to their target confident to claim that Entropy-MCMe prioritizes
distributions. Compared with Fig. 7 flat modes under limited computational budget and will eventually converge to the full posterior with adequate iterations.
18

---

Published as a conference paper at ICLR 2024
(a) SGD
(b) SGLD
(c) EMCMC (0)
(d) EMCMC (0.)
Figure 7: Synthetic experiments with sharp-mode- biased initialization.
(a) SGD
(b) SGLD
(c) EMCMC (0)
(d) EMCMC (0.)
Figure 8: Synthetic experiments with flat-mode-biase initialization.
D.2 PARAMETER SPACE INTERPOLATION
As the supplement for Fig. show additional interpolation results to demonstrate some interest- ing findings about the model 0 and the auxiliary guiding variable 0.. The additional interpolation can be separated into the following types:
D.2.1 TOWARD RANDOM DIRECTIONS
We show the interpolation results toward averaged random directions (10 random directions) in Fig. 10 For the training loss, the auxiliary guiding variable 0 is located at a flatter local region with relatively larger loss values. For the testing error, the guiding variable 0, consistently has lower generation errors, which illustrates that the flat modes are preferable in terms of generalization.
D.2.2 BETWEEN MODEL PARAMETER AND GUIDING VARIABLE
The line between the model parameter f and the guiding variable 0, special direction in the parameter space. The NLL and testing eπror are both much lower than random directions, which shown in Fig. II Besides, this special direction is biased toward the local region of 0., with averagely lower testing errors. This finding justifies the setting of adding 9, the sample set S, since the generalization performance of 0g is better.
D.3 CLASSIFICATION ON CORRUPTED CIFAR
We list the detailed classification results on corrupted CIFAR (Hendrycks & Dietterich, 2018 in Fig. [2 where each corruption type is evaluated at a corresponding subfigure. For the majority of corruption types, our method outperforms other baselines under all severity levels, and is superior especially under severe corruption.
D.4 CONFIDENCE CALIBRATION
The empirical results for confidence calibration are listed in Table The maximum softmax proba- bilities are adopted confidence scores (Hendrycks & Gimpel 2016), which are then used to predict misclassification over the test set. We use Area under ROC Curve (AUROC) (McClish) 1989) and Expected Calibration Error (ECE) (Naeini al. 2015) for evaluation. The results indicate that EM- CMC has good calibration capability compared with other baselines, and does not suffer from the overconfidence problem.
19

---

Published as a conference paper at ICLR 2024
2
-2
2
-2
0
8
(a)
(b) 0
Figure 9: Synthetic experiments with sufficient iterations. Both O and 0, will fully characterize their target distributions with enough iterations.
0.77
0.260
0.76
0,
0,
Training NLL
Testing Error
0.250
0.74
0.73
0.260
0.71
0.200
0.4
1.0
0.0
Distance
Distance
(a) Random interpolation by training NLL
(b) Random interpolation by testing error
Figure 10: Interpolation toward averaged random directions on CIFAR100, comparing the model and the guiding variable 6..
E ABLATION STUDIES
We empirically discuss several important hyper-parameters and algorithm settings in this section, which justifies our choice of their values. The hyper-parameter choices are tuned via cross- validation, and the ablation studies primarily discuss the influence of these hyper-parameters empirical results.
E. VARIANCE TERM
We compare different choices of the variance term 1) to determine influence on the performance. The experimental results are shown in Fig. 13 Generally, setting to be 10-3 for CIFAR10 and 10- -2 for CIFAR 100 will induce outstanding test accuracy. This also implies that the energy land- scapes of CIFAR 10 and CIFARI00 may be different.
E.2 STEP SIZE SCHEDULES
We compare different types of stepsize schedules in Table Specifically, we assume the initial and final stepsize to be do and an respectively. T is the total number of epochs and is the current
20

---

Published as a conference paper at ICLR 2024
0.29
Training NL
0.90
Testing Error
0.27
0.85
0.25
0.80
0.75
0
8
0.25
0
0.
0.70
0.22
-5.
5.0
7.5
-10.9
-2.5
0.0
2.5
5.0
10.0
Distance
Distance
(a) Interpolation between 0 and 0. by training NLL
(b) Interpolation between e and 0. by testing error
Figure 11: Interpolation between the model and the guiding variable 0₂ in terms of training NL. and testing error on CIFARI00.
Table 5: Confidence calibration results on CIFAR. EMCMC has good calibration capability to avoid the overconfidence problem.
CIFARTO
CIFARTO0
Method
AUROC (%)
ECE (%)
AUROC (%) 1
ECE (%)
SGD 92.90
2.94
87.85
6.41
Entropy-SGD
93.59
2.46
87.09
5.75
SAM
94.91
1.45
88.94
5.81
SGLD
94.42
0.78
87.47
2.15
Entropy-SGLD
93.80
0.33
87.59
1.66
EMCMC
93.76
0.52
87.69
3.59
epoch. The detailed descriptions of stepsize schedules are listed in Table The eyclical stepsize is best among all stepsize schedules.
Table 6: Formulas or descriptions of different stepsize schedules.
Name
Formula/Description
constant
(xl) = Okg
linear
(t)) = T
T-t (ap a1) + aq
exponential step
(1) oro (a1/00)*/7
Remain the same stepsize within one "step", and decay between "steps".
cyclical
Follow Eq. 1 in Zhang et al. (2020b).
E.3 COLLECTING SAMPLES
Due to the introduction the auxiliary guiding variable O., the composition of sample set S has multiple choices: only collect samples of 0, only collect samples of 0,, collect both samples. We conduct the comparison of all choices and the results are reported in Table g shows that using samples from both 0 and 6,, gives the best generalization accuracy.
E.4 NORMALIZATION LAYERS
During testing, the usage bath normalization layers (BN) in the model architecture induces a problem regarding the mini-batch statisties. The mean and variance of batch need calculated through at least one forward pass, which is not applicable for the guiding variable 0, since it is updated by the distance regularization during training. We try different solutions for this problem,
21

---

Published as a conference paper at ICLR 2024
THANE AFT PV)
TraRe ACK CO
RDD LNCXC
Soneity
Scntty
(a) brightness
contrast
(c) defocus blur
(d) elastic transform
5CD
SGD
5GD
SOLD
Tocke
S0LD
IMCNC
Sourto
Scricty
Scoiitt
Scndty
(e) fog
(l) frost
g gaussian blur
(h) gaussian noise
SOLD
SOLD
$GLD IVOW
DXOMC
2OOK
ii glass_blur blur
(j) impulse.nois
(k) jpeg comppossressinn
motion_ blur
Km
To ACCO
50LD 3HOC
Seatly
Serech
Swerdy
Sourdy
(m) pixclate
(n) saturate
(0) shot_noise
(p) snow
(q) spatter
(r) specceee.nooe noise
coom_blur
Figure 12: Classification accuracies under different severity levels corrupted CIFAR. The results are shown per corruption type. Our method outperforms the compared baselines on most of corrup- tion types, especially under high severity levels.
including one additional forward pass and the Filter Response Normalization (Singh & Krishnan, 2020). The comparison is listed in Table where simply adding one additional forward pass during testing can achieve promising accuracy with negligible computational overhead.
22

---

Published as a conference paper at ICLR 2024
792
Testing ACC (%)
Testing ACC (%)
700
1335
91.9
16
18"
107
14
10
n
CIFARIO
(b) CIFARI00
Figure 13: Comparison of different variance levels for Entropy algorithm. 1/ 10 is appropriate for CIFAR1O, and = 10-2 is appropriate for CIFAR100.
Table 7: Comparison of stepsize schedules CIFARIO0. Entropy-MCMM
The cyclical stepsize is the best for
Learning Rate Schedule
constant
linear 87.89
exponential 87.75
step
cyclical 89.93
Testing ACC (%) f
88.04
89.59
E.5
SGD BURN-IN
We also try SGD burn-in in our ablation studies, by adding the random noise term only to the last few epochs to ensure the fast convergence. We evaluate different settings of SGD burn-in epochs in Table TO We find that adding 40 burn-in epochs per 50 epochs is the best choice.
23

---

Published as a conference paper at ICLR 2024
Table 8: Ablation study the composition of sample set S on CIFARIO. With samples from both and a., the Bayesian marginalization can achieve the hest accuracy.
ACC (%)
95.58
95.64
95.65
Table 9: Comparison of different normalization layers CIFARIO. Simply adding additional
forward pass during testing with standard batch normalization is the best solution.
Normalization Layer
ACC (%)
Time (h)
BN
95.40
1.8
BN (one additional forward)
95.47
1.9
FRN (Singh & Krishnan 2020
93.92
2.5
Table 10: Comparison of different SGD burn-in epochs CIFARIO, In 50-epoch round, using SGD burn-in in the first 40 epochs the best choice.
SGD Burn-in Epoch
0
40
47
Test ACC (%)
95.61
95.62
95.57
95.67
95.72
95.41
24


Reviews:
[{"review_id": "vgovVff9ZP", "reviewer": "Reviewer_CmLw", "summary": "The authors define the augmented model where they aim to perform the inference over variable $\\theta_a$ in addition to model parameters.\n\nThe authors show the marginal posterior of the augmented model over model parameters give the \"right\" posterior distribution.\n\nThe authors provide several theoretical statements justifying the algorithm and the empirical evaluation of the algorithm.", "questions": "See the weaknesses section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the idea of the augmented model is interesting and the paper is a nice read.\n\nThese considerations are novel to the best of my knowledge.\n\nThe manuscript is mostly very clear.\n\nThe proposed method achieves good empirical performance.", "weaknesses": "I think the authors should clearly specify their model: prior distributions and likelihood and only after that move to the inference part to improve the clarity of the paper.\n\nI understand that when priors are uniform, the RHS of Eqn 4 effectively defines the likelihood of the augmented model the authors want to consider.\n\nThe notation using $f(\\theta)$ is confusing (e.g. because of no dependence on data) and should be replaced by substituting the definition above the Eq 3.\n\nNo line numbers in the manuscript.\n\nIs the variance $\\eta$ a free variable? Why does $\\eta$  in Eqn 21 disappear during integration?\n\nIt is not clear how data augmentation influences the empirical performance of the method.\n\nIt would be nice if the authors provided some arguments why the introduced concepts are mathematically well-defined, e.g. (4) integrates to a finite quantity (seems straightforward but the finiteness of all quantities should be ensured).\n\nWhy do authors consider the assumption 1? I.e. it's fair to say that authors what to perform the inference in the augmented model, specify the prior/likelihood and there's no need to introduce the assumption 1 (make it a remark).\n\nIt's difficult to judge the practical utility of theoretical statements.\n\nIf my questions are addressed I'm willing to increase my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eEHZ7Hug3K", "reviewer": "Reviewer_QrYa", "summary": "This paper studies the flatness aware optimization using MCMC with an entropy-adjusted loss function. The overall goal is to adapt SGD with information about the local flatness of the optimization landscape in order to find local minima in flatter regions of the loss landscape, which are widely believed to have better generalization properties compared to sharp local minima. The key technical insight is that the computational difficulties encountered applying local entropy optimization in previous works can be avoided by performing joint MCMC on the energy term inside the local entropy integral, whose marginal distributions correspond to the original Bayesian posterior distribution and the local entropy posterior distribution respectively. Joint sampling allows the local entropy weight path to facilitate movement of the original weight path. A theoretical analysis shows the proposed method has more favorable convergence rates than prior local entropy optimization methods in the strongly convex setting. Experiments show that the proposed method finds local minimum which exhibit greater flatness according to the eigenspectrum and interpolation/extrapolation experiments, and that the proposed method boosts validation accuracy compared to SGD and related entropy-based optimization methods on CIFAR-10, CIFAR-100, and ImageNet.", "questions": "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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper is easy to understand and well written.\n* The method resolves a key computational limitation in the local entropy approach from Chaudhari et. al 2019 and allows local entropy optimization with no inner loop, and essentially the same computational cost as SGD. This could make local entropy optimization much more appealing to practitioners than current methods.\n* Synthetic dataset experiments and measurement of flatness metrics corroborate the claims of the method's ability to focus its trajectory on flat minima.\n* Experiments on classifier training show improved performance relative to SGD and existing local entropy optimization methods.", "weaknesses": "* 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.\n* The SGD baselines for the classification experiments are somewhat weak. It would be interesting to see if the proposed method can push the performance of models with state of the art scores, or at least much closer to state of the art. Maybe fine-tuning rather than full training could help alleviate costs.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "drflvgQ7R9", "reviewer": "Reviewer_nUGC", "summary": "The paper describes an \"augmentation\" technique for Bayesian neural networks. \nThere is ample experimental evidence and some theoretical support that neural networks \n(pointwise estimations) obtained as local minimizers of the training loss which belong to flat \nbasins generalize better than minimizers in sharp minima.\nTherefore, in an approximate Bayesian setting where full posterior exploration is prohibitive,\nthey introduce an auxiliary replica, so that the two replicas measure has as marginals the original\nposterior and the local entropy weight (from Baldassi et al. PRL '15). \nThe argument is that while the marginal distribution is still the same, the MCMC dynamics\nof the replicated system allows early exploration of the wider modes of the posterior, thus obtaining\nbetter generalization performance.", "questions": "- Can the author comment on the possibility of enhancing the effect of attraction toward flat minima. \n  E.g. in the framework of Baldassi et al. PNAS '16 this would mean increasing the number of replicas used. \n  Can more replicas be used by also having one of the marginal equal to the origin measure?\n\n- In some of the experiments (or maybe all) experiments the author collect samples of both theta and theta_a. \n  Since the spirit of the paper is to perform Bayesian sampling, shouldn't the collect only theta samples?\n\n- Maybe mores remands to the appendices, e.g. to appendix B for proofs, are needed in the main text.\n\n- Fig. 6 shows a quite irregular dependence on eta of the test error, in particular the presence of very sharp peaks.\n  Do the authors have any intuition of why there is this peak?\n\n- Related to the previous question, did the author consider some adaptive tuning scheme for eta, such as the one proposed in the\n \"Focusing\" section of Pittorino et al ICLR '21?\n\n- A correction for the \"Related Work\" section: The concept of Local Entropy has been introduced in https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.115.128101. \n  Entropy-SGD (Chaudhari et al., 2019) is an algorithmic implementation. (Baldassi et al., 2016) use the Local Entropy framework to derive \"replicated\" algorithmic approaches. \n  Both Entropy-SGD and Replicated-SGD algorithms have been further investigated in https://openreview.net/forum?id=xjXg0bnoDmS\n\n- typo in eq. 16, there should be a +\n\n- I wonder if after the burn-in the samples are collected at each iteration.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well organized, well written, and presents some advances in the field. \n\n- The 2 replicas framework, although very similar (actually a specific instance with y=1 replicas) of the Robust Ensemble (RE)\n  introduced in Baldassi et al. PNAS '16, has the advantage over the generic RE of preserving an unbiased marginal measure,\n  while in the RE with y>1 replicas the resulting marginals are tilted. \n\n- There is good experimental coverage showing consistent (although small) improvements over competing techniques.\n\n- They provide a favorable convergence bound for log-concave target distribution.", "weaknesses": "- There is a lack of novelty with respect to Baldassi et al. PNAS '16, only partly justified by the focus on the Bayesian setting.\n\n- All experiments seem to be performed using a temperature of T=1e-4, instead of the T=1 of the purely Bayesian setting. This\n  makes the Entropy-MCM framework even more similar to the optimization setting of Baldassi et al. PNAS '16 and Pittorino et al ICLR '21.\n  Since table 4 shows only minimal performance decrease when setting the temperature to 1, I suggest to present all results \n  with T=1 only. As an alternative, they could follow the protocol of Zhang et al. ICLR 2020b setting T=0 in the burn-in phase and T=1 in\n  the sample collection phase.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YfD9NHPFVK", "reviewer": "Reviewer_JBLt", "summary": "This paper proposes entropy-MCMC (EMCMC), a method to bias posterior sampling from Bayesian neural networks towards flat basins to improve generalization performance. EMCMC works by introducing an auxiliary guiding variable with a smoothed posterior that favors flat basins. The paper presents some theoretical analysis on convergence rates, and demonstrates the effectiveness with experiments on multiple benchmarks.", "questions": "In Table 1(a), while EMCMC performs better, it seems entropy-SGLD is almost consistently worse than SGLD. Why is that? For comparison Entropy-SGD seems to consistently outperform SGD.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper tackles an important problem, naming find modes in flat basins to improve generalization. \n* The paper is clearly written, and easy to follow.\n* The core idea of the paper is quite simple: instead of relying on an inner loop to estimate the gradients of the local entropy as done in Entropy-SGD, this paper proposes to remove the integral and introduce an auxiliary variable, and sample from the new joint distribution which has a simple form and can be done in a computationally efficient way. Despite its simplicity, the idea is interesting and quite natural. Theoretical analysis and empirical results both demonstrate the effectiveness of the proposed approach.\n* The theoretical results look interesting but I did not carefully check the math.", "weaknesses": "This paper presents the method as a sampling/MCMC approach, which seems a bit confusing to me. When we talk about sampling we usually aim to sample from the true underlying distribution, including both the flat basins and the sharp nodes. However, it seems in the context of this paper, it is perfectly fine to be stuck at a good flat basin without ever exploring around the sharp nodes. Moreover, Theorem 1 presents a convergence result, implying that the marginal distribution of \\theta should eventually converge to the true posterior distribution, but this means it would have to eventually visit the sharp modes. Moreover, if the initial samples are concentrated in the flat basins as claimed, doesn't that mean later there would be a period where the sampler stays around the sharp nodes? Is this because of the strongly convex assumption for Theorem 1? If this assumption is violated, to what extent does the sampler still converge to the true posterior distribution? Should I take the theoretical results as establishing the local convergence rates once we are near a mode?\n\nIt would be helpful if the authors can clarify the above issue. I am also curious to see a simple new experiment where we run EMCMC on a distribution like the one shown in Figure 2 for a really long time to see whether EMCMC would come back to the sharp mode and stay there for a while.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uktPGzwCZK", "reviewer": "Reviewer_Prwy", "summary": "The paper introduces innovative sampling schemes tailored for Bayesian deep learning, aiming to bias the obtained samples during the training process towards flat basins, which helps to obtain better generalization. Rather than focusing on the original posterior, the authors augment it with an auxiliary variable.This augmentation results in a joint distribution with marginals corresponding to the original posterior and a smoothed version of it, achieved through convolution with a Gaussian kernel. The joint distribution's particular structure facilitates straightforward sampling using the stochastic gradient Langevin algorithm. Experimental results demonstrate that acquiring samples from both marginals during the test phase leads to a diverse set of samples that effectively capture the characteristics of flat basins.", "questions": "- In section 6.1, the definition of the ridge between the two modes of the mixture needs clarification. It is also important to explain why choosing an initial point on this ridge does not favor one mode over the other.\n\n- Regarding the MCMC algorithm used in the synthetic example, details about variance parameter $\\eta$, number of MCMC iterations, and step sizes are necessary.  Results from Figure 2, 7, and 8 seem to suggest that both SGLD and EMCMC does not mix-well. \n\n- And it would be better to elaborate why the result is independent to the choice of $\\eta$? (To my understanding, if $\\eta \\to 0$, the smoothed target will be identical to the exact target, so 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 choice of $\\eta$) on the slice of (x, x).\n- This could be due to my limited understanding to Bayesian NN: For the image classification experiement, how do you compare the results obtained from sampling-based and optimization-based algorithm? Specifically, do you obtain certain point estimates of NN's weights from MCMC samples, and compute metrics on test dataset using NN with the estimated weights?   \n- Is there a systematic way of choosing the variance parameter $\\eta$?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The connection between the  proposed posterior augmentation and flatness-aware optimization is quite novel.\n- The proposed methodology is very simple, yet effectively improves the performance of BNN.\n- The theoretical comparison between EMCMC and entropy-SGD and entropy-SGLD is insightful. \n- The reviewer appreciate the extensive set of experiments demonstrating the effectiveness of the method.", "weaknesses": "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). In my opinion, 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].\n\n[1] Wang, et.al. 2021, Analyzing the Generalization Capability of SGLD Using Properties of Gaussian Channels", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)