SAM operates far from home: eigenvalue regularization as a dynamical phenomenon

Previous experimental work suggested that decreasing batch size causes SAM to display both stronger regularization and better generalization (Andriushchenko & Flammarion, 2022). This analysis also suggested that SAM may induce more sparsity.

A recent theoretical approach studied SAM close to a minimum, where the trajectory oscillates about the minima and provably decreases the largest eigenvalue (Bartlett et al., 2022). A contemporaneous approach studied the SAM algorithm in the limit of small learning rate and SAM radius, and quantified how the implicit and explicit regularization of SAM differs between full batch and batch size $1$ dynamics (Wen et al., 2023).

In contrast to other theoretical approaches, we study the behavior of SAM far from minima. We find that SAM regularizes the eigenvalues throughout training through a dynamical phenomenon and analysis only near convergence cannot capture the full picture. In particular, in simplified models we show:

• •

  Near initialization, full batch SAM provides limited suppression of large eigenvalues (Theorem 2.1).

• •

  SAM induces a modified edge of stability (EOS) (Theorem 2.2).

• •

  For full batch training, the largest eigenvalues stabilize at the SAM-EOS, at a smaller value than pure gradient descent (Section 3).

• •

  As batch size decreases, the effect of SAM is stronger and the dynamics is no longer controlled by the Hessian alone (Theorem 2.3).

We then present experimental results on realistic models which show:

• •

  The SAM-EOS predicts the largest eigenvalue for WideResnet 28-10 on CIFAR10.

Taken together, our results suggest that SAM can operate throughout the learning trajectory, far from minima, and that it can use non-linear, discrete dynamical effects to stabilize large curvatures of the loss function.

We consider a quadratic regression model (Agarwala et al., 2022) which extends a linear regression model to second order in the parameters. Given a $P$-dimensional parameter vector $\bm{\theta}$, the $D$-dimensional output is given by $\mathbf{f}(\bm{\theta})$:

Here, $\mathbf{y}$ is a $D$-dimensional vector, $\mathbf{G}$ is a $D\times P$-dimensional matrix, and $\bm{\mathsfit{Q}}$ is a $D\times P\times P$- dimensional tensor symmetric in the last two indices - that is, $\bm{\mathsfit{Q}}(\cdot,\cdot)$ takes two $P$-dimensional vectors as input, and outputs a $D$-dimensional vector $\bm{\mathsfit{Q}}(\bm{\theta},\bm{\theta})_{\alpha}=\bm{\theta}^{\top}\bm{\mathsfit{Q}}_{\alpha}\bm{\theta}$. If $\bm{\mathsfit{Q}}=\bm{0}$, the model corresponds to linear regression. $\mathbf{y}$, $\mathbf{G}$, and $\bm{\mathsfit{Q}}$ are all fixed at initialization.

Consider optimizing the model with under a squared loss. More concretely, let $\mathbf{y}_{tr}$ be a $D$-dimensional vector of training targets. We focus on the MSE loss

We can write the dynamics in terms of the residuals $\mathbf{z}$ and the Jacobian $\mathbf{J}$ defined by

The loss can be written as $\mathcal{L}(\bm{\theta})=\frac{1}{2}\mathbf{z}\cdot\mathbf{z}$. The full batch gradient descent (GD) dynamics of the parameters are given by

which leads to

The $D\times D$-dimensional matrix $\mathbf{J}\mathbf{J}^{\top}$ is known as the neural tangent kernel (NTK) (Jacot et al., 2018), and controls the dynamics for small $\eta||\mathbf{J}^{\top}\mathbf{z}||$ (Lee et al., 2019).

We now consider the dynamics of un-normalized SAM (Andriushchenko & Flammarion, 2022). That is, given a loss function $\mathcal{L}$ we study the update rule

We are particularly interested in small learning rate and small SAM radius. The dynamics in $\mathbf{z}-\mathbf{J}$ space are given by

to lowest order in $\eta$ and $\rho$.

From Equation 7 we see that for small $\eta||\mathbf{z}||$ and $\rho||\mathbf{z}||$, the dynamics of $\mathbf{z}$ is controlled by the modified NTK $(1+\rho\mathbf{J}\mathbf{J}^{\top})\mathbf{J}\mathbf{J}^{\top}$. The factor $1+\rho\mathbf{J}\mathbf{J}^{\top}$ shows up in the dynamics of $\mathbf{J}$ as well, and we will show that this effective NTK can lead to dynamical stabilization of large eigenvalues. And note that when $\rho=0$, these dynamics coincide with that of gradient descent.

It has been noted that the effects of SAM have a strong dependence on batch size (Andriushchenko & Flammarion, 2022). While a full analysis of SGD is beyond the scope of this work, we can see some evidence of stronger regularization for SGD in the quadratic regression model.

Consider SGD dynamics, where a random fraction $\beta=B/D$ of the training residuals $\mathbf{z}$ are used to generate the dynamics at each step. We can represent the sampling at each step with a random projection matrix $\mathbf{P}_{t}$, and replacing all instances of $\mathbf{z}_{t}$ with $\mathbf{P}_{t}\mathbf{z}_{t}$. Under these dynamics, we can can prove the following:

