Paper: Quadratic models for understanding catapult dynamics of neural networks (ICLR-2024)

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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?

Reviewer comment (Theory):
Regardless of the new title, the paper still needs to clarify why or why not chaos is present in the model and clarify its connection to the studies in this line.

Author response:
We note that the catapult dynamic considered in our work is not chaotic. See, e.g., X. Chen, et al. "From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regression"  arXiv:2310.01687 (2023) for a nice discussion (we will add the ref. to the paper).

We hope we have addressed your concerns. If so, please consider raising your score.

Reviewer comment (Experiments):
The neural network in Eq (9) is initialized as $u_i \sim \mathcal{N}(0, I_d)$ 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?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Theory):
$p_1(t)$ is the top eigenvector of $K(t)$. 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)^TR_K(t)p_1(t)$. Could you please explain in more detail?

Author response:
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.

Reviewer comment (Theory):
In Eq(12), $R_\lambda(t)$ is defined without the minus sign. In the following paragraph, you mention that "$R_\lambda(t)$ stays positive and results in monotonic decrease of kernel" which makes sense since $\lambda(t)=\lambda(t-1)-R_\lambda(t)$. But in the next paragraph, you write down $\lambda(t)=\lambda(0)+\sum_{\tau=0}^{t} R_\lambda(\tau)$ which suggests that $R_\lambda(t)$ should include the minus sign. Please clarify.

Author response:
Thanks for pointing it out. We have fixed the typo.

Reviewer comment (Theory):
Related to above, I think it should be $\lambda(t+1)=\lambda(0)+\sum_{\tau=0}^{t} R_\lambda(\tau)$ or $\lambda(t)=\lambda(0)+\sum_{\tau=0}^{t-1} R_\lambda(\tau)$.

Author response:
Thanks for pointing it out. We have fixed the typo.

Reviewer comment (Theory):
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.

Author response:
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.

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

Author response:
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)$.

Reviewer comment (Theory):
Page 27 in the Appendix, it should be $\Pi_1 \mathcal{L}(t)=K_1(t)\Pi_1 \mathcal{L}(t-1)$.

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

Reviewer comment (Theory):
I have some concerns about the proof of lemma 1 in appendix D: you use a proof by induction, and still use O, o notations. But summing $O$ terms remains $O$ only when the number of iterations is controlled. To still have $O$ in 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 \frac{\log(m)}{\sqrt{m}}$. Indeed $u(t)$ goes from $O(1/m)$ to $O(\log(m)/m)$. In that case, $T \sim \frac{\log(\log(m))\sqrt{m}}{\log(m)}$ which is not $O(1)$.

Author response:
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.

Reviewer comment (Theory):
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. For example $(1+\delta))^2<(1+\delta-(-3\delta^2/\log(m)))^2$. I think this statement in its current form would perhaps need an additional study of the constants, their sign or if they are zero.

Author response:
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.

Reviewer comment (Theory):
The non-linearity that is used in this paper is a ReLU: it allows to compute easily the second derivative of the neural network function (it deletes the diagonal terms of the hessian of the neural network function). My only concern is regarding the generalization to non-linearities that are not piecewise linear and hence which make another term appear in the Hessian of the neural network function, which corresponds to the second derivative of the non-linearity (correct me if I'm wrong). 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?

Author response:
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.

Reviewer comment (Novelty):
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 (cf the paper "Second-order regression models exhibit progressive sharpening to the edge of stability" that is not cited but seems to me very related to your setting). It seems interesting because both studies explore the influence of a non-linear quadratic term on the coupled dynamics of the sharpness and the learning rate.

Author response:
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.

Reviewer comment (Writing):
I find the results interesting and they provide good contribution. I would increase my score if the authors address my concerns about the clarity of $O,o,\Omega$ notations in the proof, especially the proof by induction.

Author response:
We have revised the proof of Lemma 1 by using explicit constants to replace the $O,o,\Omega$ notation. Our results hold as the dependence on the induction step $t$ will be absorbed into the summation of the geometric sequence.  

The logic of Lemma 1 is as follows: we consider the interval $[0,T]$ where $T$ is the first time that $u(T)$ increased to $\Theta(\delta^2/\log m)$. To prove Lemma 1, we only need to show that the ratio $u(t+1)/u(t):=\kappa(t) >1$ for each $t\in [0,T]$ so that $u(t)$ keeps increasing until reaching to the order $\Theta(\delta^2/\log m)$. Note that, in this whole procedure,  the accumulation of $O$ terms across $T$ steps does not present a concern. To show $\kappa(t)>1$, we need to bound the change of $\kappa(t)$ from $\kappa(0)$ to be small, so that $\kappa(t)$ remains above $1$ by triangle inequality. Note that the change of $\kappa(t)$ from $\kappa(0)$ is controlled by  $u(t)$ and $v(t)$. As $u(t)$ increases geometrically and $u(t) = O(\delta^2/\log m)$ for all $t\in[0,T]$,  we can derive a bound of the sum of $u(t)$ independent of $t$. A similar bound can be done for $v(t)$. Therefore, we derive a bound for the change of $\kappa(t)$ from $\kappa(0)$ independent of $t$.   Consequently, even if the inductive step $t$ may depend on $m$, as it will not appear in the bound, our results hold.

Please see more details in the revised proof.