Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory

Summary

This paper discovers and studies a new phenomenon that occurs when training neural nets on 1 data point with large enough learning rates. The phenomenon concerns the end-stage evolution of the tuple (residual, sharpness), where the residual (defined as "prediction - target") is a scalar because there is only 1 data point. If one fixes the sharpness and considers the evolution of the residual, and if one linearizes the network, then the dynamical system that governs the evolution of the residual has a stable period-2 orbit whose value depends on the fixed sharpness, i.e. residual $\to \pm$ function(sharpness). The main finding of the paper is that the residual and sharpness rapidly evolve onto the curve "residual $=\pm$ function(sharpness)" and then remain on this curve for the end stage of training. This finding is made empirically on some real architectures and is also proved rigorously in two highly simplified settings. This is interesting because it is a richer characterization of the training dynamics than has been previously known (previous knowledge is limited to the behavior of the sharpness and says nothing about the residual).

Strengths

- I think the identification of the trajectory alignment phenomenon is a solid contribution. - It's great that the authors can rigorously establish this phenomenon in several simplified settings

Weaknesses

- The main weakness is that the trajectory alignment phenomenon seems mostly limited to the setting of training on one data point. The authors do briefly try to extend it to the setting of multiple data points, but the experimental evidence is not highly convincing (especially the $\ell_\infty$ plot) and I do not really see why the phenomenon _would_ be true. For example, I do not see why the generalized sharpness metric (which is kind of like trace(Hessian) / n) would be the right generalization. - Another weakness is that the EOS setting seems different than the one from Cohen et al '21. In that paper, training with gradient flow makes the sharpness continually increase, whereas my understanding is that in this paper, training with gradient flow does not make the sharpness continually increase. Is my understanding correct? If so, I would prefer if the paper were clearer about the differences from Cohen et al '21, so as not to add confusion to the literature.

Questions

- Could the authors please cite these two papers in the related work section? https://arxiv.org/abs/1807.05031, https://arxiv.org/abs/2002.09572. These papers gave evidence prior to Cohen et al '21, that the step size affects the sharpness along the optimization trajectory. - I wonder if the reason why the trajectory alignment phenomenon occurs on wide networks but not narrow ones is that the linearization of the network is more accurate for wide networks. Therefore, I wonder if one could also see trajectory alignment on narrow nets if one adopted the scaling trick from Chizat and Bach (https://arxiv.org/abs/1812.07956) which makes any network closer to its linearization. - I believe Lemma 2.1(ii) has a typo, it should be r_hat(q)

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Please see the second point in the 'weaknesses' section.

Summary

This paper studies the Edge of Stability phenomenon in training neural networks. The authors demonstrate that, on some reparameterization of the loss function, for a small enough stepsize, the trajectories of gradient descent with different initializations align on a bifurcation diagram. They show this for two classes of functions: (1) Two-layer fully connected linear neural networks trained on a single data point with a convex and Lipschitz loss. (2) Single-neuron nonlinear network trained on a single data point with the squared loss.

Strengths

1. The analysis presented by the authors is novel, and it improves upon prior work in some settings (e.g. tightening the gap on sharpness upper and lower bounds from $\mathcal{O}(\eta)$ to $\mathcal{O}(\eta^3)$ for two-layer linear neural networks). 2. The paper imports a lot of ideas from the literature on control theory to studying the optimization of neural nets. This might be useful more generally.

Weaknesses

1. Why is the whole "trajectory alignment" phenomenon is interesting in the first place? What new things does it tell us? 2. It seems very obvious that for high-dimensional, deep networks, gradient descent doesn't converge to the same minimizers for different initializations. Is this whole phenomenon a product of small $d$ and toy examples? You fix the Xavier initialization for all experiments, and as far as I can tell the widest network used has width 256. This is a far cry from modern neural networks. 3. The conditions you have on the initialization in Thm 4.2 are very restrictive. Do they hold in practice?

Questions

See the weaknesses section.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

N/A.

Summary

This paper investigates a trajectory alignment phenomenon during the edge of stability (EoS) regime in gradient descent (GD): the trajectories from different GD dynamics tend to align on a bifurcation diagram. The paper first empirically observes the trajectory alignment in deep neural network training in a canonically reparametrized space that can reveal the bifurcation. With a problem setting of a two-layer linear neural network and single neuron nonlinear network training on a single data point, the alignment is theoretically analyzed, explaining the GD trajectories in the gradient flow regime, two phases of the EoS regime, and the progressive sharpening.

Strengths

- Based on simple models, this paper clearly characterizes the gradient flow regime, two phases of the EoS regime, and progressive sharpening using the bifurcation theory. This seems a novel and original way to analyze the phenomenon. The introduction of canonical reparametrization greatly helps capture the crux of the EoS phenomenon. This work will provide useful insight to study the dynamics in the EoS regime of neural network training, which is recently getting a huge interest in the community. - This paper is in general very well written. The theories seem to be developed in a mathematically sound manner.

