Paper: Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning (ICLR-2024)

Reviewer comment (Experiments):
Although Lewkowycz et al. (2020) also studied catapults only for the squared loss, it is still concerning that this work does the same. Given that the logistic/cross-entropy loss is so often used in deep learning, at least some of the experiments should involve this setting. Since many of the tasks in this work are classification tasks, it is unrealistic to use the squared loss as it is extremely rare to train a classification model this way. This raises the question of whether catapults can be observed or induced when using the logistic or cross-entropy loss, and whether the observations and claims in this paper carry over to this more realistic setting.

Author response:
*Concern 1.a: Given that the logistic/cross-entropy loss is so often used in deep learning, at least some of the experiments should involve this setting. Especially since many of the tasks in this work is on classification, it's unrealistic to use the squared loss as it's extremely rare to train a classification model this way....whether one can observe or induce catapults at all when using the logistic or cross-entropy loss, and whether the observations and claims in this paper can be carried in this more realistic setting.*

To the best of our knowledge, there is no clear definition of catapults with other loss functions.  Indeed, all previous work on catapults has focused on the squared loss (see [1]).  We added evidence in Figure 15 that is suggestive of catapults under cross-entropy loss but a thorough investigation of this phenomenon is beyond the scope of our work.  We additionally note that the squared loss is competitive with cross-entropy across a variety of different classification tasks and architectures (see extensive experimental results in [2, 3]).  

1.  A. Lewkowycz, et al. "The large learning rate phase of deep learning: the catapult mechanism." arXiv preprint arXiv:2003.02218 (2020).

2. L. Hui, M. Belkin "Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks." ICLR 2021.

3. K. Janocha, and C. Wojciech M.  "On loss functions for deep neural networks in classification." arXiv preprint arXiv:1702.05659 (2017).

Reviewer comment (Evaluation):
The authors repeatedly emphasize that they show "catapults mainly occur in the subspace spanned by the top eigenvectors of the tangent kernel," but this point remains unclear even after reading the whole paper. While the experiments show this happens, it is not explained why this should be surprising or significant. The top eigenspace of the NTK corresponds to training set examples that matter most in terms of prediction, likely the "support vectors," so if catapults occur, it seems expected that the projection of the residual also spikes for these examples. More elaboration is needed on the significance of this observation and why the NTK matrix matters in the context of catapults. Furthermore, this part feels disconnected from the latter part of the paper, as it is unclear how "catapults show up in the top eigenspace of the NTK" relates to AGOP alignment. Without further elaboration, it seems like a standalone observation with uncertain significance.

Author response:
*Concern 2.a: Why is the observation "catapults mainly occur in the subspace spanned by the top eigenvectors of the tangent kernel" significant? The top eigenspace of the NTK corresponds to the examples in the training set that matter the most in terms of prediction, so they are likely the "support vectors". So if catapults were to occur, it's not surprising that the projection of the residual also spikes for these examples....more elaboration about this observation is needed.*

We are the first to observe that catapults mainly occur in the subspace spanned by the top eigenvectors of the tangent kernel. In our opinion, it is surprising as the catapult dynamics of complicated neural networks, including Wide ResNets and Vision Transformer, indeed occur in a low dimensional subspace. 

We are not sure we understood your comment about the ''support vectors''. Could you please clarify how you see the connection between top eigenvectors of the NTK and the individual training examples?  

*Concern 2.b: Why should we care about the NTK matrix when it comes to catapults?*

Catapults are related to the loss function, which is controlled by the Hessian. The NTK and the Hessian of the loss are closely related. This can be seen as follows: for MSE  $\mathcal{L}(W;X) = \frac{1}{n}\sum_{i=1}^n (f(W;x_i)-y_i)^2$, we can compute its $H_\mathcal{L}$ by the chain rule:
\begin{align}
 H_\mathcal{L}(w) = \frac{2}{n}(A(w) +B(w)),
\end{align}
where $A(w) = \sum_{i=1}^n\left(\frac{\partial f(w;x_i)}{\partial w}\right)^T \frac{\partial f(w;x_i)}{\partial w}$, and $B(w) = \sum_{i=1}^n  (f(w;x_i)-y_i)\frac{\partial^2 f(w;x_i)}{\partial w^2}$. 

