Quadratic models for understanding catapult dynamics of neural networks

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

W1: The paper covered too many topics that it feels that it does not achieve any point satisfactorily. For example, the paper feels mistitled. The main focus of the paper is the catapult mechanism -- which does appear in deep learning but cannot represent all types of "neural network dynamics." In my opinion, the catapult mechanism is a rather special / specific type of dynamics and the title is an overclaim. If the authors change the title to a more proper one, the paper could be much easier to evaluate.

We feel the title should not preclude the evaluation of the paper. Nevertheless, we have changed the title to "Quadratic models for understanding catapult dynamics of neural networks" to emphasize connections to catapult dynamics. We hope this addresses your concern and will result in a more favorable evaluation and score.

W2: Lack of discussion of a highly relevant problem. Essentially, within the framework of quadratic models, it appears to me that the catapult mechanism is nothing but what the academia traditionally calls "chaos." For example, consider a width-1 quadratic model, and compare the dynamics of GD to that of a logistic map -- the dynamics is essentially identitical -- the loss of local stability leads to chaos in the logistic map, and the same here in the quadratic model. The authors need to discuss this point in my opinion.

The catapult dynamics considered in our paper is not chaotic, and as such does not seem to be connected with the dynamics of the logistic map. Would you be able to share a reference so we understand your point better?

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

W1.a: The architectures and datasets used are still relatively limited... A CNN and transformer experiment would be particularly welcome.

Indeed we have experiments with CNNs as shown in Figure 4 and 16. We experimented with various datasets, including vision (CIFAR10, SVHN, MNIST), speech (FSDD) and language (AG NEWS). As a primarily theoretical paper, we believe this is a good scope of experiments.

W1.b: It would also be good to highlight what the incremental value of this paper is over prior work on quadratic models.

The related paper [1] introduced and analyzed the NQM (called quadratic models in their book), but our work analyzed the training dynamics with large learning rates, which was not discussed in [1].

1. Roberts, Daniel A., Sho Yaida, and Boris Hanin. The principles of deep learning theory. Cambridge, MA, USA: Cambridge University Press, 2022.

W2: From a cursory literature review, I have found this paper which has relatively substantial overlap in topic, and I believe should at least be cited (Meltzer and Liu https://arxiv.org/abs/2301.07737). A more thorough set of references on recent work around the catapult effect and edge of stability would also benefit the paper.

Thanks for pointing out this related work. In fact, the first version of our work was prior to [Meltzer and Liu], and we weren't aware of it. We added a discussion in the updated draft.

As for the Edge of Stability (EoS), we have not seen a clear connection between EoS and catapult dynamics. Although there seems to be some similarities between Edge of Stability (EoS) and catapult dynamics at first glance, we noticed a few key differences between the two: 1) the training loss can increase to a large value, a order of $\Theta(m)$ with $m$ being the network width, in the catapult dynamics, while it does not observed to increase significantly in EoS, 2) there is a single loss spike in the catapult dynamics of full-batch gradient descent, while there are often multiple loss oscillations in EoS, 3) the catapult phase phenomenon often improve generalization across various tasks [Lewkowycz et al. 2020], while it is unclear if EoS has any effect on the generalization. The original EoS paper [1] conjectures that the training loss oscillations could be micro-catapults, however, there is no direct evidence for the connection between EoS and catapult dynamics.

1. J. Cohen, et al. "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability." International Conference on Learning Representations. 2020.

W3: From empirical work, I've also seen that quadratic model's validation curve only tracks the NN's curve in the early stages of training, but then diverges from it. It would be nice to give an example of the limitations of the quadratic model for understanding generalization. Being clear about potential limitations of quadratic models to fully model neural networks would be welcome.

In our experiments, the generalization performance of neural networks seems to perform better than their corresponding NQMs when the learning rate is large, i.e., catapult occurs. One possible reason is that NQMs are the second-order approximation of neural networks, while they can exhibit the catapult phase phenomenon the same as the neural networks, they may not be able to learn the same features as the neural networks. That is to say, to learn certain features, higher-order terms are needed. This is a complex question and a direction for future work.

