SAM AS AN O PTIMAL R ELAXATION OF B AYES
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
Sharpness-aware minimization (SAM) (Foret et al., 2021) and related adversarial methods (Zheng et al., 2021; Wu et al., 2020;
Kim et al., 2022) have been shown to improve generalization, calibration, and robustness in various applications of deep learning (Chen et al., 2022; Bahri et al., 2022), but the reasons behind their success are not yet fully understood.
The original proposal of SAM was geared towards biasing the training trajectories towards ﬂatter minima, and the effectiveness of such minima has various Bayesian explanations, for example, those relying on the optimization of description lengths (Hinton & Van Camp, 1993; Hochreiter & Schmidhuber, 1997),
PAC-Bayes bounds (Dziugaite & Roy, 2017; 2018; Jiang et al., 2020; Alquier, 2021), or marginal likelihoods (Smith & Le, 2018).
However, SAM is not known to directly optimize any such Bayesian criteria, even though some connections to PAC-Bayes do exist (Foret et al., 2021).
The issue is that the ‘max-loss’ used in SAM fundamentally departs from a Bayesian-style ‘expected-loss’ under the posterior; see Fig. 1(a).
The two methodologies are 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. 1(a)).
From this,
SAM can be seen as optimizing the relaxation of Bayes to ﬁnd the mean of an isotropic Gaussian posterior while keeping the variance ﬁxed.
Higher variances lead to smoother objectives, which biases the solution towards ﬂatter regions (Fig. 1(b)).
Essentially, the result connects SAM and
Bayes through a Fenchel biconjugate that replaces the expected loss in Bayes by a maximum loss.
What do we gain by this connection?
The generality of our result makes it possible to easily combine the complementary strengths of SAM and Bayes.
For example, we show that the relaxedBayes objective can be used to learn the variance parameter, which yields an Adam-like extension of SAM (Alg. 1).
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 uncertainty estimates to the best Bayesian methods.
We also discuss connections of the relaxed-Bayes objective to proximal algorithms (Parikh & Boyd, 2013; Chaudhari et al., 2018; 2019).
Our work complements similar extensions for SGD, RMSprop, and Adam from the Bayesian deeplearning community (Gal & Ghahramani, 2016;
Mandt et al., 2017; Khan et al., 2018; Osawa et al., 2019).
So far there is no work on such connections between SAM and Bayes, except for a recent empirical study by Kaddour et al.
Husain & Knoblauch (2022) give an adversarial interpretation of Bayes, but it does not cover methods like SAM.
Our work here is focused on connections to approximate Bayesian methods, which can have both theoretical and practical impact.
We discuss a new path to robustness by connecting the two ﬁelds of adversarial
[Figure]
[Figure]
Expected-LossBiconjugateLoss
SAM minimizer5 -
SAM minimizer
SAM minimizer Variance = 0.
Variance = Variance =
B
A
C
B
A
C
B
A
C
[Figure]
Our main result in Theorem 2 connects the two by using the optimal convex relaxation (Fenchel biconjugate) of the expected loss.
It shows that the role of ρ and σ are exactly the same, and a
SAM minimizer obtained for a ﬁxed ρ can always be recovered from the relaxation for some σ .
An example is shown in Panel (b) for a 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 smoother (red line).
Higher σ give smoother objectives where sharp minima A and C slowly disappear.
The
SAM minimizer is shown with a red star which matches the minimizer of the relaxation.
SAM AS AN O PTIMAL R ELAXATION OF B AYES
In SAM, the training loss (cid:96) ( θ ) is replaced by the maximum loss within a ball of radius ρ > 0 around the parameter θ ∈ Θ ⊂ R P , as shown below for an (cid:96) 2 -regularized problem,
[Equation]
with δ > 0 as the regularization parameter.
The use of ‘maximum’ above differs from Bayesian strategies that use ‘expectation’, for example, the variational formulation by Zellner (1988),
[Equation]
where q ( θ ) is the generalized posterior (Zhang, 1999; Catoni, 2007), p ( θ ) is the prior, and D
KL ( · (cid:107) · ) is the Kullback-Leibler divergence (KLD).
For an isotropic Gaussian posterior q ( θ ) =
N ( θ | m , σ 2 I ) and prior p ( θ ) =
N ( θ | 0 , I /δ ) , the objective with respect to the mean m becomes
[Equation]
which closely resembles Eq. 1, 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 ﬁnally rewriting θ = m + (cid:15) (cid:48) .
There are some resemblances between the two objectives which indicate that they might be connected.
For example, both use local neighborhoods: the isotropic Gaussian can be seen as a ‘softer’ version of the ‘hard-constrained’ ball used in SAM.
The size of the neighborhood is decided by their respective parameters ( σ or ρ in Fig. 1(a)).
But, apart from these resemblances, are there any deeper connections?
Here, we answer this question.
We will show that Eq. 1 is equivalent to optimizing with respect to m an optimal convex relaxation of Eq.
3 while keeping σ ﬁxed.
See Theorem 2 where an equivalence is drawn between σ and ρ , revealing similarities between the mechanisms used in SAM and Bayes to favor ﬂatter regions.
Fig. 1(b) shows an illustration where the relaxation is shown to upper bound the expected loss.
We will now describe the construction of the relaxation based on Fenchel biconjugates.
2. F ROM EXPECTED - LOSS TO MAX - LOSS VIA F ENCHEL BICONJUGATES
We use results from convex duality to connect expected-loss to max-loss.
Our main result relies on
Fenchel conjugates (Hiriart-Urruty & Lemar´echal, 2001), also known as convex conjugates.
Con-
[Figure]
Biconjugate
Expected Loss
Original negative loss B
A
C
[Figure]
Biconjugate
Expected Loss
Original negative
[Figure]
sider a function f ( u ) , deﬁned over u ∈ U and taking values in the extended real line.
Given a dual space v ∈ V , the Fenchel biconjugate f ∗∗ and conjugate f ∗ are deﬁned as follows,
[Equation]
where (cid:104)· , ·(cid:105) 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 to connect the max-loss in SAM to the expected-loss involving regular, minimal exponential-family distributions (Wainwright et al., 2008) of the following form,
[Equation]
with natural parameter λ ∈ Ω , sufﬁcient statistics T ( θ ) , log-partition function A ( λ ) , and a base measure of 1 with Ω := { λ : A ( λ ) < ∞} .
Dual spaces naturally arise for such distributions and can be used to deﬁne conjugates.
For every λ ∈ Ω , there exists a unique ‘dual’ parameterization in terms of the expectation parameter µ =
E θ ∼ q λ [ T ( θ )] .
The space of all valid µ satisﬁes µ = ∇ A ( λ ) , and all such µ form the interior of a dual space which we denote by M .
We can also write λ in terms of µ using the inverse λ = ∇ A ∗ ( µ ) , and in general, both µ and λ are valid parameterizations. We summarize the notations in a table in App.
J for the readers convenience.
In what follows, we will always use µ to parameterize the distribution, denoting q µ ( θ ) .
As an example, for a Gaussian N ( θ | ω, v ) over scalar θ with mean ω and variance v > 0 , we have,
[Equation]
Because v > 0 , the space M is the epigraph of parabola ( ω, ω 2 ) , and Ω is a negative half space because the second natural parameter is constrained to be negative; see Fig. 2(a).
The dual spaces can be used to deﬁne the Fenchel conjugates.
We will use the conjugates of the ‘expected negative-loss’ acting over µ ∈ interior ( M ) ,
[Equation]
For all µ / ∈ interior ( M ) , we set f ( µ ) =
+ ∞ .
Using Eq.
4, we can obtain the biconjugate,
[Equation]
where we use λ (cid:48) and µ (cid:48)(cid:48) to denote arbitrary points from Ω and M , respectively, and avoid confusion to the change of coordinates µ (cid:48)
= ∇ A ( λ (cid:48) ) .
We show in App.
A that, for an isotropic Gaussian, the biconjugate takes a form where max-loss automatically emerges.
We deﬁne µ to be the expectation parameter of a Gaussian N ( θ | ω , v I ) , and λ (cid:48) to be the natural parameter of a (different) Gaussian N ( θ | m , σ 2 I ) .
The equation below shows the biconjugate with the max-loss included in the second term,
[Equation]
The max-loss here is in a proximal form, which has an equivalence to the trust-region form used in
Eq.
1 (Parikh & Boyd, 2013, Sec. 3.4).
The ﬁrst term in the biconjugate arises after simpliﬁcation of the (cid:104) λ (cid:48) , µ (cid:105) term in Eq.
8, while the second term is obtained by using the theorem below to switch the maximum over µ (cid:48)(cid:48) to a maximum over θ , and subsequently over (cid:15) by rewriting θ = m
+ (cid:15) .
Theorem 1.
For a Gaussian distribution that takes the form (5) , the following is true
[Equation]
assuming that (cid:96) ( θ ) is lower-bounded, continuous and majorized by a quadratic function.
Moreover, there exists λ (cid:48) for which f ∗ ( λ (cid:48) ) is ﬁnite and, assuming (cid:96) ( θ ) is coercive, dom( f ∗ ) ⊆
Ω .
A proof of this theorem is in App.
B, where we also discuss its validity for generic distributions.
The red surface in Fig. 2(b) shows the convex lower bound f ∗∗ to the function f (shown with the blue mesh) for a Gaussian with mean ω and variance v (Eq. 6).
Both functions have a parabolic shape which is due to the shape of M .
Fig. 2(c) shows the same but with respect to the standard mean-variance parameterization ( ω, v ) .
In this parameterization, convexity of f ∗∗ is lost, but we see more clearly the negative (cid:96) ( ω ) appear at v → 0 (shown with the thick black curve).
At the limit, we get f ( µ ) =
f ∗∗ ( µ ) = − (cid:96) ( ω ) , but for higher variances, both f and f ∗∗ are smoothed versions of (cid:96) ( ω ) .
For example, the 3 pairs of blue and red lines show f and f ∗∗ respectively, for v = 0 . 5 , 2 , and 4.
The slices at these variances are visualized in Fig. 1(b) along with the original loss.
The function f ∗∗ ( µ ) is a ‘lifted’ version of the one-dimensional (cid:96) ( θ ) to a two-dimensional µ -space.
Such liftings also recently appeared in the optimization literature (Bauermeister et al., 2022), and our work connects them to exponential families and probabilistic inference.
The convex relaxation preserves the shape at the limit v → 0 , while giving smoothed convex lower bounds for all v > 0 .
For a high enough variance the local minima at A and C disappear, biasing the optimization towards ﬂatter region B of the loss.
We will now use this observation to connect to SAM.
2. SAM AS AN OPTIMAL RELAXATION OF THE B AYES OBJECTIVE
We will now show that a minimizer of the SAM objective in Eq.
1 can be recovered by minimizing a relaxation of the Bayes objective where expected-loss f is replaced by its convex lower bound f ∗∗ .
The new objective, which we call relaxed-Bayes, always lower bounds the Bayes objective of Eq. 2,
[Equation]
[Equation]
where in the second line we substitute f ∗∗ ( µ ) from (9), and merge E θ ∼ q µ [ 12 σ 2 (cid:107) θ − m (cid:107) 2 ] with the
KLD term.
Z denotes the normalizing constant of the second distribution in the KLD.
The relaxedBayes objective has a speciﬁc difference-of-convex (DC) structure because both f ∗∗ and KLD are convex with respect to µ .
Optimizing the relaxed-Bayes objective relates to double-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 to SAM.
Due to the DC structure, it is possible to switch the order of maximization with respect to µ and ( m , σ ) , and eliminate µ by a direct maximization (Toland, 1978;
Tao & An, 1997).
The maximum occurs when the KLD term is 0 and the two distributions in its arguments are equal.
A detailed derivation is in App.
C, which gives us a minimization problem over ( m , σ )
[Equation]
where δ (cid:48)
= 1 / ( σ 2 + 1 /δ 0 ) and we use a Gaussian prior p ( θ ) =
N ( θ | 0 , I /δ 0 ) .
This objective uses max-loss in the proximal form, which is similar to the SAM loss in Eq.
1 with the difference being that the hard-constraint is replaced by a softer penalty term.
The theorem below is our main result which shows that, for a given ρ , a minimizer of the SAM objective in Eq.
1 can be recovered by optimizing Eq.
12 with respect to m while keeping σ ﬁxed.
Theorem 2.
For every ( ρ, δ ) , there exist ( σ, δ (cid:48) )
such that
[Equation]
[Equation]
assuming that the SAM-perturbation satisﬁes (cid:107) (cid:15) (cid:107)
=
ρ 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 a special case of relaxed-Bayes, and shows that ρ 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, ﬂatter regions are preferred for higher σ because such values give rise to smoother objectives (Fig. 1(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 σ through the relaxed-Bayes objective.
Another way is to exploit the DC structure in µ -space to devise new optimization procedures to improve SAM.
In Sec. 3, we will demonstrate both of these improvements of SAM.
In App.
F, we discuss connections to proximal methods (Parikh & Boyd, 2013) and Entropy-SGD (Chaudhari et al., 2019) and similar improvements might also be possible for those variants.
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 to get a good performance in deep learning while optimizing for both mean and variance (despite some recent success (Osawa et al., 2019)).
Often, non-Bayesian methods are used to set the variance parameter (Plappert et al., 2018;
Fortunato et al., 2018; Bisla et al., 2022).
The reasons behind the bad performance remain unclear.
It is possible that a 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 our result can be used to borrow algorithmic techniques from SAM to improve Bayes, and we hope to explore this in the future.
B AYESIAN -SAM TO O BTAIN U NCERTAINTY E STIMATES FOR SAM
In this section, we will show an application of our result to improve an aspect of SAM.
We will derive a new algorithm that can automatically obtain reasonable uncertainty estimates for SAM.
We will do so by optimizing the relaxed-Bayes objective to get σ .
A straightforward way is to directly optimize Eq.
12, but we will use the approach followed by Khan et al.
(2018), based on the naturalgradient method of Khan & Rue (2021), called the Bayesian learning rule (BLR).
This approach has been shown to work well on large problems, and is more promising because it leads to Adam-like algorithms, enabling us to leverage existing deep-learning techniques to improve training (Osawa et al., 2019), for example, by using momentum and data augmentation.
To learn variances on top of SAM, we will use a multivariate Gaussian posterior q µ ( θ ) =
N ( θ | ω , V ) with mean m and a diagonal covariance V = diag( s ) − 1 , where s is a vector of precision values (inverse of the variances).
The posterior is more expressive than the isotropic Gaussian used in the previous section, and is expected to give better uncertainty estimates.
As before, we will use a Gaussian prior p ( θ )
=
N ( θ | 0 , I /δ 0 ) .
[Table]
[Table]
[Equation]
{ Improved variance estimate } 9:
[Equation]
←
[Equation]
{ Scale the gradient by s instead of √ s } 10: end while 11: Return the mean m and variance σ 2 ← 1 / ( N · s )
The BLR uses gradients with respect to µ to update λ .
The pair ( µ , λ ) is deﬁned similarly to Eq.
We will use the update given in Khan & Rue (2021, Eq. 6) which when applied to Eq.
11 gives us the following update,
[Equation]
[Equation] where α > 0 is the learning rate and λ
= (0 , − δ
I / 2) is the natural parameter of p ( θ ) .
This is obtained by using the fact that ∇ µ D KL ( q µ ( θ ) (cid:107) p ( θ ))
=
λ − λ 0 ; see Khan & Rue (2021, Eq. 23).
For α = 1 , the update is equivalent to the difference-of-convex (DC) algorithm (Tao & An, 1997).
The BLR generalizes this procedure by using α < 1 , and automatically exploits the DC structure in µ -space, unlike the naive direct minimization of (12) with gradient descent.
The update (14) requires ∇ f ∗∗ , which is obtained by maximizing a variant of Eq.
9, that replaces the isotropic covariance by a diagonal one.
For a scalable method, we solve them inexactly by using the “local-linearization” approximation used by 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 to previous applications of the BLR, see Khan et al.
; Osawa et al.
; Meng et al.
The details are in App.
E, and the resulting algorithm, which we call Bayesian-SAM (or bSAM) is shown in Alg. 1.
bSAM differs slightly from a direct application of Adam (Kingma & Ba, 2015) on the SAM objective, but uses the same set of hyperparameters.
The differences are highlighted in Alg. 1, where bSAM is obtained by replacing the blue boxes with the red boxes next to them.
For SAM with
Adam, we use the method discussed in Bahri et al.
Kwon et al.
(2021) for language models.
The main change is in line 5, where the perturbation (cid:15) is adapted by using a vector s , as opposed to SAM which uses a scalar (cid:107) g (cid:107) .
Modiﬁcations 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, Sec 3.4).
Finally, the modiﬁcation in line 9 is due to the Newton-like updates arising from
BLR (Khan & Rue, 2021) where s is used instead of √ s to scale the gradient as explained in App.
E.
Due to their overall structural similarity, the bSAM algorithm has a similar computational complexity to Adam with SAM.
The only additional overhead compared to SAM-type methods is the computation of a random noise perturbation which is negligible compared to gradient computations.
[Figure]
N UMERICAL E XPERIMENTS
In this section, we present numerical results and show that bSAM brings the best of the two worlds together.
It gives an 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
DL and Bayesian DL literature: SGD, Adam (Kingma & Ba, 2015), SWAG (Maddox et al., 2019), and VOGN (Osawa et al., 2019).
We also compare with two SAM variants: SAM-SGD and SAMAdam.
Both are obtained by inserting the perturbed gradients into either SGD or Adam, as suggested in Foret et al.
; Bahri et al.
Kwon et al.
We compare these methods across three different neural network architectures and on ﬁve 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 (NLL), expected calibration error (ECE) (Guo et al., 2017) (using 20 bins) and area-under the ROC curves (AUROC).
For an explanation of these metrics, we refer the reader to (Osawa et al., 2019, Appendix H).
For Bayesian methods, we report the performance of the predictive distribution p ( y | D ) =
(cid:82) p ( y | D , θ ) q ( θ ) d θ , which we approximated using an average over 32 models drawn from q ( θ ) .
All further details about hyperparameter settings and experiments are in App. H.
4. I LLUSTRATION ON
[Figure]
Fig.
3 compares the predictive probability on a simple Bayesian logistic regression problem (Murphy, 2012, Ch. 8.4).
For exact Bayes, we use numerical integration, which is feasible for this toy 2D problem.
For the rest, we use diagonal Gaussians.
White areas indicate probabilities around 0 .
5 , while red and blue are closer to 0 and 1 .
We see that the bSAM result is comparable to the exact posterior and corrects the overconﬁdent predictions of SAM.
Performing a Laplace approximation around the SAM solution leads to underconﬁdent predictions.
The posteriors are visualized in Fig. 5(c) and other details are discussed in App. G.1.
We show results on another small toy example (“two moons”) in App.
I.
We can also use this toy problem to see the effect of using f vs f ∗∗ in the Bayesian objective.
Fig.in App.
G.2 shows a comparison of optimizing Eq.
11 to an optimization of the original Bayesian objective Eq.
For both, we use full covariance to avoid the inaccuracies arising due to a diagonal covariance assumption.
We ﬁnd that the posteriors are very similar for both f and f ∗∗ , indicating that the gap introduced by the relaxation is also small.
4. R EAL D ATASETS
We ﬁrst 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 is useful in applications where it is not easy to do data-augmentation (for example, tabular or time series data).
The results are shown in Table 1, and the proposed bSAM method overall performs best with respect to test accuracy as well as uncertainty metrics.
Gains get larger as the models get larger.
In Table 2 we show results with data augmentations (random horizontal ﬂipping and random cropping).
We consider CIFAR-10, CIFAR-100 and TinyImageNet, and still ﬁnd bSAM to perform favorably, although the margin of improvement is smaller with data augmentation.
We note that, for
[Table] [Table]
[Figure]
SGD on CIFAR-10 in Table 2, the accuracy (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 architecture with ﬁlter response normalization (FRN), which has 0 .
27 million parameters.
Experiments with a larger ResNet are provided in App.
G.7. bSAM not only improves uncertainty estimates and test accuracy, it also reduces the sensitivity of
SAM to the hyperparameter ρ .
This is shown in Fig.
4 for SAM-SGD, SAM-Adam and bSAM, where bSAM shows the least sensitivity to the choice of ρ .
Plots for other metrics are in App.
G.3.
To understand the reasons behind the gains obtained with bSAM, we do several ablation studies.
First, we study the effect of Gaussian sampling in line 3 in Alg. 1, which is not used in SAM.
Experiments in App.
G.4 show that Gaussian sampling consistently improves the test accuracy and uncertainty metrics.
Second, in App.
G.5, we study the effect of using posterior averaging over the quality of bSAM’s 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 et al., 2021) and Bayesian methods (Osawa et al., 2019).
In App.
G.6, we show that larger values of m lead to better performance for bSAM too.
In our experiments, we used m = 8 .
For completeness, we show how m -sharpness is implemented in bSAM in App. H.
D ISCUSSION
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, we obtain a relaxed
Bayesian objective, which involves a maximum loss, similar to SAM.
We showed that SAM can be seen as optimizing the mean while keeping variance ﬁxed.
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 in 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 to 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 to improve theoretical understanding of SAM-like methods.
It is possible to use the results from convex-optimization literature to study and improve adversarial as well as Bayesian methods.
Our work may open such directions and provide a new path for research in improving robustness. We expand on future work and limitations of our contributions in ?? .
More broadly speaking, expectations with respect to distributions arise naturally in many other problems, and often sampling is used to approximate them.
Our paper gives a different way to approximate such expectations using convex optimization rather than sampling.
Such problems may also beneﬁt from our results.
A PPENDIX A D ERIVATION OF THE F ENCHEL B ICONJUGATE
Here, we give the derivation of the Fenchel Biconjugate given in Eq. 9.
For an isotropic Gaussian, we have T ( θ ) =
( θ , θθ (cid:62) ) and λ (cid:48) =
( m / σ 2 , − 1 / 2 σ 2 I ) .
We will use the following identity to simplify,
[Equation]
Using this, the ﬁrst term in Eq.
8 becomes
[Equation]
Similarly, for the second term, we use Theorem 1 to and use Eq. 15 again, sup
[Equation]
Subtracting (17) from (16), we arrive at Eq. 9.
B P ROOF OF T HEOREM
Theorem 1.
For a Gaussian distribution that takes the form (5) [Equation] the following is true
[Equation]
assuming that (cid:96) ( θ ) is lower-bounded, continuous and majorized by a quadratic function.
Moreover, there exists λ (cid:48) for which f ∗ ( λ (cid:48) ) is ﬁnite and, assuming (cid:96) ( θ ) is coercive, dom( f ∗ ) ⊆
Ω .
Proof.
We ﬁrst pull out the expectation in the ﬁrst equality in Eq.
[Equation]
We will denote the expression inside the expectation by J ( θ ) , and prove that the supremum over µ (cid:48)(cid:48) is equal to the supremum over θ .
We will do so for two separate cases where the supremum
J ∗ = sup θ ∈ Θ J ( θ ) < ∞ is ﬁnite, and inﬁnite respectively.
When J ∗ is ﬁnite: For this case, our derivation proceeds in two steps, where we show that both the following two inequalities are true at the same time,
[Equation]
This is only possible when the two are equal, which will establish the result.
To prove the ﬁrst inequality, we take a sequence { θ t } ∞ t =1 for which J ( θ t )
→ J ∗ when sending t → ∞ .
By the deﬁnition of supremum, there always exists such a sequence.
From the sequence { θ t } ∞ t =1 we will now construct a sequence { µ (cid:48)(cid:48) t } ∞ t =1 for which the expectation of J ( θ ) also converges to J ∗ .
This will establish that the supremum over µ (cid:48)(cid:48) is larger or equal than the supremum over θ .
It is only larger or equal, since hypothetically there could be another sequence of { µ (cid:48)(cid:48) t } ∞ t =1 which achieves higher value.
We will see soon that this cannot be the case.
We use the sequence { µ (cid:48)(cid:48) t } ∞ t =1 where µ (cid:48)(cid:48) t
= ( θ t , Σ t + θ t θ (cid:62) t ) with { Σ t } ∞ t =1 ⊂ S P × P + is taken in the space of symmetric positive-deﬁnite matrices with the sequence such that Σ t → 0 .
Now, we can see that the expectation E θ ∼ q µ (cid:48)(cid:48) t [ J ( θ )] converges to the supremum as follows:
[Equation]
In the above steps, we used the fact that J is Gaussian-integrable and continuous.
A sufﬁcient condition for the integrability is that it is bounded by quadratics.
With this, we have established that the supremum over µ (cid:48)(cid:48) is at least as large as the one over θ .
To prove the second inequality, we will show that the supremum over µ (cid:48)(cid:48) is also less or equal than the supremum over θ .
To see that, we note that J ( θ (cid:48) )
≤ J ∗ for all θ (cid:48) ∈ Θ .
This implies that for any distribution q µ (cid:48)(cid:48) , we have that E θ (cid:48) ∼ q µ (cid:48)(cid:48)
[ J ( θ (cid:48) )] ≤ J ∗ .
Since the inequality holds for any distribution, we also have that sup µ (cid:48)(cid:48) ∈M E θ (cid:48) ∼ q µ (cid:48)(cid:48)
[ J ( θ (cid:48) )] ≤ J ∗ .
Having established both inequalities, we now know that the suprema in (10) agrees whenever the supremum over θ is ﬁnite.
When J ∗ is inﬁnite: For such cases, we will show that sup µ (cid:48)(cid:48) ∈M E θ ∼ q µ (cid:48)(cid:48)
[ J ( θ )] = ∞ , which will complete the proof.
Again, by deﬁnition of supremum, there exists a sequence { θ t } ∞ t =1 for which J ( θ t ) → ∞ when t → ∞ .
This means that for any M > 0 there is a parameter θ t ( M ) in the sequence
such that J ( θ t ( M ) ) > M .
Now from the previous part of the proof, we know that we can construct a sequence q µ (cid:48)(cid:48) t of Gaussians whose expectation converges to J ( θ t ( M ) ) .
Since the above holds for any M > 0 , the supremum over the expectation is also + ∞ .
Conditions for existence of λ (cid:48) for which J ∗ is ﬁnite: The ﬁnal step of our proof is to give the conditions when there exist at λ (cid:48) for which J ∗ is ﬁnite, and therefore the result is non-trivial.
Since (cid:96) ( θ ) is upper-bounded by a quadratic function, there exists a constant c ∈ R , vector b ∈ R P and matrix A ∈ S P × P + such that the following bound holds:
[Equation]
Then for the candidate λ (cid:48) =
( − b , − 12 A )
we know that J ( θ ) =
(cid:96) ( θ ) + (cid:104) λ (cid:48) , T ( θ ) (cid:105) ≤ c has a ﬁnite upper bound for any θ .
Hence, the supremum is ﬁnite for this choice of λ (cid:48) ∈ Ω .
Domain of f ∗ ( λ (cid:48) )
: For a Gaussian distribution, valid ( λ (cid:48) 1 , λ (cid:48) 2 ) ∈ Ω satisfy the constraint λ (cid:48) 2 ≺ 0 .
We now show that if λ (cid:48) / ∈ Ω , then f ∗ ( λ (cid:48) ) =
+ ∞ .
Since λ (cid:48) 2 ⊀ 0 , there exists a direction d (cid:54) = {+
+} such that d (cid:62) λ (cid:48) 2 d ≥ 0 and λ (cid:48) 1 (cid:62) d ≥ 0 .
Taking the sequence { θ t } ∞ t =0 with θ t = t · d we see
[Equation]
, since coercivity of (cid:96) ( θ ) implies that (cid:96) ( θ t ) → ∞ for any sequence { θ t } ∞ t =1 with (cid:107) θ t (cid:107) → ∞ .
Some additional comments on the assumptions used in the theorem are in order.
(i)
Any reasonable loss function that one would want to minimize is lower-bounded, so this assumption is always satisﬁed in practice.
(ii)
The existence of an upper-bounding quadratic holds for functions with Lipschitz continuous gradient, which is a standard assumption in optimization known as the descent lemma.
In practice, our conjugates are also applicable to cases where this assumption is not satisﬁed.
In such a case, one can consider the conjugate function only in a 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.
(iii)
We assumed that (cid:96) ( θ ) is continuous since virtually all loss functions used in deep learning are continuous.
Interestingly, if (cid:96) ( θ ) is not continuous, the conjugate function f ∗ will involve its semicontinous relaxation.
(iv) The coercivity assumption on (cid:96) ( θ ) is another standard assumption in optimization.
For noncoercive losses, there could exist quadratic upper-bounds which are ﬂat in some direction.
This would lead to a dual variable λ (cid:48) at the boundary of Ω .
The above Theorem 1 also holds for other exponential family distributions such as scalar Bernoulli and Gamma distributions, but we 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 ﬁxed covariance.
The above proof does not work there, because the closure (see e.g.,
Csisz´ar & Mat´uˇs (2005); Malag`o et al.
(2011)) of the family of distributions does not contain all
Dirac delta distributions.
Intuitively, the requirement is that we should be able to approach any delta as a limit of a sequence of members of the family, which excludes the ﬁxed-covariance case.
C D ERIVATION OF THE E QUIVALENT R ELAXED O BJECTIVE (12)
Inserting the biconjugate (9) into the relaxed Bayes objective (11) yields
[Equation]
where we expanded the KLD as D KL ( q µ ( θ ) (cid:107) p ( θ ))
=
E θ ∼ q µ [log q µ ( θ )] − E θ ∼ q µ [log p ( θ )] .
Interchanging the order of maximization, and noticing the ﬁrst term does not depend on µ , we can ﬁnd a closed-form expression for the maximization in µ as follows:
[Equation]
In the last step
[Equation]
we used the fact that the product of two isotropic Gaussians is again a isotropic
Gaussian.
The family described by M contains all isotropic Gaussians, and therefore inﬁmum over the KLD is attained at zero for two identical distributions.
Finally
[Equation]
we compute Z as
[Equation]
where in the last step we inserted the normalization constant for the product of the two Gaussians
N ( θ | m , σ 2 I ) and N ( θ | 0 , 1 δ 0 I ) as for example given in Rasmussen & Williams (2006, App. 2).
This gives the following expression
[Equation]
where δ (cid:48)
= 1 / ( σ 2 + 1 /δ 0 ) .
Switching to a minimization over m , σ 2 and inserting the result for the exact µ -solution (18) leads to the ﬁnal objective
[Equation]
which is the energy function in (12).
D P ROOF OF T HEOREM
Theorem 2.
For every ( ρ, δ ) , there exist ( σ, δ (cid:48) )
such that
[Equation]
assuming that the SAM-perturbation satisﬁes (cid:107) (cid:15) (cid:107)
=
ρ at a stationary point.
Proof.
We will show that any stationary point of E SAM is also a stationary point of E relaxed when σ and δ (cid:48) 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 ﬁrst write the two objectives in terms of perturbation (cid:15) (cid:48) and (cid:15) respectively, [Table] [Equation] [Table]
[Equation]
[Table]
[Equation] [Table] [Equation] where i {(cid:107) · (cid:107) ≤ ρ } is the indicator function which is zero inside the (cid:96) 2 -ball of radius ρ and + ∞ 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., 2020, Prop. 3).
At such points, the ﬁrst derivatives in θ and (cid:15) (cid:48) (or (cid:15) ) vanish.
Taking the derivatives and setting them to 0, we get the following conditions for the two problems,
[Equation]
for some constant multiplier µ > 0 .
Here, we used the assumption that the constraint (cid:107) (cid:15) ∗ (cid:107) ≤ ρ is active at our optimal point.
Since the constraint is active, the negative gradient is an element of the (non-trivial) normal cone, which gives us the multiplier µ in the second line above.
The two equations are structurally equivalent.
Any pair ( m ∗ , (cid:15) (cid:48)∗ ) satisfying the ﬁrst line also satisﬁes the second line for σ 2 = ρ/µ and δ (cid:48) =
δ .
The opposite is also true when using the pair ( θ ∗ , (cid:15) ∗ ) in the ﬁrst line, which proves the theorem.
Finally, let us comment on the assumption we used in the proof.
We assumed that the constraint in the perturbation is active, i.e., (cid:107) (cid:15) ∗ (cid:107)
= ρ . This assumption would be violated if there exists a local maximum within a ρ -ball around the parameter θ ∗ .
However, such a local maximum is unlikely to exist since the parameter θ ∗ is determined by minimization and tends to lie inside a ﬂat minimum.
E D ERIVATION OF THE B SAM A LGORITHM
In this section of the appendix, we show how the natural gradient descent update (14) will lead us to the proposed bSAM algorithm shown in Alg. 1.
As mentioned in the main text, for bSAM we consider a generalized setting where our Gaussian distribution q µ ( θ ) =
N ( θ | ω , V ) has diagonal covariance.
This corresponds to the following setting:
[Equation]
Here, ω ∈ R P is the mean and diag ( s ) =
V − 1 denotes the entries of the inverse diagonal covariance matrix.
All operations (squaring, multiplication) are performed entrywise.
E. T HE GRADIENT OF THE BICONJUGATE
First, we discuss how to compute that gradient of the Fenchel biconjugate, as the BLR update (14) requires it in every iteration.
The diagonal Gaussian setup leads to a slightly generalized expression for the Fenchel biconjugate function:
[Equation]
where Σ − 1 = diag ( b ) denotes the diagonal covariance matrix with 1 / b on its diagonal.
It follows from results in convex analysis, for instance Rockafellar & Wets (1998, Proposition 11.3), that ∇ f ∗∗ ( µ ) =
( Σ − 1 ∗ m ∗ , − 12 b ∗ ) , i.e. 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 (cid:15) , m , and Σ (through b ).
For the ﬁrst two, we will make very similar approximation to what is used for SAM (Foret et al., 2021), and simply add an additional optimization over b .
To simplify computation, we will make an additional approximation.
Instead of an iterative joint minimization in m and b , we will perform a single block-coordinate descent step where we ﬁrst minimize in m for ﬁxed b
= s and then afterwards in b for ﬁxed ω
= m .
This will give rise to a simple algorithm that works well in practice.
Optimization with respect to (cid:15) .
To simplify the supremum in (cid:15) we will use SAM’s technique of local linearization (Foret et al., 2021) with one difference.
Instead of using the current solution to linearize, we will sample a random linearization point θ ∼ N ( θ | ω , diag( s ) − 1 ) .
This is preferable from a Bayesian perspective and is also used in the other Bayesian variants (Khan et al., 2018).
Linearization at the sampled θ then gives us,
[Equation]
Using this approximation in (20), we get a closed form solution for (cid:15) ∗ =
Σ ∇ (cid:96) ( θ ) .
Optimization with respect to m .
To get a closed-form solution for the optimization in m , we again consider a linearized loss: (cid:96) ( m + (cid:15) ) ≈ (cid:96) ( ω + (cid:15) ) + (cid:104)∇ (cid:96) ( ω + (cid:15) ) , m − ω (cid:105) .
Moreover, we set (cid:15) =
V ∇ (cid:96) ( θ ) , corresponding to a ﬁxed b = s .
This reduces the optimization of (20) to
[Equation] [Equation]
which has closed form solution m ∗ = ω − Σ ∗ ∇ (cid:96) ( ω + (cid:15) ) which we here leave to depend Σ ∗ .
Optimization with respect to Σ (via b ).
Inserting (cid:15) ∗ into (20) with the linearized loss, the minimization over b while ﬁxing m = ω also has a simple closed-form solution as shown below,
[Equation]
[Equation]
where θ ∼ N ( θ | m , diag( s ) − 1 ) and g = ∇ (cid:96) ( θ ) .
E. S IMPLIFYING THE NATURAL - GRADIENT UPDATES
Having a way to approximate the biconjugate-gradient, we now rewrite and simplify the BLR update (14).
The derivations are largely similar to the ones used to derive the variational online Newton approaches, see Khan et al.
(2017; 2018); Khan & Rue (2021).
Inserting the parametrization (19) and our expression ∇ f ∗∗ ( µ ) =
( Σ − 1 ∗ m ∗ , − 12 b ∗ ) into the BLR update (14), we get the following updates
[Equation]
where we also used the prior λ 0 = (0 , − 12 δ 0 I ) .
By changing the order of the updates, we can write them in an equivalent form that resembles an adaptive gradient method:
[Equation]
[Equation] Now, inserting the solutions Σ − 1 ∗ ( ω − m ∗ ) = ∇ (cid:96) ( θ + (cid:15) ) and b ∗ = |∇ (cid:96) ( θ ) | · √ s from Eqs. 21 and 22 into these updates, we arrive at [Equation]
[Equation]
(26) where (cid:15) = V ∇ (cid:96) ( θ ) and θ ∼ N ( θ | ω , V ) .
E. F URTHER MODIFICATIONS FOR LARGE - SCALE DEEP LEARNING
Now we will consider further modiﬁcations to the updates (25) – (26) to arrive at 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 to deal with larger problems.
To that end, we now brieﬂy describe the general learning setup.
The loss is then written as a sum of individual losses, (cid:96) ( θ ) =
(cid:80) Ni =1 (cid:96) i ( θ ) where N is the number of datapoints.
The N datapoints are disjointly partitioned into K minibatches { B i } Ki
=1 of B examples each.
We then apply a stochastic variant to the bound N 1 K (cid:80) Ki
=
f ∗∗ i ( µ ) ≤ f ∗∗ ( µ ) , where the individual functions are f i ( µ ) =
1 B (cid:80) j ∈ B i E θ ∼ q µ [ − (cid:96) j ( θ )] .
Essentially, the variant computes the gradient on an incrementally sampled f i ∗∗ rather than considering the full dataset.
With s = N ˆ s , δ
=
Nδ this turns (25) – (26) into the following updates:
[Equation] [Equation]
(28) where g = 1 B (cid:80) j ∈M ∇ (cid:96) j ( θ ) , θ ∼ N ( θ | ω , diag ( σ 2 )) , σ 2 = ( N ˆ s ) − 1 and M denotes a randomly drawn minibatch.
Moreover, we have g (cid:15) = 1 B (cid:80) j ∈M ∇ (cid:96) j ( ω + (cid:15) ) with (cid:15) =
g / ( N ˆ s )
=
ρ g / ˆ s where we introduce the step-size ρ > 0 in the adversarial step to absorb the 1 /N factor and to gain additional control over the perturbation strength.
Following the practical improvements of Osawa et al.
(2019), we introduce a damping parameter γ > 0 and seperate step-size β 2 > 0 on the precision estimate, as well as an exponential moving average for the gradient β 1 > 0 .
Under these changes, and dividing out the N -factors in (27) – (28) we arrive at the following method:
[Equation]
[Equation]
31) with g and g (cid:15) deﬁned as above.
In practice, the method performed better without the √ N -factor in the ˆ s -update.
Renaming ˆ s as s , the updates (29) – (31) are equivalent to bSAM (Alg. 1).
F C ONNECTION TO THE P ROXIMAL P OINT A LGORITHM
In this section, we will see how the BLR update on the relaxed Bayesian objective can be specialized to a proximal point algorithm.
This will provide us with connections to Entropy-SGD (Chaudhari et al., 2018) and the Lookahead optimizer (Zhang et al., 2019).
Moreover, it will furnish our understanding in which situations the method can avoid sharp minima.
To get the proximal point method, we have to consider the BLR update (23) – (24) for ﬁxed precision s = b ∗ =
t − 1 1 , δ 0 = 0 and an exact minimization for the biconjugate gradient computation.
Under these settings, we get:
[Equation]
(32) where the proximal mapping (Parikh & Boyd, 2013) of h with step-size t > 0 is given by
[Equation]
(33) with h ( m ) =
sup (cid:15) ∈ R P (cid:96) ( m + (cid:15) ) − 12 t (cid:107) (cid:15) (cid:107) 2 .
The proximal-mapping (33) is due to the optimization problem (20), which computes a gradient of the biconjugate.
The above update is the relaxed proximal point algorithm of Eckstein & Bertsekas (1992, pp. 9) with relaxation parameter α ∈ (0 , 1] .
When using an iterative scheme to compute inexact/approximate solutions to (33), this leads to a “double-loop” method in similar style to the Lookahead optimizer (Zhang et al., 2019) and
Entropy-SGD (Chaudhari et al., 2019).
However, here the inner-loop is applied to our SAM-like loss function h ( m ) .
Similar to our method which requires a solution of an optimization problem to compute the gradient of the biconjugate in every iteration, proximal point methods also have seemingly “intractable” updates due to the inner minimization.
Nevertheless these methods are very useful, both in practice and theory.
We refer the interested reader to the review by Drusvyatskiy (2017) for a discussion on this matter.
We conjecture the BLR on the biconjugate loss to be of similar conceptual and guiding nature as the proximal point method, and believe it may aid the design of future learning and optimization algorithms when other (perhaps more accurate) approximations to the gradient of the biconjugate function are considered.
The proximal mapping (33) corresponds to a gradient descent step with step-size t on the Moreau envelope, see (Parikh & Boyd, 2013, Eq. 3.3).
Therefore, the update (32) can be viewed as gradient descent on the Moreau envelope of h , again with regularization parameter and step-size t > 0 .
This interpretation has been used in Entropy-SGD (Chaudhari et al., 2018) to argue why such a double-loop scheme prefers ﬂat minima: the Moreau smoothing widens all minima, and the hope is that eventually sharp ones are eliminated; we refer the interested reader to (Chaudhari et al., 2018,
Fig. 1), where the Moreau-smoothing corresponds to the “non-vicious Hamilton-Jacobi” case in that ﬁgure.
However, the function h used in (32) is the negative Moreau envelope of the negative loss with smoothing parameter t .
As we have seen, the function h eliminates sharp minima in the loss function.
This can be seen as a complementary type of smoothing to the proximal one, as the proximal point update widens the remaining ﬂat minima in h .
Therefore, for ﬁxed precision, the BLR on the relaxed Bayesian objective naturally combines the two different types of regularization proposed in earlier literatures: the one from SAM (Foret et al., 2021) and the one from Entropy-SGD
(Chaudhari et al., 2018) (non-viscious HJ setting), to avoid sharp minima.
We expect a faithful implementation, perhaps using inner iterations, to combine the advantages of both worlds.
Finally, let us understand in which situations such a strategy works.
The Moreau envelope of h is a double-envelope, as we take the Moreau envelope twice.
Such a construction corresponds to the (negative) proximal hull of the negative loss (Rockafellar & Wets, 1998, Ex. 1.44).
If the loss is ω -semiconcave and t − 1 ≥ ω , then by (Rockafellar & Wets, 1998, Ex. 11.26d)
this proximal hull coincides with the original function.
Therefore, in that case, the update (32) actually reduces to plain gradient descent on the original loss with step-size tα .
In other words, the method cannot avoid minima which are ﬂatter than the precision of the variational distribution.
Finally, we remark that as described in (Chaudhari et al., 2018, Sec. 6.1), semiconcavity can be a way of measuring the width of a local minimum.
G A DDITIONAL D ETAILS AND E XPERIMENTS G. D ETAILS ON THE L OGISTIC R EGRESSION E XPERIMENT
Fig. 5(a) shows the binary classiﬁcation data-set (red vs blue circles) we adopted from (Murphy, 2012, Ch. 8.4).
The data can be linearly seperated by a linear classiﬁer 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. 5(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 classiﬁers shown in the left ﬁgure.
In Fig. 5(c), we show the solutions found by SAM, bSAM and an approximate Bayesian posterior.
On this example, the posterior approximation found by bSAM is similar to the Bayesian one.
Performing a Laplace approximation at the SAM solution leads to the covariance shown by the green dashed ellipse.
It overestimates the extend of the posterior as it is based only on local information.
[Figure]
G
H OW TIGHT IS OUR LOWER - BOUND ?
In this section, we empirically check how tight our proposed lower-bound (11) is in practice.
To that end, we compute a near exact solution to the relaxed problem as well as the original problem (2) for a full Gaussian variational family.
For small problems, gradients of the biconjugate function can be easily approximated using convex optimization.
To see this, notice that:
[Equation]
[Equation]
This is a optimization problem with linear cost function and convex constraints.
To implement the constraints, we approximate the loss by a piecewise-linear function at L points θ i ∼ q µ drawn from the current variational distribution.
This leads to L constraints for i ∈ { 1 , . . .
, L } as follows:
[Equation]
[Equation]
The convex program (34) with convex constraints (35) is then solved using the software package
CVXPY (Diamond & Boyd, 2016) to obtain the gradient of the biconjugate function.
Having the gradients available, our posterior is obtained by the Bayesian learning rule of Khan &
Rue (2021),
[Equation]
relaxed
[Equation]
[Equation]
Fig.
6 compares both solutions obtained by iterating (36) on the logistic regression problem described in App.
G.1.
We can see that the relaxation induces a lower-bound as expected [Figure] but the relaxed solution is still a reasonable approximation to the true posterior
[Figure]
[Figure]
maxima of the curves.
The proposed bSAM method adapts to shape of the loss and is overall more robust to misspeciﬁed parameter
ρ while also giving the overall best performance for all four metrics (Accuracy ↑ , NLL ↓ , ECE ↓ , AUROC ↑ ).
G. S ENSITIVITY OF B SAM TO THE CHOICE OF ρ
We show the sensitivity of all considered SAM variants to the hyperparameter ρ in Fig.
As the bSAM algorithm learns the variance vector, the method is expected to be less sensitive to the choice of ρ .
We conﬁrm this here for a ResNet–20 on the CIFAR-100 dataset without data augmentation.
G. A BLATION STUDY : E FFECT OF G AUSSIAN S AMPLING (L INE 3 IN A LG . 1)
In the derivation of the bSAM algorithm in Sec.
3 we performed a local linearization in order to approximate the inner problem.
We claimed that performing this local linearization not at the mean but rather at a sampled point is preferable from a Bayesian viewpoint.
The following Table 3 shows that the noisy linearization helps also in practice. [Table] [Table] ResNet-
[Table]
FRN/CIFAR-
[Table]
G. A BLATION STUDY : E FFECT OF B AYESIAN MODEL AVERAGING
In Table 4 we show how the predictive performance and uncertainty of bSAM depends on the number of MC samples used to approximate the integral in the Bayesian marginalization
[Table]
G. A BLATION STUDY : D EPENDANCE OF SAM AND B SAM ON “ M - SHARPNESS ”
It has been observed in Foret et al.
(2021) that the performance of SAM depends on the number m of subbatches the minibatch is split into.
This phenomenon is referred to as m -sharpness .
From a theoretical viewpoint, a larger value of m corresponds to a 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 Table 5. [Table]
the effective minibatch size decreases, all performance metrics (accuracy and uncertainty) tend to improve.
G. E XPERIMENTS ON L ARGER N EURAL N ETWORKS
The following Table 6 compares SGD, SAM-SGD and the proposed bSAM method on a larger
ResNet with 11 million parameters. This network has 4 – 5 times more parameters than the “ResNet20-FRN” used in the main paper and uses standard batch normalization instead of ﬁlter response normalization (FRN) [Table] As can be seen in the Table
[Table] our proposed bSAM method again signiﬁcantly improves over SAM and SGD in all considered metrics. This indicates that bSAM scales well to larger architectures.
[Table]
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 on each dataset.
For all experiments, the hyperparameters were selected using a grid-search over a moderate amount of conﬁgurations to ﬁnd the best validation accuracy.
Our neural network outputs the natural parameters 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 a batch-size of B = 128 .
For SAM-SGD, SAM-Adam and bSAM we split each minibatch into m = 8 subbatches, for VOGN we set m = 16 and consider independently computed perturbations for each subbatch.
Finally, to demonstrate that bSAM does not require an excessive amount of tuning and parameters transfer well, we used the same hyper-parameters found on CIFAR10 for the CIFAR-100 experiments.
We summarize all hyperparameters in Table 7.
MNIST.
The architecture is a 784 – 500 – 300 – 9 fully-connected multilayer perceptron (MLP) with ReLU activation functions.
All methods were trained for 75 epochs.
We used a cosine learning rate decay scheme, annealing the learning rate to zero.
The hyperparameters of the individual methods are shown in Table 7.
For SWAG, we ran SGD for 60 epochs (with the same parameters).
Then collected SWAG statistics with ﬁxed α = 0 .
01 , β 1 = 0 .
For all SWAG runs, we set the rank to 20 , and collected a sample every 100 steps for 15 epochs.
FashionMNIST.
The architecture is a LeNet–5 with ReLU activations.
We used a cosine learning rate decay scheme, annealing the learning rate to zero.
We trained all methods for 120 epochs.
For
SWAG, we run SGD for 105 epochs and collect 15 epochs with α = 0 .
001 and β 1 = 0 . 8 .
CIFAR-10/100.
The architecture is a ResNet-20 with ﬁlter response normalization nonlinearity adopted from Izmailov et al.
The ResNet-20 has the same number of parameters as the one used in the original publication (He et al., 2016).
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 forepochs with α = 0 . 0001 , β 1 = 0 .
H. E XPERIMENTS IN T ABLE
The bSAM method depends on the number of data-samples N .
To account for the random cropping and horizontal ﬂipping, we rescale this number by a factor 4 which improved the performance.
This corresponds to a tempered posterior, as also suggested in Osawa et al.
In all experiments, all methods use cosine learning rate schedule which anneals the learning rate across 180 epochs to zero.
The batch size is again chosen as B = 128 .
For an overview over all the hyperparameters, please see Table 8. [Table]
[Table]
H. E XPERIMENTS IN T ABLE
We trained
[Table]
H. B SAM A LGORITHM WITH m -S HARPNESS
The bSAM algorithm parallelizes well onto multiple accelerators.
This is done by splitting up a minibatch into m parts and computing independent perturbations for each part in parallel. The algorithm remains exactly the same as Alg. [Table] with lines 3 – 6 replaced by the following lines 1 – 8 :
1: Equally partition M into {M 1 , . . .
, M m } of B examples each 2: for k = 1 . . .
m in parallel do 3:
[Equation]
[Equation]
[Equation]
←
[Equation]
6: g (cid:15),k ← (1 /B )
(cid:80) i ∈M k ∇
[Equation]
[Table]
[Figure]
I C OMPARISON ON A L OW -D IMENSIONAL C LASSIFICATION P ROBLEM
In Fig.
8, we compare predictive uncertainties of SAM and bSAM on the two-moons classiﬁcation problem.
We ﬁt a small network ( 2 − 24 − 12 − 12 − 1 ).
For SAM+Laplace, we consider a diagonal GGN approximation to the Hessian (which was indeﬁnite at the SAM solution) and use the linearized model for predictions, as suggested by Immer et al.
J T ABLE OF N OTATIONS
The following App.
J provides an overview of the notations used throughout the paper.
Multidimensional quantities are written in bold-face , scalars and scalar-valued functions are regular face.
Symbol Description (cid:104)· , ·(cid:105)
An inner-product on a ﬁnite-dimensional vector-space.
(cid:107) · (cid:107) The (cid:96) 2 -norm.
f ∗ , f ∗∗ Fenchel conjugate and biconjugate of the function f .
θ
Parameter vector of the model, e.g., the neural network.
T ( θ ) Sufﬁcient statistics of an exponential family distribution over the parameters.
µ , µ (cid:48) , µ
(cid:48)(cid:48)
Expectation (or mean/moment) parameters of exponential family.
λ , λ (cid:48) , λ (cid:48)(cid:48)
Natural parameters of an exponential family.
Ω Space of valid natural parameters.
M Space of valid expectation parameters.
q µ Exponential family distribution with moments µ .
q λ Exponential family distribution with natural parameters λ .
A ( λ ) Log-partition function of the exponential family in natural parameters.
A ∗ ( µ ) Entropy of q µ , the convex conjugate of the log-partition function.
[Table]