Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning

Concern 2.c: This part of the paper also seems to be somewhat disconnected from the latter part - it's unclear how "catapults show up in the top eigenspace of the NTK" is related to the AGOP alignment.

The observation that "catapults show up in the top eigenspace of the NTK" is not directly related to the AGOP alignment in the paper. The contributions of our work include two areas: optimization, where we propose an explanation for the spikes in the training loss using SGD; and generalization, where we explain how catapults improve generalization. The observation that catapults show up in the top eigenspace of the NTK serves as evidence that the loss spikes in SGD are catapults (contribution in optimization). AGOP alignment is the feature learning that explains how catapults improve generalization (contribution in generalization). Both contributions are about the impact of catapults in neural networks.

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.

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.

We thank the reviewer for the detailed and insightful comments. We will address your concerns and questions below.

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).

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:

$

H_\mathcal{L}(w) = \frac{2}{n}(A(w) +B(w)),

$

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.

Q1: Why does the critical step size only depend on the initialization and not varying throughout training?

The critical step size can definitely vary through training, which is exactly what we leveraged to induce multiple catapults in GD (see Figure 3). We can observe multiple catapults with the same step size in SGD due to fluctuations in the critical learning rate due to the stochasticity of the batch.

Q2: Paragraph under Claim 1: shouldn't it be $\mathcal{PL}_n$ instead of $\mathcal{PL}_5$ that corresponds to the spike in the training loss? Is even $\mathcal{PL}_5$ plotted somewhere that I missed?

In Figure 5, we plotted the loss corresponding to the complement of the top-5 eigendirections of the tangent kernel, i.e., $\mathcal{PL}_5^\bot$ which decreases almost monotonically. Therefore, $\mathcal{PL}_5$, which is $\mathcal{L}-\mathcal{PL}_5^\bot$ corresponds to the spikes in the training loss.

Q3.a: Choice of top eigenspace dimension $s$: The first sentence after this - I'm confused by what this sentence is trying to convey, especially the second part. What linear dynamics?

For a clarifying example, consider a linear model $h(w;x) = w^Tx$ where $w$ is the parameter. For $n$ data points $(X,y)$, by minimizing the squared loss $L(w) = \frac{1}{n}||{h(w;X) - y}||^2$ its training dynamics of GD with learning rate $\eta$ satisfies the equation:

$

h(w_{t+1};X) - y = (I - \eta \frac{2}{n}XX^T) (h(w_t;X) - y).

$

Then we project the residual $h(X)-y$ to the eigenspace of $XX^T$, which we denote by $u_i\in\mathbb{R}^n$ and we get:

$

u_i^T( h(w_{t+1};X) - y) = (1 - \eta\frac{2\lambda_i}{n})u_i^T( h(w_{t};X) - y),

$

where $\lambda_i$ is the i-th eigenvalue of $XX^T$.

Note that if $\eta > n/\lambda_i$, the loss will increase along the eigen-direction $u_i$.

As wide neural networks are close to their linear approximations before the spike, they follow approximately linear dynamics. The sample matrix in this case is the tangent kernel. Then a larger learning rate will lead to the increase of the loss in more eigen-directions.

Q3.b: Do you have an explanation to why small $s$ contradicts to what you expect?

We find this observation interesting since the dramatic change of the loss turns out to occur in a low-dimensional subspace. Moreover, this phenomenon is quite consistent among our experiments covering a broad class of neural networks.

Q3.c: Is there a way to verify $s$ corresponds to support vectors?

Could you please clarify what you mean by the connection between $s$ and ''support vectors''? We are not sure we understood your comment about the ''support vectors''.

Q4: $f(X_{batch})$ has the same length as the number of examples in the batch. But the projection operator acts on prediction vectors of length $n$.

We further clarify this point. The projection operator is defined based on the tangent kernel corresponding to the mini-batch, hence it acts on the prediction vectors of length $b$. We have made it clear in the revision. Specifically, we edited the beginning of section 3.1: we define the network output $f$ first then define the corresponding NTK and the training loss. In this way, the projection operator depends on the length of the output.

Q5: Do you need to recompute $\eta_{crit}(X_{batch})$ for each batch? How does it work?

We recompute $\eta_{crit}(X_{batch})$ for each batch. For a mini-batch $X_{batch}$, we can compute the batch loss $L(X_{batch})$ and the corresponding Hessian of the loss $H_L(X_{batch})$. The critical learning rate is computed as $2/||H_L(X_{batch})||_2$.

Q6: Does high AGOP alignment imply high generalization in low-rank matrix sensing problem?

This is a very interesting example for understanding AGOP alignment. In fact, this is a problem we are currently working on but the results are out of scope of this submission.

We thank you for the minor suggestions. We have revised our draft according to your suggestions.

Minor suggestions: ...To me it's more like Cohen et al. "observed" numerous catapults at EoS, but conjectured that EoS and catapults have the same underlying cause.

Indeed in [Cohen et al. 2020] they observed numerous loss spikes but they did not provide any evidence that these are catapults. We have made it clear in our revised draft.