I thank the authors for their input. The distinction the authors' made in their rebuttal between studying EoS dynamics and catapult dynamics was helpful in being able to understand the relative novelty of this paper. The paper by Agarwala, Pedregosa, and Pennington: "Second-order regression models exhibit progressive sharpening to the edge of stability" tackles the EoS phenomenon for quadratic models. I'd encourage the authors to add a sentence of two about the distinction between these effects in the final draft to remove confusion between their contribution and earlier ones. Thanks also for including the discussion of the Meltzer and Liu paper. These additions will substantially improve the literature review.
After some thought, I have decided to raise my score. This is due to the set of experiments pointed out to me, together with the clear introduction and general importance of the problem being studied. My confidence in my understanding of the proofs is still relatively low, but there are no glaring errors that I have found. Also, the empirical results give credence to the claims.

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

W1: The proposed term Neural Quadratic Models seems to be unnecessary as for linear models we don't call them Neural Linear Models.

The term "quadratic model" can refer to any model that generically has a quadratic form. In this paper, we study the special instance: the quadratic model that best approximates the neural network. Hence, we think the term "Neural Quadratic Model” describes this specific model the best. This is similar to the situation of "Neural tangent kernel" (NTK). NTK is a specific kernel associated with neural networks}

We believe it is important to highlight the close relation between the quadratic models and neural networks since they are the second-order Taylor expansion of the networks.

Some works, e.g., [1] indeed refer to the linear approximation of neural networks as neural tangent models, to highlight their close relation.

1. B. Ghorbani, et al. (2019). Limitations of lazy training of two-layers neural network. NeurIPS.

Q1: Some recent findings such as the deep double descent are high-dependent on the size of the neural networks and other hyper-parameters. Is the "catapult phase" here sensitive to other hyper-parameters apart from learning rates? I don't see many conditions for the theorem presented in the paper.

The width $m$ of the network affects the magnitude of the catapult dynamics. As we have shown in the paper, the loss value at the peak of the catapult scales as $O(m/\log m)$. We don't observe dependence on other hyper-parameters.

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

W1: The neural network in Eq (9) is initialized as $u_i \sim \mathcal{N}(0,I)$ and $v_i \sim Unif(-1,1)$. Can you explain why $v_i$ is initialized different from $u_i$, also different from the typical inverse fanin/fanout initialization in practice?

This initialization/parameterization strategy adopted in our paper is the so-called NTK parameterization [1]. It is standard in theoretically analyzing neural networks with large width, e.g., [1,2,3]. In the "typical inverse fan-in/fan-out initialization in practice", NTK scales as $O(m)$, which blows up for large width. However, in the NTK parameterization, NTK scales as $\Theta(1)$.

The reason why we use the uniform distribution of $v_i$, instead of the normal distribution $\mathcal{N}(0,1)$, is to make sure $v_i$ is bounded, which simplifies the analysis. In the case of $v_i \sim \mathcal{N}(0,1)$, similar results will hold with extra $log(m)$ factors in some $O$ terms, with a high probability that depends on $m$.

Drawing $v_i$ from a bounded distribution is a common setting in literature, for example, [2,3].

1. A. Jacot, et al. “Neural tangent kernel: Convergence and generalization in neural networks”. NeurIPS, 2018.

2. S. Du S., et al. "Gradient Descent Provably Optimizes Over-parameterized Neural Networks." International Conference on Learning Representations. 2018.

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

W2: The assumption that width ($m$) is larger than data size ($n$) which is assumed to be a small constant is unrealistic. While in the limit where width goes to $inf$, this would be the case but for any finite width network, this does not hold in practice. Is this assumption crucial, can you still assume that $n/m$ does not necessarily go to zero?

Having $m>n$, i.e., over-parameterization is a common setting in theoretical research of recent years. For example, [1] assumes $m = \Omega(n^6)$, [2] assumes $m = \Omega(n^4)$, [3] assumes $m = \Omega(n^{24})$. Many interesting findings and results in this large width setting can often be experimentally observed in practical settings. In practice, the catapult dynamics still occur when the width and the size of datasets are comparable. For example, in our experiment for FSDD dataset, the catapult still occurs when the width is $1,000$ and the size of the dataset is $256$.

