SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {\boldsymbol {a}} \ell (\boldsymbol {\theta}) = \ell^ {\prime} (\boldsymbol {\theta}) \cdot \sigma (\boldsymbol {W} \boldsymbol {x}), \quad \nabla_ {\boldsymbol {w} _ {j}} \ell (\boldsymbol {\theta}) = \ell^ {\prime} (\boldsymbol {\theta}) \cdot a _ {j} \sigma^ {\prime} (\langle \boldsymbol {w} _ {j}, \boldsymbol {x} \rangle) \boldsymbol {x}.
$$
Then a direct computation gives us the following expression for the full gradient norm:
$$
\| \nabla \ell (\boldsymbol {\theta}) \| _ {2} = | r | \cdot \| \nabla f (\boldsymbol {\theta}) \| _ {2} = | \langle \boldsymbol {a}, \sigma (\boldsymbol {W} \boldsymbol {x}) \rangle - y | \sqrt {\| \sigma (\boldsymbol {W} \boldsymbol {x}) \| _ {2} ^ {2} + \| \boldsymbol {x} \| _ {2} ^ {2} \cdot \| \boldsymbol {a} \odot \sigma^ {\prime} (\boldsymbol {W} \boldsymbol {x}) \| _ {2} ^ {2}},
$$
where  $\odot$  denotes element-wise multiplication and  $r$  denotes the residual  $f(\theta) - y$ . Then the update of  $\boldsymbol{w}_j$  for neuron  $j$  on each step of SAM with step size  $\eta$  can be written as:
$$
\boldsymbol {w} _ {j} := \boldsymbol {w} _ {j} - \eta (\nabla \ell (\boldsymbol {\theta}) + \rho \nabla \| \nabla \ell (\boldsymbol {\theta}) \| _ {2}) + \mathcal {O} \left(\rho^ {2}\right)
$$
$$
\boldsymbol {w} _ {j} := \boldsymbol {w} _ {j} - \underbrace {\eta r \left(1 + \rho \| \nabla f (\boldsymbol {\theta}) \| _ {2} / \sqrt {\ell (\boldsymbol {\theta})}\right) a _ {j} \sigma^ {\prime} (\langle \boldsymbol {w} _ {j} , \boldsymbol {x} \rangle) \boldsymbol {x}} _ {\text {d a t a f i t t i n g t e r m}} - \underbrace {\eta \rho \sqrt {\ell (\boldsymbol {\theta})} / \| \nabla f (\boldsymbol {\theta}) \| _ {2} \sigma (\langle \boldsymbol {w} _ {j} , \boldsymbol {x} \rangle) \boldsymbol {x}} _ {\text {r e g u l a r i z a t i o n c o m p o n e n t}} + \mathcal {O} \left(\rho^ {2}\right)
$$
where we used the fact that  $\sigma'(\langle \pmb{w}_j, \pmb{x} \rangle) \sigma(\langle \pmb{w}_j, \pmb{x} \rangle) = \sigma(\langle \pmb{w}_j, \pmb{x} \rangle)$  and second-order terms are zero almost everywhere for ReLUs. The data fitting term is the same as normal gradient but with a larger effective learning rate  $\eta (1 + \rho \| \nabla f(\pmb{\theta}) \|_2 / \sqrt{\ell(\pmb{\theta})})$  instead of  $\eta$ . The most interesting term is the one coming from  $\nabla \| \nabla f(\pmb{\theta}) \|_2$  which drives pre-activations  $\langle \pmb{w}_j, \pmb{x} \rangle$  to negative values on every step of SAM which can be seen directly from the update rule:
$$
\begin{array}{l} \langle \boldsymbol {w} _ {j}, \boldsymbol {x} \rangle := \langle \boldsymbol {w} _ {j}, \boldsymbol {x} \rangle - \eta r \left(1 + \rho^ {\| \nabla f (\boldsymbol {\theta}) \| _ {2}} / \sqrt {\ell (\boldsymbol {\theta})}\right) a _ {j} \sigma^ {\prime} (\langle \boldsymbol {w} _ {j}, \boldsymbol {x} \rangle) \| \boldsymbol {x} \| _ {2} ^ {2} \\ - \eta \rho \sqrt {\ell (\boldsymbol {\theta})} / \| \nabla f (\boldsymbol {\theta}) \| _ {2} \sigma \left(\langle \boldsymbol {w} _ {j}, \boldsymbol {x} \rangle\right) \| \boldsymbol {x} \| _ {2} ^ {2} + \mathcal {O} \left(\rho^ {2}\right), \\ \end{array}
$$
where we note that  $\eta \rho \sqrt{\ell(\pmb{\theta})} / \| \nabla f(\pmb{\theta}) \|_2 \sigma (\langle \pmb{w}_j, \pmb{x} \rangle) \| \pmb{x} \|_2^2$  is always non-negative for ReLU.
We confirm empirically in Figure 15 in Appendix that using only the first-order terms in the SAM update leads to the same effect in all key metrics including the feature rank and the number of active ReLUs. Overall, this mechanism suggests how SAM can suppress redundant activations if they are not needed to fit the training data. Moreover, this effect takes place at every iteration of SGD but is stronger at the beginning of training since it is proportional to the loss  $\sqrt{\ell(\pmb{\theta})}$ . Moreover, a similar argument can be made for a multi-layer network which can explain why the low-rank effect occurs at multiple layers since the  $\| \nabla f(\pmb {\theta})\| _2$  term has activations of all layers in it. Also it suggests an intuitive picture of what a flat minimum of SAM means in terms of the learned function: a minimum that corresponds to a network sparsely activated on the training data.
5 Investigation of low-rank mechanisms on deep networks
In this section, we confirm that the low-rank mechanism described above can also occur in post-activation ResNets, although the overall rank reduction mechanism can be more complex, particularly for networks with pre-activation skip connections and self-attention layers.
Post-activation ResNet on CIFAR-10. We consider a standard, post-activation ResNet-18 (He et al., 2016a) with 4 residual superblocks trained on CIFAR-10 in the minimal setting (i.e., small step sizes, no momentum, no weight decay). Since values in the residual stream of this network (unlike for PreAct ResNets) are taken after ReLU, we can expect to see the ReLU pruning effect described in the previous section. We count a ReLU as active if at least on one training example, it achieves at least  $5\%$  of the maximum value in the feature matrix. We plot the results in Figure 6 which confirms the ReLU pruning effect: the number of active ReLUs decreases with  $\rho$  at later residual blocks. For example, for  $\rho = 0.6$ , the number of active ReLUs at block 4 reaches  $77\%$ , which directly contributes to a lower feature rank, compared to  $100\%$  of standard training. Finally, we note that this mechanism is not likely to be the only one contributing to the lower rank since the trends for the feature rank and number of active ReLUs do not perfectly agree.
Pre-activation ViT on MS-COCO. Next, we consider the ViT with 12 pre-activation residual blocks used in our experiments on MS-COCO. First, we show how the feature ranks change in the residual stream after attention and MLP subblocks in Figure 7 (left, middle). We observe that almost each attention block gradually increases the feature rank, but the increase is smaller for models with higher  $\rho$ . At the same time, MLP blocks can both increase and decrease the rank, but for higher  $\rho$ , the difference between post-MLP rank and pre-MLP rank tends to be larger. Thus, we conclude that the
Figure 6: Post-activation ResNet-18 on CIFAR-10. Both the feature rank and the number of active ReLUs tend to decrease with the  $\rho$  of SAM, particularly at later ResNet blocks.
Figure 7: ViT image encoder on MS-COCO. We show the difference in the feature ranks before and after attention subblocks (left), before and after MLP subblocks (middle), and the number of active GeLUs inside of residual branches (right).
rank reduction effect is coming primarily from the attention blocks. Moreover, we do not observe any coordinate-aligned sparsity in the residual stream. Additionally, we plot the number of active GeLUs inside of MLP subblocks in Figure 7 (right). Since GeLU can be seen as a differentiable approximation of ReLU, the analogous notion of active GeLU units also makes sense. However, we do not observe any decrease in the number of GeLUs over the  $\rho$  of SAM for these models. We show more detailed plots about the behavior of feature ranks in Figure 16 in Appendix. Moreover, in Figure 17 we show that even when the text encoder remains frozen and only a single linear layer is trained on top of it, the low-rank effect is still observed. This indicates that SAM can effectively utilize even a single matrix multiplication to achieve features of reduced rank. Overall, we conclude that the mechanism behind the low-rank effect of SAM can vary significantly, and SAM can leverage different components of the network, including activations, attention layers, and even individual linear layers.
6 Do low-rank features lead to better generalization on natural data?
We conclude the discussion on the low-rank features of SAM by demonstrating that directly inducing low-rank features does not result in improved generalization for natural data such as MS-COCO.
Setup. We induce the low-rank directly at the last layer by using a linear bottleneck layer which is also known as the Burer-Monteiro factorization (Burer and Monteiro, 2003): we parametrize the weight matrix  $W^{(L)} \in \mathbb{R}^{d_{L-1} \times d_L}$  at the last layer  $L$  as  $W^{(L)} = U^{(L)} V^{(L)}$  where  $U^{(L)} \in \mathbb{R}^{d_{L-1} \times h}$  and  $V^{(L)} \in \mathbb{R}^{h \times d_L}$ . We perform training on MS-COCO by varying the inner dimension  $h$  for the last layers of both image and text encoders, while keeping the remaining training process the same as described in Section 3.
Observations. We plot the results in Figure 8 where we observe that enforcing low-rank features alone does not consistently enhance generalization. Introducing bottleneck layers only marginally improves the imagewise retrieval error (from  $23.3\%$  to  $22.7\%$ ) but negatively affects textwise retrieval (from  $22.2\%$  to  $22.9\%$ ). In contrast, SAM demonstrates a significant improvement of  $-4.7\%$  and  $-3.0\%$  in the respective retrieval tasks. Thus, we conclude that SAM stands out in its ability to simultaneously enhance generalization and select low-rank features in a data-adaptive manner.
Figure 8: Contrastive image-text training on MS-COCO. We compare the effect of SAM and linear bottleneck layers on the retrieval error and feature ranks of the last layer. We observe that linear bottleneck layers do not improve the retrieval error at all, unlike SAM.
7 Discussion and future directions
Finally, we discuss the implications of low-rank features learned by SAM as well as future directions.
More efficient retrieval. As illustrated in Section 3, despite the lower rank of the features learned by SAM, their effectiveness for nearest neighbor retrieval is preserved or even improved. This suggests that one can achieve faster retrieval by reducing the dimensionality of the feature space by using, for example, only the top- $k$  components of PCA. Furthermore, the low-rank bias of SAM appears to be data-adaptive, meaning that it adjusts to the data structure rather than relying on a fixed predetermined dimension.
More efficient feature quantization. On a related note, having features of reduced rank can also be helpful to speed up feature quantization. This is confirmed, in particular, in the DALL-E 2 paper (Ramesh et al., 2022) which reports that this effect combined with PCA can lead to a  $3 \times$  reduction in the number of tokens predicted during inference, and also can improve the training stability.
Understanding the features learned by SAM. Developing a better understanding of the features learned by popular methods such as SAM is an important goal on its own. The low-rank feature bias can be advantageous in scenarios where the data is generated by a low-rank function. However, for realistic datasets like MS-COCO, the impact of this effect on generalization is not as significant, yet it still remains a valuable side effect.
Future directions. Our findings suggest that low-rank features alone do not enhance generalization for natural data beyond simple teacher-student setups. Consequently, the low-rank effect of SAM appears to be a useful side effect. Therefore, understanding the impact of SAM on learned features that leads to generalization improvements on natural data remains an open question. Additionally, further theoretical analysis of the low-rank effect of SAM for more complex architectures, such as those involving skip-connections and self-attention layers, would be of great interest.
We thank the anonymous reviewers for their suggestions that helped to improve the paper. The work is supported by the Google/EPFL Research Collab program. M.A. was additionally supported by the Google Fellowship and Open Phil AI Fellowship.
M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. In ICML, 2022.
M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. SGD with large step sizes learns sparse features. In ICML, 2023.
D. G. Barrett and B. Dherin. Implicit gradient regularization. In ICLR, 2021.
P. L. Bartlett, P. M. Long, and O. Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513, 2022.
T.-Y. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitnick. Microsoft COCO: Common objects in context. In ECCV, 2014.
S. Liu, D. Papailiopoulos, and D. Achlioptas. Bad global minima exist and SGD can reach them. In NeurIPS, 2019.
R. Mulayoff, T. Michaeli, and D. Soudry. The implicit bias of minima stability: A view from function space. In NeurIPS, 2021.
C. Na, S. V. Mehta, and E. Strubell. Train flat, then compress: Sharpness-aware minimization learns more compressible models. In Findings of EMNLP, 2022.
M. S. Nacson, R. Mulayoff, G. Ongie, T. Michaeli, and D. Soudry. The implicit bias of minima stability in multivariate shallow ReLU networks. In ICLR, 2023.
A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748, 2018.
V. Papyan, X. Han, and D. L. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. PNAS, 2020.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen. Hierarchical text-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06125, 2022.
S. L. Smith, B. Dherin, D. G. Barrett, and S. De. On the origin of implicit regularization in stochastic gradient descent. In ICLR, 2021.
A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and L. Beyer. How to train your ViT? data, augmentation, and regularization in vision transformers. TMLR, 2021.
N. Timor, G. Vardi, and O. Shamir. Implicit regularization towards rank minimization in relu networks. In ALT, 2023.
I. Tolstikhin, N. Houlsby, A. Kolesnikov, L. Beyer, X. Zhai, T. Unterthiner, J. Yung, A. P. Steiner, D. Keysers, J. Uszkoreit, M. Lucic, and A. Dosovitskiy. MLP-mixer: An all-MLP architecture for vision. In NeurIPS, 2021.
K. Wen, T. Ma, and Z. Li. How sharpness-aware minimization minimizes sharpness? In ICLR, 2023.
D. Wu, S.-T. Xia, and Y. Wang. Adversarial weight perturbation helps robust generalization. In NeurIPS, 2020.
C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. In ICLR, 2017.
X. Zhang, R. Xu, H. Yu, H. Zou, and P. Cui. Gradient norm aware minimization seeks first-order flatness and improves generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20247-20257, 2023.
Y. Zheng, R. Zhang, and Y. Mao. Regularizing neural networks via adversarial model perturbation. In CVPR, 2021.
- Section A: implementation details and hyperparameters of all experiments.
- Section B: additional results for classification experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet.
- Section C: additional results for contrastive image-text training on MS-COCO.
- Section D: additional results for two-layer networks.
- Section E: additional results for investigations of the low-rank mechanism on deep networks.
In this section, we specify the implementation details and hyperparameters of all our experiments.
Experiments on CIFAR-10, CIFAR-100, Tiny ImageNet. We train pre-activation ResNets-18 (He et al., 2016b) on CIFAR-10, CIFAR-100, Tiny ImageNet for experiments in Section 3 and post-activation ResNets-18 for experiments in Section 5. We train these models with batch size 256 for 200 epochs using standard augmentations (random crops and random mirroring). For the minimal setting, we use plain SGD with the learning rate 0.05. For the state-of-the-art setting, we use SGD with the learning rate 0.1 (decayed by a factor of  $10 \times$  after  $50\%$  and  $90\%$  epochs), momentum parameter 0.9, weight decay 0.0005. We train models with a grid of different parameters  $\rho$  of SAM where we adjust the grid such that it includes both values of  $\rho$  that improve generalization (small  $\rho$ , usually in the range [0.05, 0.2]) but also those which slightly degrade it (large  $\rho$ , usually larger than 0.5). We evaluate the feature rank on 10240 training examples by counting the minimal number of PCA components needed to span  $99\%$  variance.
Experiments on ImageNet-1k. We refer to the original papers Dosovitskiy et al. (2021), Tolstikhin et al. (2021), and Chen et al. (2022) for the details on how these models were trained. We evaluate the feature rank on 12800 training examples by counting the minimal number of PCA components needed to span  $99\%$  variance.
Experiments on MS-COCO. We fine-tune the  $\mathrm{R + Ti / 16}$  vision transformer from Steiner et al. (2021) and BERT from Devlin et al. (2018) on MS-COCO using the InfoNCE contrastive loss. We add a linear head to each of the encoders with the final feature dimension of 768. We use Adam (Kingma and Ba, 2014) with learning rate 0.0001 which is decayed down to 0 using a cosine decay schedule. We train these models with batch size 128 for 25 epochs without data augmentations. We evaluate the feature rank on 1280 training examples by counting the minimal number of PCA components needed to span  $99\%$  variance. To aggregate the features over different tokens, we always take the embedding of the first token since this is what is used for classification of the pretrained model.
Experiments on two-layer networks in the teacher-student setup. We use the teacher-student setup with 3 teacher neurons,  $m = 100$  student neurons, and input dimension  $d = 3$ . We consider training with stochastic gradients based on mini batch size of 1 using the learning rate of 0.1. We initialize the weights of the teacher network with the standard deviation of 1.0 and the student network with the standard deviation of 0.3. We use SAM for the first  $50\%$  of iterations and then disable it to achieve full convergence as we observed that running SAM, especially with high  $\rho$ , hinders convergence. We evaluate the feature rank on all 20 training examples by counting the minimal number of PCA components needed to span  $99.99\%$  variance. We exceptionally choose this PCA threshold (instead of  $99\%$  as in all other experiments) to obtain a smooth trend of rank reduction and clear separation between different  $\rho$ .
Computational resources. We performed all experiments on a single Nvidia A100 GPU where we used an internal cluster for all experiments except experiments on MS-COCO, for which we used a cloud provider. A typical training run on CIFAR-10 and CIFAR-100 takes around 2 hours, on Tiny ImageNet around 5 hours, and on MS-COCO around 7 hours.
B Additional results for classification
In Section 3, we focused on the original SAM as introduced in Foret et al. (2021) since it is still the most popular SAM variant in the community and it is implemented without any further approximations. However, to provide a comprehensive comparison of SAM variants, we also include results on more recent variants of SAM such as ASAM (Adaptive Sharpness-Aware Minimization) (Kwon et al., 2021) and GAM (Gradient norm Aware Minimization) (Zhang et al., 2023). Similarly to our main experiments in Section 3, we also use ResNets on CIFAR-10. We select the default settings given in their code repositories (which includes a smaller ResNet for ASAM compared to GAM) and vary only the perturbation radius  $\rho$ . We report the results in Table 2 that confirm that the low-rank observation also extends to other recent SAM variants.
ρ of ASAM	0.0	0.5	1.0	2.0	4.0
Test error	7.29%	6.53%	6.38%	7.12%	10.64%
Feature rank	5048	4801	4699	4578	4383
ρ of GAM	0.0	0.2	0.4	0.8	1.6
Test error	4.04%	3.65%	3.64%	3.81%	4.81%
Feature rank	7633	7381	7303	6897	6927
Table 2: Results for more recent variants of SAM: ASAM (Kwon et al., 2021) and GAM (Zhang et al., 2023).
In Figure 9, we plot feature ranks computed on the training and test sets of CIFAR-10 and CIFAR-100 for post-activation and pre-activation ResNets-18, respectively, taken at the last epoch. We select two different architectures to check the generality of our findings. We observe that the feature ranks stay very close, no matter if we evaluate them on the training or test sets. Thus, we conclude that the low-rank effect of SAM generalizes also to unseen samples.
Post-activation ResNet-18 on CIFAR-10
Pre-activation ResNet-18 on CIFAR-100
Figure 9: ResNets-18 on CIFAR-10 and CIFAR-100. Feature ranks computed on the training and test sets are very similar to each other. We conclude that the low-rank effect of SAM generalizes also to unseen samples.
In Figure 10, we plot feature ranks for  $95\%$ ,  $99\%$ ,  $99.9\%$  PCA variance. Feature ranks computed with different PCA variance thresholds exhibit very similar trends. We observe that the overall trend and relative ranking between models trained with different  $\rho$  is unaffected by the choice of the PCA variance threshold.
In Figure 11, we plot metrics like in Figure 1 but for larger batch sizes (512 and 1024 instead of 256). We see that the low-rank trend is clearly visible there as well.
Finally, in Figure 12, we show feature ranks computed at the last block of ResNet (just before the linear head). We see that the neural collapse phenomenon occurs: the feature rank converges to 10 (i.e., the number of classes) even for  $\rho = 0$ . The feature rank with SAM is lower at the beginning but then increases to a slightly higher value than 10 towards the end of training.
CIFAR-10, minimal setting
CIFAR-10, state-of-the-art setting
CIFAR-100, minimal setting
CIFAR-100, state-of-the-art setting
Tiny ImageNet, minimal setting
Tiny ImageNet, state-of-the-art setting
Figure 10: ResNet18 on CIFAR-10, CIFAR-100, Tiny ImageNet. Feature ranks computed with different PCA variance thresholds exhibit very similar trends.
95% PCA variance
99% PCA variance
(Reported in the main part)
99.9% PCA variance
Batch size 512, minimal setting (small learning rate, no momentum, no weight decay)
Batch size 512, state-of-the-art setting (large learning rate, momentum, weight decay)
Batch size 1024, minimal setting (small learning rate, no momentum, no weight decay)
Figure 11: ResNet-18 on CIFAR-10 with different batch sizes. Similarly to the experiments in Figure 1 that were done with batch size 256, SAM with larger batch sizes (512 and 1024) also improves test error (left), leads to more generalizable features (middle), and noticeably reduces the feature rank at the intermediate ResNet block (right).
Batch size 1024, state-of-the-art setting (large learning rate, momentum, weight decay)
Figure 12: ResNets-18 on CIFAR-10. Feature ranks computed at the last block of ResNet (just before the linear head). We see that the neural collapse phenomenon occurs: the feature rank converges to 10 (i.e., the number of classes) even for  $\rho = 0$ . The feature rank with SAM is lower at the beginning but then increases to a slightly higher value than 10 towards the end.
C Additional results for contrastive image-text training
In Figure 13, we plot the retrieval error and feature rank for models with a frozen (i.e., not updated during training) BERT text encoder trained with different  $\rho$  of SAM. We observe that the low-rank effect in the final text features (as well as image features) occurs even when a single layer is attached to a fixed BERT text encoder. Moreover, as a side note, SAM also improves the retrieval error in this setting as well.
Figure 13: Contrastive image-text training on MS-COCO (with a frozen text encoder). The low-rank effect observed in the final text features even when a single layer is attached to a fixed BERT text encoder.
D Additional results for two-layer networks
In Figure 14, we plot the same metrics as in Figure 5 but, in addition, for tanh and absolute value activations. We can see that these models, when trained with higher  $\rho$  of SAM, have a smaller feature rank, smaller number of active tanh/abs units, and higher weight norm. This illustrates the fact that the same low-rank mechanism can be at play for many activations, not just for ReLU.