Here $A(w)$ shares the same eigenvalue with NTK with extra $p-n$ $0$ eigenvalues ($w\in\mathbb{R}^p$). $||B(w)||$ is small due to $|f_i-y_i| = O(1)$ and $||\frac{\partial^2 f(w;x_i)}{\partial w^2}|| = \tilde{O}(1/\sqrt{m})$ during training in the kernel regime [1].  Therefore, the spectrum of the tangent kernel is close to the spectrum of $H_\mathcal{L}(w)$ in the kernel regime.  As the critical learning rate is determined by the top eigenvalue of the Hessian, we can estimate the critical learning rate by NTK which greatly improves the time efficiency as $p\gg n$. This approximation of critical learning rates is validated in catapult dynamics as well. Please see Appendix A for experiments.

Furthermore, the catapult was originally characterized as a spike in the training loss and a decrease in the spectral norm (i.e., largest eigenvalue) of the NTK [Lewkowycz et al. 2020].

[1]C. Liu, et al. "On the linearity of large non-linear models: when and why the tangent kernel is constant." Advances in Neural Information Processing Systems 33 (2020): 15954-15964.

Reviewer comment (Experiments):
The paper claims to connect three phenomena in deep learning, the third one being "small batch SGD leads to better generalization." However, the paper by Geiping et al., 2022 ("Stochastic Training is Not Necessary for Generalization") points out that this is not always the case if tuning is done well. In Figures 7 and 18–20, only one step size is given for each task (page 22), and it is assumed to be the same across batch sizes. It is unclear how these step sizes are selected and why they are the same for all batch sizes, raising skepticism about the setup, especially in light of Geiping et al., 2022.

Author response:
*Concern 3.a: "small batch SGD leads to better generalization"  is not always the case given the paper by Geiping et al., 2022*

Thanks for the interesting reference. It is certainly widely believed that small batch sizes generally lead to better generalization[1,2,3,4].  We have no specific opinion on the prevalence of this behavior but rather focus on showing how catapult can lead to such behavior.  

1. I. Kandel, and C. Mauro. "The effect of batch size on the generalizability of the convolutional neural networks on a histopathology dataset." ICT express 6.4 (2020): 312-315.
2. D. Masters, and L. Carlo. "Revisiting small batch training for deep neural networks." arXiv preprint arXiv:1804.07612 (2018).
3. N. Keskar S., et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).
4. S. Smith L., et al. "On the Origin of Implicit Regularization in Stochastic Gradient Descent." International Conference on Learning Representations. 2020.

*Concern 3.b: It seems like in Figure 7, 18-20, only one step size is given for each task (page 22), so I'm assuming they are the same across batch sizes as well. It's unclear how these step sizes are selected and why they are the same for all batch sizes. I'm a bit skeptical about this setup, especially given the results in Geiping et al., 2022.*

We used the same step size across batch sizes for each task in Figure 7. The step size is chosen as $\frac{1}{2}\eta_{crit}$ corresponding to the whole training set. For large datasets e.g., SVHN, we estimate $\eta_{crit}$ using a subset of the whole training set. We have clarified the details in the appendix.

The reason we use the same learning rate for different batch sizes is as follows. Consider the update rule of SGD: $w_{t+1} = w_t - \frac{\eta}{b} \frac{\partial }{\partial w} \sum_{j\in X_{batch}}\ell_j(w_t) $ where  $\ell_j(w_t) = (f(w_t;x_j)-y_j)^2$ and $b = |X_{batch}|$.  If we assume $(x_j,y_j)$ are i.i.d., then $\mathbb{E} \left[\frac{1}{b}\frac{\partial }{\partial w} \sum_{j\in X_{batch}}\ell_j(w_t)\right] = \mathbb{E} \left[\frac{\partial }{\partial w} \ell_j(w_t)\right] $. Note that a larger batch size will make the gradient closer to the expectation, and a smaller batch size can have a higher variance in the gradient. To ensure the magnitude of the norm of weight update is similar, we use the same learning rate for different batches.

Reviewer comment (Evaluation):
Claim 2 is about changes in the (projected) batch loss, but experiments on SGD show the full loss. This raises the question of whether Claim 2 is truly significant if the full loss is what we actually care about.

Author response:
*Concern 4: Claim 2 is about changes in the (projected) batch loss, but the experiments on SGD all show the full loss. Is Claim 2 really significant if what we truly care about is the full loss?*

Claim 2 shows that the mechanism of the catapults in GD, that it occurs when the learning rate is larger than the critical one, also holds in SGD. Since the gradient of the loss is calculated based on the mini-batch in SGD, the mechanism also manifests in the change of the batch loss. Together with the observation on the behavior of the training loss and the NTK, it serves as a strong support that catapults also occur in SGD.