1. S. Du S., et al. "Gradient Descent Provably Optimizes Over-parameterized Neural Networks." International Conference on Learning Representations. 2018.
2. S. Du S., et al. "Gradient descent finds global minima of deep neural networks." International conference on machine learning. PMLR, 2019.
3. Z. Allen-Zhu, et al. "A convergence theory for deep learning via over-parameterization." International conference on machine learning. PMLR, 2019.

W3: Given that $p_1(t)$ is not necessarily equal to $p_1(t+1)$, it is not clear how you derived $\lambda_1(t+1) = \lambda_1(t) - p_1(t)^T R_K(t) p_1(t)$...

We suppose you were referring to Appendix N.3. This is a typo: the equation should be $p_1(t)^T K(t+1)p_1(t) = \lambda_1(t) - p_1(t)^T R_K(t) p_1(t)$. The experiments in Appendix N.3 show that the training dynamics confined to the top eigenspace of the tangent kernel i.e., $p_1(t)$ can almost capture the catapult dynamics. We have fixed it in the revision.

W4: In Eq(12), $R_\lambda(t)$ is defined without the minus sign....which suggests that $R_\lambda(t)$ should include the minus sign.

W5: Related to above, I think it should be $\lambda(t) = \lambda(0) - \sum_{\tau = 0}^{T-1} R_\lambda(\tau)$ ...

Thanks for pointing it out. We have fixed the typo.

W6: You mention in the decreasing phase section in page 6 that decrease in $v(t)$ would cause a decrease in $\kappa(t)$. But in Eq (13), reducing $v(t)$ would increase inside of square which should lead to an increase in $\kappa(t)$; unless $u(t) + w(t)<0$ which is not clear if it holds. Please clarify.

Note that at initializaiton $v(0) > 2$ and $u(0) + w(0) = O(1/m)$. Therefore $\kappa(0) = (1-v(0) + O(1/m))^2 >1$ since $1-v(0) < -1$. As $v(t)$ decreases, $|1-v(t)|$ will also decrease therefore $\kappa(t)$ decreases.

Additionally, the fact that $u(t) + w(t)>0$ and keeps increasing will further make $\kappa(t)$ decrease.

W7: In Eq.(10), $1/\sqrt{d}$ is missing from $\sigma(u_{0,i}^T x).$ It is present in Appendix A

In the theoretical analysis, we only consider ReLU activation function which is homogenous, hence we can extract the scaling factor $1/\sqrt{d}$ out of the function $\sigma(\cdot)$.

W8: Page 27 in the Appendix, it should be $\Pi_1 L_1(t) = K_1(t) \Pi_1 L(t-1$

We denote the scaling factor by $\kappa_1(t)$, consistent with the notation used in the single training sample scenario.

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

W1: I have some concerns about the proof of lemma 1 in appendix D ... But summing $O$ terms remains $O$ only when the number of iterations is controlled. To still have $O$ at the end of the summation with respect to $m$, it should be checked that the summation index $T$ is itself $O(1)$. However I am not sure such a result is proved (correct me if I'm wrong). Especially it seems to me that it might not be the case by considering the fact that T must satisfy a relation of the form $(1+\delta)^T \sim \log m$ with $\delta \sim \log m/\sqrt{m}$. Indeed $u(t)$ goes from $O(1/m)$ to $O(\log m /m)$. In that case, $T \sim \log (\log m)\sqrt{m}/\log m$ which is not $O(1)$.

In the proof of Lemma 1, $u(t)$ grows geometrically from $O(1/m)$ to $O(\log m /m)$ with rate $\kappa(t) \geq (1+\delta - O(\delta^2/\log m))^2>1$ for large enough $m$. Note that the summation of a sequence that grows geometrically with a ratio $(1+\delta)^2>1$ can be controlled by the last term, i.e., $\sum_{t=1}^T u(t) = \Theta(u(T)/\delta)$. Hence, even if $T$ is not $O(1)$, our results still hold. We updated the draft to clarify this point.

W2: Another ambiguity about the $O$ notation is for example when it is stated: $\kappa(0)>(1+\delta - O(1/m))^2 > (1+\delta - O(\delta^2/\log m))^2$ in appendix D. In full generality it seems wrong for general sequences as it depends on the constants in the $O$ itself and their sign. I think this statement in its current form would perhaps need an additional study of the constants, their sign or if they are zero.

We apologize for the confusion. In our analysis, the $O$-terms are always positive, and we explicitly write out their sign before $O$. For this inequality $\kappa(0)>(1+\delta - O(1/m))^2 > (1+\delta - O(\delta^2/\log m))^2$, we have $m$ large enough to make sure that each term within the brackets is positive. We have updated the draft to clarify this.

W3: My only concern is regarding the generalization to non-linearities that are not piecewise linear... Could the author comment about how to handle this and if they expect the analysis to be the same and the results to still hold?

While our analysis is specific to ReLUs, our experiments for tanh and swish activation functions in Appendix N.8 also verify that our results extend to such smooth activation functions. In the literature, the analysis for smooth non-linear activation functions is often handled differently from ReLUs.

In addition, ReLU is so widely used in both theory and practice that many important theoretical results are specific to neural networks with ReLU activations. See for example:[1,2,3]

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. D. Zou, et al. (2020). Gradient descent optimizes over-parameterized deep ReLU networks. Machine learning, 109, 467-492.

Q1. I am curious if you could clarify the links, if they exist, between catapult dynamics and recent works on edge of stability analysis of neural networks...

Although there seem to be some similarities between Edge of Stability (EoS) and catapult dynamics at first glance, we noticed a few key differences between the two: 1) the training loss can increase to a large value, an order of $\Theta(m)$ with $m$ being the network width, in the catapult dynamics, while it is not observed to increase significantly in EoS, 2) there is a single loss spike in the catapult dynamics of full-batch gradient descent, while there are often multiple loss oscillations in EoS, 3) the catapult phase phenomenon often improve generalization across various tasks [Lewkowycz et al. 2020], while it is unclear if EoS has any effect on the generalization. The original EoS paper [1] conjectures that the training loss oscillations could be micro-catapults, however, there is no direct evidence for the connection between EoS and catapult dynamics.

