Benign Oscillation of Stochastic Gradient Descent with Large Learning Rate

Q2: Why only train one layer? Essentially, we understand that training with one layer is very different from training multiple layers simultaneously. For example, see [1].

A2: We only train the hidden layer because for our model since this procedure is sufficient to characterize the key difference between the dynamics and the generalization properties between small learning rate training and large learning rate training, which is the main focus of our paper.

Moreover, we point out that the result for one-layer setup in [1] you mentioned does not apply to our problem, since the so called training with one layer setup in [1] considers the neural tangent kernel (NTK) regime with linear models. In contrast, our model focuses on the feature learning perspective with quadratic activation (nonlinear), so that we are able to show the following three types of training dynamics:

• when the learning rate $\eta$ is small, the dynamic converges smoothly to a nearby local minimum;
• when the learning rate $\eta$ is reasonably larger, the oscillation happens during which the model gradually learns all useful patterns from data by learning the weak signal;
• when the learning rate $\eta$ is too large, despite that we have not directly shown in our paper, it is obvious to see the dynamic will diverge.

Then this three phase dynamic is quite similar to the non-linear dynamic shown in [1] for the two layer neural quadratic model case. Intuitively, the reason is that the $\rm ReLU^2$ activation we adopted has already provided the source of nonlinearity we want, which gives a locally varying loss Hessian and enables the benign oscillation under reasonably large learning rate training. With this in hand, we have been able to show the benefits of large learning rate training and the mechanism behind it

It is indeed interesting to consider training a multiple layer neural network. But that complicates the analysis and may not change the key message we would like to convey. So we consider only training the hidden layer in this work. We will add the related work [1] and more discussions on this in the revision.

References:

[1] Zhu, L., Liu, C., Radhakrishnan, A., & Belkin, M. (2022). Quadratic Models for Understanding Neural Network Dynamics. arXiv preprint arXiv:2205.11787.

Q3: The authors refer to the model as a "CNN", which just does not feel right to me. I doubt if the majority of readers would agree that this model is indeed a CNN.

A3: This specific style of model has been widely adopted and named as "CNN" in a sequence of pivotal works, to name a few, see [1, 2, 3, 4, 5, 6, 7]. This specific model is a surrogate of the real-world CNN architecture and is convenient for theoretical study. More specifically, each $\mathbf w\_{j,r}$ in the hidden layer $\sigma(\langle\mathbf{w}\_{j,r},\mathbf{x}\rangle)$ can be regarded as a filter of a CNN, and different patches of the data share the same filter because of the "weight sharing" principle of the CNN architecture.

References:

[1] Du, Simon S., Jason D. Lee, and Yuandong Tian. "When is a Convolutional Filter Easy to Learn?" International Conference on Learning Representations, 2018.

[2] Du, Simon, et al. "Gradient Descent Learns One-hidden-layer CNN: Don’t Be Afraid of Spurious Local Minima." International Conference on Machine Learning. PMLR, 2018.

[3] Cao, Yuan, and Quanquan Gu. "Tight Sample Complexity of Learning One-hidden-layer Convolutional Neural Networks." Advances in Neural Information Processing Systems 32, 2019.

[4] Shen, Ruoqi, Sébastien Bubeck, and Suriya Gunasekar. "Data Augmentation as Feature Manipulation." In International conference on machine learning, pp. 19773-19808. PMLR, 2022.

[5] Allen-Zhu, Zeyuan, and Yuanzhi Li. "Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning." The Eleventh International Conference on Learning Representations, 2023.

[6] Zou, Difan, et al. "Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization." The Eleventh International Conference on Learning Representations, 2023.

[7] Cao, Yuan, et al. "Benign Overfitting in Two-layer Convolutional Neural Networks." Advances in neural information processing systems 35, 2022.

Thanks very much for your review! We now address all your concerns and questions for our work in the following.

Q1: What is a $\rm ReLU^2$ activation function? Is $\rm ReLU$ more essential to the proof or is the quadratic effect more essential to the proof? Does the result hold for $\rm ReLU$? Does the result hold for the quadratic activation? This is unanswered in the paper. To me, this activation neither feels theoretically appealing nor practically relevant.

