Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos

We thank the reviewer for their careful reading and encouraging comments. We will consider their suggestion to shorten Section 5.2 in the future version of the manuscript.

What actionable insights can be derived from the analysis?

Sharpness dynamics serves as a practical way to assess the initialization scale. When sharpness reduces during training, it implies that the model initialization was large in some sense. Instead, we could have picked a flatter initialization or performed a learning rate warmup to facilitate training at higher learning rates, which typically results in better performance. Additionally, the correlation between sharpness and weight norm dynamics offers a simpler way of assessing the initialization scale without the need to measure sharpness. If the weight norm decreases during training, it would suggest that the initialization was large.

The UV model analysis suggests models initialized in the progressive sharpening region approach the zero-loss solution with low sharpness. This suggests using small initializations or parameterizations, which exhibit progressive sharpening from the start. Common practices such as setting the last layer to zero at initialization and using $\mu$P in the large-width limit regime result in zero model output at initialization, eliminating sharpness reduction during early training (see Figure 19 in the updated manuscript).

Does Section 6 rely on the trace of the Hessian, or does it analyze its top eigenvalue? ...

In Section 6, we analyze the top eigenvalue of the Hessian. The trace of Hessian is only used in the UV model for mathematical simplifications. In Appendix C.5, we have verified that the trace of Hessian in the UV model is a good proxy for its sharpness.

I was under the impression that the role of loss landscape sharpness and generalization was called into question by Dinh et al. (2017)....

We thank the reviewer for pointing out this important reference. We have added it to the first paragraph of the introduction.

In Appendix E, the examples used to validate the hypothesis about the weight norms is a 3-layer FCN, while...

Yes, the weight norms were consistent with the two-layer FCNs results. We analyzed the weight norm of models with varying depths to verify that the correlation was not limited to two-layer networks.

The authors mention that fixed point analysis is a generally powerful approach ..

As we mentioned in Sec. 7, the limitation of applying this method is the closure of the dynamical equations describing the model. For deep linear networks, this model can be generalized by introducing more function space variables. For non-linear DNNs, the function space dynamics are hard to close with a few variables. We have included this discussion in Section 7.

I think putting references in parentheses would make the paper a little easier to read..

We thank the reviewer for their suggestion, which we will incorporate in the future version of the manuscript.

We thank the reviewer for the detailed comments and highly constructive suggestions. In response to the reviewer's comments, we have included a thorough discussion on the limitations of the UV model in Section 7. Below, we address these concerns individually.

...comparing Fig. 1 and Fig. 7 (FCN vs. UV model), there seem to be some differences in behavior of the loss and $\lambda$...How truly "similar" are the dynamics in the UV model and the FCNs used to motivate this paper?...

As correctly noted by the reviewer, the UV model does not capture two aspects of the training dynamics.

First, it does not capture the non-monotonic decrease in sharpness at EoS. This is because the dynamics is fully characterized by two variables and the EoS condition $\lambda \sim 2/\eta$ puts one more constraint on the dynamics. Resultantly, the dynamics is described by only a single variable, and loss and $\lambda$ oscillates at EoS. This is in contrast to real-world models where the training converges along the stable directions at EoS. Therefore, the UV model should be thought of as a minimal model of the top eigenspace oscillations at EoS.

Second, the loss may not catapult for the $\mu$P case shown in Figure 7(c, f) as pointed out by the reviewer. This arises because the initialization $\lambda_0 = 2$ is comparable to the maximum trainable learning rate $\eta_{\text{max}} = 1.0$ (for loss increase at initialization, we require $\eta > 2 / \lambda_0 = 1$). Such situations are also observed in practice. Such as Pre-LN Transformers trained with SGD, as shown in Figure 14(a, c) of the paper.

..it seems that some parameters are much more important than others, e.g. for large enough value of $s$ any value of $c$ suffices to see EoS...

As pointed out by the reviewer, parameters such as $s$ are more important in observing EoS and the UV model does not make any such predictions. We believe this is due to the non-linear relationship between the initial sharpness and $\sigma^2_w$, $s$ and $c$. While the UV model does not capture these complex relationships, it does predict that these parameters are crucial in determining EoS.