1. J. Cohen, et al. "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability." International Conference on Learning Representations. 2020.

Q2: It is written in the text that the results hold with high probability over initialization but I'm not sure if this is written in the statement of the theorems. Perhaps it could be added in the theorems themselves.

Thanks for the suggestion. We have updated it in our revision.

I thank the authors for their response which I have the feeling make me now understand better theorem 4 and lemma 1.

I thank them for having also changed the proof of theorem 4 with explicit constants which makes it more understandable in my point of view. However I have still concerns on the proof of proposition 3 still partly because of the the $O,o,\Theta,\Omega$ notations.

Minor concerns:

1. there is a different definition of $\delta$ before the "increasing phase" in the main paper as in the proof of lemma 2 in the appendix. Could you clarify?

2. I get the intuition but do not understand the meaning of proposition 1,2,3 because I feel there lacks something of the form "for m large enough...". For example in proposition 3, you state "v(t) keeps increasing...". I feel it would be good to specify it. (or correct me if I am wrong)

Bigger concerns: 3) I think that T_2 and T_1 depend both on m. I think it is important to recall it, or else I would be happy if the authors could explain me why it does not depend on m. Following this remark, I do not understand the sentence "Therefore, for $t\geq T_2$ such that u(t) is of the order greater than $O(1/m)$. Indeed, you cannot work with fixed $t$ if you set the condition $t\geq T_2$ with $t$ depending on $m$. Perhaps it would be good to specify this point.

