SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
W _ {1} = \left[ \sqrt {r _ {i} | y _ {i} |} w _ {i} / \epsilon \right] _ {i} ^ {\top}, b _ {1} = \left[ \sqrt {r _ {i} | y _ {i} |} \left(- w _ {i} ^ {\top} x _ {i} + \epsilon\right) / \epsilon \right] ^ {\top}, W _ {2} = \left[ \operatorname {s i g n} (y _ {i}) \sqrt {| y _ {i} | / r _ {i}} \right] _ {i}.
$$
Here  $r_i = (\| x_i\|^2 + 1)^{1/2}$  and  $\epsilon$  is a sufficiently small positive number. Then it holds that  $\mathrm{relu}(W_1x_i + b_1) = \sqrt{r_i|y_i|} e_i$ , where  $e_i$  is the  $i$ -th coordinate vector. This shows there is a one-to-one correspondence between the input data and the neurons. It is easy to verify that the model interpolates the training dataset. Furthermore, for  $\mathcal{P}_{\mathrm{xor}}$  and sufficiently small  $\epsilon$ , for any input  $x \notin \{x_i\}_{i \in [n]}$ , it holds that  $\mathrm{relu}(W_1x + b_1) = 0$ . Hence the model will output the same label 0 for all the data points outside the training set. This indicates Proposition 4.1.
To show the memorization solution has minimal sharpness, we need the following lemma that relates the sharpness and the Jacobian of the model.
Lemma 4.1. For mean squared error loss  $l_{mse}$ , if model  $f_{\theta}$  is differentiable and interpolates dataset  $\{(x_i, y_i)\}_{i \in [n]}$ , then  $\mathrm{Tr}(\nabla^2 L(\theta)) = \frac{2}{n} \sum_{i=1}^{n} \| \nabla_{\theta} f_{\theta}(x_i) \|^2$ .
Figure 3: Visualization of Memorization Solutions. This is an illustration of the memorizing solutions constructed in Theorem 4.1. Here the input data points come from a unit circle and are marked as dots. The shady area with the corresponding color represents the region where the corresponding neuron is activated. One can see that the network can output the correct label for each input data point in the training set as long as the weight vector on the corresponding neuron is properly chosen. Further, the network will make the same prediction 0 for all the input data points outside the shady area and this volume can be made almost as large as the support of the training set by choosing  $\epsilon$  sufficiently small. Hence the model can interpolate the training set while generalizing poorly.
Proof of Lemma 4.1. By standard calculus, it holds that,
$$
\begin{array}{l} \mathrm {T r} \left(\nabla^ {2} L (\theta)\right) = \frac {1}{n} \sum_ {i = 1} ^ {n} \mathrm {T r} \left(\nabla_ {\theta} ^ {2} \left[ (f _ {\theta} (x _ {i}) - y _ {i}) ^ {2} \right]\right) \\ = \frac {2}{n} \sum_ {i = 1} ^ {n} \operatorname {T r} \left(\nabla_ {\theta} ^ {2} f _ {\theta} (x _ {i}) \left(f _ {\theta} (x _ {i}) - y _ {i}\right) + \left(\nabla_ {\theta} f _ {\theta} (x _ {i})\right) \left(\nabla_ {\theta} f _ {\theta} (x _ {i})\right) ^ {\top}\right) \\ = \frac {2}{n} \sum_ {i = 1} ^ {n} \operatorname {T r} \left(\left(\nabla_ {\theta} f _ {\theta} \left(x _ {i}\right)\right) \left(\nabla_ {\theta} f _ {\theta} \left(x _ {i}\right)\right) ^ {\top}\right) = \frac {2}{n} \sum_ {i = 1} ^ {n} \| \nabla_ {\theta} f _ {\theta} \left(x _ {i}\right) \| _ {2} ^ {2}. \tag {1} \\ \end{array}
$$
The first equation in Equation (1) use  $\forall i, f_{\theta}(x_i) = y_i$ . The proof is then complete.
After establishing Lemma 4.1, one can then explicitly calculate the lower bound of  $\| \nabla_{\theta}f_{\theta}(x_i)\|^2$  condition on  $f_{\theta}(x_i) = y_i$ . For simplicity of writing, we will view the bias term as a part of the weight matrix by appending a 1 to the input data point. Precisely, we will use notation  $x_{i}^{\prime}\in \mathbb{R}^{d + 1}$  to denote transformed input satisfying  $\forall j\in [d],x_i'[j] = x_i[j],x_i'[d + 1] = 1$  and  $W_{1}^{\prime} = [W_{1},b_{1}]\in \mathbb{R}^{m\times (d + 1)}$  to denote the transformed weight matrix.
By the chain rule, we have,
$$
\begin{array}{l} \left\| \nabla_ {\theta} f _ {\theta} (x _ {i}) \right\| ^ {2} = \left\| \nabla_ {W _ {1} ^ {\prime}} f _ {\theta} (x _ {i}) \right\| _ {F} ^ {2} + \left\| \nabla_ {W _ {2}} f _ {\theta} (x _ {i}) \right\| _ {F} ^ {2} \\ = \| \left(W _ {2} \odot \mathbf {1} \left[ W _ {1} ^ {\prime} x _ {i} ^ {\prime} > 0 \right]\right) x _ {i} ^ {\prime \top} \| _ {F} ^ {2} + \| \operatorname {r e l u} \left(W _ {1} ^ {\prime} x _ {i} ^ {\prime}\right) \| _ {2} ^ {2}. \\ = \| W _ {2} \odot \mathbf {1} \left[ W _ {1} ^ {\prime} x _ {i} ^ {\prime} > 0 \right] \| _ {2} ^ {2} \| x _ {i} ^ {\prime} \| ^ {2} + \| \operatorname {r e l u} \left(W _ {1} ^ {\prime} x _ {i} ^ {\prime}\right) \| _ {2} ^ {2}. \tag {2} \\ \end{array}
$$
Then by Cauchy-Schwarz inequality, we have
$$
\begin{array}{l} \left\| \nabla_ {\theta} f _ {\theta} \left(x _ {i}\right) \right\| ^ {2} = \left\| W _ {2} \odot \mathbf {1} \left[ W _ {1} ^ {\prime} x _ {i} ^ {\prime} > 0 \right] \right\| _ {2} ^ {2} \left\| x _ {i} ^ {\prime} \right\| ^ {2} + \left\| \operatorname {r e l u} \left(W _ {1} ^ {\prime} x _ {i} ^ {\prime}\right) \right\| _ {2} ^ {2} \\ \geq 2 \| x _ {i} ^ {\prime} \| \left| \left(W _ {2} \odot \mathbf {1} \left[ W _ {1} x _ {i} > 0 \right]\right) ^ {\top} \operatorname {r e l u} \left(W _ {1} ^ {\prime} x _ {i} ^ {\prime}\right) \right| = 2 \| x _ {i} ^ {\prime} \| | y _ {i} |. \tag {3} \\ \end{array}
$$
In Equation (3), we use condition  $f_{\theta}(x_i) = y_i$ . Finally, notice that the lower bound is reached when
$$
W _ {2} \odot \mathbf {1} \left[ W _ {1} ^ {\prime} x _ {i} ^ {\prime} > 0 \right] = \operatorname {r e l u} \left(W _ {1} ^ {\prime} x _ {i} ^ {\prime}\right) / \left\| x _ {i} ^ {\prime} \right\|. \tag {4}
$$
Condition Equation (4) is clearly reached for the memorization construction we constructed, where both sides of the equation are equal to  $\sqrt{|y_i| / \|x_i'\|e_i}$ . This completes the proof of Theorem 4.1.
However, the memorization network is not the only parameter that can reach the lower bound. For example, for distribution  $\mathcal{P}_{\mathrm{xor}}$ , if parameter  $\theta$  satisfies,
$$
\forall i, j \in \{0, 1 \}, W _ {1, 2 i + j + 1} = r [ (- 1) ^ {i}, (- 1) ^ {j}, \dots , 0 ], b _ {1} [ 2 i + j + 1 ] = - r, W _ {2} [ 2 i + j ] = (- 1) ^ {i + j} / r. \tag {5}
$$
$$
\forall k > 4, W _ {1, k} = [ 0, \dots , 0 ], b _ {1} [ k ] = 0, W _ {2} [ k ] = 0,
$$
(b) 1-SAM
Figure 4: Scenario II. We train a 2-layer MLP with ReLU activation with Bias using gradient descent with weight decay and 1-SAM on  $\mathcal{P}_{\mathrm{xor}}$  with dimension  $d = 30$  and training set size  $n = 100$ . One can clearly observe a distinction between the two settings. The minimum reached by 1-SAM is flatter but the model fails to generalize and the generalization performance even starts to degenerate after 4000 epochs. The difference between Figures 1b and 4b indicates a small change in the architecture can lead to a large change in the generalization performance.
(a) Baseline
(b) 1-SAM
Figure 5: Scenario II with Softplus Activation. We train a 2-layer MLP with Softplus activation  $(\text{SoftPlus}(x) = \log(1 + e^x))$  with bias using gradient descent with weight decay and 1-SAM on  $\mathcal{P}_{\mathrm{xor}}$  with dimension  $d = 30$  and training set size  $n = 100$ . We observe a similar phenomenon as Figure 4.
with  $r = (d^2 + 1)^{1/4}$ . then for any  $x \in \{-1, 1\}^d$ , it holds that  $\mathrm{relu}(W_1 x + b_1) = r e_{5/2 - x[1] - x[2]/2}$  and  $f_\theta(x) = x[1] \times x[2]$ . Hence it is possible for Equation (5) to hold while the model has perfect generalization performance.
4.2 SAM empirically finds the non-generalizing solutions
In this section, we will show that in multiple settings, SAM can find solutions that have low sharpness but fail to generalize compared to the baseline full batch gradient descent method with weight decay. This proves that flat minimization can hurt generalization performance. However, one should note that Question 2 is not denied for the current architectures.
Converged models found by SAM fail to generalize. We perform experiments on data distribution  $\mathcal{P}_{\mathrm{xor}}$  in Figure 4. We apply small learning rate gradient descent with weight decay as our baseline and observe that the converged model found by SAM has a much lower sharpness than the baseline. However, the generalization performance of SAM is much worse than the baseline. Moreover, the generalization performance even starts to degenerate after 4000 epochs. We conclude that in this scenario, sharpness minimization can empirically hurt generalization performance.
1-SAM may fail to generalize with other activation functions. A natural question is whether the phenomenon that 1-SAM fails to generalize is limited to ReLU activation. In Figure 5, we show empirically that 1-SAM fails to generalize for 2-layer networks with softplus activation trained on the same dataset, although there is no known guarantee for the existence of memorizing solutions.
(a) Standard Training
(b) Projected Training
Figure 6: Scenario III. We train two-layer ReLU networks with simplified LayerNorm on data distribution  $\mathcal{P}_{\mathrm{xor}}$  with dimension  $d = 30$  and sample complexity  $n = 100$  using 1-SAM. In Figure 6a, we use standard training. In Figure 6b, we restricted the norm of the weight and the bias of the first layer as 10, to avoid minimizing the sharpness by simply increasing the norm. We can see that in both cases, the models almost perfectly generalize.
5 Scenario III: Both Flattest Generalizing and Non-generalizing Models Exist, and SAM Finds the Latter
5.1 Both generalizing and non-generalizing solutions can be flattest
Despite the surprising contrary between Theorems 3.1 and 4.1, experiments show that Question 2 consistently hold. However, we will provide a counterexample in this section. Specifically, we will consider data distribution  $\mathcal{P}_{\mathrm{xor}}$  and 2-layer ReLU MLP with simplified LayerNorm. One can first show both generalizing and non-generalizing solutions exist similar to Theorem 4.1 and propositions 4.1 and 4.2.
Theorem 5.1. If the input data points  $\{x_{i}\}$  of the training set form a set of extreme points (Definition 4.1), for sufficiently small  $\epsilon$ , then there exists a width- $n$  2-MLP-Sim-LN with hyperparameter  $\epsilon$  that is a memorizing solution (Definition 4.2) for the training dataset and has minimal sharpness over all the interpolating solutions.
Proposition 5.1. For data distribution  $\mathcal{P}_{\mathrm{xor}}$ , for sufficiently small  $\epsilon$ , for any number of samples  $n$ , there exists a width- $n$  2-MLP-Sim-LN with hyperparameter  $\epsilon$  that memorizes the training set as in Theorem 4.1, reaches minimal sharpness over all the interpolating models and has generalization error  $\max \{1 - n / 2^d, 0\}$  measured by zero one error.
Proposition 5.2. For data distribution  $\mathcal{P}_{\mathrm{xor}}$ , for sufficiently small  $\epsilon$ , for any number of samples  $n$ , there exists a width- $n$  2-MLP-Sim-LN with hyperparameter  $\epsilon$  that interpolates the training dataset, reaches minimal sharpness over all the interpolating models, and has zero generalization error measured by zero one error.
The construction and intuition behind Theorem 5.1 and propositions 5.1 and 5.2 are similar to that of Theorem 4.1 and propositions 4.1 and 4.2. The proof is deferred to Appendix B.
5.2 SAM empirically finds generalizing models
Notice in Section 5.1 our theory makes the same prediction as in Section 4. However, strikingly, the experimental observation is reversed (Figure 6). Now running SAM can greatly improve the generalization performance till the model perfectly generalizes. This directly denies Question 2 as now we have a scenario in which sharpness minimization algorithms can improve generalization till perfect generalization while there exists a flattest minimizer that will generalize poorly.
We present theoretical and empirical evidence for (1) whether sharpness minimization implies generalization subtly depends on the choice of architectures and data distributions, and (2) sharpness minimization algorithms including SAM may still improve generalization even when there exist flattest models that generalize poorly. Our results suggest that low sharpness may not be the only cause of the generalization benefit of sharpness minimization algorithms.
ACKNOWLEDGEMENTS
The authors would like to thank the support from NSF IIS 2045685.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-aware minimization leads to low-rank features. arXiv preprint arXiv:2305.16292, 2023a.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flammarion. A modern look at the relationship between sharpness and generalization. arXiv preprint arXiv:2302.07011, 2023b.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09745, 2022.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513, 2022.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an Ornstein-uhlenbeck like process. arXiv preprint arXiv:1904.09080, 2019.
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint arXiv:1804.10200, 2018.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers, 2021.
Lijun Ding, Dmitriy Drusvyatskiy, and Maryam Fazel. Flat minima generalize for low-rank matrix recovery. arXiv preprint arXiv:2203.03756, 2022.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pp. 1019-1028. JMLR.org, 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv:1703.11008, 2017.
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21: 136, 2020.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
Khashayar Gatmiry, Zhiyuan Li, Ching-Yao Chuang, Sashank Reddi, Tengyu Ma, and Stefanie Jegelka. The inductive bias of flatness regularization for deep matrix factorization. arXiv preprint arXiv:2306.13239, 2023.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In Advances in Neural Information Processing Systems, pp. 6151-6159, 2017.
Jeff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.08680, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Jean Honorio and Tommi Jaakkola. Tight bounds for the expected risk of linear classifiers and pac-bayes finite-sample guarantees. In Artificial Intelligence and Statistics, pp. 384-392. PMLR, 2014.
Colin Wei and Tengyu Ma. Data-dependent sample complexity of deep neural networks via lipschitz augmentation. In Advances in Neural Information Processing Systems, pp. 9722-9733, 2019a.
Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification via an all-layer margin. arXiv preprint arXiv:1910.04284, 2019b.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems, pp. 9709-9721, 2019.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating Turing machines with transformers. Advances in Neural Information Processing Systems, 35:12071-12083, 2022.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729, 2022.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. arXiv preprint arXiv:2002.09277, 2020.
Lei Wu and Weijie J Su. The implicit regularization of dynamical stability in stochastic gradient descent. arXiv preprint arXiv:2305.17490, 2023.
Lei Wu, Chao Ma, et al. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530, 2016.