A1: First, we clarify that the $\rm ReLU^q$-style activation (or ReLU with a polynomial smoothing) is actually widely adopted in existing works on deep learning theory, see e.g., [1, 2, 3, 4, 5, 6, 7] and the references therein. We followed the setup of this line of works for our study on large learning rate NN training. Therefore, we respectfully disagree with your comment that this activation is not theoretically appealing.

Moreover, we need the $\rm ReLU^2$ activation to make the loss Hessian vary locally throughout the training process, which is important to enable the benign oscillation phenonemon we aim to study. This is consistent with the well-known edge-of-stability work [8, 9, 10], where they show that SGD with different learning rates tend to find the solutions with different Hessians. In contrast, with $\rm ReLU$ activation, we get only a locally quadratic objective and the loss sharpness will remain locally unchanged for our model, then the SGD/GD may lack the ability to explore unreached region but will be stuck at a local region (i.e., it will either smoothly converge or oscillate locally).

Finally, our analysis can still work if we use a pure quadratic activation and the key proof techniques will be similar. The major reason of using ReLU activations is to be aligned with the existing line of related research.

References:

[1] Zhong, Kai, et al. "Recovery Guarantees for One-hidden-layer Neural Networks." International Conference on Machine Learning. PMLR, 2017.

[2] Lyu, Kaifeng, and Jian Li. "Gradient Descent Maximizes the Margin of Homogeneous Neural Networks." International Conference on Learning Representations, 2019.

[3] Fan, Fenglei, Jinjun Xiong, and Ge Wang. "Universal Approximation with Quadratic Deep Networks." Neural Networks 124 (2020): 383-392.

[4] Cao, Yuan, et al. "Benign Overfitting in Two-layer Convolutional Neural Networks." Advances in neural information processing systems 35, 2022.

[5] Shen, Ruoqi, Sébastien Bubeck, and Suriya Gunasekar. "Data Augmentation as Feature Manipulation." In International conference on machine learning, pp. 19773-19808. PMLR, 2022.

[6] Allen-Zhu, Zeyuan, and Yuanzhi Li. "Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning." The Eleventh International Conference on Learning Representations, 2023.

[7] Zou, Difan, et al. "Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization." The Eleventh International Conference on Learning Representations, 2023.

[8] Cohen, Jeremy, et al. "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability." International Conference on Learning Representations, 2020.

[9] Arora, Sanjeev, Zhiyuan Li, and Abhishek Panigrahi. "Understanding Gradient Descent on the Edge of Stability in Deep Learning." International Conference on Machine Learning. PMLR, 2022.

[10] Wang, Zixuan, Zhouzi Li, and Jian Li. "Analyzing Sharpness Along GD Trajectory: Progressive Sharpening and Edge of Stability." Advances in Neural Information Processing Systems 35, 2022.

Dear reviewer yxwb, we thank you again for your time reviewing our paper! We have dealt with all your concerns regarding our problem setup including the model architecture and the training protocol, and we have added new references to our draft according to Q2 and Q4 (see Appendix A). As the end of author-review discussion period is approaching, we would appreciate it if you can kindly let us know whether your previous concerns have been fully addressed. We are happy to answer any further questions. As you also mentioned that our work conveys sufficiently novel message and insight for large learning rate NN training, if you think our work is worthy of a higher score, we would be immensely grateful!

