SAM as an Optimal Relaxation of Bayes
Authors: Thomas Möllenhoff
Mohammad Emtiyaz Khan
Venue: ICLR 2023 notable top 5% (2023)
Keywords: bayesian deep learning
sharpness-aware minimization
variational bayes
convex duality

Abstract: Sharpness-aware minimization (SAM) and related adversarial deep-learning methods can drastically improve generalization, but their underlying mechanisms are not yet fully understood. Here, we establish SAM as a relaxation of the Bayes objective where the expected negative-loss is replaced by the optimal convex lower bound, obtained by using the so-called Fenchel biconjugate. The connection enables a new Adam-like extension of SAM to automatically obtain reasonable uncertainty estimates, while sometimes also improving its accuracy. By connecting adversarial and Bayesian methods, our work opens a new path to robustness.  

Published as a conference paper at ICLR 2023
SAM AS AN OPTIMAL RELAXATION OF BAYES
Thomas Möllenhoff & Mohammad Emtiyaz Khan
RIKEN Center for Advanced Intelligence Project
Tokyo, Japan
hthomas.e. www mmnollennoff. www.eennoffmiiiaaattieei khan) ip
ABSTRACT
Sharpness-aware minimization (SAM) and related adversarial deep-learning methods can drastically improve generalization, but their underlying mechanisms are not yet fully understood. Here, we establish SAM as relaxation of the Bayes objective where the expected negative-loss is replaced by the optimal convex lower bound, obtained by using the so-called Fenchel biconjugate. The connection en- ables a new Adam-like extension of SAM to automatically obtain reasonable un- certainty estimates, while sometimes also improving its accuracy. By connecting adversarial and Bayesian methods, our work opens a new path to robustness.
INTRODUCTION
Sharpness-awee minimization (SAM) (Foret al. 2021) and related adversarial methods (Zheng jet al. 2021 Wu et al.) 2020)) Kim et al. 2022) have been shown to improve generalization, calibra-
tion, and robustness in various applications of deep learning (Chen et al. 2022, Bahri et al.. 2022). but the reasons behind their success are not fully understood. The original proposal of SAM was geared towards biasing training trajectories towards flat minima, and the effectiveness of such min- ima has various Bayesian explanations, for example, those relying on the optimization of description lengths (Hinton & Van Camp 1993 Hochreiter & Schmidhuber, 1997). PAC-Bayes bounds (Dziu- gaite & Roy] [017] 2018]] Jiang et al. 2020; | or marginal likelihoods (Smith & Le 2018). However, SAM not known to directly optimize any such Bayesian criteria, even though some connections to PAC-Bayes do exist (Foret et al. (021). The issue is that the 'max-loss' used in
SAM fundamentally departs from a Bayesian-style "expected-loss" under the posterior; see Fig. T(a) The two methodologies distinct, and little is known about their relationship.
Here, we establish a connection by using a relaxation of the Bayes objective where the expected negative-loss is replaced by the tightest convex lower bound. The bound is optimal, and obtained by Fenchel biconjugates which naturally yield the maximum loss used in SAM (Fig. I(a).. From this, SAM can be seen as optimizing the relaxation of Bayes to find the mean of an isotropic Gaussian posterior while keeping the variance fixed. Higher variances lead to smoother objectives, which biases the solution towards flatter regions (Fig. [(b)). Essentially, the result connects SAM and Bayes through Fenchel biconjugate that replaces the expected loss in Bayes by maximum loss.
What do gain by this connection? The generality of our result makes possible to easily com- bine the complementary strengths of SAM and Bayes For example, we show that the relaxed- Bayes objective can be used learn the variance parameter, which yields an Adam-like extension of SAM (Alg. [ The variances are cheaply obtained from the vector that adapts the learning rate, and SAM's hyperparameter is adjusted for each parameter dimension via the variance vector. The extension improves the performance of SAM on standard benchmarks while giving comparable un- certainty estimates to the best Bayesian methods.
Our work complements similar extensions for SGD, RMSprop, and Adam from the Bayesian deep- learning community (Gal & Ghahramani 2016 Mandt et al., 2017 7 Khan et al., 2018 Osawa et al.) 2019). So far there is work on such connections between SAM and Bayes, except for recent empirical study by Kaddour et al, 2002). Husain & Knoblauch (202)) give an adversarial interpre- tation of Bayes, but it does not cover methods like SAM. Our work here is focused connections to approximate Bayesian methods, which can have both theoretical and practical impact. We discuss new path robustness by connecting the two fields of adversarial and Bayesian methodologies.

---

Published as a conference paper at ICLR 2023
SAM:
6(0 + c)
Variance
Variance
Variance
le|<p
Expectec-Lose Biconjugate Lose
SAM minimizer
Fenchel Biconjugate
SAM minimizer
SAM minimizer
Bayes: Em (((++))))
(a) SAM vs Bayes
(h) Higher variances give smoother objectives
Figure 1: Panel (a) highlights the main difference between SAM and Bayes. SAM uses max-loss (the red dot), while Bayes expected-loss (the blue dot shows {(O + e) for an (e |0, o o²))- Our main result in Theorem 2 connects the two by using the optimal convex relaxation (Fenchel biconjugate) of the negative expected loss. shows that the role p and a are exactly the same, and SAM minimizer obtained for fixed can always be recovered from the relaxation for some σ. An example shown Panel (b) for loss (gray line) with 3 local minima indicated by A, B, and C. The expected loss is smoother (blue line) but the relaxation, which upper bounds it. is even more smooth (red line). Higher give smoother objectives where sharp minima and slowly disappear. The SAM minimizer is shown with red star which matches the minimizer of the relaxation.
SAM AS AN OPTIMAL RELAXATION OF BAYES
In SAM, the training loss ((0) is replaced by the maximum loss within a ball radius p > 0 around the parameter 0 E c R.P, as shown below for an fz-regularized problem,
(0 sup {(O + +0)+)++ 2010 ||
(1)
with 5 as the regularization parameter. The use of 'maximum' above differs from Bayesian strategies that use 'expectation', for example, the variational formulation by Zellner (988).
[(g) = Ea~g [-{ DEL (4(0) [|))))
where (0) is the generalized posterior (Zhang) 1999 Catoni 2007), p(0) is the prior, and
DRL (- II-) the Kullback-Leibler divergence (KLD). For isotropie Gaussian posterior (0)
N m, oI) and prior p(0) = N 0, 1/5), the objective with respect to the mean m becomes
Bayes m 6) = (e [e(m + e') + 2 |m||
which closely resembles Eq. I but with the maximum replaced by an expectation. The above objective is obtained by simply plugging in the posterior and prior, then collecting the terms that depend on m, and finally rewriting 0 m e'.
There are some resemblances between the two objectives which indicate that they might be con- nected. For example, both use local neighborhoods: the isotropic Gaussian can be seen "softer' version of the 'hard-constrained" ball used SAM. The size of the neighborhood is decided by their respective parameters (a or p in Fig. ((a).. But, apart from these resemblances, are there any deeper connections? Here, we answer this question.
We will show that Eq. Α is equivalent optimizing with respect to optimal convex relaxation of Eq. 3 while keeping a fixed. See Theorem 2 where an equivalence drawn between and revealing similarities between the mechanisms used in SAM and Bayes to favor flatter regions. Fig. shows an illustration where the relaxation shown to upper bound the expected loss. We will now describe the construction of the relaxation based on Fenchel biconjugates.
2. FROM EXPECTED-LOOSS TO MAX-LOSS VIA FENCHEL BICONIUGATES
We use results from convex duality to connect expected- to max-loss. Our main result relies on Fenchel conjugates (Hiriart-Uratt & Lemaréchal 2001), also known convex conjugates. Con-
2

---

Published as a conference paper at ICLR 2023
V
Original negative
Expected Loss
M
VA
Ω
(a) Duality of exponential family
Biconjugate
+
W
(b) Biconjugate of expected loss
(c) Biconjugate and expected loss in (w,, pa-pprametteritioonn
Figure 2: Panel (a) shows the dual-structure of ayponential-family for Gaussian defined in Eq. The space M lies above the parabola, while 52 half space. Panel (b) shows f(u) and its bi- conjugate (*) with the blue mesh and red surface respectively. We use the loss from Fig. [(b) The f function perfectly matches the original loss at the boundary where v - (black curve), clearly showing the regions for three minima
B and C. Panel (c) shows the same but with spect to the standard mean-variance (w,,) v) parameterization. Here, convexity of f is lost, but more clearly the negative E(w) appear
(black curve). At this limit, also get f(u) (p) -f(w), but for higher variances, both f and f** are smoothed version of f(w).
The slices 0.5, 2, and give rise to the curves shown in Fig. [(b)
sider function f(u), defined over U and taking values in the extended real line. Given dual space v V, the Fenchel biconjugate fo* and conjugate f are defined as follows,
f (u) sup (v' u)) f'
f*(v)) = sup (u",v) v') - f(u"),
(4)
v'ey
n"
where (.-) denotes the inner product. Note that f need not be convex, but both f and f are
convex functions. Moreover, the biconjugate is the optimal convex lower bound: f(u) f (u).
We will use the conjugate functions connect the max-loss in SAM to the expected-loss involving regular, minimal exponential-famin distributions (Wainwright et al. 2008) of the following forn
(0) exp ((A, T(0)) A())),
(5)
with natural parameter X E S. sufficient statistics T(Ø), log-partin partition function A()), and a base measure of with (A A()) )o). Dual spaces naturally arise for such distributions and can be used to define conjugates. For every X E Ω. there exists unique 'dual' parameterization in terms of the expectation parameter μ = Ee [T))). The space all valid μ satisfies μ = V((A), and all such ya form the interior of dual space which we denote by M. We can also write À in terms of μ using the inverse À = V (μ), and in general, both μ and are valid parameterizations.
In what follows, will always use to parameterize the distribution, denoting qu (0). As an example, for a Gaussian (1 w, v) over scalar 9 with mean w and variance v >> O, we have,
T(#) = (0, 8²), = (w,w² + X = (w/v, - 1/(2w))
66)
Because v 0, the space M is the epigraph of parabola (w, w²), and is negative half space because the second natural parameter is constrained to be negative; see Fig. 2(a) The dual spaces can be used define the Fenchel conjugates.
We will use the conjugates of the 'expected negative-loss' acting μ E interior(M).
f(1) Ee-q, [- [(()]
'The notation used in the paper summarized in Table 10 in App-

---

Published as a conference paper at ICLR 2023
For all μ & interior(M).. we set f(+) +o0. Using Eq. obtain the biconjugate,
(1) sup (A, - sup x'
(n", X²) Ro-q," (-((())
(8)
EM
where we use X' and μ² denote arbitrary points from Q and M,, respectively, and avoid confusion to the change of coordinates μ = VA()).
We show in App. A that, for isotropic Gaussian, the biconjugate takes form where max-loss automatically emerges. We define to be the expectation parameter of a Gaussian ( w, eI), and X to be the natural parameter of (different) Gaussian (( m, I2).. The equation below shows the biconjugate with the max-loss included in the second term,
0 m | sup f(m + +c) 20
1 20
(1) = sup Ea m..7
||E
(9)
The max-loss here is a proximal form, which has an equivalence to the trust-region form used in Eq. M (Parikh & Boyd 2013 Sec. 3.4). The first term in the biconjugate arises after simplification of the (X, (1) term in Eq. 8 while the second term is obtained by using the theorem below to switch the maximum over ı² to maximum over 0, and subsequently over e by rewriting 0 t.
Theorem 1. For Gaussian distribution that takes the form
(). the conjugate f (X') can be ob-
tained by taking the supremum over instead of M, that is,
(x) "Eawa, | [() sup #MEM
sup
-
(10)
OE
assuming that E(0) is lower-bounded, continuous and majorized by a quadratic function. Moreover, there exists X for which f'X)) is finite and, assuming {(0) is coercive, dom(f*) S Ω.
proof of this theorem is in App. where we also discuss its validity for generic distributions.
The red surface in Fig. 2(b shows the convex lower bound for to the function (shown with the blue mesh) for a Gaussian with mean w and variance o (Eq. 6)_ Both functions have a parabolic shape which is due to the shape of M. Fig. [[c] shows the same but with respect the standard mean-variance parameterization (w, v). In this parameterization, convexity of f is lost, but we see more clearly the negative f(w) appear at u + (shown with the thick black curve). At the limit, we get f(1)) = (") = -ww), but for higher variances, both f and f are smoothed versions of E(w). Fig. 2(c) shows three pairs of blue and red lines. The blue line in each pair is slice of the function f. while the red line is slice of f The values of for which the functions are plotted are 0.5, 2, and The slices at these variances are visualized in Fig. [(b]along with the original loss.
The function n (14) is 'lifted' version of the one-dimensional (()) to two-dimensional f4-space. Such liftings have been used the optimization literature (Bauermeister et al. 2022). Our work connects them to exponential families and probabilistic inference. The convex relaxation preserves the shape at the limit v + O, while giving smoothed lower bounds for all v > For a high enough
variance the local minima A and C disappear, biasing the optimization towards flatter region B of the loss. We will now use this observation to connect to SAM.
2.2 SAM AS AN OPTIMAL RELAXATION OF THE BAYES OBJECTIVE
We will now show that a minimizer of the SAM objective in Eq, D can be recovered by minimizing relaxation of the Bayes objective where expected-loss f is replaced by its convex lower bound f The new objective, which we call relaxed- -ayees, always lower bounds the Bayes objective of Eq
sup C(gu) > sup HEM
(4) (g (()))
(11)
uEM
1 20 || - DKL. qu (n) ll m0(e)e 2 10/ + log
= sup PEM m.o
sup
sup {(m + +)
where in the second line substitute f (44) from 99) and merge Eg 2 1 m|| with the
KLD term. Z denotes the normalizing constant of the second distribution in the KLD. The relaxed- Bayes objective has specific difference-of-convex (DC) structure because both for and KLD
4

---

Published as a conference paper at ICLR 2023
convex with respect to μ.. Optimizing the (elaxed-Bayee objective relates to double-loop loop algorithms where the inner-loop uses max-loss to optimize for (m, σ), while the outer loop solves for μ.. We will use the structure to simplify and connect SAM.
Due to the DC structure, it is possible to switch the order of maximization with respect to μl and (m,o), o), and eliminate ja by direct maximization (Toland 1978). The maximum occurs when the KLD term and the two distributions in its arguments are equal. A detailed derivation is in App. which gives us a minimization problem over (m,
1
8'
P
Erelaxed (m, (m 0;)
sup {(m + +c) 202
+
|m||
2
log(o²8"),
(12)
log
where 8' = +1(6g) and we use a Gaussian prior p(0) N 0. 0/6g) This objective uses max-loss in the proximal form, which is similar to the SAM loss in Eq. [ with the difference that the hard-constraint is replaced by a softer penalty term.
The theorem below is our main result which shows that, for a given p. minimizer of the SAM objective in Eq. [ can be recovered by optimizing Eq. with respect to m while keeping a fixed.
Theorem 2. For every (p. 8), there exist (a, 8') such that
arg min EsAM p. 6) arg min BEB
Erelased m, Ø; 8')
(13)
mee
assuming that the SAM- AAM-perturbatton satisfies ||e]| p at a stationary point.
A proof is given in App. D and essentially follows from the equivalences between proximal and trust-region formulations. The result establishes SAM as special case of relaxed- -Byyes, and shows that p and play the exactly same role. This formalizes the intuitive resemblance discussed at the beginning of Sec. 2 and suggests that, in both Bayes and SAM, flatter regions are preferred for higher because such values give rise to smoother objectives (Fig. [(b).
The result can be useful for combining complementary strengths of SAM and Bayes. For example, uncertainty estimation in SAM can now be performed by estimating a through the relaxed-Bay Bayed objective. Another way is to exploit the DC structure in ju-space to devise new optimization proce- we will demonstrate both of these improvements of SAM. dures improve SAM. In Sec. 3
It might also be possible to use SAM to improve Bayes, something we would like to explore in the future. For Bayes, it still is challenging get a good performance deep learning while
optimizing for both mean and variance (despite some recent success (Osawa et al. 2019)). Often, non- -ayyBayssian methods are used set the variance parameter (Plappert al.) 2018
Fortunato al., 2018 Bisla et al. 2022). The reasons behind the bad performance remain unclear. It is possible that simplistic posterior is to blame (Foong et al. 2020 Fortuin et al. 2022 Coker et al.. 2022). but our result suggests to not disregard the opposite hypothesis (Farquhar et al. 2020). SAM works well, despite using a simple posterior, and it is possible that the poor performance of Bayes is due to poor algorithms, not simplistic posteriors. We believe that result can be used to borrow algorithmic techniques from SAM to improve Bayes, and we hope to explore this in the future.
3 BAYESIAN-SAM TO OBTAIN UNCERTAINTY ESTIMATES FOR SAM
In this section, we will show an application of our result to improve an aspect of SAM. We will derive new algorithm that can automatically obtain reasonable uncertainty estimates for SAM. We will do so by optimizing the relaxed- relaxed-Bayes objective to get σ. A straightforward way is to directly optimize Eq. T2 but we will use the approach followed by Khan et al. (2018), based on the natural- gradient method of Khan & Rue (202)), called the Bayesian learning rule (BLR). This approach has been shown to work well on large problems, and is more promising because it leads Adam-like algorithms, enabling us to leverage existing deep-learning techniques to improve training (Osawal al. 2019) for example, by using momentum and data augmentation.
To learn variances on top of SAM, we will use a multivariate Gaussian posterior (0) (e w, V) with mean m and a diagonal covariance V = diag(s) -1, where s is a vector of preci- sion values (inverse of the variances). The posterior is more expressive than the isotropic Gaussian
5

---

Published as a conference paper at ICLR 2023
Algorithm 1 Our Bayesian-SAM (bSAM) is simple modification of SAM with Adam Just add the red boxes, or use them replace the blue boxes next them. The advantage of bSAM is that it automatically obtains variances o²² through the vector that adapts the learning rate (see line 11). Two other main changes are in line 3 where Gaussian sampling is added, and line 5 where the vector s is used to adjust the perturbation e. The change in line 8 improves the variance by adding the regularizer ó and a constant 7 (just like Adam), while line uses Newton-like update where s is used instead of vs. We denote by b, a/b and (a), the element-wise multiplication, division and absolute value, respectively. N denotes the number of training examples.
Input: Learning rate a; 91,, T--regularizer 8, SAM parameter p. A3 = 0.999, 7 10-$
I: Initialize mean ~ +- (NN weight init), scale + 0 and momentum vector gm +- 0
2: while not converged do
3: 0 ee, where Niee ((e)0,
a 1/(N
(ample to improve the varianoe
4: g B Σiem Vl((0) where M is minibatch of B examples
5:
E
&-
{Rescale by vector instead of scaler Igll used in sax]
6:
g. - B EiEM V4
7:
8m t 3igm + (1 31) (g. + )w)
8: 8 s B2 s + (1 - 32)
(g. + وwسة
(vs. g +8+
[mmprovdd varlance est inate)
9:
a
g
gm s
{scole the gredient by s instead of vE}
Vs+>
10: end while
11: Return the mean and variance の² <- 1/(N s)
used in the previous section, and is expected to give better uncertainty estimates. As before, we will use Gaussian prior p(0)= (0 0, 1/6)).
The BLR uses gradients with respect to μ to update
A. The pair (μ, A) is defined similarly to Eq. We will use the update given in Khan & Rue 202T Eq. 6) which when applied to Eq. gives us.
A (1 - a)
V V (4) + Xu)_
(14)
where a > is the learning rate and A (0. - 122 the natural parameter of p(0 This obtained by using the fact that V,LDLL(qu (0) (()) x -oo: see Khan & Rue 2021 Eq. 23). For a = 1, the update is equivalent to the differeecee-f-cconves (DC) algorithm (Tao & An 1997). The BLR generalizes this procedure by using < 1, and automatically exploits the DC structure in μu-space, unlike the naive direct minimization of (22 with gradient descent.
The update (14) requires Vf* which is obtained by maximizing variant of Eq. that replaces the isotropic covariance by diagonal one. For a scalable method, we solve them inexactly by using the 'ocaal-ineerizationn approximation used in SAM, along with an additional approximation to do single block-coordinate descent step to optimize for the variances. The rest of the derivations is similar the previous applications of the BLR; see Khan al. (2018); Osawa et al. (2019). The details are in App. E. and the resulting algorithm, which we call Bayesian-SAM is shown in Alg.
bSAM differs slightly from direct application of Adam (Kingma & Ba 2015) on the SAM ob- jective, but uses the same set of hyperparameters. The differences are highlighted in Alg. T where bSAM is obtained by replacing the blue boxes with the red boxes next them. For SAM with Adam, we use the method discussed in Bahri et al. (2022); Kwon et al. (2021) for language models. The main change is line 5, where the perturbation 6 is adapted by using a vector 8, as opposed to SAM which uses scalar ||g]]. Modifications in line 3 and 8 are aimed to improve the variance estimates because with just Adam we may not get good variances; see the discussion in Khan et al. 2018 See 3.4). Finally, the modification in line 9 is due the Newton-liks updates arising from BLR Khan & Rue 2021) where s is used instead of / to scale the gradient as explained App.
Due to their overall structural similarity, the bSAM algorithm has similar computational com- plexity to Adam with SAM. The only additional overhead compared to SAM-type methods is the computation of random noise perturbation which is negligible compared to gradient computations.

---

Published as a conference paper at ICLR 2023
1.00
0.75
(a) Exact Bayes
(b) bSAM
(c) SAM (point est.)
(d) SAM (Laplace)
Figure 3: For 2D logistic regression, bSAM gives predictive uncertainties that are similar those obtained by an exact Bayesian posterior. White areas show uncertain predictive probabilities around 0.5. SAM's point estimate gives overconfident predictions, while Laplace leads underconfidence.
Model/ Data
Accuracy 1
NLL +
ECE
AUROC 1 (higher better)
Method
(higher betler)
better)
better)
SGD
(6.55
0.56(0.014)
0.0839 (0.003)
0.878 (a.005)
SAM-SGD
87.49
(.4990
0.0710
0.891 (a.004)
SWAG
(6.80
(.53()
(.0774(u.ou1)
0.880 (0.00)))
ResNet-20-FRN
VOGN
(0.30
0.38 (0.00)
0.0315 0003
0.890 003)
CIFAR-10 10
Adam
80.85(0.98)
0.83
0.1317
0.820 (0.013)
SAM-Adam
85.26
0.46
0.0228 (0.002)
0.874
bSAM (ours)
(8.72(0.24)
(0.05)
0.0163 (0.002)
0.903 (0.003)
SGD
55.82
(o.
0.1695
0.811
SAM-SGD
(8.58
(o.0().022)
(.0999((o.00))
(0.07(0.003)
SWAG
(6.43
1.86(0.018)
0.1604(o.004)
0.81 ((o.004)
ResNet-20-FRN
VOGN
(9.83
(..4(0.011)
0.0756
0.830 (0.002)
CIFAR-100
Adam
(9.77
2.29 (o.04)
0.0295 (0.018)
0. (.77 004))
SAM-Adam
(3.25
1.71 (0.05)
0.0401 (0.005)
0.0401
0.818 0.818 (a.005)
ASA (ours)
62.64(0.33)
(o.00)
0.0311 (0.003)
0.841 (0.004)
Table I: Comparison without data augmentation. The results are averaged over 5 random seeds, with standard deviation shown in the brackets. The best statistical significant performance is shown in bold. bSAM (gray row) gives either comparable better performance than the rest. Small scale experiments on MNIST are in Table 6 App. F.8
NUMERICAL EXPERIMENTS
In this section, we present numerical results and show that bSAM brings the best of the two worlds together. It gives improved uncertainty estimate similarly to the best Bayesian approaches, but also improves test accuracy, just like SAM. We compare performance to many methods from the deep learning (DL.) and Bayesian DL. literature: SGD, Adam (Kingma & Ba 2015), SWAG (Maddox et al. (019). and VOGN (Osawa et al.) 2019). We also compare with two SAM variants: SAM-SGD and SAM-Adam. Both are obtained by inserting the perturbed gradients into either SGD Adam, as suggested in Foret et al. 2 Bahri et al. (2022); Wwon et al. (202T). We compare these methods across three different neural network architectures and on five datasets of increasing complexity.
The comparison is carried out with respect to four different metrics evaluated on the validation set. The metrics are test accuracy, the negative log-likelihood (NIL.), expected calibration error (ECH) (Guo et al. (017) (using 20 bins) and area-under the ROC curves (AUROC). For an explana- tion of these metrics, we refer the reader tsJaa et al. (2019) Appendix H). For Bayesian methods, report the performance of the predictive distribution p(y DD) p(v D, 0) g()) d0, which we approximate using an average over 32 models drawn from (0).
4.1
ILLUSTRATION TOY EXAMPLE
Fig.
compares the predictive probability simple Bayesian logistic regression problem (Mur-
phy 2012] Ch. 8.4). For exact Bayes, we use numerical integration, which is feasible for this toy
21)
problem. For the rest, we use diagonal Gaussians. White areas indicate probabilities around

---

Published as a conference paper at ICLR 2023
Model
Accuracy
NLL
ECE &
AUROC t
Method
Dataset
(higher is better)
(lower is better)
(lower is better)
Ohigher is better)
SGD
91.68(0.26)
(.29(o.008)
0.0397 0.0397(0.002)) (0.002)
0.915
ResNet-20-FRN
SAM-SGD
92. (0.29(0.39)
0.0266 0.25:0.004
0.920(o.003)
CIFAR-10
Adam
89.97 (0.97)
0.41 (0.01)
0.0610
0.900
SAM-Adam
91.37 (0.21)
0.26 (0.004)
0.0329 (0.002)
0.918 (.918(0.001)
bSAM (ours)
92.16
0.23 0.23( (0.008)
0.0057 (0.002)
0.925 0.925 (0.001)
SGD
66.48
0.0524 (.20
0.846
ResNet-20-FRN/
SAM-SGD
(0.22))
(0.133
0.0481 (0.001)
0.848
CIFAR-100
Adam
(1.76(p.67)
0.158220.000) 1.66)(0040)
0.826
SAM-Adam
65.34 (0.32)
(0.02)
2.01666(0000)
0.847 (.847(0.002)
bSAM (ours)
68. 22 (0.44)
1.10
0.0258 (0.003)
0.857 (0.004)
SGD
52.01 (0.36)
(0.8))
0.0330
0.832
ResNet-20-FRNN
SAM-SGD
(2.25
0.0155 (0.033)
(0.002)
(0827(0.005)
TinylmageNet
Adam
(9.04
2.14(0.024)
0.0502(0.000
0.820
SAM-Adam
51.17 (0.45)
(v.044
(o.0160(0.004))
(0..8
SSAM (ours)
52.90
1.94 (0.009)
0.0199
(0.008)
0.831 (0.00)
ResNet-18/
SGD
(1.76
(v.00))
0.0304
0.926
CIFAR-10
SAM-SGD
(5.72(0.14)
0.14 (0.004)
(.0134
0.949
bSAM (ours)
96.15(0.08)
0.1 0.12 (0.02)
0.0044
0.954 001)
ResNet-18/
SGD
(o.54
0.98
0.0501 (0.002)
0.869
CIFAR-100
SAM-SGD
(8.74
0.79 (0.007)
0.0445
0.887 (0.83)
bSAM (ours)
80.22 (0.28)
0.70 (0.08)
0.031 (0.003)
(0.92
Table 2: Comparison with data augmentation. Similar to Table 1 the shaded row show that bSAM consistently improves over the baselines and the overall best method.
0.5, while red and blue are closer to and respectively. We see that the bSAM result compara- ble to the exact posterior and corrects the overconfident predictions of SAM. Performing Laplace approximation around the SAM solution leads to underconfident predictions. The posteriors are vi- sualized in Fig. 5(c) and other details are discussed in App. F.I We show results another small toy example (" ('two moons') )) in App. F2
We can also use this toy problem in App. F3 shows a comparison of optimizing Eq. to to an optimization of the original Bayesian objective Eq. 2 For both, use full covariance to avoid the inaccuracies arising due to diagonal covariance assumption. We find that the posteriors are very similar for both f and f** indicating that the gap introduced by the relaxation is also small.
the effect of using vs f** in the Bayesian objective. Fig.
4.2 REAL DATASETS
We first perform a set of experiments without any data augmentation, which helps us to quantify the improvements obtained by the regularization induced by our Bayesian approach. This useful in applications where it is not easy to do data- aata-augmentation (for example, tabular or time series data). All further details about hyperparameter settings and experiments are in App. G The results are shown in Table [ and the proposed bSAM method overall performs best with respect to test accuracy as well as uncertainty metries. bSAM also works well when training smaller networks on MNIST-like datasets, see Table 6 App. F.8
In Table 2 we show results with data augmentations (random horizontal flipping and random crop- ping). We consider CIFAR-10, CIFAR-100 and TinylmageNet, and still find bSAM to perform favorably, although the margin of improvement is smaller with data augmentation. We note that, for SGD CIFAR-10 10 in Table 2 the accuracy (91 (91.68% accuracy) is slightly better than the 91.25% reported in the original ResNet-paper (He et al. 2016 Table 6). by using a similar ResNet-20 archi- tecture with filter response normalization (FRN), which has 0.27 million parameters.
8

---

Published as a conference paper at ICLR 2023
ResNet20-FRN CIFAR-100
ResNe120-FRN CIFAR-100
Accuracy (bigher is betler)
SAM-SGD
NLL (lower better)
SAM-Adam
bSAM
SAM-
SAM-Adam
DSAM
Figure 4: bSAM less sensitive to the choice of p than both SAM-SGD and SAM-Adam. We show accuracy and NLL. Plots for the other metrics are in App. F.4
bSAM not only improves uncertainty estimates and test accuracy, it also reduces the sensitivity of SAM to the hyperparameter p. This is shown in Fig. [ for SAM-SGD, SAM-Adam and bSAM, where bSAM shows the least sensitivity to the choice of p. Plots for other metrics in App. F.
To understand the reasons behind the gains obtained with bSAM, we do several ablation studies. First, study the effect of Gaussian sampling in line 3 in Alg. [ which is not used in SAM. Experiments in App. F.5 show that Gaussian sampling consistently improves the test accuracy and uncertainty metrics. Second, in App. [6 we study the effect of using posterior averaging over the quality of bSA predictions, and show that increasing the number of sampled models consistently improves the performance. Finally, we study the effect of "m-sharpness", where we distribute the mini-batch across several compute nodes and use a different perturbation for each node (m refers to the number of splits). This has been shown to improve the performance of both SAM (Foret al., 2021) and Bayesian methods (Osawa et al. 2019). In App- F.7 we show that larger values of m lead to better performance for bSAM too. our experiments, we used m = 8. For completeness, we show how m-sharpness is implemented in bSAM in App.
DISCUSSION
In this paper, we show a connection between SAM and Bayes through a Fenchel biconjugate. When the expected negative loss in the Bayes objective is replaced by the biconjugate, obtain relaxed Bayesian objective, which involves a maximum loss, similar to SAM. We showed that SAM can be seen optimizing the mean while keeping variance fixed. We then derived an extension where the variances are optimized using a Bayesian learning algorithm applied to the relaxation. The numerical results show that the new variant improves both uncertainty estimates and generalization, as expected of a method that combines the strengths of Bayes and SAM.
We expect this result to be useful for researchers interested designing new robust methods. The principles used in SAM are related to adversarial robustness which is different from Bayes. We show that the Fenchel conjugate is the bridge between the two, and we used it obtain an extension of SAM. Since SAM gives an excellent performance on many deep-learning benchmark, the techniques used there will be useful in improving Bayesian methods. Finally, the use of convex duality can be useful improve theoretical understanding of SAM-like methods. It is possible to use the results from convex-optimizan Our work may open such directions and provide new path for research in improving robustness. Expectations with respect to distributions arise naturally in many other problems, and often sampling is used to approximate them. Our paper gives different way approximate such expectations using convex optimization rather than sampling. Such problems may also benefit from our results.
literature to study and improve adversarial as well as Bayesian methods.
ACKNOWLEDGEMENTS
We would like to thank Pierre Alquier for his valuable comments and suggestions. This work is supported by the Bayes duality project, JST CREST Grant Number JPMICR2112.

---

Published as a conference paper at ICLR 2023
AUTHOR CONTRIBUTIONS STATEMENT
List of Authors: Thomas Möllenhoff (T.M.), Mohammad Emtiyaz Khan (M.E.K.).
T.M. proposed the Fenchel biconjugates, their connections to SAM (Theorems I and 2 and derived the new bSAM algorithm. M.E.K. provided feedback simplify these. T.M. conducted all the experiments with suggestions from M.E.K. Both authors wrote the paper together.
REFERENCES
Pierre Alquier. User-friendly introduction to PAC-Bayes bounds.
2021.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Association for Computational Linguistics (ACL), 2022.
1 6
Hartmut Bauermeister, Emanuel Laude, Thomas Möllenhoff, Michael Moeller, and Daniel Cre- mers. Lifting the convex conjugate in Lagrangian relaxations: tractable approach for continu- ous Markov random fields. SIAM J. Imaging Sci., 15(3):1253-1281. 2022.
Devansh Bisla, Jing Wang, and Anna Choromanska. Low-pass filtering SGD for recovering flat optima in the deep learning optimization landscape. International Conference on Artificial Intel- ligence and Statistics (AISTATS). 2022. [
Olivier Catoni. PAC-Bayesian Supervised Classification. Number 56. Institute of Mathematical Statistics Lecture Notes Monograph Series, 2007. 2
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform ResNets without pretraining or strong data augmentations. In International Conference on Learning Rep- resentations (ICLR), 2022.
Beau Coker, Wessel P Bruinsma, David R Burt, Weiwei Pan, and Finale Doshi-Velez Wide mean- field Bayesian neural networks ignore the data. International Conference on Artificial Intelli- gence and Statistics (AISTATS), 2022. 5
Imre Csiszár and František Matús. Closures of exponential families. Annals of probability, pp. 582-600, 2005.
Steven Diamond and Stephen Boyd. CVXPY: A Python-embedded modeling language for convex optimization. J. Mach. Learn. Res. (JMLR), 17(1)290992913, 2016. 20
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Conference on Uncertainty in Artificial Intelligence (UAI), 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Eatropy-SGD optimizes the prior of a PAC-Bayes bound: Generalization properties Entropy-SGD and data-dependent priors. In International Conference on Machine Learning (ICML), 2018.
Sebastian Farquhar, Lewis Smith, and Yarin Gal. Liberty depth: Deep bayesian neural nets do not need complex weight posterior approximations. Advances in Neural Information Processing
Systems, 33:4346-43577 2020. 5
Andrew Foong, David Burt, Yingzhen Li, and Richard Turner. On the expressiveness approximate inference in Bayesian neural networks. Advances in Neural Information Processing Systems, 33, 2020. 5
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations (ICLR), 2021.
22
Vincent Fortuin, Adrià Garriga-Alonse Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. In Interna- fional Conference on Learning Representations (ICLR), 2022.
10

---

Published as a conference paper at ICLR 2023
Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Ian Osband, Alex Graves, Vlad Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, Charles Blundell, and Shane Legg. Noisy networks for exploration. In International Conference Learning Representations (ICLR), 2018. 5
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model
uncertainty in deep learning. In International Conference on Machine Learning (ICML), 2016.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning (ICML), 2017. 7
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2016. 23
Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Conference on Learning Theory (COLT), 1993. |
Jean-Baptir Ban--Batistt Hiriart- Hiriart-Uruty and Claude Lemaréchal. Fundamentals of Convex Analysis. Springer, 2001. 2
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Hisham Husain and Jeremias Knoblauch. Adversarial interpretation of Bayesian inference. In International Conference on Algorithmic Learming Theory, 2022. [
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021. 19
Pavel Izmailov, Sharad Vikram, Matthew D. Hoffman, and Andrew Gordon Wilson. What are bayesian neural network posteriors really like? In International Conference Machine Learning
(ICML), 2021. 23
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where find them. In International Conference Learning Representations (ICLR), 2020. [
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What local optimality in nonconvex-nonconnan minimax optimization? In International Conference Machine Learning (ICML), 2020. 16
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner. A fair comparison of two pop- ular flat minima optimizers: Stochastic weight averaging vs. sharpness-aware minimization.
2.Xiv:22022.0661.. 2022. [
Mohammad Emtiyaz Khan and Haavard Rue. The Bayesian learning rule. arXiv:2107.04562. 2021. 6
Mohammad Emtiyaz Khan, Wu Lin, Voot Tangkaratt, Zuozhu Liu, and Didrik Nielsen. Variational adaptive-newton method for explorative learning. In NIPS Workshop Advances in Approximate Bayesian Inference, 2017. 17
Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivas- tava. Fast and scalable Bayesian deep learning by reight-perturbation in Adam. In International Conference on Machine Learning (ICML), 2018. 5
Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning (ICML), 2022. I
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference Learning Representations (ICLR), 2015. 6
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Con- ference Machine Learning (ICMI.), 2021. 6

---

Published as a conference paper at ICLR 2023
Yong Liu, Siqi Mai, Minhao Cheng, Xiangning Chen, Cho-Jui Hsich, and Yang You. Random sharpness-aware minimization. Advances in Neural Information Processing Systems (NeurIPS), 2022. 21
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Luigi Malago, Matteo Matteucei, and Giovanni Pistone. Towards the geometry of estimation of distribution algorithms based on the exponential family. In Foundations of Genetic Algorithms (FOGA), 2011.
15
Stephan Mandt, Matthew Hoffman, and David Blei. Stochastic gradient descent as approximate Bayesian inference. J. Mach. Learn. Res. (JMLR), 18:1 8:1135, 2017.
Kevin Patrick Murphy. Machine learning: a probabilistic perspective. MIT Press, 2012.
19
Kazuki Osawa, Siddharth Swaroop, Anirudh Jain, Runa Eschenhagen, Richard Turner, Rio Yokota, and Mohammad Emtiyaz. Khan. Practical deep learning with Bayesian principles. Ad- vances Neural Information Processing Systems (NeurIPS), 2019.
6
23
Neal Parikh and Stephen Boyd. Proximal algorithms. Foundations and Trends in Optimization, 1: 123-231, 2013.
16
Matthias Plappert, Rein Houthooft, Prafulla Dhariwal, Szymon Sidor, Richard Chen, Xi Chen, Tamim Asfour, Pieter Abbeel, and Marcin Andrychowicz. Parameter space noise for exploration. In International Conference on Learning Representations (ICLR), 2018. 5
Carl Edward Rasmussen and Christopher K. I. Williams, Gaussian processes for machine learning. MIT Press, 2006. 15
Ralph Tyrelll Rockafellar and Roger J.-B Wets. Variational Analysis. Springer, 1998.
17
Samuel L. Smith and Quoc V. Le. A Bayesian perspective generalization and stochastic gradient
descent. In International Conference on Learning Representations (ICLR), 2018. |
Pham Dinh Tao and Le Thi Hoai An. Convex analysis approach to DC programming: Theory, algorithms and applications. Acta Math. Vietnam., 22(1):289-355, 1997.
6
John F Toland. Duality in nonconvex optimization. Math. Anal. Appl., 66(2):399-415, 1978. S
Martin J Wainwright, Michael Jordan, et al. Graphical models, exponential families, familiis and variational inference. Foundations and Trends@ in Machine Learning, 1(1-2):1-305, 2008. 3
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust gener- alization. In Advances in Neural Information Processing Systems (NeurIPS), 2020. []
Aroold Zellner. Optimal information processing and Bayes's theorem. The American Statistician, 42(4):278-288) 1988.
Tong Zhang. Theoretical analysis of a class of randomized regularization methods. In Conference Learning Theory (COLT), 1999. 2
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2021. |
12

---

Published as a conference paper at ICLR 2023
APPENDIX
A DERIVATION OF THE FENCHEL BICONJUGATE
Here, give the derivation of the Fenchel Biconjugate given in Eq. For an isotropic Gaussian, have T(0) (0,00') and X (m/a², -1/(22)))) We will use the following identity to simplify.
( T())) = m
0 - Tr
00
1
m
01 + 20
m 12
(15)
o²
20²
20²
Using this, the first term in Eq. becomes
() Eo IX', T(O)) E
01
ml|
(16)
20
20
Similarly, for the second term in Eq. 8 we use Theorem and Eq. 15 to simplify,
sup x² EM ,,,,,)) x²) Ea-gg. [-((0) = sup (X', T(0)) + ((0)
sup e
2a² 0 m ml
+
m
1| E())
20
sup {(m + E) 20
||e||
+
m
12,
(17)
20
where in the last step we have performed change of variables 0 m Subtracting 16
from
arrive at Eq. 9
B PROOF OF THEOREM
Theorem 1. For Gaussian distribution that takes the form tained by taking the supremum over instead of M, that is,
5) the conjugate f f*(X) can be ob-
(p" EErag [()] sup (T(O). -1- EMEM
sup
(10)
oEe
assuming that (()) is lower-bourded, continuous and majorized by a quadratic function. Moreover, there exists X for which f*x(') is finite and, assuming {(0) is coercive, dom(J") S.
Proof. We first pull out the expectation in the left-hand side of Eq. 10
(n". X') Ee-gg [[( "MEM
sup
sup Egr (X,T(0)) 3²EM
(())
We will denote the expression inside the expectation by J(0), and prove that the supremum
μu" is equal to the supremum over 0. We will do so for two separate cases where the supremum
J. supece J(0) co finite, and infinite respectively.
When J. is finite:
For this case, our derivation proceeds in two steps, where we show that both the following two inequalities are true the time,
sup Ea-g.. [J(0)] < Ja, #
sup Ee-g,M (J(O) 2 J.
#MEM
This is only possible when the two equal, which will establish the result.
To prove the first inequality, we take a sequence [0.) for which J(,)) - J. when sending - By the definition of supremum, there always exists such sequence. From the sequence [0 we will now construct sequence {μr] for which the expectation of ((0) also converges to J.. This will establish that the supremum over 1² is larger equal than the supremum over 0. It only larger or equal, since hypothetically there could be another sequence of {u") co which
achieves higher value. We will see soon that this cannot be the case.
13

---

Published as a conference paper at ICLR 2023
We use the sequence {Mr [Mi hi=1 where μnر = (0,, E, 0,0,) with {E,) SPxPP is taken in the space of symmetric positive-definit matrices with the sequence such that Σ, + Now, we can that the expectation Ex-awy []))] converges the supremum as follows:
lim Eg Eenawy [J(O)] lim Ba-V( [] (0.+ =2,²2 4000
t-000
lim (0, E/ {
10000
= sup J(0). OERP
In the above steps, we used the fact that J is Gaussian Gaussian-integrable integrable and continuous. A sufficient condition for the integrability that E(O) lower-bounded and majorized by quadraties. With this, we have established that the supremum over μ² is at least as large as the one over O.
To prove the second inequality, we will show that the supremum over fu" also less or equal than the supremum over 0. To see that, we note that ((0) < J. for all 0' E e. This implies that for any distribution yu², we have that Egg-a [J ] I J.. Since the inequality holds for any distribution, we also have that SUP Eor [J(O') I J.. Having established both inequalities, we now know that the suprema in TO0 agrees whenever the supremum over finite.
When J. is infinite: For such cases, will show that suP Ee [J(0)] oc, which will complete the proof. Again, by definition of supremum, there exists sequence {0. ) for which ((,)) + oo when + co. This means that for any M > O there is a parameter (c(ar) in the sequence such that J(O+(M))) > M. Now from the previous part the proof, we know that we can construct a sequence que Gaussians whose expectation converges to ((0) Since the above holds for any M > 0, the supremum over the expectation is also +oc.
Conditions for existence of X' for which J. is finite: The final step of our proof is to give the conditions when there exist at X for which J, is finite, and therefore the result non-trivial.
Since {() is majorized by a quadratic function, there exists a constant e E IR, vector b E R and matrix A s.x such that the following bound holds:
{(0) S c b 0 1 A0.
Then for the candidate X' (-b, AA ) we know that J(Ø) = ((0) (X, T(0)) I c has a finite upper bound for any 0. Hence, the supremum is finite for this choice of X' f2.
Domain of (x()): For a Gaussian distribution, valid (X X)) )) E S satisfy the constraint X, < 0.
We now show that if X/ S2 then f (x( +00. Since X₂ A 0, there exists direction d * 0 such
that dl 2 2 0 and X d 2 0. Taking the sequence {0, -0 with 0, = we see
00, ((O,) tX d d 2²d + l(dd) 2 f(td) + 3o
since coercivity of E(0) implies that ((O.) + for any sequence (0. J with |0, co.
Some additional comments on the assumptions used in the theorem are in order.
1. Any reasonable loss function that one would want to minimize lower-hounded, so this assumption is always satisfied in practice.
2. The existence of an upper-bounding quadratic holds for functions with Lipschitz continu- ous gradient, which is a standard assumption optimization known as the descent lemma. In practice, conjugates
also applicable to cases where this assumption is not satis- fied. In such case, one can consider the conjugate function only in local neighbourhood around the current solution by restricting the set of admissible in the supremum. It may not matter that the bounding quadratic will eventually intersect the loss far away outside of our current region of interest.
3. We assumed that ((0) is continuous since virtually all loss functions used in deep learning are continuous. Interestingly, if (()) is not continuous, the conjugate function f" will involve its semicontinous relaxation.
14

---

Published as a conference paper at ICLR 2023
4. The coercivity assumption E()) is another standard assumption in optimization. For non-coercive losses, there could exist quadratic upper-bounde which are flat in some direc- tion. This would lead to dual variable X' the boundary of f?.
The above Theorem aalso holds other exponential family distributions such as sealar Bernoulli and Gamma distributions, but do not give more details here because they are not relevant for connecting to SAM. There are also cases where the result does not hold. The simplest case is a Gaussian with fixed covariance. The above proof does not work there, because the closure (see e.g., Csiszár & Matús (200);; Malagò et al. (2011) of the family of distributions does not contain all Dirac delta distributions. Intuitively, the requirement is that we should be able approach any delta as limit of a sequence of members of the family, which excludes the fixed-covariance case.
C DERIVATION OF THE EQUIVALENT RELAXED OBJECTIVE (22
Inserting the biconjugate ( into the relaxed Bayes objective
yields
sup sup
sup {(m + c) 20 ||e| Ee
0 ml + log (ogp(0) log 9u (0)
2o
M
where we expanded the KLD as DKL( (0) (()) = Eo-g, (log qu (0)) Bo log (())
Interchanging the order of maximization and noticing that the first term does not depend on μ., we can find a closed-form expression for the maximization in μe as follows:
Ea-qy M
sup
0
m||
+ log p(0) log qu (0)
20
1 |
sup Eawgn MEM
log exp
mll2
p(0)
-log (p)
log
20
1
= log Z
inf DKL
(u(())
exp
10 m
p(0)
μEM
Z
20
= log Z.
In the last step, the KLD is attained at zero with (0)
2 exp
(- 2
m m² p(0). 11
Finally, we compute 2 as
2= = exp -
20 110 [00|]] p(0)d0
1
= (2102) (2TT02)-P7 1/22
exp
2 00
m 2
p(0)
de
m||²
= (02) 1/56) -P/2 cxp
2 + 2)
where in the last step we inserted the normalization constant for the product of the two Gaussians
N(O m, I')) and NO 0, 1) as for example given in Rasmussen & Williams 20066 App. 2).
This gives the following expression
P
P log(o* o² 1/80)
1 2(02 + 1 m 2
log Z log(o²)
8'
log(o?) ) + 2
log 0')
|mml
(18)
2
where 8' = 1/ 1/60).
Switching to minimization over m, x² and inserting the result for the exact ju-solution the final objective
(8
leads
1
8'
P
P
inf m
sup {(m E)
||||
+
m
log
log(6"),
202
2
which is the energy function in 12
15

---

Published as a conference paper at ICLR 2023
D PROOF OF THEOREM
Theorem 2. For every (p. 8), there exist (a, 8') such that
arg min BEB
EsAM 00 p. 8) arg min meo
Erelased (m, 0; 8'),
(13)
assuming that the SAM- SAM-perturbation satisfies | = at a stationary point.
Proof. We will show that any stationary point of &SAM is also stationary point of Erelaxed when a and are chosen appropriately. The correspondence between the stationary points follows directly from the equivalences between proximal and trust-region formulations (Parikh & Boyd 2013 Sec. 3.4).
To show this, we first write the two objectives in terms of perturbation e' and e respectively,
8'
min sup {(m +c)+ 20 ||''|12
2 ||
e'
min sup 8(e + c) +f s p) + |
where ifll p} is the indicator function which zero inside the f2- of radius p and +oo otherwise.
Since these are min-max problems, we need an appropriate notion of "stationary point", which we consider here to be a local Nash equilibirum (Jin et al. [020)] Prop. 3). At such points, the first derivatives 0 and (or c) vanish. Taking the derivatives and setting them to we get the following conditions for the two problems.
8' D'm. Vi(m, + E,,,,
é, = -2 Vl(m + +e,),
6, P V0(. + p
50. VE(O, +.):
for some constant multiplier μ > 0. Here, we used the assumption that the constraint ||ex|| S is active at optimal point. Since the constraint is active, the negative gradient is element of the (non-trivial) normal cone, which gives us the multiplier in the second line above.
The equations structurally equivalent. Any pair (m., c..) satisfying the first line also satisfies the second line for o² = p/ and 8' = 6. The opposite is also true when using the pair (0.,e,) e.) in the first line, which proves the theorem. We assume that the constraint the perturbation is active, that is, ||c_|| around the parameter e.. However, such local maximum is unlikely to exist since the parameter 0, is determined by minimization and tends to lie inside flat minimum.
This assumption would be violated if there exists local maximum within p-ball
E DERIVATION OF THE BSAM ALGORITHM
In this section of the appendix, we show how the natural gradient descent update (44 leads us to the proposed bSAM algorithm shown Alg. [] As mentioned in the main text, for bSAM we consider generalized setting where our Gaussian distribution gr (0) -N w. V) has diagonal covariance. This corresponds to the following setting:
T(0) (0,00)); ) μ (w, 2 diag(s) A (s w. diaag()))
(19)
Here, w E P is the mean and diag(s) = V denotes the entries the inverse diagonal covariance matrix. All operations (squaring, multiplication) are performed entrywise.
E THE GRADIENT OF THE BICONJUGATE
First, we discuss how to compute that gradient of the Fenchel biconjugate, as the BLR update
[44
requires it in every iteration. The diagonal Gaussian setup leads slightly generalized expression
16

---

Published as a conference paper at ICLR 2023
for the Fenchel biconjugate function:
(1L)
min mERF DER
sup CERP
8(m + +6)
- -
+ +Bruqu
| (0 m)l
min mERP .ERR
sup CER"
f(m + o) - - 1
{| 2
E 1 (w m)ll -v(--V), (00)
where Σ = diag(1/ /b) denotes the diagonal covariance matrix with 1/b on its diagonal.
follows from results in convex analysis, for instance from Rockafellar & Wets (1998] Propo-
sition 11.3), that (n) = (Σ,'m. m. - b, that is, the gradient of the biconjugate can be constructed from the optimal solution (m., b.) of the optimization problem (20).
In (20) there are three optimization steps: over E, m. and b. For the first two, we will make very similar approximation to what is used for SAM (Foret et al. (021). and simply add an additional optimization over To simplify computation, we will make an additional approximation. Instead of an iterative joint minimization in m and b, we will perform single block-coordinate descent step where we first minimize in m for fixed b = s and then afterwards in b for fixed w m. This will give rise to a simple algorithm that works well in practice.
Optimization with respect to 6. To simplify the supremum in E we will use SAM's technique of local lincarization (Foret et al. (222)) with one difference. Instead of using the current solution to linearize, we will sample a random linearization point 0 ~ N w. diag(s) 11). This is preferable from Bayesian perspective and is also used in the other Bayesian variants (Khan et al. 2018). Linearization at the sampled 0 then gives us,
f(m + E) x {(0) + (VE(O), m 0).
Using this approximation in (20). we get closed form solution for e, = EV((O).
Optimization with respect to m. To get a closed-form solution for the optimization in m, we again consider a linearized loss: f(m + +) ) ~ f(w + (Vl(w + +), m w). Moreover, we set VVE(0), corresponding
fixed b
This reduces the optimization of (20)
m, = arg min (Ve(w + E), m) + 2 |ز 1/21/2(mm w)1l?
(21)
mERP
which has closed form solution m. = w Σ. Vi(w E) which we here leave to depend Σ.
Optimization with respect to via b. Inserting c, into 20) with the linearized loss, the mini- mization over b while fixing m w also has simple closed-form solution as shown below,
b
b. = arg min
/2 /2(00)
g2 lgl.. Vs
(22)
bER
where O ~ N/( m, diag(s)-1) and g VE(().
E.2 SIMPLIFYING THE NATURAL-GRADIER UPDATES
Having a way to approximate the biconjugate-gradiin
we now rewrite and simplify the BLR up-
date (1).. The derivations are largely similar to the used to derive the variational online Newton approaches, see Khan et al. [2017] 2018); Khan & Rue (202I).
"m., Inserting the parametrization 19) and our expression V (F'( = 'E.m,,, b2.b) update (14), we get the following updates
into the BLR
(1 a)s w + a mm..
5 + (1 a)s + abb, + do),
where we also used the prior Ao = (0. - 10)1) By changing the order of the updates, we can write them in equivalent form that resembles an adaptive gradient method:
- (1 - a)s + a(b. + do),
(23)
V Σ (w m.) + 8w
(24)
17

---

Published as a conference paper at ICLR 2023
Now, inserting the solutions Σ. (w "w and 22 into these updates, we arrive at
m.) Vε(0 e) and b. [V((0)] vs from Eqs. 21
s + (1 - a)s + a ( v0+
(25)
w+ aV (Vi(w + é) )ow),
(26)
where 6 = VV and ~ ~O W, V).
E.3 FURTHER MODIFICATIONS FOR LARGE- SCALE DEEP LEARNING
Now we will consider further modifications to the updates 25
226 to arrive our bSAM method
shown in Alg. | These are similar to the ones used for VOGN (Osawa et al. 2019).
First, we consider a stochastic approximation deal with larger problems. To that end, we now briefly describe the general learning setup. The loss is then written as a sum of individual losses, (()) = Σi= (o( where N is the number of datapoints. The N datapoints are disjointly parti- tioned into K minibatches {B: of B examples each.
We then apply stochastic variant to the bound N Σ (1) S (4), where the individual functions are (1( Eee,, (-))) Essentially, the variant computes the gradient an incrementally sampled fi" ** rather than considering the full dataset.
With s = Ns, 60 N8 this turns 25
26
into the following updates:
NS - (1-a)
N8 + NVNS |g|
(27)
w - w [No + Ng. (NVS)
(28)
where g= g EEEM Ve,(O), N(0 )) 0² (Ns) and M denotes randomly
drawn minibatch. Moreover, we have g. = B EjEM Ve + c) with g/(Ns) gg/s
where we introduce the step-size 0 in the adversarial step to absorb the I/N factor and to gain additional control over the perturbation strength.
Following the practical improvements of Osawa et al. (019)), we introduce damping parameter > 0 and seperate step-size B2 > on the precision estimate, as well as an exponential moving
average for the gradient 31 > 0. Under these changes, and dividing out the N N-factors in 27 we arrive at the following method:
28
8m + Bigm + B1) (ow + ge),
(29)
- Bys + (1 - 32) + + v Ns. g
(30)
w - agm//s,
(31)
with g and g. defined as above. In practice, the method performed better without the v -actor in the S- s-update. Renaming å as s, the updates (29 (31) are equivalent to bSAM (Alg
E.4 BSAM ALGORITIM WITH m- MI-SHARPNESS
The bSAM algorithm parallelizes well onto multiple accelerators. This is done by splitting up minibatch into parts and computing independent perturbations for each part in parallel. The algorithm remains exactly the same as Alg. [ with lines 3 6 replaced by the following lines 8:
1: Equally partition M into (Mi...., of examples each
2: for A 1... in parallel do
3: 02 t w + ck. Ck ~ NN(c 0", 12+1/(N...
4: g - (1/B) EiEM, (1,(O,)
5: 6 + 81/s
6: Bi, + (1/B) EiEM, Vliw+ee
7: end for
8: g. (1/m) Em Benk: g (1/m) Σm= 8x
18

---

Published as a conference paper at ICLR 2023
F ADDITIONAL DETAILS AND EXPERIMENTS
F.I DETAILS ON LOGISTIC REGRESSION EXPERIMENT
Fig. ((a) shows the binary classification data-set (red vs blue circles) adopted from (Murphy 2012 Ch. 8.4). The data can be linearly seperated by a linear classifier with two parameters whose decision boundary passes through zero. We show the decision boundaries obtained from the MAP solution and from samples of the exact Bayesian posterior in Fig. ((a)
Fig. 5(b) shows the 2D probability density function of the exact Bayesian posterior over the two parameters, where the black cross indicates the mode (MAP solution), and the orange triangles the weight-samples corresponding to the classifiers shown the left figure. Fig. [(c) we show the solutions found by SAM, bSAM and an approximate Bayesian posterior. On this example, the posterior approximation found by bSAM similar to the Bayesian one. Performing Laplace approximation at the SAM solution leads to the covariance shown by the green dashed ellipse. It overestimates the extend of posterior as based only on local information.
5)
Figure 5: Posterior shapes for the logistic regression problem shown in Sec. 4.1
F.2 COMPARISON ON A LOW-DIMENSIONAL CLASSIFICATION PROBLEM
In Fig. 6 we compare predictive uncertainties of SAM and bSAM on the two-moons classification
problem. We fit small network (2 24 12 12 1). For SAM+Laplace, consider
diagonal GGN approximation to the Hessian (which was indefinite at the SAM solution) and use the linearized model for predictions, suggested by Immer et al. 2021..
98
98
40g
bSAM
SAM
SAM Laplace
Figure 6: bSAM improves SAM's predictive uncertainty. Using Laplace approximation at the SAM solution can lead to an overestimation of the predictive uncertainty.
F.3 EFFECT OF USING f* INSTEAD
In this section, empirically check the effect of using f in the Bayesian objective instead of f. To that end, we compute an exact solution to the relaxed problem as well as the original problem for full Gaussian variational family.
19

---

Published as a conference paper at ICLR 2023
For small problems, gradients of the biconjugate function can be easily approximated using convex optimization. To see this, notice that:
(H) = arg min f*(x) - (X' m
X é a
arg min c- (x) s.t. c2 fX)).
(32)
This is a optimization problem with linear cost function and convex constraints. To implement the constraints, we approximate the loss by maximum over linear functions, which is always possible for convex losses. The linear functions are first-order approximations of the loss points 0; qu
drawn from the current variational distribution. Under this approximation, the constraint in (32 can
be written as L seperate constraints for i (1...., L} as follows:
sup og
0 1 0 + ((.,) (Ve(O,); 0.) Sc
0
1
x2 + V((0.)) cs (VE(O.), E(O,).
(33)
The convex program (32) with convex constraints (33) then solved using the software package
CVXPY (Diamond & Boyd 2016) to obtain the gradient of the biconjugate function.
Having the gradients of the biconjugate available, our posterior is obtained by the Bayesian learning rule of Khan & Rue (2021).
No VS(μ), 2o + Vf (p). for relaxed-Baye Bayes
for Bayes
A - (1 a ta
(34)
the logistic regression problem de- scribed in App. ET We can see that the relaxation induces lower-bound as expected, but the relaxed solution is still reasonable approximation to the true posterior.
Fig. compares both solutions obtained by iterating 34
Bayes
relaxed Bayes
ELRO
Bayas
relaxed-Brig Bayee
20
60
60
100
iterations
(a)
(b)
Figure 7: Fig. (aa compares the posterior approximations obtained by solving (2) and our lower- bound 11 for full Gaussian variational family. Fig. 7(b) shows the gap in the evidence-lower bound (HLBO) induced by the relaxation.
F,4 SENSITIVITY BSAM TO THE CHOICE OF P
We show the sensitivity of all considered SAM variants to the hyperparameter in Fig. 8 As the bSAM algorithm learns the variance vector, the method expected be less sensitive to the choice of p. We confirm this here for ResNet-20 the CIFAR- 100 dataset without data augmentation.
F.5 ABLATION STUDY: EFFECT OF GAUSSIAN SAMPLING (LINE 3 IN ALG.
In the derivation the bSAM algorithm Sec. B we performed a local linearization order to approximate the inner problem. There, claim that performing this local linearization not the
20

---

Published as a conference paper at ICLR 2023
ResNe120-FRN CIEAR-
ResNerZ0 CIPAR-
Accuracy (higher is better)
SAM-SGD
64)
NLL (lower is better)
SAM-Adam
50
bSAM
SAM-SGD
SAM-Adam
6SAM
10-4
10P
10
311-1
107
ResNec20-FRN CIFAR-100
0.4
0.90
SAM-SGD
AUROC (higher is botter)
ECE (lower is better)
0.85
0.3
SAM Adaan
0.80
0.75
0.70
SAM-SGD
SAM-Adam
0.65
bSAM
10
10
10
10 10a
10
16
10
10
P
Figure 8: Sensitivity of SAM-variants to p. For bSAM, we plot the markers 100 p to roughly align the minima/muxinu of the curves. The proposed bSAM method adapts to shape of the loss and overall more robust to misspecified parameter p while also giving the overall best performance for all four metrics (Accuracy 1, NLL ECE 4, AUROC f).
mean but rather at sampled point is preferable from a Bayesian viewpoint. A concurrent work by Liu et al. (2022) also shows that combining SAM with random sampling improves the performance. The following Table confirms that the noisy linearization helps in practice.
Model /
Accuracy t (higher better)
NLL (lower bettet)
ECE +
AUROC 1 (higher better)
Method
Data
(lower better)
MLP/
bSAM w/o Noise
98.67
0.041
0.0030 0007)
0.981
MNIST
bSAM (Alg. T
98.78 (8.06)
0.038
0.0024 (0.0044)
0.982
Improvement:
+00.11
+ +0.003
0.0006
+0.001
LeNet-5
bSAM w/o Noise
92.01 (a.11)
(.23
0.0128 0021))
0.918(0.002)
FMNIST
bSAM (Alg. T
92.08
0.22((.005)
0.0066 0022)
0.920
Improvement:
000.07
+0.01
0.0062
+0.002
ResNet-20-FRN
bSAM w/o Noise
87.74(0.28)
(0.09)
1.0251(o.000)
0.896;0.004)
/CIFAR-10
bSAM (Alg. T
(8.72
(.34)p.o0s)
0.0163
0.903
Improvement:
+00.98
+ +0.04
+0.0088
+0.007
ResNet-20-FRN /CFARR-100
bSAM w/o Naise
(0.3)
1.44(0.000)
0.0226
0.832(0.002)
bSAM (Alg I
62.64(0.34)
(o.2()
0.0311 (0.002)
0.841 (0.003)
Improvement:
02233
+0.12
-0.0085
+0.008
Table 3: Performing the mos-llinearizatin the noisy point rather than the mean ("bSAM w/o Noise") typically improves the results with respect to all metrics.
F.6 ABLATION STUDY: EFFECT OF BAYESIAN MODEL AVERAGING
In Table 4 we show how the predictive performance and uncertainty of bSAM depends on the num- ber of MC samples used to approximate the integral in the Bayesian marginalization.
21

---

Published as a conference paper at ICLR 2023
Accuracy t (higher is better)
NLL 4
ECE
AUROC (higher better)
Model / Data
Samples
(lower is better)
(lower better)
0 (mean)
(8.34().22))
0.392 007)
0.0538 (0.003)
(o.00)))
MossNet-20-FRA /
8
(0.2)
0.349 (0. 006)
0.0212 (2.12(000000 (0.002)
0.9038 (0.005)
CIFAR-10
64
88.59 (8.59(
0.339 to. 005)
0.0186 (0.001)
0.9056
512
88.69 (0.20)
0.338 (0.000)
0.0173 (0.0022
0.9047 (0.0u6)
2048
(8.63(o.20)
0.338 (0.00))
0.0173 (0.003)
0.9058 (o.005)
Table 4: Increasing the number of MC samples to approximate the marginal Bayesian predictive distribution tends to improve the performance terms of NLL and ECE.
F.7 ABLATION STUDY: DEPENDANCE OF SAM AND BSAM ON "M-SHARPNESS"
It has been observed in Foret et (02T)) that the performance of SAM depends the number m of subbatches the minibatch is split into. This phenomenon is referred to as m-sharpness. From theoretical viewpoint, larger value of m corresponds to looser lower-bound for bSAM, as the sum of biconjugates is always smaller than the biconjugate of the sum.
Here we investigate the performance of SAM and bSAM on the parameter m for a ResNet-20-FRN on CIFAR-100 without data augmentation. The other parameters are chosen as before, and the batch-size is set to 128. The results are summarized in the following T
Accuracy (bigher is better)
NLL
ECE 4
AUROC t (higher is better)
Algorithm
m
(lower better)
(lower is better)
1
61.388 (7)
1.39().007)
0.0264(0.002)
0.836
2
61.47 16)
1.38 002)
0.0233
(.837(0.003)
bSAM
4
62.35
1.34(0.014)
0.0208
0.837 (0.33(00.00))
8
62.81 (0.57)
1.31(0.017)
0,0327 (0.04)
0.842 (0.003)
16
63.31 (0.32)
1.30
0.0592
0.840 (o.00))
1
55.88 (0.81)
1.82
0,1501
(0.00)00000) 004)
0.820 (o.00))
2
56.66(0.74)
1.76().041)
0.1421 (0.006)
0.81 0.819(o.oos)
SAM-SGD
56.53
(.78)).058)
0.1289 0060))
0.810
8
58 (8.62
1.60
0.1017 (0.005)
0.825(0.003)
16
59.51 (0.67)
1.51 (0.040)
0.0681 (0.006)
0.829 (.829(o.004)
Table 5: Effect of "m-sharpness" mini-batch size for SAM-SGI and bSAM for ResNet-20-FRN on CIFAR-100 without data augmentation. As the effective minibatch size decreases, all perfor- mance metrics (accuracy and uncertainty) tend to improve. This boost in performance is not captured by our theory, and understanding it is interesting direction for future work.
F.8 EXPERIMENTS ON MNIST
In Table compare bSAM to several baselines for smaller neural networks MNIST and FashionMNISST. Also in these smaller settings, bSAM performs competitively to state-of-the-ard
G CHOICE OF HYPERPARAMETERS AND OTHER DETAILS
G.1 EXPERIMENTS IN TABLE AND TABLE 6
In the following, we list the details of the experimental setup. First, we provide some general remarks, followed by tables of the detailed parameters used each dataset.
For all experiments, the hyperparameters are selected using a grid-search over a moderate amount of configurations to find the best validation accuracy. Our neural network outputs the natural param-
22

---

Published as a conference paper at ICLR 2023
Model / Data
Accuracy f (higher better)
NLL 4
BCB
AUROC1
Method
(lower is better)
(lower better)
(higher is better)
SGD
98.63(0.06)
0.044(0.00)
0.0028 (0.0070))
0.979
SAM-SGD
98 (0.03)
0.039 (0.0005)
0.0031(0.0002)
0.978(0.003)
SWAG
98. (5.66)
0.044 0002)
0.0028
0.976
MLP MNIST
VOGN
98.54
0.046 (0.0008)
0.0033 (0.0007)
0.980 (0.0022
Adam
98.41 (0.06)
(o.50
0.0036 (.0036(o.0007)
(o.02)
SAM-Adam
(8.58
0.046
0.0044(o.ouo0)
0.980 (0.002)
bSAM (ours)
98. 08.7700
0.0 0.0388j.m (338g00.01)) (1.0011)
0.0024 0.0024(0.0000 (0.0004)
0.982
SGD
(0.37(
(o.20)
0.0429 (0.0015)
0.897 (0.004)
SAM-SGD
91.95
0.22 (0.004)
0.0062 0.0062 (0.0007)
0.91 7 (0.002)
SWAG
(1.38(0.27)
0.31 (o.00)
0.0397 (0.0026)
0.901 (0.005)
LeNet-5
VOGN
91.24(0.25)
0.0071 (.24().004)
(0.0012)
0.916 (0.00)
FMNIST
Adam
91.14(0.25)
(0.30)
0.0450(0.0000)
0.897 (0.005)
SAM-Adam
91.66(0.19)
(.2550.004)
0.0225
0.913
bSAM (ours)
92. 12.10
0.22
0.0066
0.920
Table 6: Comparison on small datasets.
eters of the categorical distribution as a minimal exponential family (number of classes minus one output neurons). The loss function is the negative log-likelihood
We always use batch-size of B 128. For SAM-SGD, SAM-Adam and bSAM we split each minibatch into m = subbatches, for VOGN we set m 16 and consider independently computed perturbations for each subbatch. Finally, to demonstrate that bSAM does not require excessive amount of tuning and parameters transfer well, we use the same saper-parameters found on CIFAR- 10 for the CIFAR-100 100 experiments.
We summarize all hyperparameters in Table
The architecture is 784 500 - 300 -9 fully-connected multilayer perceptron (MLP) with ReLU activation functions. All methods trained for 75 epochs. We use cosine learning rate decay scheme, annealing the learning rate to zero. The hyperparameters
MNIST.
the individual methods SGD for 60 epochs (with the same parameters). Then collect SWAG statistics with fixed a = 0.01, 31 = 0.95. For all SWAG rums, we set the rank to 20, and collect sample every 100 steps for 15 epochs.
shown in Table 7 For SWAG,
The architecture is LeNet-5 with ReLU activations. We cosine learning rate decay scheme, annealing the learning rate to zero. We train all methods for 120 epochs. For SWAG, we run SGD for 105 epochs and collect 15 epochs with 0.001 and 0.8.
FashionMNIIST.
CIFAR-10/100. The architecture is ResNet-20 with filter response normalization nonlinearity adopted from Izmailov et al. (202T). A. The ResNet-20 has the same number of parameters as the one used in the orgginal publication (He et al.) (016). We train for 180 epochs and decay the learning rate by factor 10 at epoch 100 and 120. For SWAG, we run SGD for 165 epochs and collect for 15 epochs with 0.0001, 31 0.9.
G.2 EXPERIMENTS IN TABLE
The bSAM method depends on the number of data-samples N. To account for the random cropping and horizontal flipping. rescale this number by factor 4 which improved the performance. This corresponds to a tempered posterior, as also suggested in Osawa et al. (2019).
In all experiments, all methods cosine learning rate schedule which anneals the learning rate across 180 epochs to zero. The batch size is again chosen as B 128.
For overview over all the hyperparameters, please Table
23

---

Published as a conference paper at ICLR 2023
Configuration
Method
B1
N
8
82
SGD
0.1
0.95
24
X
X 0.05
VOGN
0.0005
0.95
25
0.999
MNIST
Adam
0.0005
0.8
C
MLP
SAM-SGD
0.1
0.95
30
0.05
X
SAM-Adam
0.0005
0.8
24
0.02
bSAM
0.1
0.8
15
0.01
SGD
0.01
0.8
60
X
X 0.05
VOGN
0.0001
0.95
25
0.99
FashionMNIST/
Adam
0.001
0.9
60
O
LeNet-5
SAM-SGD
0.05
0.8
60
0.05
X
SAM-Adam
0.001
0.9
60
0.02
bSAM
0.1
0.95
50
0.002
SGD
0.1
0.9
25
X
x 0.999
X 0.01
VOGN
0.001
0.9
25
CIFAR-I0/
Adam
0.001
0.9
250
ResNet-20-FRN
SAM-SGD
0.05
0.9
0.05
X
SAM-Adam
0.0003
0.9
25
0.1
bSAM
0.1
0.9
25
0.001
SGD
0.1
0.9
30
X
X
x 0.01
VOGN
0.001
0.9
25
0.999
CIFAR- /
Adam
0.001
0.9
250
O
ResNet-20-FRN
SAM-SGD
0.05
0.9
0.05
X
SAM-AN A-MAdam
0.0003
0.9
25
0.1
ASAM
0.1
0.9
25
0.001
Table 7: Hyperparameters used in the experiments without data augmentation. "x" denotes that this method does not use that hyperparameter. "O" indicates fixed hyperparameter across all datasets.
Dataset
Method
81
N
5
82
SGD
0.1
0.95
10
X
x
X
Adam
0.005
0.7
2.5
O
CIFAR-10
SAM-SGD
0.03
0.95
25
0.01
SAM-Adam
0.001
0.8
bSAM
1.0
0.95
0.01
SGD
0.03
25
Adam
0.005
0.9
x
O
CIFAR-100
SAM-SGD
0.2
0.8
10
0.02
X
SAM-Adam
0.005
0.7
0.2
bSAM
1.0
0.95
0.01
SGD
0.1
0.95
20
Adam
0.002
0.9
20
O
TinylmageNet
SAM-SGD
0.1
0.95
50
0.01
OOX
X
0000
SAM-Adam
0.001
0.95
10
0.1
hSAM
0.1
0.9
25
0.00005
Table 8: Hyperparameters used in the experiments with data augmentation shown in Table 2 "x" denotes that this method does not use that hyperparameter. "O" indicates fixed hyperparameter across all datasets.
24

---

Published as a conference paper at ICLR 2023
G.3 RESNET- 18 EXPERIMENTS IN TABLE
We trained all methods over 180 epochs with a cosine learning rate scheduler, annealing the learning rate The batch-size set to 200, and both SAM and bSAM m--saarpesss with 8.
Dataset
Method
81
N
B2
Y
SGD
0.03
0.9
25
X
X
CIEAR-I0
SAM-SGD
0.03
0.9
0.05
X
bSAM
0.5
0.9
0.01
O
O
SGD
0.03
X
X
CIFAR-100
SAM-SGD
0.03
0.9
25
0.05
X
X
bSAM
0.5
0.9
10
0.01
O
O
Table 9; Hyperparameters used in the ResNet-18 experiments of Table method does not use that hyperparameter. "O" indicates fixed hyperparameter across all datasets.
"x" denotes that this
H TABLE OF NOTATIONS
The following Table 10 provides an overview of the notations used throughout the paper. Multidi- mensional quantities are written in bold-face, scalars and scalar-valued functions are regular face.
Symbol
Description
( II
An inner-product on finte-ddimensional vector-space.
The to-nr..
Fenchel conjugate and biconjugate of the function f-
0
Parameter vector of the model, e.g., the neural network.
T(Ø)
Sufficient statistics of an exponential family distribution over the parameters.
p. u. p"
Expectation (or mcan/moment) parameters of exponential family. Natural parameters of an exponential family. Space of valid natural parameters.
A. X.X X"
Ω
M
Space of valid expectation parameters. Exponential family distribution with moments Exponential family distribution with natural parameters
4r
4
A(A)
nog-partition function of the exponential family in natural parameters.
(²)
Entropy of qμ. the convex conjugate of the log-partition function.
w.
Mean and variance of normal distribution, corresponding μ and X.
m. 22
Mean and variance of another normal, corresponding to fı' and X'
N N(elm,o"))
follows normal distribution with mean m and covariance I Precision of isotropic Gaussian prior.
5
Table 10: Notation.
25


Reviews:
[{"review_id": "uKde2ZFpXLh", "reviewer": "Reviewer_Qr1q", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Sharpness-aware minimization (SAM) is a learning objective that improves upon maximum likelihood in terms of generalization, calibration and robustness of the resulting models. Intuitively, SAM drives the optimizer toward flatter minima, which are known to be desirable from the Bayes-optimal learning perspective. However, so far SAM has not been shown to directly correspond to any Bayesian learning objectives. In this paper authors derive this connection, showing that SAM is equivalent to minimizing a convex upper bound on a variational Bayesian learning objective. Authors then use this connection to propose an extension of SAM (Bayesian-SAM) that, in addition to the point-estimate of the parameters, estimates the parameter uncertainty, and leads to improved predictive performance and calibration in the numerical experiments.", "summary_of_the_review": "I strongly recommend accepting this paper. It's a rare example of a paper that theoretically connects two important ideas, _and_ uses this connection to derive a novel practical method that works well on real data. The theoretical connection is non-trivial, and the method is presented and evaluated well (albeit the evaluation could be made even stronger, see my suggestions above).", "strength_and_weaknesses": "Strengths:\n- A novel connection between two important ideas (SAM and Bayesian learning) that required a non-trivial amount of complex mathematical reasoning.\n- A well-motivated, novel practical method driven by the theoretical results.\n- Strong experimental results on a wide set of numerical benchmarks with a strong set of baseline methods.\n- A clear, polished write-up with consistent notation. A precise comparison to the baseline SAM-Adam method with algorithmic differences highlighted line-by-line.\n- Experiments to understand the sensitivity of the method to $\\rho$, as well as additional ablation studies in the appendix.\n\nWeaknesses:\n- The toy experiment could be more informative: have authors considered doing non-linear classification with a NN instead? This might accentuate the differences between uncertainty estimates.\n- I would recommend replacing the MNIST datasets in Table 1 with alternative, more challenging datasets. MNIST-based datasets are likely too easy for meaningful comparisons between methods, as indicated by more ambiguous results on these.\n- A (brief) discussion on the computation cost of the method and how it compares to SAM/SAM-Adam would be useful.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "X5EDrmp3iOu", "reviewer": "Reviewer_VM1e", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a loss function that looks similar to SAM loss, but is a lower bound of Bayesian objective.\nThe authors state that the proposed pipeline can be used to obtain reasonable uncertainty estimates for NNs.", "summary_of_the_review": "There is little empirical evidence that the proposed approach is better than other. It has only limited theoretical justification. So, I lean towards the reject the paper in its current state.", "strength_and_weaknesses": "Strengths:\n- Numerical experiments, to some extent, justify the method proposed in the paper.\n- The idea is nice and can lead to some insights about how SAM works (but sadly, more on that is in the weakness part)\n\nWeakness:\n- Actually, losses (12) and (1) are very different from each other, as in (1) we take the supremum over a neighbourhood and (12) we look at all space of possible values $\\varepsilon$. So, the related term $l(m + \\varepsilon)$ in (12) can be arbitrarily larger than the corresponding term in (1), making the whole approach inspired by Fenchel biconjugates but not equivalent to it. We also note that in Fenchel biconjuates, we search for a convex analogue for the function we optimize. This makes little sense in general deep learning, as the loss is very far away from convex with many symmetries and local optima observed.\n- So, we can't directly optimize (12).\n- The accuracy of this approximation is, thus, should be the main concern in the paper. This is not true now. The experiments follow a common pipeline for DL papers comparing to other approaches and making reasonable ablation studies.\n- what is the difference between ResNet-20-FRN and ResNet-20? The paper [1] reports similar results for ResNet-20 in Tables 2 and 3 for different variants of SAM for CIFAR10 and CIFAR100. For CIFAR10, they obtain 93.82% accuracy (c.t. bSAM 92.16% in the paper). For CIFAR100 they obtain 71.40% (c.t. bSAM 68.22% in the paper). What is the difference in the protocol that provides such a decrease in quality for your approach compared to the baselines?\n\nSuggestions:\n- I suppose that the biconjugate for the Exponential family is worth deeper examination, as it leads to convex problems, and thus, everything should be correct and can lead to interesting theoretical insights and analysis (also, I suppose that you can generalize Theorems to the paper in this more general case). \n\n\n[1] Kwon, Jungmin, et al. \"Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.\" International Conference on Machine Learning. PMLR, 2021.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "MmVjeYnpwpd", "reviewer": "Reviewer_FsSv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper sheds light on the theoretical underpinnings of Sharpness Aware Minimization (SAM) by taking the biconjugation of the PAC-Bayes loss. The result is three-fold:\n\n1. SAM can be understood as an optimal convex relaxation of the PAC-Bayesian loss.\n2. This relaxation of the PAC-Bayes loss retains the logical sense sense in which it measures epistemic uncertainty. (The new loss being a looser bound.)\n3. The PAC-Bayes relaxation is computed using convex optimization rather than Monte Carlo approximations.\n\nUsing the above insight, the authors modify SAM+Adam to fit a posterior mean and variance for Bayesian Neural Networks on standard image benchmarks.", "summary_of_the_review": "This reviewer strongly endorses this paper be accepted. The authors present a simple and enlightening connection between (vanilla) SAM and Bayesian methods. We believe such connections are critical toward getting ever closer to understanding how/why modern DNN techniques perform so well and aligning those empirical successes with the theoretical underpinnings of Bayesian methods.\n\nWhile we believe the lack of attention to m-sharpness was disappointing, we nevertheless recognize this work as a significant step toward better understand SAM and hopefully, ultimately leading to even more significant gains.\n\nWe thank the authors for their remarkably well written presentation and their resolve to offer insight into one of the recent empirical success stories, SAM.", "strength_and_weaknesses": "SAM is a SOTA optimizer which is poorly understood; any thoughtful story with appropriate empirical validation which explains how it works is a welcome contribution to the community. This paper easily exceeds that bar.  The mathematical details appeared correct and the prose was free from flaws or mistakes.\n\nI only see one weakness in the paper which I share in the interest of constructive feedback and which in no way undermines my enthusiasm for the paper.\n\nI would like to see much more study of the m-sharpness aspect of SAM, perhaps in a future work. M-sharpness was (seemingly?) only mentioned in the appendix. This is unfortunate because:\n\n1. m-sharpness _further_ weakens the bound\n2. m-sharpness leads to considerable test improvement.\n\nThe fact that weakening the bound improves results undermines the otherwise elegance of the paper's story. That is, we must now square the fact that an optimal convex relation of an otherwise defensible loss (von Neumann-Morgenstern and Savage utility theory / PAC-Bayes) is itself still somehow not optimal.  This is a hard pill to swallow because now that the loss is convex it is harder to lay blame to algorithmic concerns.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "Kw-Cu9DGgIA", "reviewer": "Reviewer_SPsx", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors establish a connection between sharpness-aware minimization (SAM) and a Bayes objective by using a relaxation of the latter, such that the expected negative-loss is replaced by a tight convex lower bound, which is shown to be equivalent to SAM. Moreover, they propose a Bayesian extension of SAM that improves both generalization and uncertainty estimation.", "summary_of_the_review": "This paper provides strong theoretical results, and extensive experimental results. However, some of the main questions remain unanswered, and some of the experimental results are not convincing enough due to weak baselines.", "strength_and_weaknesses": "Strength\n\nUsing the Fenchel biconjugate to bridge SAM and the Bayes objective is an interesting and insightful idea, and the resulting relaxation of the Bayes objective has shown good empirical performance. Moreover, the authors have conducted a variety of experiments and ablation studies on both toy and real datasets, providing deeper understanding of the proposed method.\n\nWeaknesses\n1. While the equivalence between SAM and the relaxed Bayes objective is interesting, it doesn't explain why the relaxed Bayes objective (or SAM) has better generalization performance than the original Bayes objective. In other words, it seems the \"maximum loss\" performs better than the \"expected loss\" despite the established connection, the reason of which is unclear.\n2. The baselines used in the experiments are too weak. For instance, test accuracies around 95% and 75%, respectively, on CIFAR-10 and CIFAR-100 would be more convincing.\n3. The readability of Sec. 2 needs improvement.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%