Minor suggestions: When defining the AGOP alignment in page 4, the "true model" $f$ should be defined more rigorously...

We actually made it clear that when the true model is not available, we use a SOTA model as a substitute in footnote 1 at the end of page 2. We will mention it again when we define true AGOP to avoid confusion.

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.

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.

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.

Comment 8: The convergence speed of GD is determined by the condition number of the sample matrix $XX^T$. Do you mean $X^T X$? And even for linear regression, the Hessian and the NTK are different (they are not even of the same size). How are the eigenspaces comparable?

For a linear model $h(w;x) = w^Tx$ where $w$ is the parameter. For $n$ data points $(X,y)$ where $X\in \mathbb{R}^{n\times d}$, $y\in\mathbb{R}^n$, consider minimizing the squared loss $L(w) = \frac{1}{n}||{h(w;X) - y}||^2 = \frac{1}{n}|| Xw - y||^2$.

The loss Hessian takes the form $H_L = \nabla^2 L(w) = X^T X \in \mathbb{R}^{d\times d}$, and the NTK takes the form $K = \frac{d h(w;X)}{dw} \frac{d h(w;X)}{dw}^T = XX^T \in \mathbb{R}^{n\times n}$.

$XX^T$ and $X^TX$ share the same non-zero eigenvalues.

Comment 9: 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. It'd be helpful to make this more formal or at least include this in the paper (along with the NTK parameterization descriptions).

We have updated the draft to add a description of the quick loss drop in the introduction and NTK parameterization in Appendix F.

Comment 10: 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. I think I've already commented on this earlier, but I still don't understand why you used the NTK parameterization other than someone else also used it, and why does it even matter if model is approximately linear under NTK with a small learning rate, when you're in fact using a large learning rate for catapult to occur.

We also used Pytorch default initialization in our work (see Figure 12, 18, and 20 in the appendix). The original catapult paper [1], as well as the analysis of the catapults on simplified models (e.g., [2,3]) suggest that catapults occur using NTK parameterization with a large learning rate. Our work used this well-established setting to study the impact of the catapults.

Furthermore, the analyses in [1,2,3] also suggest that models that are approximately linear with a small learning rate (including Pytorch default initialization [4]) can have catapults with a large learning rate. Therefore the kernel regime is closely related to the catapults.

Additionally, we note that small and large learning rates are generally only different by a factor of two (whether it is smaller or larger than the critical learning rate).

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).

4. G. Yang, and H. Edward J.. "Feature learning in infinite-width neural networks." arXiv preprint arXiv:2011.14522 (2020).

Comment 11: The learning rate is increased to be roughly $2/\eta_{crit}$. So how do you increase it to be roughly that value? Could you point me to where it's described in the appendix? (There's quite a lot of stuff scattered around in the appendix and I don't think I was able to find it). This is very important for reproducibility.

The learning rates and the iterations are specified in Appendix F( Experimental details), in the paragraphs corresponding to Figure 3 and Figure 6.

(Sorry, I don't know why it's not showing as part of a thread...)

Thank you for the reply. The addition of Figure 15 is nice, but since it's ran with SGD, it's unclear to me whether the catapults are due to stochasticity in the sampling or due to having a large step size. How was the step size selected here? And does more catapults also lead to higher accuracy? From this plot I can only tell that for the logistic loss, there's also a correlation between catapults and drops in $\|K\|_2$. Does it also correlate with higher accuracy?

What I mean by support vectors is that - since you are doing classification with the squared loss under some neural network, the model represents some nonlinear separator of the data into classes. The examples that are closest to this separator can be thought of as the support vectors, just like in the linear case (SVM). So if we take a large step such that a catapult is to occur for the training loss, then it ought to be these "difficult examples" that are affected the most. No each entry in the NTK measures the similarity between two examples when it comes to how much their prediction changes by a change in the model. And so it's natural for me to think that when a catapult occurs, the loss along the eigendirections corresponding to these "support vectors" will also catapult, hence why I find it somewhat unsurprising. Regarding your comment:

as the catapult dynamics of complicated neural networks, including Wide ResNets and Vision Transformer, indeed occur in a low dimensional subspace

The NTK matrix is $n\times n$ where $n$ is the number of training examples. How exactly is a subspace spanned by the eigenvectors of the NTK related to the dimension of the neural network? I think it's more related to a subset of the examples.

I agree that the NTK is closely related to the Hessian of the loss. Can you explain why

...shares the same eigenvalue with NTK with extra ...

What if $n>p$? And why is it ok to assume $f_i-y_i=O(1)$, especially when you are studying catapults where this quantity is large? Please correct me if I'm wrong, but afaik the kernel regime is when the step size is tiny so that the model doesn't move much and hence the linear approximation is accurate. Doesn't it contradict to what it requires for catapults to occur?

We thank the reviewer for the positive feedback and insightful comments. We will address your concerns and questions below.

W1: it does not describe how NTK is changing...one could expect the NTK to decrease more along its top eigenvalues for example, since these eigenvalues are highly related to the spikes.

This is a nice point and we believe this is indeed correct. We conducted an experiment by training a two-layer neural network on 128 CIFAR10 data points. The change in the eigenvalues of NTK is as follows (we report the average of 3 independent runs):

