Rethinking Sharpness-Aware Minimization as Variational Inference
Abstract
Sharpness-aware minimisation (SAM) aims to improve the generalisation of gradient-based learning by seeking out flat minima.
In this work, we establish connections between SAM and meanfield variational inference (MFVI) of neural network parameters.
We show that both these methods have interpretations as optimizing notions of flatness, and when using the reparametrisation trick, they both boil down to calculating the gradient at a perturbed version of the current mean parameter.
This thinking motivates our study of algorithms that combine or interpolate between SAM and MFVI.
We evaluate the proposed variational algorithms on several benchmark datasets, and compare their performance to variants of SAM.
Taking a broader perspective, our work suggests that SAM-like updates can be used as a drop-in replacement for the reparametrisation trick.
1. Introduction
The training loss of overparametrized models often has several equally low minima, which may differ widely in their performance at test time.
The most salient effort to characterise which minima lead to better generalisation is the one that links flatness to generalisation [3–5, 9–11].
Motivated by this, SAM aims to directly bias optimization towards flat regions leading to substantial improvement in generalisation while only doubling the computational cost per iteration.
Other methods, such as MFVI and Evolution Strategies (ES) also have flatness-seeking interpretations.
The pursuit of flat minima was originally motivated by a Bayesian, minimum description length (MDL) argument [9].
Thus, it is not surprising that minimizing a variational Bayesian objective, itself an estimate of the description length of training data [8], favours flat minima.
This connection between variational Bayesian deep learning and SAM has remained largely unexplored.
In this work we show that the connections between MFVI and SAM are more than philosophical.
Concrete implementations of SAM and MFVI (with the reparametrization trick) both calculate gradients at a perturbed point around the current parameter.
While SAM takes the worst-case perturbation, MFVI uses a random Gaussian perturbation.
In this sense, the SAM update can be thought of as a biased deterministic approximation replacing the unbiased but high-variance single-sample
Monte Carlo gradient in MFVI.
Exploring this further, we make the following contributions:
We establish connections between SAM and MFVI, characterising and comparing their flatnessseeking inductive biases (Appendix A, B).
We propose the algorithm VariationalSAM, which combines aspects of SAM and MFVI (C).
We compare all methods on benchmark datasets (CIFAR-10 and CIFAR-100).
2. Sharpness-Aware Minimization and Mean-Field Variational Inference 2.1. Sharpness-Aware Minimization (SAM)
In this section, we present SAM with a slight change in notation relative to original works in order to make connections to variational methods more apparent later.
Denoting the parameter by µ , and the number of parameters by p , the idealized loss function SAM aims to minimize can be written as
[Equation]
Σ is any positive definite matrix.
Setting Σ
= ρ 2 p
I , where ρ is a constant parameter recovers the standard SAM [6], while setting Σ as a function of µ , recovers newer variants e.g. Adaptive SAM (ASAM) [15] and Fisher SAM (FSAM)
The parameter α > 0 controls L 2 regularization.
Computing gradients of (1) with respect to µ is intractable in general, hence an approximation is used.
Taking a first-order Taylor expansion of L we can approximate the optimal ϵ [6, 12, 15]:
[Equation]
[Equation]
where ∇ ⊤ θ L ( µ ) denotes the gradient of the loss evaluated at µ .
For reasons that will be clear later, we introduce new notation η ( µ ) and re-express ϵ ∗ from η as follows:
[Equation]
where
[Equation]
The SAM update direction is obtained by plugging this estimate back into (1) and ignoring the dependence of η on µ when differentiating the resulting expression:
[Equation]
[Equation] Note that computing this update direction requires backpropagation twice - first to calculate η ( µ ) , and second to evaluate ∇ ⊤ θ L - thus doubling the computational cost.
2.2. Mean-Field Variational Inference (MFVI)
We can also take an approximate Bayesian approach to inferring model parameters from data.
We start with a prior distribution, p ( θ ) over the parameters θ , which we choose to be N (0 , σ 0 I ) .
We then attempt to infer the posterior distribution of weights using the Bayes’ rule, p ( θ |D )
∝ e − NL ( θ ) p ( θ ) , where N denotes the number of data points and we assumed that the loss function L calculates the average negative log likelihood of i. i .d.
observations.
This operation being intractable, we use an approximate posterior q ( θ ) , chosen to be Gaussian with mean µ and covariance Σ .
Minimizing the KL-divergence between this approximate q and the true posterior yields the following objective function for µ and Σ :
[Equation]
When Σ is chosen to be diagonal, the minimization of (5) is called mean-field variational inference (MFVI).
Here, we liberally extend this name to algorithms where Σ is non-diagonal.
Dropping the KL term from (5) yields variational optimization (VO) [22] and evolution strategies [7, 19, 23].
We can estimate the gradients of L MFVI using the reparametrization trick [13].
We also expand the KL divergence term, and substitute α := [Equation] to obtain:
[Equation]
We approximate this by a single-sample Monte Carlo estimate and get the update direction:
[Equation]
where η
[Equation]
Note the similarity between Eqns.
(4) and (7).
They are of precisely the same form except for the perturbation η : in SAM, η is chosen to be the worst-case perturbation calculated deterministically from µ by taking a gradient, while in MFVI, η is drawn from a standard normal.
Throughout this work, we will refer to the variant of MFVI where we fix Σ
= ρ 2 d I as RandomSAM , or L RSAM ( µ ) .
2.3. Exploring the relationship
[Table]
mentioned in this work.
All methods can be related to MFVI or
SAM by changing the perturbation type, shape of Σ , or penalty terms. Random SAM is a variation of MFVI with fixed Σ = ρ 2 p I and Variational SAM is the version of MFVI, where we learn Σ . SAM as an upper bound on the Variational Objective:
The similarity between the SAM and
MFVI updates is not surprising given that L SAM can be considered as a loose upper bound on
L MFV I when the variance Σ is sufficiently small, and the number of parameters, p is sufficiently large.
This is because in high dimensions, samples from N ( µ, Σ) concentrate around the ellipsoid ( x − µ ) ⊤ Σ − 1 ( x − µ ) =
p .
Thus, any expectation over the N ( µ, Σ) can be upper bounded by the maximum value within the ellipsoid i. e. max ( x − µ ) ⊤ Σ − 1 ( x − µ ) ≤ p L ( θ )
This upper bound relationship is exploited in the proof of theorems used to justify SAM[6],
ASAM[15] and FSAM [12].
Those proofs therefore imply an even stronger theoretical justification for MFVI as a means to achieve good generalisation.
Implicit regularization
As discussed, both SAM and MFVI have flatness-seeking inductive biases.
But can we compare these biases, or connect them to notions of sharpness we understand better?
Here, we characterize the implicit regularization towards flat minima in terms of higher order derivatives, borrowing techniques from [1, 20, 21].
We summarise these in the following four propositions (we set α = 0 ignoring the L 2 or KL penalties for readability):
In the first row
[Equation]
(what it says on the box)
[Equation]
(what the algorithm does)
[Equation]
we approximate the idealized loss functions of SAM (Eqn.
and
MFVI (Eqn.
in terms of higher order derivatives.
SAM approximately penalizes the norm of the gradient, while MFVI penalizes the trace of the Hessian.
[-Proposition 
-]Near minima, Tr [Σ H ( µ )] is a good measure of sharpness, but it can take negative values around saddle points and local maxima.
In the second row we approximate the implicit regularization properties of stochastic gradient descent with SAM (Eqn.
(4)) or MFVI gradients (Eqn. (7)), with small but finite learning rate δ .
The δ 4 ∥ g ( µ ) ∥ term represents the implicit bias of GD with finite step size as in as in [21].
SAM’s flatness penalty is the same as the idealized, but for MFVI the penalty contains the trace of the Hessian squared, which is now always positive. These results show us the regularizing properties of both SAM and MFVI. The detailed theorems and proof can be found in Appendix A, B.
In Figure 1 we illustrate the sharpness-avoidance of [Figure] L SAM and L RSAM on a 2D toy example originally introduced in [12].
We can see that both transformations are effective at reducing the attraction basin of the sharp minimum.
However, it is still a question which translates better to high-dimensional problems we encounter in deep learning, and which one leads to better generalization.
The above connections between MFVI and SAM motivate our two main research questions which we investigate in the rest of this report:
All things being equal, is SAM or RandomSAM (MFVI with fixed Σ
= ρ 2 p I more effective at finding minima that generalise
well?
It is possible to learn Σ in SAM via gradient descent along with µ as in MFVI.
This results in an algorithm we call VariationalSAM .
Does this have advantages over SAM where Σ is fixed.
(Details of the algorithm are in Appendix C, E with a PAC-Bayes.)
3. Image Classification Experiments
In this section
[Table]
we empirically assess the generalisation performance of the mentioned algorithms: vanilla SGD, SAM, RandomSAM (RSAM) and VariationalSAM (VSAM).
Our implementations of these methods is open source
[Table]
and can be accessed in this repository.
Details of the parameters of the experiments can be found in Appendix D.
[Table]
SAM visibly outperforms all other methods in all experiments.
However, on CIFAR-10, RandomSAM performs almost as well as SAM.
Given that RandomSAM is a much simpler method operating with random noise and only one backpropagation, this result somewhat weakens the advantages of SAM.
It raises the question whether the superior results of SAM on larger datasets could be reproduced by other methods using random perturbations in the parameter space.
On CIFARwe see a significant gap in performance, however this may be due to our limited computing budget not allowing for more extensive grid search of parameters.
Our adaptive version of SAM, Variational SAM, does not appear to work better than SAM.
This may be because during training, Σ starts to increase, and often reaches a magnitude where the
Taylor approximation underlying SAM no longer holds.
This raises the question whether random
MFVI would have an advantage over SAM when more flexible covariance structures are used.
Although the evaluation of Adaptive SAM and Fisher SAM are not in the focus of this paper, we have included a limited set of results in Appendix G for comparison.
4. Summary and Future Work
In this work, we have provided a novel interpretation of SAM, from the angle of Variational Inference.
This led to the comparison of SAM with methods taken from Variational Inference: RandomSAM
(Variational Optimization) and Variational SAM.
The latter method performed rather unpromisingly.
Interestingly, RandomSAM performed similarly to SAM on the CIFAR-10 dataset.
This raises questions about optimality of SAM in biasing the optimisation against sharp minima, which we would like to explore in the future.
For example, one could construct an algorithm that interpolates between SAM and RandomSAM by adding random noise to the SAM perturbation.
In most experiments, a limitation of computational power has prevented extensive grid-searches on the hyperparameters in our models.
Therefore, the reported test accuracies can likely be improved with further parameter calibration.
5. Appendix A. Implicit regularization of SAM
Proposition [-
-]1 The following approximation holds for the SAM objective [Equation] Σ
[Equation]
Furthermore, for general Σ , this takes the form
[Equation]
Proof Plugging the solution ϵ ∗ in (Eqn. (2)) into the SAM objective and performing a first-order Taylor approximation, we arrive at
[Equation]
where we have used that Σ is a symmetric matrix.
In order to show that the error term is controlled, let us use express ϵ ∗ as in (Eqn. as ϵ ∗
= Σ 12 η .
Then
[Equation]
From (Eqn. (3)), we see that ∥ η ∥ 2 = √ p , which means that as long as the largest (in magnitude) eigenvalue of Σ , λ max < K 1 p for some K constant, we have η ⊤ Σ η < K .
In the special case of SAM, we have Σ = ρ 2 p
I , hence the error term scales with ρ 2 .
Proposition 2 Full batch gradient descent using the SAM step follows a path that is closest to the exact continuous path given by ˙ µ = −∇ µ ˜ L SAM ( µ ) , where ˜ L SAM ( µ ) is given by
[Equation]
11) where δ is the stepsize of the gradient descent algorithm. Proof In the proof, we follow [1, 21] in finding a modified loss surface, along which the exact path of gradient flow is closer to the discrete steps of gradient descent on the approximated gradient of the SAM objective (Eqn. The general formula in [1, 21] of the modified loss surface for loss function E ( µ ) : = L SAM ( µ ) is
[Equation]
Now we can use the approximation in Proposition 1.
In the following, we concentrate on the general case (general Σ ).
For readability, we use the notation ∇ θ L SAM ( µ ) =
g ( µ ) .
[Equation]
We can approximate the rightmost term as
[Equation]
since the first two of the remaining terms scale with √ p Σ
[Equation]
which can be made smaller than ρ for sufficiently small Σ .
The last remaining term scales with p Σ , which can be made smaller than ρ 2 .
Since we also have the scaling factor δ in (Eqn. (13))
[Equation]
we may neglect these terms.
B. Sharpness avoidance of variational optimization
Proposition
The following approximation holds for the RandomSAM with constant Σ = ρ 2 p I :
[Equation]
[Equation]
where H ( µ ) denotes the Hessian.
For MFVI with general (but symmetric positive definite) Σ we have the following approximation:
[Equation]
Proof Recall the objective of MFVI
[Equation]
where, similarly as in (Eqn. (1)), the term in brackets can be interpreted as a sharpness penalty.
Note that we have fixed Σ
= σ 2 I for a constant σ .
Using a second-order Taylor expansion around µ , where H ( µ ) is the Hessian at µ ,
[Equation]
[Equation]
If Σ = ρ 2 p I , we recover the first part of the theorem. Since the MFVI objective in (Eqn. (17)) is not typically available in closed form, in practice one uses a single-sample Monte Carlo estimate based on the reparametrization trick.
This yields a stochastic objective as follows.
[Equation]
We can follow the method of [ ? ] to see that SGD on the single-sample Monte Carlo estimates of the gradient display additional implicit regularisation towards wider minima. This is carried out in
Proposition 4. Proposition 4 Gradient descent with MFVI step follows a path that is closest to gradient flow path on ˜ L MFVI ( µ, Σ) , where ˜ L MFVI ( µ, Σ) is the following [Equation]
where g ( µ + Σ 12 η ) = ∇ θ L ( µ + Σ 12 η ) . Proof By [ ? ], we have the first approximation, i.e.
[Equation] [Equation] [Equation] Now with Taylor expansion we get
[Equation]
(27) and
[Equation]
Now by using (26), (27) and (28) we obtain ˜ [Equation]
This shows that Mean Field Variational Inference implicitly regularizes the trace of the square of the Hessian, and hence it also regularizes the magnitude of its eigenvalues.
A way of defining sharpness is via the local curvature of the loss function around the minimum given that it is a critical point [4].
Since local curvature is encoded in the Hessian eigenvalues, this means that MFVI penalizes a notion of sharpness at critical points of the loss landscape
C. Algorithms
Algorithm
Input : Training set S
= { ( x i , y i ) } , parameter σ
Initialize: µ for t = 1 , 2 , · · · do Sample batch B ∼ S η N
−
[ ∇ end
Algorithm Input : Training set S
= { ( x i , y i ) } , parameter σ
Initialize: Σ and µ for t = 1 , 2 , · · · do (1) Sample batch B ∼ S (2) Take a sample of η ∼ N (0 , 1) (
Compute
B i η
Compute
B i. σ I )]
Update µ Σ µ ←− µ − η 1 [ ∇ µ L B ] µ +Σ 12 η Σ ←− Σ − η 2 (cid:16) ∇ Σ [ L B ] µ +Σ 12 η + ∇ Σ KL [( µ, Σ) || (0 , σ 0 I )] (cid:17) end
Input
Initialize: Σ and µ for t = 1 , 2 , · · · do
D. Details of Experiments
We use WideResNets [24] on the CIFAR-10/100 datasets [14]. Following prior work [6, 12, 15] we calibrate the SGD optimiser with momentum 0 . 9 , weight decay 0 . 0005 and initial learning rate 0 .
We use stepwise decreasing learning rate scheduling as we have found this more effective than cosine learning rate scheduling.
Using batch size 128 , we train the optimizers requiring two backpropagations per step (SAM, VSAM, MixSAM) for up to 200 epochs, while those with only one backpropagation (SGD and RandomSAM) are trained for up to 400 epochs.
For CIFAR-10, we employ label smoothing [18] with factor 0 .
We follow [6] in setting ρ = 0 . 05 in the SAM optimizer.
Random SAM requires the standard deviation of the noise σ be specified.
After exploration of the hyperparameter space, we have set this to 5 e − 5 .
In VariationalSAM, we used a learning rate of 0 .
01 on Σ .
We set values for ρ and the penalty coefficients in order to match the KL-divergence of the prior and posterior distributions.
The matrix Σ was taken to be diagonal and was initialized in order to render the perturbation size around 0 .
For details see Appendix E. The results are summarized in Table 3.
E. Information theoretic motivation for VariationalSAM
The SAM, ASAM and FSAM algorithms all modify the loss function to penalise the maximum loss value within a small neighbourhood around the current weights.
The considered neighbourhood is an Euclidean ball in SAM, a weight dependent ellipsoid in ASAM, and an ellipsoid defined by the
Fisher information matrix in FSAM.
What we suggest is a generalization of these approaches: our method omits any constraint on the ellipsoid defining the neighbourhood, and treats it as an object to learn.
Specifically, besides µ the algorithm also optimizes a symmetric positive-definite matrix Σ , because for such matrix ϵ ⊤ Σ ϵ ≤ p describes an ellipsoid.
Following the ideas of the previously mentioned algorithms, the modified loss function would be
[Equation]
Our intention is to do coordinate descent on L VSAM ( µ, Σ) w.r.t.
µ and Σ .
However, minimizing the loss in Σ would lead to the null matrix, which is undesirable, therefore we would like to add additional terms to the loss.
The idea is motivated by the Variational Bayesian methods, specifically the maximization of the Evidence Lower Bound (ELBO) and equivalently, the minimization of the negative ELBO.
Consider a prior p ( Z ) , a likelihood p ( X | Z ) , and an arbitrary distribution q θ ( Z ) , then the posterior p ( Z | X ) and the evidence p ( X ) can be computed, and the minimization of the negative ELBO has the form of argmin
[Equation]
When the loss function is defined as the negative log-likelihood, the expectation can be considered as the expectation of the loss.
[Equation]
Choosing p ( Z ) =
N (0 , σ 20 I ) and q θ ( Z ) =
N ( µ, Σ) as k -dimensional Gaussians, the KL-divergence can be rewritten as
[Equation]
By slightly rephrasing E z ∼ q µ L ( z ) as E ϵ ∼N (0 , Σ) L ( µ + ϵ ) we get argmin µ
[Equation]
The expectation above can be bounded with max ϵ ⊤ Σ − 1 ϵ ≤ p L ( µ + ϵ ) in a similar way to the derivations in SAM and FSAM, if ρ is sufficiently large.
This motivates the minimization of
[Equation]
We note that the derivations of SAM, ASAM, FSAM also use the same approach, but they bound the KL term further.
Instead, we found these terms interesting to keep.
Using Appendix A, we can approximate the loss as
[Equation]
We take gradient descent steps w.r.t.
µ and Σ alternately.
F. PAC-Bayes bound for VariationalSAM
Theorem 5 For a parameter space M (with later described properties) and for any ( µ, Σ) ∈ M × Θ we have with probability at least 1 − δ , that
[Equation]
where L D is the generalization loss, and max ϵ T Σ − 1 ϵ ≤
[Equation]
2 L S ( µ + ϵ ) is the empirical SAM loss, with the geometry provided by Σ , i.e. the VSAM loss without the KL divergence term and γ = m (1 + (cid:112) log n/m )
Proof To build our PAC-Bayes bound we follow the steps laid out in the PAC-Bayes bound for
Fisher SAM [12], and we omit some details which can be found there.
Let’s take the parameter spaces µ ∈ M ⊂ R m and Σ ∈ Θ , where Θ is the subset of the diagonal, positive definite matrices.
Also assume, that diam ( M ) ≤ M and 1 /λ ≤ | Σ k,k | ≤ λ for any Σ ∈ Θ and k .
We will take the set of ellipsoids P 1 , 1 , · · · , P t,s , where
[Equation]
where c k ∈ [ − c, c ] ∀ k } (45)
There exists a finite set of these ellipsoids, that cover M × Θ .
Let’s take the following subset of R m
[Equation]
thus t = O ( M m /r m )
= , so log(
[Equation]
=
O ( m ) .
We also assumed 1 λ ≤ Σ k,k ≤ λ .
Let’s think about Θ as a subset of R m .
Then as in part (1) we can define
[Equation]
thus s =
O (( λ + 1 λ ) m / (2 c ) m ) , so log( s ) =
O ( m )
Now we use the PAC-Bayes Theorem of McAllester [17] as follows.
For any prior distribution
P ( µ, Σ) , posterior distribution Q ( µ, Σ) and training set S we have with probability at least 1 − δ
[Equation]
L D ( µ, Σ) and L S ( µ, Σ) are generalization and empirical losses.
We choose Q ( µ, Σ) =
N ( µ 0 , Σ 0 ) , where ( µ 0 , Σ 0 ) ∈
M × Θ
and choose our set of priors as P 1 , 1 , · · · , P t,s , where P i,j = N ( µ i , Σ j ) .
Then as in (46) for each i, j we have with probability 1 − δ i,j , that ∀ Q
[Equation]
µ, Σ)
[Equation]
In the intersection (47) holds for every P i,j .
So by Union bound theorem the intersection is at least 1 − (cid:80) i,j δ i,j , so if we take δ i,j = δst , we have with probability at least 1 − δ
[Equation]
If we choose the prior closes to Q ( µ, Σ) we get the following:
[Equation]
From our assumptions Tr (Σ
[Equation]
0 | | Σ j | ≤ (cid:80) k log (1 + c k ) ≤ (cid:80) k c k ≤ mc and ( µ 0 − µ i ) T Σ − 1 j ( µ 0 − µ i ) ≤ r 2 .
Thus we get
[Equation]
Lets rephrase 48 as we plug in µ + ϵ instead of µ where ϵ N (0 , Σ) .
Thus we get:
[Equation]
Let u = Σ − 1 / 2 ϵ , so u ∼ N (0 , 1) using the result from Laurent & Massart [16] we get that with probability at least 1 − 1 √ n
[Equation]
Let’s partition the space into two parts.
One where (53) holds, where we take the maximum loss and on where (53) does not hold, where we choose the loss bound L max .
So we have
[Equation]
Plugging (54) into (52) we get:
[Equation]
Note, that similar bounds can be build for Random SAM and MFVI.
G. Further Experiments
[Table]
For these experiments, we have set additional hyperparameters to the value reported in [15] and [12].
Namely, in ASAM we use γ = 0 .
5 , η = 0 .
01 for CIFAR-10, γ = 1 . 0 , η = 0 .
1 for CIFARand in FSAM γ = 0 .
1 , η = 1 . 0 for both datasets.