Reviewer comment (Writing):
On page 2, the sentence "This explains why the loss drops quickly to pre-spike levels from the peak of the spike" is problematic. First, it is unclear how catapult behavior in the eigenspace explains this. Second, the sentence is vague and handwavy. The meaning of "pre-spike levels" is unclear—does it refer to the iteration right before the spike or much earlier? It is also unclear what drop rate is considered quick and what it is being compared to.

Author response:
*Concern 5.a: I don't see how catapult in the eigenspace explains the quick loss drop.*

As a simple clarifying example, consider linear regression. The convergence speed of GD is determined by the condition number of the sample matrix $(XX^T)$ since the learning rate is determined by its largest eigenvalue. That is to say, the convergence speed is fast along the top eigendirections but is slow along the remaining directions. The neural networks are close to their linear approximations locally by Taylor's theorem with the tangent kernel as the sample matrix.  Therefore to result in such a quick loss drop, the change of the loss must be in the top eigendirections of the tangent kernel.  In contrast, if the catapults also occurred in the non-top eigendirections, i.e., the loss corresponding to the non-top eigendirections was large at the peak, since the learning rate is determined by the top eigenvalue, the convergence speed would be slow for non-top eigendirections hence the loss would decrease slowly. 

*Concern 5.b: I find this entire sentence  "This explains why the loss drops quickly to pre-spike levels from the peak of the spike" vague and handwavy. It's unclear what the authors mean by pre-spike levels. Is it the loss level at the iteration right before the spike occurs, or 50 iterations before the spike occurs? What drop rate is considered quick? Quick compared to what?*

This is an informal qualitative description. We observed that the loss is much smaller at the iteration right before the spike occurs. Quick means in just a few iterations (typically less than 5). The decrease of the loss at the peak of the spike is much quicker than the decrease from the loss that is not in the spikes.

Reviewer comment (Reproducibility):
More experiment details should be provided clearly rather than scattered through the main text and appendix. For example, the paper frequently mentions that the NTK parameterization is used, but it is unclear exactly what this entails, particularly to unfamiliar readers. The authors should describe explicitly the model architecture, initializations (crucial to NTK parameterization), and why NTK parameterization is used—is it crucial to producing these results? Another example is the step size used for Figure 3: while it may be inferred from a figure, exact numbers are not given, nor is the method for choosing them. The step size is said to start small ($\eta<\eta_{\mathrm{crit}}\approx\eta/\lambda_{\max}(K)$) and increase to $\eta>\eta_{\mathrm{crit}}$, but the mechanism of this change is unclear—what does "≈" mean, and is the increase multiplicative? This information does not seem to be explicitly provided.

Author response:
*Concern 6.a: what exactly is NTK parameterization is unclear...I recommend that the authors explicitly describe what the model architecture is, what the initializations are (crucial to the NTK parameterization), etc...*

In NTK parameterization, each weight parameter is i.i.d. drawn from the standard normal distribution, instead of having width-dependent variance. There will be a scaling factor $1/\sqrt{m}$ to control the scale of each neuron, where $m$ is the width. NTK parameterization is commonly used to analyze over-parameterized neural networks in literature, e.g.,[1,2,3].
 
 We described the NTK parameterization in preliminaries, which follows from the original definition in [Jacot et al. 2018]. All the models considered in our work are standard architecutres hence should be easy to find corresponding information. We have added more information in the appendix.

1. S. Du S., et al. "Gradient Descent Provably Optimizes Over-parameterized Neural Networks." International Conference on Learning Representations. 2018.
2. Z. Ji, and T. Matus. "Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks." International Conference on Learning Representations. 2019.
3. J. Lee, et al. "Wide neural networks of any depth evolve as linear models under gradient descent." Advances in neural information processing systems 32 (2019).

*Concern 6.b: Why is the NTK parameterization used? Is it crucial that you must use this to produce the results?*

The NTK parameterization is not crucial for our results, and similar results can be produced using the Pytorch default parameterization (see Figure 12, 18, and 20 in the appendix).  We follow the parameterization used in [Lewkowycz et al. 2020], under which setting the training dynamics for wide neural networks is approximately linear with a small learning rate and catapults occur with a large learning rate. With NTK parameterization, each entry of the NTK does not scale with the width otherwise it becomes infinity for infinite width.

*Concern 6.c: how did you choose the the step size used for Figure 3?*