Weaknesses

- The claim that the trajectory alignment on a bifurcation diagram is independent of network architecture and training data seems somewhat excessive. It does not corroborate well with some experimental results in the main text and appendix (see Figures 5, 10, 13, 14, 15), which shows that the degree of alignment depends greatly on the width size and the number of data. Also, the analyses are only performed for specific network architecture and loss function combinations. - The theory does not cover some widely used losses, e.g., the cross-entropy loss. Is the loss function necessary to be symmetric for this alignment to occur?

Questions

In addition to the questions posed in the **Weaknesses** section, I have the following questions: - Is any further analysis possible for the period-4 (or more) oscillation shown in Figure 1 (b) and (c)? - Can the EoS for GD with momentum also be explained by this theory? - How do the experimental conditions differ for each color in Figure 2? - In Figure 5 of the appendix, why there seems to be more alignment to the bifurcation diagram for the mean and $L_\infty$-norm cases than the $L_1$, $L_2$-norm cases? - What happens if $q_0 < \frac{1}{2}$ in line 221? - How is the sharpness anticipation for progressive sharpening in **Theorem 4.4** accurate for neural networks larger than those in Figure 1 (a)? - Can the progressive sharpening also be interpreted in the bifurcation diagram?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Some limitations of the current results are discussed in **Conclusion**. For other possible limitations, please refer to the weaknesses and questions above.

Summary

This paper studies the edge of stability phenomenon and proposes a "trajectory alignment" phenomenon, which states that with different learning rates and different initializations, if one plots the residuals and the inverse squared gradient norms, these lie on a fixed curve. The paper supports this theory with experiments on small networks with log-cosh loss and rigorously proves it for a two-layer linear net and a single neuron nonlinear net, both trained on a single datapoint.

Strengths

- While the bifurcation analysis of GD at EoS is not new, this paper takes this analysis further than previous papers and rigorously characterizes the limiting curve for toy models trained on a single datapoint. - The paper also provides fine-grained analyses for each of the toy cases it considers including precise predictions for the limiting sharpness which improves over previous work.

Weaknesses

- The abstract and the introduction state that GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram. However, the (p,q) variables in the paper are only a reparameterization in the toy case when $\theta$ has dimension 2. In the general case (e.g. in section 3) it is not fair to call this a reparameterization as the variables (p,q) do not determine the GD trajectory. - The abstract claims that the paper establishes the trajectory alignment phenomenon through empirical studies. However, the experiments in the paper are extremely limited. The experiments in figure 3 use a single datapoint equal to $e_1 \in \mathbb{R}^{10}$ and the non-standard log-cosh loss function, while the experiments in figure 4 use 10 datapoints sampled from a Gaussian distribution in $\mathbb{R}^{10}$. While these settings are certainly more difficult to analyze than the settings in the rest of the paper, they are still very much toy settings. It would be interesting to see if similar phenomena occur on more realistic datasets/models with more realistic loss functions (e.g. square loss or cross entropy loss). - The paper states a couple of times that the bifurcation diagram is determined solely by the loss function, independent of the network architecture, training data, and step size. However, this strong claim is not demonstrated outside of the toy setting of a single datapoint. Figure 4 does suggest that in the general case it may be independent of the learning rate but there is no evidence for the full claim. - It seems that all of the results in terms of the "canonical parameterization" in eq (2) are specific to assumption 2.4. In particular, the results in section 5 for square loss require defining a different parameterization (definition 5.2) which challenges the generality of the results in the rest of the paper. As the original EoS phenomenon as described in Cohen et al. was observed for square/logistic/cross entropy losses, it is unclear whether the phenomenon being studied is the same as the phenomenon observed in Cohen et al. Minor Comments: - There seems to be a clear separation of colors in Figures 2, 3, 4 but as far as I can tell these colors are never defined. Also, increasing the size of the datapoints might help with visibility, especially in Figure 3. - When discussing previous EoS papers with bifurcation diagrams, the paper should cite Chen et al. "Beyond the Edge of Stability via Two-step Gradient Updates" (it is in the list of references but not in the related work)

Questions

- Why did the authors chose to parameterize $q$ in terms of the inverse squared gradient norm instead of simply the squared gradient norm? At first glance, the reciprocal $1/q$ seems like an easier object to work with? - How is the analysis in section 5 for square loss different from the analyses in the rest of the paper?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have adequately addressed the limitations of this work.

multiple data points

It may be worth trying to see whether the trajectory alignment holds for multiple datapoints if you consider (1) the largest eigenvalue of the NTK, and (2) the projection of the residuals onto the largest eigenvector of the NTK?