• (Large learning rate case) From the update (1) regarding the strong signal $\mathbf{u}$, by using the same argument as in the illustrative deduction in Section 3, we can show that the oscillation of $f$ around $1$ will accumulate the loss derivatives under large learning rate, that is, $\sum_{t}{\ell^{\prime}}^{(t)}$ is increasing (see the first conclusion of Lemma 17 and its proof). Applying this fact to the update formula (2) regarding the weak signal $\mathbf{v}$, we are able to show that $h_2^{(t)}\cdot\sigma(\langle\mathbf{w}^{(t)},\mathbf{v}\rangle)$ can keep increasing, which follows from the same proof of the second conclusion of Lemma 17. This demonstrates the benefits of large learning rate oscillating training in helping the model to learn the weak features.
• (Small learning rate case) In contrast, training two layer $\mathrm{ReLU}$ activation NN with small learning rate cannot learn the weak signal. One can consider the same two-phase analysis as in our paper. In Phase 1, with a small initialization (the initialization scalings for $\mathbf{h}$ and $\mathbf{w}$ need to be similar to enable learning both layers), the NN is trying to learn both signals simultaneously and the growth of both signals are at least exponentially fast in early stages (note that ${{h_1^{(t)}}^2\|\|\mathbf{u}\|\|_2^2+\sigma(\langle\mathbf{w}^{(t)},\mathbf{u}\rangle)^2\ge h_1^{(t)}\sigma(\langle\mathbf{w}^{(t)},\mathbf{u}\rangle)\|\|\mathbf{u}\|\|_2}$ and similar result can hold for $h_2^{(t)}\sigma(\langle\mathbf{w}^{(t)},\mathbf{v}\rangle)$). However, due to the difference in the strength of the two signals, when Phase 1 quickly ends, the NN only learns the strong signal but makes little progress in learning the weak signal. Next, in Phase 2, we can prove that the dynamic smoothly converges to a local minima by exploiting the learned strong signal before it can make sufficient progress in learning the weak signal. In combination of Phase 1 & 2, we can show that training the two-layer NN with $\mathrm{ReLU}$ activation fails to learn the weak signal under small learning rate.

Thus, there is still a division of generalization between large learning rate training and small learning rate training if we consider training both layers with standard $\rm ReLU$ activation. This shows the potential extendibility of our proof of the benign oscillation phenomenon, at least for training two-layer $\rm ReLU$ neural networks. We will add the discussion on the potential extensions to more general settings in the revision.

References:

[1] Zhong, Kai, et al. "Recovery Guarantees for One-hidden-layer Neural Networks." International Conference on Machine Learning. PMLR, 2017.

[2] Lyu, Kaifeng, and Jian Li. "Gradient Descent Maximizes the Margin of Homogeneous Neural Networks." International Conference on Learning Representations, 2019.

[3] Fan, Fenglei, Jinjun Xiong, and Ge Wang. "Universal Approximation with Quadratic Deep Networks." Neural Networks 124 (2020): 383-392.

[4] Cao, Yuan, et al. "Benign Overfitting in Two-layer Convolutional Neural Networks." Advances in neural information processing systems 35, 2022.

[5] Shen, Ruoqi, Sébastien Bubeck, and Suriya Gunasekar. "Data Augmentation as Feature Manipulation." In International conference on machine learning, pp. 19773-19808. PMLR, 2022.

[6] Allen-Zhu, Zeyuan, and Yuanzhi Li. "Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning." The Eleventh International Conference on Learning Representations, 2023.

[7] Zou, Difan, et al. "Understanding the Generalization of Adam in Learning Neural Networks with Proper Regularization." The Eleventh International Conference on Learning Representations, 2023.

Thank you very much for your detailed review and your support of our work! We address all your concerns and questions in the following.

Q1: While the authors show that the oscillation is consistent with their hypothesis, it could be great that the authors comment a bit more of this necessity to assume this. How difficult would it be to remove this hypothesis? Is the problem to control that the iterates do not diverge? Or on the contrary that they do not converge locally? If this is the second option, could the authors add some (bounded) label noise at each iteration to show this?

A1: Thanks for your questions! We will explain the necessity of making this assumption in the following. In short, it is intuitively possible to derive a sufficient condition that enables the stable oscillation. However, proving a rigorous and clean sufficient condition for the multiple-data setting would be highly complicated. We thus resort to summarizing this specific mode of oscillation as an assumption in order to better present the key idea and finding of our work.

Now we answer your questions in detail.

We first explain why we need this assumption. As we illustrate in the single training data setup (Section 3), the oscillation (in terms of the model parameter) tend not to cancel with each other, but to have a strict positive value, which turns out to be the driving force of the weak signal learning. Intuitively, the amount of the accumulation of oscillation will determine the amount of weak signal learning:

Therefore, it boils down to mathematically characterize how fast the oscillation accumulation $(\star)$ would grow, which is not simply to only control the iterates from divergence or to prove no local convergence. What we want is a control over the magnitude of the oscillation (i.e., $\delta$) so that the summation $(\star)$ can be lower bounded.

Intuitively, for the single training data setup, one can achieve this by properly scaling the initialization and the learning rate. i.e., controlling the learning rate by using the loss smoothness parameters at the initialization and the nearby local minima. Then, as the magnitude of the oscillation can be expressed as a joint function of the loss smoothness, the learning rate, and the initialization, we can then control the magnitude via properly scaling the learning rate and the initialization so as to give a sufficient condition for certain oscillation magnitude.

However, in the general multiple data setup, the training dynamic when oscillation happens could be very difficult to precisely characterize, and the specific modes of oscillation could be quite sensitive to the choice of the learning rate and initialization (different LR and initialization can result in different modes of oscillation and may require different analysis). To this end, we took the cleanest and least troublesome way by assuming that the magnitude of the oscillation is uniformly bounded away by a some $\delta>0$, which then leads to the linear accumulation we show in the paper (e.g., Eq. (7) in Section 3). This saves us from diving into the explicit mathematical characterization of the oscillation dynamic and better explain the key message regrading the benefits of oscillation we would like to convey.

Dear reviewer ahqj, we thank you again for your time reviewing our paper and your support of our work! We have revised our draft according to Q2 (see Appendix A) and Q3 (see Appendix C). As the end of author-review discussion period is approaching, we appreciate it if you can kindly let us know whether your previous concerns have been fully addressed. We are happy to answer any further questions. Thanks!

Thank you very much for your appreciation of our work! We address all your concerns and questions in the following.

Q1: How to justify the feature-noise data general model? E.g., is this a natural assumption for real-world dataset like CIFAR-10?

A1: The original version of this multi-patch data model can be found in the pivotal work in feature learning [1] and our version is simplified to only three patches with strong/weak noises and signals. We put the answer to this problem in two folds:

1. How original data model in [1] can approximate real world image-classification data? In the multi-view feature model therein, the authors assume that the images from given class should exhibits some specific signals with high probability. For example, for an image that represents a car, it is highly possible that "wheel", "window", or "headlight" appears some where in the image. Moreover, "wheel" and "wheel" usually share high similarity between images, wheras "wheel" and "window" (or other irrelvant noise/feature) are typically uncorrelated due to the high-dimensionality. This allows us to encode the features with fixed and mutually orthogonal signal vectors and set the noise orthogonal to all the signals. The reviewers can refer to Section 2.3 in [1] for a brief introduction and Appendix A for technical specifications to the data model. We also note that the authors in [1] mentioned that this data model can be considered as the intermediate ouput of previous well-trained convolution layer.
2. How does our model reasonably simplify the multi-view feature data model? In short, we proposed several simplifications towards the data model to make our presentation easier to follow without losing the generality.
   • Simplification 1: We only consider the binary classification with three-patches image data. One can imagine that all the information that is irrelavant to the binary classificaiton of our interests is covered in the noise patch.
   • Simplification 2: Fixed position signal and orthogonal Gaussian noise. In the original multi-view data model [1], the signals appear randomly among all possible patches, which aligns with the real images. But CNN is unable to identify this positional information because of the "weight sharing" principle, thus our simplication is reasonable.
   • Simplification 3: No feature noise. The original model assume that the important features are corrupted (additively) by other irrelavant features or noises. However, in our setting, this is not essential in presenting our key mesaage. Suchg simplification can reduce some unnecessary technical issues.

References:

[1] Allen-Zhu, Z., & Li, Y. (2022, September). Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning. In The Eleventh International Conference on Learning Representations.

Dear reviewer 1fN5, we thank you again for your time reviewing our paper and your support of our work! As the end of author-review discussion period is approaching, we appreciate it if you can kindly let us know whether your previous concerns have been addressed. We are happy to answer any further questions. Thanks!