The weak correlation between performance and EoS arises from clipping $\eta \lambda^H_t / 2$ at $1$. We have updated Figure 4 to show $\eta \lambda^H / 2$ without any cap. These results show a better correlation between test accuracy and EoS.

The lack of period doubling in large networks is a very interesting point....

As observed by the reviewer, Section 6.4 demonstrates that real-world data suppresses the period doubling route to chaos. Nevertheless, we do observe a reminiscent of period doubling route to chaos in the power spectrum of sharpness trajectories of real-datasets (peaks in power spectrum Figure 5(f) at $\omega=1/2, 1/4, 1/8$).

This is a limitation of the UV model, which is trained on a single example and thus cannot capture the complex correlations between training examples. We believe highlighting this limitation and identifying its root cause is important. Nevertheless, the UV model successfully captures various other features of sharpness dynamics, such as progressive sharpening, early sharpness reduction, and late-time sharpness oscillations around $2/\eta$. The emergence of these phenomena in the simplified UV model suggests they are fundamental properties of neural network training dynamics.

I did not follow how the authors got that the loss was ...

We use the relation that on EoS manifold $\lambda = 2 \|x \| (\Delta f + y) / \sqrt{n_{\mathrm{eff}}}$, to substitute $\Delta f$ by $\lambda$ and other variables, then we pull target $y$ out of the parentheses and combine other variables into $\eta_c = \sqrt{n_{\mathrm{eff}}} / \|x\|y$. We have updated section C.9 for detailed derivations.

The title of the paper "Universal sharpness" is perhaps too strong...Does the analysis performed by the authors demonstrate "Universal" properties?..