4. I do not understand the proof of Proposition 2 in part H.2, starting "decreasing. Specifically...". Could your reformulate this part a bit? Especially: first what is meant by "order at least O(.)". Do you mean "order at least $\Omega$" or "order at most $O(.)$"? Second, are you sure about $O(1/\sqrt{log(m})$? it seems strange to me since this order is actually bigger than $1/\log(m)$ that you were discussing before. Shouldn't it be $O(1/log(m)^2)$

5. I will probably seem very attached to details but the proof by contradiction of proposition 2 is troubling me: you prove that the following statement is wrong: "$\forall t \geq T_1, v(t)=4-o(1)$". My concern is in terms of the mathematical statement of the contradiction of this. To me, the correct statement would be: "$\exists t \geq T_1, 4-v(t)$ is not $o(1)$" However, it seems not possible here since $T_1$ depends on $m$ and therefore $t$ cannot be defined. Furthermore I would like if the authors could clarify the following point: it seems to me that you use that the contradiction of being $o(1)$ is being $\Omega(1)$ but I think that it it is only true for a subsequence.

6. In the proof of proposition 2 it is written $u(T_1)=\Omega(\frac{1}{\log(m)})$. Could the authors explain more why it is true as it seems in contrast with the statement above proposition 2: $T_1=inf (t, u(t)=\Theta(\frac{\delta^2}{\log(m)}))$ with $\delta$ small.

We thank the reviewer for the reply. We have revised the statements of propositions and proofs to get rid of $O,\Omega,o$ by using explicit constants. And we will address your specific concerns below.

Concern 1: there is a different definition of $\delta$ before the "increasing phase" in the main paper as in the proof of lemma 2 in the appendix. Could you clarify?

This is a typo. We have fixed it.

Concern 2: I get the intuition but do not understand the meaning of proposition 1,2,3 because I feel there lacks something of the form "for m large enough...". For example in proposition 3, you state "v(t) keeps increasing...". I feel it would be good to specify it. (or correct me if I am wrong)

We have revised the statements of propositions and used explicit constants in the proof. Now there is an explicit lower bound for $m$ for the propositions to hold.

Concern 3: I think that $T_2$ and $T_1$ depend both on m. I think it is important to recall it, or else I would be happy if the authors could explain me why it does not depend on m. Following this remark, I do not understand the sentence "Therefore, for $t\geq T_2$ such that $u(t)$ is of the order greater than $O(1/m)$. Indeed, you cannot work with fixed $t$ if you set the condition $t \geq T_2$ with $t$ depending on $m$. Perhaps it would be good to specify this point.

We have clarified it in the proof that $T_1,T_2$ depend on $m$.

We have revised the statements, and now the proposition becomes under certain conditions, there exits $T_2$ such that $v(T_2) <3$ (see Proposition 3 in the updated draft).

Concern 4: I do not understand the proof of Proposition 2 in part H.2, starting "decreasing. Specifically...". Could your reformulate this part a bit? Especially: first what is meant by "order at least $O()$". Do you mean "order at least $\Omega$ " or "order at most $O$ "? Second, are you sure about $O(1/\sqrt{\log m})$ ? it seems strange to me since this order is actually bigger than $1/\log m$ that you were discussing before. Shouldn't it be $1/\log^2 m$?

We have revised the proof by using explicit constants. Specifically, we added Proposition 2 to clarify when $u(t)$ is larger than $\frac{4C}{m (4-v(t))^2}$, $v(t)$ will decrease.

Concern 5: I will probably seem very attached to details but the proof by contradiction of proposition 2 is troubling me: you prove that the following statement is wrong: "$\forall t\geq T_1, v(t) = 4-o(1)$ ". My concern is in terms of the mathematical statement of the contradiction of this. To me, the correct statement would be: " $\exists t\geq T_1$,$4-v(t)$ is not $o(1)$ " However, it seems not possible here since $T_1$ depends on $m$ and therefore $t$ cannot be defined. Furthermore, I would like if the authors could clarify the following point: it seems to me that you use that the contradiction of being $o(1)$ is being $\Omega(1)$ but I think that it it is only true for a subsequence.

We have revised the Proposition and used explicit constants in the proof. Now we show that there exists $T_2 \geq T_1$ ($T_2$ depends on $m$) such that $v(T_2)<3$. Note that this statement is more explicit than " $\exists t\geq T_1$,$4-v(t)$ is not $o(1)$". The idea of the proof is similar: $v(t)$ will keep decreasing if $v(t)\geq 3$ therefore $v(t)$ will be less than $3$ at a certain iteration.

Concern 6: In the proof of proposition 2 it is written $u(T_1) = \Omega(1/\log m)$ . Could the authors explain more why it is true as it seems in contrast with the statement above proposition 2: $T_1 = \inf_t\{u(t) = \Theta(\delta^2\log m)\}$ with $\delta$ small.

The proposition shows that there exits $T_2$ such that $v(T_2)$ is not so close to $4$. For the case where $\delta$ is small, i.e., $v(T_1)$ close to $2$, we can just let $T_2 = T_1$. For the case $v(T_1)$ is close to $4$, we have $\delta >1$. We have clarified it in the updated draft.