The calculations are detailed in Appendix A.2. This suggests that there are two possible sources of increased regularization for SGD: the first being the additional terms proportional to $\beta(1-\beta)$. In addition to the fact that $\beta(1-\beta)>\beta^{2}$ for $\beta<\frac{1}{2}$, we have

for left and right eigenvectors $\mathbf{v}_{\alpha}$ and $\mathbf{w}_{\alpha}$ of $\mathbf{J}_{0}$, where $\circ$ is the Hadamard (elementwise) product. This term can be large even if $\mathbf{v}_{\alpha}$ and $\mathbf{z}_{t}$ have small dot product. This is in contrast to $\beta^{2}(\mathbf{v}_{\alpha}\cdot\mathbf{z}_{0})^{2}$, which is small if $\mathbf{z}_{0}$ does not have a large component in the $\mathbf{v}_{\alpha}$ direction. This suggests that at short times, where the large eigenmodes decay quickly, the SGD term can still be large. Additionally, the onto the largest eigenmode itself decreases more slowly in the SGD setting (Paquette et al., 2021), which also suggests stronger early time regularization for small batch size.

We can explore the effects of SAM and show the SAM EOS behavior via numerical experiments on the quadratic regression model. We use the update rule in Equation 6, working directly in $\mathbf{z}$ and $\mathbf{J}$ space as in (Agarwala et al., 2022). Experimental details can be found in Appendix A.3.

For small learning rates, we see that SAM does not reduce the large eigenvalues of $\mathbf{J}\mathbf{J}^{\top}$ in the dynamics (Figure 2, left). In fact in some cases the final eigenvalue is larger with SAM turned on. The projection onto the largest eigenmodes of $\mathbf{J}\mathbf{J}^{\top}$ exponentially decreases to $0$ quicker than any other mode; as suggested by Theorem 2.1, this leads to only a small decreasing pressure from SAM. The primary dynamics of the large eigenvalues is due to the progressive sharpening phenomenology studied in (Agarwala et al., 2022), which tends to increase the eigenmodes.

However, for larger learning rates, SAM has a strong suppressing effect on the largest eigenvalues (Figure 2, middle). The overall dynamics are more non-linear than in the small learning rate case. The eigenvalues stabilize at the modified EOS boundary $\eta(\lambda_{max}+\rho\lambda_{max}^{2})=2$ (Figure 2, right), suggesting non-linear stabilization of the eigenvalues. In Appendix A.3 we conduct additional experiments which confirm that the boundary predicts the largest eigenvalue for a range of $\rho$, and that consequently generally increasing $\rho$ leads to decreased $\lambda_{max}$.

We can see this phenomenology in more general non-linear models as well. We trained a fully-connected network on the first $2$ classes of CIFAR with MSE loss, with both full batch gradient descent and SAM. We then computed the largest eigenvalues of $\mathbf{J}\mathbf{J}^{\top}$ along the trajectory. We can see that in both GD and SAM the largest eigenvalues stabilize, and the stabilization threshold is smaller for SAM (Figure 3). The threshold is once again well predicted by the SAM EOS.

Sharpness For MSE loss, edge of stability dynamics can be shown in terms of either the NTK eigenvalues or the Hessian eigenvalues (Agarwala et al., 2022). For more general loss functions, EOS dynamics takes place with respect to the largest Hessian eigenvalues (Cohen et al., 2022a; Damian et al., 2022). Following these results and the analysis in Equation 12, we chose to measure the largest eigenvalue of the Hessian rather than the NTK. We used a Lanczos method (Ghorbani et al., 2019) to approximately compute $\lambda_{max}$. Any reference to $\lambda_{max}$ in this section refers to eigenvalues computed in this way.

CIFAR-10 We conducted experiments on the popular CIFAR-10 dataset (Krizhevsky et al., 2009) using the WideResnet 28-10 architecture (Zagoruyko & Komodakis, 2016). We report results for both the MSE loss and the cross-entropy loss. In the case of the MSE loss, we replace the softmax non-linearity with Tanh and rescale the one-hot labels ${\bf y}\in\{0,1\}$ to $\{-1,1\}$. In both cases, the loss is averaged across the number of elements in the batch and the number of classes. For each setting, we report results for a single configuration of the learning rate $\eta$ and weight decay $\mu$ found from an initial cross-validation sweep. For MSE, we use $\eta=0.3,\mu=0.005$ and $\eta=0.4,\mu=0.005$ for cross-entropy. We use the cosine learning rate schedule (Loshchilov & Hutter, 2016) and SGD instead of Nesterov momentum (Sutskever et al., 2013) to better match the theoretical setup. Despite the changes to the optimizer and the loss, the test error for the models remains in a reasonable range (4.4% for SAM regularized models with MSE and 5.3% with SGD). In accordance with the theory, we use unnormalized SAM in these experiments. We keep all other hyper-parameters to the default values described in the original WideResnet paper.