The exact numbers of the step size can be found in Appendix E. The learning rate is increased to be roughly $2\eta_{crit}$ (it can vary for different tasks) to generate extra catapults.

Reviewer comment (Novelty):
The related work section is underdeveloped. It currently reads as a list of related works without sufficient discussion of how they relate to each other and to the current work. For example, at the end of the "Catapult phase" paragraph, instead of re-explaining catapult (which was already covered in the introduction), the authors should elaborate on what these works have found, how their setups differed, and how their results relate to the current work. Similar improvements are needed for the "EoS" paragraph.

Author response:
Concern 7: The related work section needs more work, in the sense that it currently presents a laundry list of related works, but how they are related to each other and to the current work is inadequately discussed....at the end of the Catapult phase paragraph, instead of explaining what catapult is, which was already done in the intro...


We have updated the draft to further clarify the contributions of related works. Due to its relevance for our work, we wanted to give a detailed definition of the catapult phenomenon and point out the difference and the relation between the catapult phase phenomenon and the kernel regime.

Reviewer comment (Theory):
The paper never answers *why* the catapults occur. The observations that the catapults happen in the top eigenspace of NTK and that it leads to a decrease in the spectral norm of NTK are not the actual mechanisms behind why catapults occur. Rather, they are direct conclusions of catapults occurring.

Author response:
Indeed we do not analyze the mechanism underlying the catapult phenomenon in this paper but rather focus on their implications.  It is a well-established phenomenon first identified in  [Lewkowycz et al. 2020]. We refer the reviewer to, e.g., works [1,2,3] for theoretical analyses in simplified models.

1. A. Lewkowycz, et al. "The large learning rate phase of deep learning: the catapult mechanism." arXiv preprint arXiv:2003.02218 (2020).

2. L. Zhu, et al. "Quadratic models for understanding neural network dynamics." arXiv preprint arXiv:2205.11787 (2022).

3. D. Meltzer, and L. Junyu. "Catapult Dynamics and Phase Transitions in Quadratic Nets." arXiv preprint arXiv:2301.07737 (2023).

Reviewer comment (Theory):
The authors discuss in detail that the number of catapults correlates strongly with AGOP, which is known to be correlated strongly with test accuracy. But to me, these are still just correlations not explicit inner workings of catapults and their effects on generalization.

Author response:
We believe there is a causal relation between catapults and generalization. This is verified in our experiments, particularly in  Fig. 6, where we induce catapults by increasing the learning rate and each consecutive catapult leads improved generalization.  We empirically show that after each catapult, AGOP alignment improves, and it is well-known that better AGOP alignment leads to improved generalization (see [1, 2, 3, 4, 5]).  

1. W. Härdle, and S. Thomas M. "Investigating smooth multiple regression by the method
of average derivatives." Journal of the American statistical Association 84.408 (1989): 986-995.
2. M. Hristache, et al. "Structure adaptive approach for dimension reduction." Annals of Statistics (2001): 1537-1566.
3. S. Trivedi,  et al. "A Consistent Estimator of the Expected Gradient Outerproduct." UAI. 2014.
4. Y. Xia, et al. "An adaptive estimation of dimension reduction space." Journal of the Royal Statistical Society Series B: Statistical Methodology 64.3 (2002): 363-410.
5. A. Radhakrishnan, et al. "Feature learning in neural networks and kernel machines that recursively learn features." arXiv preprint arXiv:2212.13881 (2022).

Reviewer comment (Experiments):
Any results on (stochastic) gradient descent with momentum?

Author response:
In Section 4, we show that the correlation between the test performance and AGOP alignment holds for SGD+momentum and other optimization algorithms. This suggests that learning AGOP is useful to improve generalization with SGD+momentum.  As to the catapults in SGD with momentum, we believe it is an important direction of future work.

Reviewer comment (Evaluation):
Although the authors did provide sufficient motivation for considering AGOP, I'm still curious about the relations to other "generalization measures" such as flatness (max eigenvalue of Hessian), tail-index of the trajectory/weight matrix [2,3]...etc.

Author response:
Connections to flat minima is an interesting and complicated issue.  There is evidence that flat minima lead to better generalization in, for example, multi-index models [1]. Yet, the same paper also provided examples for which flat minima do not always generalize.  Understanding connections between AGOP and flatness is an interesting area for future exploration. 

1. W. Kaiyue, et al. "Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization." arXiv preprint arXiv:2307.11007 (2023).