Perturbing BatchNorm and Only BatchNorm Benefits Sharpness-Aware Minimization
Abstract
We investigate the connection between two popular methods commonly used in training deep neural networks: Sharpness-Aware Minimization (SAM) and Batch
Normalization.
We find that perturbing only the affine BatchNorm parameters in the adversarial step of SAM benefits the generalization performance, while excluding them can decrease the performance strongly.
We confirm our results across several models and SAM-variants on CIFAR-10 and CIFAR-100 and show preliminary results for ImageNet.
Our results provide a practical tweak for training deep networks, but also cast doubt on the commonly accepted explanation of SAM minimizing a sharpness quantity responsible for generalization
Introduction
The generalization of deep neural networks has been linked to the flatness of the loss-surface already by [8].
In the past years, several studies attempted to better understand this relation [10, 5, 11, 16] and to leverage it for improved learning algorithms [19].
Most notably, [6] proposed SharpnessAware Minimization (SAM), an optimization technique designed towards directly minimizing a sharpness-based generalization bound alongside the training loss and achieved state-of-the-art in a variety of vision-benchmarks.
SAM minimizes an adversarially perturbed loss, where the perturbation is computed with respect to a small l 2 -ball around the current point in the optimization trajectory.
Follow-up works tried to either reduce the computational cost of the method [4], or improve its performance.
[20] minimize a different objective called surrogate gap , whereas others aim at defining a more accurate perturbation model, which is e.g. adaptive to the scale of the parameters and hence invariant to reparametrizations of the network (ASAM [14]) or respects the parameter space geometry induced by the Fisher information (Fisher-SAM [12]).
While the empirical success of SAM-like methods is commonly attributed to them finding favorable flat minima, there remain doubts on whether this explanation can provide a conclusive picture. [1] argue that the generalization bound, which provides the main theoretical justification for the
SAM algorithm, is based on average-case sharpness, but SAM with random instead of worst-case perturbations does not substantially improve over SGD.
Further, the empirical success of m-sharpness, a gradient averaging heuristic explained in Section 2, and the observation that a more accurate optimization of the perturbation model decreases SAMs generalization performance (discussed in [6]) additionally weakens this reasoning.
[1] hypothesize that some other quantity, possibly correlated with (m-)sharpness, could be responsible for generalization and the main success of SAM might instead be due to its benefitial implicit bias. Has it Trained Yet?
Workshop at the Conference on Neural Information Processing Systems (NeurIPS 2022).
We therefore take a step back from the sharpness-explanation and instead investigate the interplay between SAM and Batch Normalization (BN) [9], another technique that has shown to be crucial for training well-generalizing networks, but remains poorly understood.
[17] showed that the initial explanation of BN-layers mitigating a covariate shift in the activations is insufficient, and instead linked it to a smoothing of the loss-surface, whereas [15] understood it as an implicit regularization scheme.
[7] showed that the affine parameters of the layer are by themselves already surprisingly expressive: Training only γ and β of deep ResNets and freezing all other parameters at their random initialization yielded non-trivial accuracy on CIFAR-10 and ImageNet, which could not be achieved when training an equivalent number of other parameters.
In this work, we show that ASAM, as it was proposed in [14], relies crucially on the BatchNorm parameters and that computing the adversarial perturbation without them degrades its performance strongly.
We further demonstrate for a range of models and SAM-variants, that perturbing only the
BatchNorm parameters in the adversarial step typically boosts their performance, leading to a simple adjustment that can easily be incorporated in existing SAM-like methods.
Background: Sharpness-Aware Minimization
We recapitulate SAM, ASAM and Fisher-SAM with their respective perturbation models.
To this end, we consider a neural network f w : R d −→ R k which is parameterized by a vector w as our model.
The train dataset consists of pairs S train = { ( x 1 , y 1 ) , ...
( x n , y n ) } which are drawn from the data distribution D and we write the loss function as l : R k ×
R k −→ R
+ .
The goal is to learn a model f w with good generalization performance, i.e. low expected loss L D ( w ) =
E ( x , y ) ∼ D [ l ( y , f w ( x ))] on the distribution D .
The training loss can be written as L ( w ) = 1 n (cid:80) ni =1 l ( y i , f w ( x i )) .
Conventional
SGD-like optimization methods minimize (a regularized version of) L by stochastic gradient descent.
SAM aims at additionally minimizing the worst-case sharpness of the training loss in a neighborhood defined by an l p ball around w , i.e. max ||
ϵ
|| p <
ρ L ( w + ϵ [Equation] − L ( w ) .
This leads to the overall objective
[Equation]
In practice, SAM uses p = 2 and approximates the inner maximization by a single gradient step, yielding ϵ = ρ ∇ L ( w ) / ||∇ L ( w ) || 2 and requiring an additional forward-backward pass compared to
SGD.
The gradient is then re-evaluated at the perturbed point w + ϵ , giving the actual weight update:
[Equation]
Computing ϵ separately for each GPU in multi-GPU settings and then averaging the resulting perturbed gradients for the update step in (2) has been shown to increase SAMs performance.
This method is called m-sharpness, with m being the number of samples on each GPU.
Since the perturbation model in (1) is not invariant to rescaling of f w [3], ASAM [14], a partly scale-invariant version of SAM, was proposed, with the objective
[Equation]
where T w is a normalization operator, making the perturbation adaptive to the scale of the network parameters.
[14] choose T w to be diagonal with entries T iw = | w i | + η for weight parameters and T iw = 1 for bias parameters, and η is typically set to 0 .
Equivalently to SAM, the inner maximization is solved by a single gradient step:
[Equation]
We note that for T w being the identity matrix, this is equivalent to the SAM formulation.
Recently, [12] proposed to use a distance metric induced by the Fisher information instead of a Euclidean distance measure between parameters.
The approach can also be framed as a variant of ASAM, with T w being diagonal with entries T iw = 1 / √
+ ηf i and f i approximating the i th diagonal entry of the Fisher-matrix by the squared average batch-gradient, f i = ( ∂ w i L Batch ( w )) 2 .
For our experiments, we additionally employ layerwise normalization.
This is, we set the diagonal entries of
T w i =
|| W layer [ i ] || 2 , which corresponds to a normalization with respect to the l 2 -norm of a layer.
Method
In this paper, we focus on the interplay between Sharpness-Aware Minimization variants and BatchNorm.
BatchNorm layers transform an input x with batch mean µ B and batch variance σ 2 B according to
[Equation]
During training, µ B and σ B are computed from the current batch-statistics, and running estimates are used at test time.
In our experiments, we focus on the trainable parameters γ and β , which perform a channel-wise affine transformation.
In our first experiment, we exclude them from the adversarial
SAM-step ( no-bn ).
Taking inspiration from [7], in our main experiment we perturb only γ and β and neglect all other parameters ( only-bn ).
only-bn corresponds to setting all entries T w i = 0 , if w i is not a BatchNorm parameter, and vice versa for no-bn .
Apart from this change in T w , which leads to a change of the perturbation ϵ according to (4), we use the conventional SAM-algorithm.
Experiments
[Figure]
We first showcase that excluding γ and β from the ϵ -perturbation ( no-bn ) can lead to drastic performance drops in some variants compared to using all parameters ( all ), but has little effect on others.
We train a range of ResNet-like models with several SAM-variants on CIFAR-10 and CIFAR- [13] and monitor the performance when computing the perturbation ϵ only along the directions of γ and β ( only-bn ).
We confirm our findings with experiments on ImageNet [2].
For our CIFAR experiments, we consider a range of SAM-variants which differ either in the norm ( p ∈ { 2 , ∞} ) or in the definition of the normalization operator.
We use SGD, the original SAM with no normalization and p = 2 , Fisher-SAM and the following ASAM-variants: elementwise- l ∞ , layerwise- l 2 , and elementwise- l 2 .
For each of the ASAM-variants, we normalize both bias and weight parameters and set η = 0 .
Additionally, we employ the original ASAM-algorithm, where the bias parameters are not normalized and η = 0 .
We train all models on a single GPU for 200 epochs, and m-sharpness is not employed.
We use both basic augmentations (random cropping and flipping) and strong augmentations (basic+AutoAugment).
Like [14], adopt a learning rate of 0.1, momentum of 0.9, weight decay of 0 . 0005 and use label smoothing with a factor of 0.1.
We showcase the effect of excluding γ and β from the adversarial step for a WideResNet-28-10 on
CIFAR-100 in Figure 1.
For the elementwise- l 2 variants we observe a strong drop in accuracy, while for SAM, Fisher-SAM, layerwise- l 2 and elementwise- l ∞ there is either no or very little change.
If, in contrast, we use only the batch-norm parameters for the ϵ computation, we observe that almost all variants obtain higher accuracy.
Except for Fisher-SAM, the difference is especially pronounced for the variants which did not experience a performance drop for no-bn .
For those, the ideal ρ shifts towards larger values, indicating that the perturbation model cannot perturb the BatchNorm parameters enough when all parameters are used.
In order to confirm this observation, we train a range of models on CIFAR-10 and CIFAR-100.
For each SAM-variant and dataset, we probe a set of pre-defined ρ -values (shown in Table 4) with a ResNet-56 and fix the best-performing ρ for all other models to compare only-bn to all .
We report mean accuracy and standard deviation over 3 seeds for
CIFAR-100 in Table 1.
On average, only-bn outperforms all for all considered SAM-variants and layerwise- l 2 works particularly well.
On CIFAR-10 (Table 3 in the appendix) the results are similar, but slightly less pronounced due to the very high accuracy of all methods.
For ImageNet, we adopt the timm training script ([18]).
We train all models for 100 epochs on 8 2080Ti GPUs with m = 64 , leading to an overall batch-size of 512.
Apart from ρ , all hyperparameters are shared for all models and can be found in the appendix in Table 5.
Due to the computational cost of training on large-scale datasets like ImageNet, we can only show preliminary results.
In particular, we could not run all models with all methods, but instead selected the most promising only-bn variants and compared them against the established methods (SGD, SAM, ASAM elementwise l 2 ).
For SAM and ASAM, ρ has been tuned by [6] and [14] and we adopt those values.
For the methods involving elementwise l ∞ and layerwise l 2 with their respective all and only-bn variant, we probe two ρ values each and report the result of the better one.
The results are shown in Table 2, where we observe that the only-bn models outperform their all counterparts for elementwise l 2 and elementwise l ∞ .
For layerwise l 2 , the all variant achieves higher accuracy, which might be due to the ρ -value of the corresponding only-bn variant not being properly tuned.
Nevertheless, all only-bn variants outperform the previously established methods (SGD, SAM, ASAM).
For reference, we also show the values reported for ESAM [4] and GSAM [20], two other SAM-variants we did not include in our study.
[Table] [Table]
Conclusions and Future Work
In this paper we showed that excluding the BN-parameters in ASAM degrades its performance, whereas performing the adversarial SAM-step only for γ and β brings improvements for a variety of SAM-variants.
This provides a practical training tweak, but also casts doubt on the commonly accepted explanation of SAMs empirical success.
Since the BatchNorm parameters typically account only for a tiny fraction of the models parameters (for a WideResNet-28-10, 0 .
05% of all parameters belong to BatchNorm layers), it is unclear if the sharpness quantity used by SAM-variants could still be optimized well with only-bn .
In addition, the good performance of layerwise normalization, which is not invariant to parameter rescaling, questions the relevance of designing reparametrizationinvariant sharpness measures.
While we do not have a conclusive answer on why only-bn works, we provide an experiment in Appendix A.3 giving additional insights into the method’s impact on BN parameters.
A Appendix
The appendix is structured as follows:
In A.1 we show the results of the CIFAR-10 experiment and the ρ values that were used.
In A.2 we show the hyperparameters used for the ImageNet runs.
In A. we show an additional experiment, aimed at providing insight into the effects of only-bn on the scale of the BatchNorm parameters.
A.1 CIFAR-Experiments
Table 3 show the results for fixed ρ on CIFAR-10.
The values of ρ we considered for each method can be found in Table 4.
The networks we considered for the CIFAR-experiments are DenseNet (D100), ResNet56 (RN56), ResNeXt-29-32x4d (RnxT), and WideResNet-28-10 (WRN)
[Table]
[Table]
[Table] [Table]
l
[Table]
l
[Table]
layerwise
[Table]
marked bold.
A.2 ImageNet Experiments
Table 5 shows the hyperparameters for all variants used for ImageNet training.
For SGD, SAM and elementwise- l 2 we used the hyperparameters from [6] and [14].
For the elementwise l 2 and elementwise- l ∞ we tried 2 ρ -values per configuration and report the results of the better one (named ρ (reported) in the table).
ρ (discarded) refers to the ρ value we probed, but found to perform worse than the other one.
[Table]
A.3 Additional Experiment
In order to get a better understanding of the impact of only-bn on γ and β
[Table]
we train a WideResNet28-10 with different SAM-variants and both only-bn and all .
We show the distribution of | w i |
[Table]
i [Table] e. the parameter magnitudes, at the end of training for different layer types in figure 2.
The y -axis is shown on log-scale, since most convolutional parameters are almost zero and would make the effects on the BatchNorm parameters invisible.
For elementwise l 2 there is no visible change in the distribution of the BatchNorm parameters between all and only-bn .
For elementwise l ∞ , layerwise l 2 and SAM, however, the magnitude of the BatchNorm parameters shifts clearly towards larger values, especially for the weight parameters.
We note that this resembles a pattern we observed when comparing the optimal ρ -value for all and only-bn in figure 1: The optimal ρ of elementwise l 2 did not change much, whereas for the other considered methods, it shifted towards larger values for only-bn .
Additionally and in contrast to the other methods, the elementwise l 2 variant showed a strong performance decrease in no-bn , indicating that it implicitly focuses on perturbing the BatchNorm layers already.
We observe a similar, yet weaker effect when comparing basic to strong augmentations: The BatchNorm parameters shift slightly towards larger values (shown for SGD in figure 2).
We note that larger BatchNorm parameters do not necessarily indicate a functionally different network, since there are many reparametrization invariances in ReLU networks, some of which ASAM tries to leverage in its perturbation definition (3).
Nevertheless, the scale of the network still has an impact on the training dynamics, since other methods like e.g. weight decay depend on it.
While this does not provide a conclusive explanation for the success or the underlying mechanism of only-bn -SAM,
[Figure]
SGD [Figure] AA
FC
Conv
BN weight
[Figure]
We truncate the x -axis at x = 0 . 3 for better visualization. we think it should be taken as a starting point to investigate the impact of SAM-like methods on other parts of the training of neural networks.