As shown in Figure 4 (left), the maximum eigenvalue increases significantly throughout training for all approaches considered. However, the normalized curvature $\eta\lambda_{max}$, which sets the edge of stability in GD, remains relatively stable early on in training when the learning rate is high, but necessarily decreases as the cosine schedule drives the learning rate to $0$ (Figure 4, middle).

SAM radius drives curvature below GD EOS. As we increase the SAM radius, the largest eigenvalue is more controlled (Figure 4, left) - falling below the gradient descent edge of stability (Figure 4, middle). The stabilizing effect of SAM on the large eigenvalues is evident even early on in training.

Eigenvalues stabilize around SAM-EOS. If we instead plot the SAM-normalized eigenvalue $\eta(\lambda_{max}+\rho\lambda_{max}^{2})$, we see that the eigenvalues stay close to (and often slightly above) the critical value of $2$, as predicted by theory (Figure 4, right). This suggests that there are settings where the control that SAM has on the large eigenvalues of the Hessian comes, in part, from a modified EOS stabilization effect.

Altering SAM radius during training can successfully move us between GD-EOS and SAM-EOS. Further evidence from EOS stabilization comes from using a SAM schedule. We trained the model with two settings: early SAM, where SAM is used for the first $2500$ steps ($50$ epochs), after which the training proceeds with SGD ($\rho=0$), and late SAM, where SAM is used for the first $2500$ steps, after which only SGD is used. The maximum eigenvalue is lower for early SAM before $2500$ steps, after which there is a quick crossover and late SAM gives better control (Figure 5). Both SAM schedules give improvement over SGD-only training. Generally, turning SAM on later or for the full trajectory gave better generalization than turning SAM on early, consistent with the earlier work of Andriushchenko & Flammarion (2022).

Plotting the eigenvalues for the early SAM and late SAM schedules, we see that when SAM is turned off, the normalized eigenvalues lie above the gradient descent EOS (Figure 5, right, blue curves). However when SAM is turned on, $\eta\lambda_{max}$ is usually below the edge of stability value of $2$; instead, the SAM-normalized value $\eta(\lambda_{max}+\rho\lambda_{max}^{2})$ lies close to the critical value of $2$ (Figure 5, right, orange curves). This suggests that turning SAM on or off during the intermediate part of training causes the dynamics to quickly reach the appropriate edge of stability.

Networks with cross-entropy loss behave similarly. We found similar results for cross-entropy loss as well, which we detail in Appendix C.1. The mini-batch gradient magnitude and eigenvalues vary more over the learning trajectories; this may be related to effects of logit magnitudes which have been previously shown to affect curvature and general training dynamics (Agarwala et al., 2020; Cohen et al., 2022a).

Minibatch gradient norm varies little. Another quantity of interest is the magnitude of the mini-batch gradients. For SGD, the gradient magnitudes were steady during the first half of training and dropped by a factor of $4$ at late times (Figure 6). Gradient magnitudes were very stable for SAM, particularly for larger $\rho$. This suggests that in practice, there may not be much difference between the normalized and un-normalized SAM algorithms. This is consistent with previous work which showed that the generalization of the two approaches is similar (Andriushchenko & Flammarion, 2022).

Much like the study of EOS before it, our analysis of SAM suggests that sharpness dynamics near minima are insufficient to capture relevant phenomenology. Our analysis of the quadratic regression model suggests that SAM already regularizes the large eigenmodes at early times, and the EOS analysis shows how SAM can have strong effects even in the large-batch setting. Our theory also suggested that SGD has additional mechanisms to control curvature early on in training as compared to full batch gradient descent.

The SAM schedule experiments provided further evidence that multiple phases of the optimization trajectory are important for understanding the relationship between SAM and generalization. If the important effect was the convergence to a particular minimum, then only late SAM would improve generalization. If instead some form of “basin selection” was key, then only early SAM would improve generalization. The fact that both are important (Andriushchenko & Flammarion, 2022) suggests that the entire optimization trajectory matters.

We note that while EOS effects are necessary to understand some aspects of SAM, they are certainly not sufficient. As shown in Appendix A.3, the details of the behavior near the EOS have a complex dependence on $\rho$ (and the model). Later on in learning, especially with a loss like cross entropy, the largest eigenvalues may decrease even without SAM (Cohen et al., 2022a) - potentially leading the dynamics away from the EOS. Small batch size may add other effects, and EOS effects become harder to understand if multiple eigenvalues are at the EOS. Nonetheless, even in more complicated cases the SAM EOS gives a good approximation to the control SAM has on the eigenvalues, particularly at early times.

This work provides additional evidence that understanding some regularization methods may in fact require analysis of the optimization dynamics - especially those at early or intermediate times. This is in contrast to approaches which seek to understand learning by characterizing minima, or analyzing behavior near convergence only. A similar phenomenology has been observed in evolutionary dynamics - the basic $0$th order optimization method - where the details of optimization trajectories are often more important than the statistics of the minima to understand long-term dynamics (Nowak & Krug, 2015; Park & Krug, 2016; Agarwala & Fisher, 2019).