We chose `Universal sharpness' as the title because several prior works were operating in different sharpness regimes. For instance, (Lewkowyz et al. 2020, Zhu et al. 2022) operated in the sharpness reduction regime, whereas (Jastrzebski et al. 2020, Cohen et al. 2022) operated in the progressive sharpening regime. Our work resolves this discrepancy by providing a unifying analysis for different sharpness regimes, motivating the chosen title. Furthermore, these sharpness trends are observed across architectures, including small CNNs, ResNets, and Transformers, where sharpness computation is feasible.

We thank the reviewer for their detailed comments.

..the theoretical approach is hugely similar to existing work (Lewkowyz, 2020), though many additional insights revealed

While we use a similar model compared to Lewkowycz et al. (2020), our analysis differs in two key aspects:

1. Our work recognizes the crucial role of non-zero target $y$ and parameterization. Lewkowycz et al. (2020) set $y = 0$ (Eqs. 6-7 of their paper), which results in a model that cannot exhibit progressive sharpening (and consequently EoS) for any initialization or parameterization. In comparison, we observe progressive sharpening and EoS, precisely because of non-zero target $y$ and $\mu$P.

2. Furthermore, their analysis relies on specific approximations to simplify the dynamical equations to analyze catapult dynamics, particularly leveraging the existence of $\mathcal{O}(1/n)$ terms that can be ignored in the NTK regime. This approach cannot be generalized to $\mu$P ($n_{\text{eff}} = 1$) where such approximations are infeasible. In contrast, our generalized UV model dynamics comprehensively captures different sharpness phenomena across parameterizations.

Section 6.1 claims that $\mu$P networks begin training in a flat region of the landscape where gradients point to increased sharpness...

We thank the reviewer for highlighting this potential source of confusion.

For the same variance ($\sigma^2_w = 2.0$), the two parameterizations show completely different dynamics. Excluding large learning rates resulting in loss and sharpness catapults, models in NTP/SP experience sharpness reduction, whereas $\mu$P exhibits progressive sharpening after one step of reduction. From the UV model, if we set $f_0 = 0$, then the model is necessarily initialized in the progressive sharpening regime, which is satisfied by $\mu$P at large width. However, due to finite width corrections, training can be initialized in the sharpness reduction regime and may experience a slight sharpness reduction. To validate this, we set $f_0 = 0$ for models in Figure 1, by setting the last layer to zero at initialization. Figure 19 shows that setting the last layer to zero completely eliminates sharpness reduction in all cases. This result suggests that the initial one-step sharpness reduction in $\mu$P networks observed in Figure 1(f) is likely due to finite width effects.

We note that this distinction extends beyond the choice between NTP/SP and $\mu$P; the initialization scale $\sigma_w$ also influences the initial conditions, potentially leading to different sharpness trajectories. In Figure 1 (a, d), we deliberately chose a small enough $\sigma_w$ value with SP that exhibits all four training regimes, including the initial sharpness reduction. If we consider an even smaller value of $\sigma_w$, then training exhibits an even smaller sharpness reduction, and the dynamics aligns more with the $\mu$P case.

These outcomes align with our UV model's predictions.

Section 6.2 suggests that real-world deep neural networks (sometimes) undergo an interim decrease in $\lambda$...

We request the reviewer to clarify the following concern: ``Section 6.2 suggests that real-world deep neural networks (sometimes) undergo an interim decrease in $\lambda$ and the weight norm, but the real-world models presented in Fig. 1 and App. E exhibit a decrease in weight-norm.''

Regarding initialization in regime R1, can the reviewer also clarify what they imply by " initialization in R1 can exist for real-world networks is limited,"?

While we await your clarification, we would like to take the opportunity to clarify the initialization in the R1 region for NTP/SP and $\mu$P models in a real-world setting. NTP/SP models are unlikely to be initialized in the R1 regime in real-world scenarios because in high dimensions the initialized function will be concentrated around the variance for a zero-mean initialization with high probability. In comparison, the initial variance of the $\mu$P models scales as $1/n$, and for sufficiently large widths, they can be initialized in the R1 regime. The UV model does not capture these high-dimensional cases which is a limitation.

We thank the reviewer for their careful reading and detailed comments.

... The UV model studied in the paper is a toy model and is only a slight generalization of the well-known model from Lewkowycz et al. (2020)

While we use a similar model compared to Lewkowycz et al. (2020), our analysis differs in two key aspects:

1. Our work recognizes the crucial role of non-zero target $y$ and parameterization. Lewkowycz et al. (2020) set $y = 0$ (Eqs. 6-7 of their paper), which results in a model that cannot exhibit progressive sharpening (and consequently EoS) for any initialization or parameterization. In comparison, we observe progressive sharpening and EoS, precisely because of non-zero target $y$ and $\mu$P.

2. Furthermore, their analysis relies on specific approximations to simplify the dynamical equations to analyze catapult dynamics, particularly leveraging the existence of $\mathcal{O}(1/n)$ terms that can be ignored in the NTK regime. This approach cannot be generalized to $\mu$P ($n_{\text{eff}} = 1$) where such approximations are infeasible. In contrast, our generalized UV model dynamics comprehensively captures different sharpness phenomena across parameterizations.

All the regimes discussed in the paper - the catapult mechanism, edge of stability, etc. - have already been demonstrated repeatedly in very similar models.

While we acknowledge that certain training regimes (such as catapult dynamics, progressive sharpening, EoS, etc.) have been demonstrated in prior works, our study is unique in providing a unified analysis encompassing all four regimes (especially the sharpness reduction regime) within a single model, which does not exist in the literature. Additionally, our work analyzes practical conditions under which different sharpness phenomena, such as early sharpness reduction, progressive sharpening, and EoS, are not observed.

Where exactly is this shown in Yang, Hu (2021)? Table 1 and Figure 2 there actually show SP and NTP as different parameterizations

The $abc$ parameterization introduced by Yang & Hu (2021) exhibits a symmetry (Section 3.2). For any scalar $\theta \in \mathbb{R}$, the optimization trajectory remains invariant under the following simultaneous transformations: $a^l \rightarrow a^l + \theta$, $b^l \rightarrow b^l - \theta$, and $c^l \rightarrow c^l - 2\theta$. Starting from Standard Parameterization, if we consider $\theta = 1/2$, then from Table 1, we observe that NTP and SP are equivalent except for the first layer. This first layer difference arises because in this table, input dimension $d_{\text{in}}$ is considered to be $\mathcal{O}(1)$ wrt width. If we consider 'constant width' networks ($d_{\text{in}} = \mathcal{O}(n)$), then NTP and SP are equivalent parameterizations. We acknowledge that the phrase 'constant width' is confusing. Accordingly, we have clarified this in the updated version of the manuscript. The difference between NTP and SP in Fig. 2 of Yang & Hu (2021) perhaps arises from the difference in the first layer.