The reasons we only show the change of the spectral norm (i.e., the top eigenvalue) of the tangent kernel are the following:

1. the catapult phase is characterized by the decrease of the spectral norm of the tangent kernel in the original paper [Lewkowycz et al. 2020] for GD. We use this characterization to show that catapults occur in SGD as well.

2. the top eigenvalue relates directly to the optimization properties since the critical learning rate is determined by the top eigenvalue.

W2: The only explanation that is proposed is that the AGOP alignment increases, but it is not explained why these spikes should increase AGOP alignment.

The mechanism of how these spikes increase AGOP alignment can be seen as follows. In the literature, e.g., [3], the change of the NTK is closely related to feature learning. Note that with $0$ catapult, i.e., NTK regime, due to the small weight change, AGOP is not able to update a lot to align with the true AGOP. In work[1,2], it was proved that NTK models cannot effectively learn features. However, with catapults, the weights are allowed to change a lot (as the gradient of the loss i.e., $dL/dw = (f-y)\frac{df}{dw}$ scales with the value of the loss, which is huge when catapults occur). Consequently, it leads to a significant change in the AGOP. Please see our experiments of multiple catapults in GD (Figure 6), where a greater number of catapults in GD leads to a higher (better) AGOP alignment and smaller (better) test loss/error. This mechanism is similar to the feature learning process with near-zero initialization[3], in which case the AGOP is also allowed to change a lot during training (validated by our experimental results in Figure 16). Note that to get a better generalization/AGOP alignment, it is necessary to have a significant AGOP change when the true AGOP is low-rank (e.g., multi-index models), compared to the full-rank AGOP at initialization. And this can be realized by having catapults. Our experiments on real-world datasets such as SVHN and CelebA suggest that this is the case of how neural networks learn features in practice. A more detailed explanation is in the paragraph " Intuition of the connection between catapults and AGOP alignment" after Figure 6.

1. B. Ghorbani, et al. "Limitations of lazy training of two-layers neural network." Advances in Neural Information Processing Systems 32 (2019).
2. C. Wei, et al. "Regularization matters: Generalization and optimization of neural nets vs their induced kernel." Advances in Neural Information Processing Systems 32 (2019).
3. G. Yang, and H. Edward J.. "Feature learning in infinite-width neural networks." arXiv preprint arXiv:2011.14522 (2020).

W3: ...it is unlikely that AGOP alignment fully captures feature learning. The authors could have considered other measures of feature learning, such as the alignment between the NTK and the task. This could have been especially interesting ...

Q1: I would prefer if you did not present the AGOP as the mechanism throught which NN learn features... While the AGOP alignement seems well suited to low-index tasks, it seems ill suited to measure feature learning when the learned features depend nonlinearly on the inputs, such as in (https://arxiv.org/abs/2305.19008).

Thank you for the suggestion and the reference. We note that AGOP has been shown to accurately capture features learned in a broad class of neural network architectures, including both nonlinear fully connected and convolutional networks (see [1, 2]). Given this line of work and our results in this paper, we believe that a major contribution to neural feature learning comes from AGOP. Nevertheless, we feel connecting AGOP alignment and other measures of generalization is an important direction of future work. We have added a discussion of the suggested references to our draft.

1. A. Radhakrishnan, et al. "Feature learning in neural networks and kernel machines that recursively learn features." arXiv preprint arXiv:2212.13881 (2022).

2. D. Beaglehole, et al. "Mechanism of feature learning in convolutional neural networks." arXiv preprint arXiv:2309.00570 (2023).

We thank the reviewer for the positive feedback and insightful comments. We will address your concerns and questions below.

W1: No rigorous theoretical analysis for spikes in neural network training loss.

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).

W2: Not sure if the linearized approximation between two spikes holds in real-world situations.

It looks like the reviewer's comment refers to our Remark 2, which does appear to be cryptic. We clarify that we are not assuming the linear dynamics hold between the spikes. We only used the fact that wide neural networks are locally linear in a small neighborhood to estimate the critical learning rate (please see Claim 2 where the critical learning rate depends on iteration $t$). We have updated our draft to clarify this point.

We thank the reviewer for the positive feedback and insightful comments. We will address your concerns and questions below.

W1: The paper never answers why the catapults occur...are not the actual mechanisms behind why catapults occur. Rather, they are direct conclusions of catapults occurring.

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).

W2: The correlation between the number of catapults and AGOP is not explicit inner workings of catapults and their effects on generalization.

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).

Q1: Any results on (stochastic) gradient descent with momentum?

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.

Q2: ...relations of AGOP to other "generalization measures" such as flatness (max eigenvalue of Hessian), tail-index of the trajectory/weight matrix [2,3]...etc

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).

Q3: Any connection to recently uncovered saddle-to-saddle hopping dynamics [3,4]? Can the hoppings be equated to catapults?

In saddle-to-saddle hopping dynamics, there are no spikes in the training loss and these dynamics typically happen with small initialization. Thus, we believe there are no direct connections between these phenomena.