SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The effect of initialization scale. In Figures 2a and 2b, we examine the effect of the initialization scale  $\alpha$  on GD trajectories in a two-layer fully-connected linear network with a width of  $m = 256$ . In Figure 2a, when the weights are initialized with a smaller scale ( $\alpha = 5$ ), the initial value of  $q$  is greater than 1, and it converges towards the minimum with only a small change in  $q_{t}$  until convergence. In this case, the limiting sharpness is relatively smaller than  $2 / \eta$ , and the EoS phenomenon does not occur. This case is referred to as the gradient flow regime (Ahn et al., 2023). On the other hand, in Figure 2b, when the weights are initialized with a larger scale ( $\alpha = 10$ ), the initial value of  $q$  is less than 1, and we observe convergence towards the point (close to)  $(p,q) = (0,1)$ . This case is referred to as the EoS regime. We note that choosing larger-than-standard scale  $\alpha$  is not a necessity for observing EoS; we note that even with  $\alpha = 1$ , we observe the EoS regime when  $\eta$  is larger.
Trajectory alignment on the bifurcation diagram. In order to investigate the trajectory alignment phenomenon on the bifurcation diagram, we plot the bifurcation diagram  $q = r(p) = \frac{\ell'(p)}{p}$  and observe that GD trajectories tend to align with this curve, which depends solely on  $\ell$ . Figure 2b clearly demonstrates this alignment phenomenon. Additionally, we analyze the evolution of  $\frac{q_t}{r(p_t)}$  and  $p_t$  in Figure 2c. We observe that the evolution of  $\frac{q_t}{r(p_t)}$  follows two phases. In Phase I,  $\frac{q_t}{r(p_t)}$  approaches to 1 quickly. In Phase II, the ratio remains close to 1. Notably, the convergence speed of  $\frac{q_t}{r(p_t)}$  towards 1 is much faster than the convergence speed of  $p_t$  towards 0. In Sections 4 and 5, we will provide a rigorous analysis of this behavior, focusing on the separation between Phase I and Phase II.
The effect of width and depth. In Figure 3, we present the GD trajectories of tanh-activated networks with different widths and depths  $(\alpha = 5)$ . All three cases belong to the EoS regime, where GD converges to a point close to  $(p,q) = (0,1)$ , resulting in a limiting sharpness near  $2 / \eta$ . However, when comparing Figures 3a and 3b, we observe that the trajectory alignment phenomenon is not
(a)  $P(z) = \frac{1}{n}\sum_{i = 1}^{n}z_{i}$
(b)  $P(z) = \| z\| _2$
(c) Loss and  $q$  vs.  $t$
Figure 4: Training on a subset of CIFAR-10. The GD trajectories trained on a 50 example subset of CIFAR-10 under the generalized canonical reparameterization (5) with different choices of  $P$  and network architecture. Each color corresponds to a single run of GD. Top row: fully-connected neural network with tanh activation. Bottom row: CNN with tanh activation. Full implementation details and further experimental results are given in Appendix A.4.
observed for the narrower network with  $m = 64$ , whereas the GD trajectories for the wider network with  $m = 256$  are clearly aligned on the bifurcation diagram. This suggests that network width plays a role in the trajectory alignment phenomenon, which is reasonable since wide networks are well approximated by their linearized models, hence Eq. (3) is more accurate. Furthermore, we note that the trajectory alignment phenomenon is also observed for a deeper network with  $L = 10$ , as depicted in Figure 3c.
Multiple training data points. In our trajectory alignment analysis, we have primarily focused on training with a single data point. However, it is important to explore the extension of this phenomenon to scenarios with multiple training data points.
To investigate this, we train a neural network on a dataset  $\{(\pmb{x}_i,y_i)\}_{i = 1}^n$ , where  $\pmb{x}_i\in \mathbb{R}^d$  and  $y_{i}\in \mathbb{R}$  by minimizing the empirical risk  $\mathcal{L}(\Theta)\coloneqq \frac{1}{n}\sum_{i = 1}^{n}\ell (f(\pmb{x}_i;\Theta) - y_i)$ . Defining  $\pmb {X}\in \mathbb{R}^{n\times d}$  as the data matrix and  $\pmb {y}\in \mathbb{R}^n$  as the label vector, we introduce a generalized canonical reparameterization:
$$
(p, q) := \left(P (f (\boldsymbol {X}; \boldsymbol {\Theta}) - \boldsymbol {y}), \frac {2 n}{\eta \| \sum_ {i = 1} ^ {n} (\nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x} _ {i} ; \boldsymbol {\Theta})) ^ {\otimes 2} \| _ {2}}\right), \tag {5}
$$
where  $P:\mathbb{R}^n\to \mathbb{R}$  can be a function such as a mean value or a specific vector norm.
In Figure 4, we consider training on a 50 example subset of CIFAR-10 with only 2 classes and vary the network architecture. We use three-layer fully-connected network with tanh activation and convolutional network (CNN) with tanh activation. Under the generalized canonical reparameterization (5) for various choices of  $P$ , including the mean and the  $\ell_2$  norm, we observe the trajectory alignment phenomenon throughout all settings, indicating a common alignment property of the GD trajectories. However, unlike the single data point case, the alignment does not happen on the curve  $q = \frac{\ell'(p)}{p}$ . The precise characterization of the coinciding curve is an interesting direction for future research.
4 Trajectory Alignment of GD: A Theoretical Study
In this section, we study a two-layer fully-connected linear network defined as  $f(\pmb{x}; \Theta) \coloneqq \pmb{v}^{\top} \pmb{U} \pmb{x}$ , where  $\pmb{U} \in \mathbb{R}^{d \times m}$ ,  $\pmb{v} \in \mathbb{R}^{m}$ , and  $\Theta$  denote the collection of all parameters  $(\pmb{U}, \pmb{v})$ . We consider training this network with a single data point  $\{(x, 0)\}$ , where  $x \in \mathbb{R}^{d}$  and  $\| x \|_{2} = 1$ . We run GD with step size  $\eta$  on the empirical risk
$$
\mathcal {L} (\boldsymbol {\Theta}) := \ell (f (\boldsymbol {x}; \boldsymbol {\Theta}) - 0) = \ell (\boldsymbol {v} ^ {\top} \boldsymbol {U} \boldsymbol {x}),
$$
where  $\ell$  is a loss function satisfying Assumption 4.1. We note that our assumptions on  $\ell$  is motivated from the single-neuron linear network analysis ( $d = m = 1$ ) by Ahn et al. (2023).
Assumption 4.1. The loss  $\ell$  is convex, even, 1-Lipschitz, and twice differentiable with  $\ell''(0) = 1$ .
The canonical reparameterization (Definition 2.1) of  $\Theta = (\pmb{U},\pmb{v})$  is given by
$$
(p, q) := \left(\boldsymbol {v} ^ {\top} \boldsymbol {U} \boldsymbol {x}, \frac {2}{\eta \left(\| \boldsymbol {U} \boldsymbol {x} \| _ {2} ^ {2} + \| \boldsymbol {v} \| _ {2} ^ {2}\right)}\right).
$$
Under the canonical reparameterization, the 1-step update rule of GD can be written as
$$
p _ {t + 1} = \left[ 1 - \frac {2 r \left(p _ {t}\right)}{q _ {t}} + \eta^ {2} p _ {t} ^ {2} r \left(p _ {t}\right) ^ {2} \right] p _ {t}, \quad q _ {t + 1} = \left[ 1 - \eta^ {2} p _ {t} ^ {2} r \left(p _ {t}\right) \left(2 q _ {t} - r \left(p _ {t}\right)\right) \right] ^ {- 1} q _ {t}, \tag {6}
$$
where we define the function  $r$  by  $r(p) \coloneqq \frac{\ell'(p)}{p}$  for  $p \neq 0$  and  $r(0) \coloneqq 1$ . Note that the sequence  $(q_t)_{t=0}^{\infty}$  is monotonically increasing if  $q_0 \geq \frac{1}{2}$ , which is the case our analysis will focus on.
We have an additional assumption on  $\ell$  as below, motivated from Lemma 2.1.
Assumption 4.2. The function  $r(p) = \frac{\ell'(p)}{p}$  corresponding to the loss  $\ell$  satisfies Assumption 2.4.
We now present our theoretical results on this setting, and defer the proofs to Appendix B.
4.1 Gradient flow regime
We first consider the gradient flow regime, where  $q$  is initialized with  $q_0 > 1$ .
Theorem 4.1 (gradient flow regime). Let  $\eta \in (0, \frac{2}{33})$  be a fixed step size and  $\ell$  be a loss function satisfying Assumptions 4.1 and 4.2. Suppose that the initialization  $(p_0, q_0)$  satisfies  $|p_0| \leq 1$  and  $q_0 \in \left(\frac{2}{2 - \eta}, \min\left\{\frac{1}{16\eta}, \frac{r(1)}{2\eta}\right\}\right)$ . Consider the GD trajectory characterized in Eq. (6). Then, the GD iterations  $(p_t, q_t)$  converge to the point  $(0, q^*)$  such that
$$
q _ {0} \leq q ^ {*} \leq \exp (C \eta^ {2}) q _ {0} \leq 2 q _ {0},
$$
where  $C = 8q_0\left[\min \left\{\frac{2(q_0 - 1)}{q_0},\frac{r(1)}{2q_0}\right\} \right]^{-1} > 0.$
Theorem 4.1 implies that in gradient flow regime, GD with initialization  $\Theta_0 = (U_0, v_0)$  and step size  $\eta$  converges to  $\Theta^*$  which has the sharpness bounded by:
$$
(1 - C \eta^ {2}) (\| \boldsymbol {U} _ {0} \boldsymbol {x} \| _ {2} ^ {2} + \| \boldsymbol {v} _ {0} \| _ {2} ^ {2}) \leq \lambda_ {\max } (\boldsymbol {\Theta} ^ {*}) \leq (\| \boldsymbol {U} _ {0} \boldsymbol {x} \| _ {2} ^ {2} + \| \boldsymbol {v} _ {0} \| _ {2} ^ {2}).
$$
Hence, for small step size  $\eta$ , if the initialization satisfies  $\|U_0\pmb{x}\|_2^2 + \|v_0\|_2^2 < \frac{2}{\eta} - 1$ , then the limiting sharpness is slightly below  $\|U_0\pmb{x}\|_2^2 + \|v_0\|_2^2$ . Note that we assumed the bound  $|p_0| \leq 1$  for simplicity, but our proof also works with the assumption  $|p_0| \leq K$  for any positive constant  $K$  modulo some changes in numerical constants. Moreover, our assumption on the upper bound of  $q_0$  is  $1/\eta$  up to a constant factor, which covers most realistic choices of initialization.
4.2 EoS regime
We now provide rigorous results in the EoS regime, where the GD trajectory aligns on the bifurcation diagram  $q = r(p)$ . To establish these results, we introduce additional assumptions on the loss  $\ell$ .
Assumption 4.3. The function  $r(z) = \frac{\ell'(z)}{z}$  is  $C^4$  on  $\mathbb{R}$  and satisfies
(i)  $z\mapsto \frac{r'(z)}{r(z)^2}$  is decreasing on  $\mathbb{R}$
(ii)  $z \mapsto \frac{zr'(z)}{r(z)}$  is decreasing on  $z > 0$  and increasing on  $z < 0$ ,
(iii)  $z \mapsto \frac{zr(z)}{r'(z)}$  is decreasing on  $z > 0$  and increasing on  $z < 0$ .
We note that both the log-cosh loss and the square root loss satisfy Assumptions 4.1, 4.2, and 4.3.
Theorem 4.2 (EoS regime, Phase I). Let  $\eta$  be a small enough step size and  $\ell$  be a loss function satisfying Assumptions 4.1, 4.2, and 4.3. Let  $z_0 \coloneqq \sup_z \left\{ \frac{z r'(z)}{r(z)} \geq -\frac{1}{2} \right\}$  and  $c_0 \coloneqq \max \{ r(z_0), \frac{1}{2} \}$ . Let  $\delta \in (0, 1 - c_0)$  be any given constant. Suppose that the initialization  $(p_0, q_0)$  satisfies  $|p_0| \leq 1$  and  $q_0 \in (c_0, 1 - \delta)$ . Consider the reparameterized GD trajectory characterized in Eq. (6). We assume that for all  $t \geq 0$  such that  $q_t < 1$ , we have  $p_t \neq 0$ . Then, there exists a time step  $t_a = \mathcal{O}_{\delta, \ell}(\log (\eta^{-1}))$ , such that for any  $t \geq t_a$ ,
$$
\frac {q _ {t}}{r (p _ {t})} = 1 + h (p _ {t}) \eta^ {2} + \mathcal {O} _ {\delta , \ell} (\eta^ {4}),
$$
where  $h(p) \coloneqq -\frac{1}{2}\left(\frac{pr(p)^3}{r'(p)} + p^2 r(p)^2\right)$  for  $p \neq 0$  and  $h(p) \coloneqq -\frac{1}{2r''(0)}$  for  $p = 0$ .
One can check that for log-cosh and square-root losses, the ranges of  $h$  are  $(0,3/4]$  and  $(0,1/2]$ , respectively. Theorem 4.2 implies that in the early phase of training  $(t \leq t_a = \mathcal{O}(\log(\eta^{-1})))$ , GD iterates  $(p_t, q_t)$  approach closely to the bifurcation diagram  $r(p) = q$ , which we called Phase I in Section 3. In Phase II, GD trajectory aligns on this curve in the remaining of the training  $(t \geq t_a)$ . Theorem 4.3 provides an analysis on Phase II stated as below.
Theorem 4.3 (EoS regime, Phase II). Under the same settings as in Theorem 4.2, there exists a time step  $t_b = \Omega((1 - q_0)\eta^{-2})$  such that  $q_{t_b} \leq 1$  and  $q_t > 1$  for any  $t > t_b$ . Moreover, the GD iterates  $(p_t, q_t)$  converge to the point  $(0, q^*)$  such that
$$
q ^ {*} = 1 - \frac {\eta^ {2}}{2 r ^ {\prime \prime} (0)} + \mathcal {O} _ {\delta , \ell} (\eta^ {4}).
$$
Theorem 4.3 implies that in EoS regime, GD with step size  $\eta$  converges to  $\Theta^{*}$  with sharpness
$$
\lambda_ {\max} (\boldsymbol {\Theta} ^ {*}) = \frac {2}{\eta} - \frac {\eta}{| r ^ {\prime \prime} (0) |} + \mathcal {O} _ {\delta , \ell} (\eta^ {3}).
$$
Note that Ahn et al. (2023) study the special case  $d = m = 1$  and prove that the limiting sharpness is between  $2 / \eta -\mathcal{O}(\eta)$  and  $2 / \eta$ . Theorem 4.3 provides tighter analysis on the limiting sharpness in more general settings, reducing the gap between the upper bound and lower bound to only  $\mathcal{O}(\eta^3)$ . Also, our result is the first to prove that the limiting sharpness in the EoS regime is bounded away from  $2 / \eta$  by a nontrivial margin.
We also study the evolution of sharpness along the GD trajectory and prove that progressive sharpening (i.e., sharpness increases) occurs during Phase II.
Theorem 4.4 (progressive sharpening). Under the same setting as in Theorem 4.2, let  $t_a$  denote the obtained iteration. Define the function  $\tilde{\lambda} : \mathbb{R}_{>0} \to \mathbb{R}$  given by
$$
\tilde {\lambda} (q) := \left\{ \begin{array}{l l} \left(1 + \frac {\hat {r} (q) r ^ {\prime} (\hat {r} (q))}{q}\right) \frac {2}{\eta} & \text {i f q \leq 1 , a n d} \\ \frac {2}{\eta} & \text {o t h e r w i s e .} \end{array} \right.
$$
Then, the sequence  $\left(\tilde{\lambda}(q_t)\right)_{t=0}^{\infty}$  is monotonically increasing. Moreover, for any  $t \geq t_a$ , the sharpness at  $GD$  iterate  $\Theta_t$  closely follows the sequence  $\left(\tilde{\lambda}(q_t)\right)_{t=0}^{\infty}$  by satisfying
$$
\left| \lambda_ {\max} (\boldsymbol {\Theta} _ {t}) - \tilde {\lambda} (q _ {t}) \right| \leq 1 + \mathcal {O} _ {\ell} (\eta).
$$
The gap between  $\lambda_{\max}(\Theta_t)$  and  $\tilde{\lambda}(q_t)$  is bounded by a numerical constant, which becomes negligible compared to  $2/\eta$  for small  $\eta$ . In Figure 1a, we perform numerical experiments on a single-neuron case and observe that  $\tilde{\lambda}(q_t)$  closely approximates the sharpness.
Note that Cohen et al. (2021) observe an increase in sharpness during training in the gradient flow regime, while our work reveals that sharpness increases when exhibiting oscillations in the EoS regime. This distinction may be linked to the selection of the loss function, as our study focuses on Lipschitz convex losses, while Cohen et al. (2021) examine the squared loss.
5 EoS in Squared Loss: Single-neuron Nonlinear Network
Our canonical reparameterization has a limitation in explaining the EoS phenomenon under squared loss  $\ell(p) = \frac{1}{2}p^2$ , as the function  $r(p) = \frac{\ell'(p)}{p} = 1$  does not satisfy Assumption 2.4. However, empirical studies by Cohen et al. (2021) have observed the EoS phenomenon in GD training with squared loss. In this section, we analyze a simple toy model to gain insight into the EoS phenomenon and trajectory alignment of GD under squared loss.
We study the GD dynamics on a two-dimensional function  $\mathcal{L}(x,y)\coloneqq \frac{1}{2} (\phi (x)y)^2$ , where  $x, y$  are scalars and  $\phi$  is a nonlinear activation satisfying Assumption 5.1 below.
Assumption 5.1 (sigmoidal activation). The activation function  $\phi : \mathbb{R} \to \mathbb{R}$  is odd, increasing, 1-Lipschitz and twice continuously differentiable. Moreover,  $\phi(0) = 0$ ,  $\phi'(0) = 1$ ,  $\lim_{x \to \infty} \phi(x) = 1$ , and  $\lim_{x \to -\infty} \phi(x) = -1$ .
One good example of  $\phi$  satisfying Assumption 5.1 is tanh. For this section, we use an alternative reparameterization defined as below.
Definition 5.2. For given step size  $\eta$ , the  $(p,q)$  reparameterization of  $(x,y) \in \mathbb{R}^2$  is defined as
$$
(p, q) := \left(x, \frac {2}{\eta y ^ {2}}\right).
$$
Under the reparameterization, the 1-step update rule can be written as
$$
p _ {t + 1} = \left(1 - \frac {2 r \left(p _ {t}\right)}{q _ {t}}\right) p _ {t}, \quad q _ {t + 1} = \left(1 - \eta \phi \left(p _ {t}\right) ^ {2}\right) ^ {- 2} q _ {t}, \tag {7}
$$
where the function  $r$  is given by  $r(z) \coloneqq \frac{\phi(z)\phi'(z)}{z}$  for  $z \neq 0$  and  $r(0) \coloneqq 1$ .
We can observe a notable resemblance between Eq. (7) and Eq. (6). Indeed, our theoretical findings for a single-neuron nonlinear network closely mirror those of the two-layer linear network discussed in Section 4. Due to lack of space, we summarize our theorems in this setting as the following:
Theorem 5.1 (informal). Under suitable assumptions on  $\phi$ , step size, and initialization,  $GD$  trained on the squared loss  $\mathcal{L}(x,y) \coloneqq \frac{1}{2} (\phi(x)y)^2$  exhibits the same gradient flow,  $EoS$  (Phase I, II), and progressive sharpening phenomena as shown in Section 4.
In Theorem C.3, we prove that in the EoS regime, the limiting sharpness is  $\frac{2}{\eta} - \frac{2}{|r''(0)|} + \mathcal{O}(\eta)$ . For formal statements of the theorems and the proofs, we refer the reader to Appendix C.
In this paper, we provide empirical evidence and rigorous analysis to demonstrate the interesting phenomenon of GD trajectory alignment in the EoS regime. Importantly, we show that different GD trajectories, under the canonical reparameterization, align on a bifurcation diagram independent of initialization. This discovery is notable due to the intricate and non-convex nature of neural network optimization, where the algorithm trajectory is heavily influenced by initialization choices. Our theoretical analysis not only characterizes the behavior of limiting sharpness but also establishes progressive sharpening of GD. One immediate future direction is to understand the trajectory alignment behavior when trained on multiple data points. Lastly, it will be interesting to extend our analysis to encompass squared loss for general neural network, going beyond the toy single-neuron example.
Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-12 Jul 2020. URL https://proceedings.mlr.press/v125/woodworth20a.htm1.
Jingfeng Wu, Vladimir Braverman, and Jason D Lee. Implicit bias of gradient descent for logistic regression at the edge of stability. arXiv preprint arXiv:2305.11788, 2023.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W. Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE International Conference on Big Data (Big Data), pages 581-590, Dec 2020. doi: 10.1109/BigIntData50022.2020.9378171.
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=ZsZM-4iMQkH.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with a minimalist example. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=p7EagBsMAEO.
1.1 Our contributions 2
1.2 Related works 3
2.1 Problem settings 4
2.2 Canonical reparameterization 4
2.3 Bifurcation analysis 4
3 Trajectory Alignment of GD: An Empirical Study 6
4 Trajectory Alignment of GD: A Theoretical Study 7
4.1 Gradient flow regime 8
4.2 EoS regime 8
5 EoS in Squared Loss: Single-neuron Nonlinear Network 10
A.1 Experimental setup 15
A.2 Training on a single data point 16
A.3 Training on multiple synthetic data Points 17
A.4 Training on real-world dataset 21
B Proofs for the Two-Layer Fully-Connected Linear Network 24
B.1 Proof of Theorem 4.1 24
B.2 Proof of Theorem 4.2 25
B.3 Proof of Theorem 4.3 33
B.4 Proof of Theorem 4.4 34
C Proofs for the Single-neuron Nonlinear Network 38
C.1 Formal statements of Theorem 5.1 38
C.2 Proof of Theorem C.1 39
C.3 Proof of Theorem C.2 41
C.4 Proof of Theorem C.3 48
C.5 Proof of Theorem C.4 49
In this section, we present additional empirical evidence demonstrating the phenomenon of trajectory alignment, which supports the findings discussed in Section 3 of our main paper.
Objective function. We run gradient descent (GD) to minimize the objective function defined as
where  $\Theta$  represents the parameters,  $\ell : \mathbb{R} \to \mathbb{R}$  is a loss function,  $f: \mathbb{R}^d \to \mathbb{R}$  is a neural network, and  $\{(x,y)\}$  denotes a single data point with  $x \in \mathbb{R}^d$  and  $y \in \mathbb{R}$ . We also consider training on multiple data points  $\{(x_i, y_i)\}_{i=1}^n$  with  $x_i \in \mathbb{R}^d$  and  $y_i \in \mathbb{R}$  for each  $1 \leq i \leq n$ , where we minimize the objective function
$$
\mathcal {L} (\Theta) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (f (\boldsymbol {x} _ {i}; \Theta) - y _ {i}).
$$
In our experiments, we primarily focus on the log-cosh loss function  $\ell_{\mathrm{log - cosh}}(p) = \log (\cosh (p))$  but we also investigate the square root loss  $\ell_{\mathrm{sqrt}}(p) = \sqrt{1 + p^2}$ .
Model architecture. We train a fully-connected  $L$ -layer neural network, denoted as  $f(\cdot ;\Theta):\mathbb{R}^d\to \mathbb{R}$ . The network is defined as follows:
where  $\phi : \mathbb{R} \to \mathbb{R}$  is an activation function applied entry-wise,  $W_{1} \in \mathbb{R}^{m \times d}$ ,  $W_{l} \in \mathbb{R}^{m \times m}$  for  $2 \leq l \leq L - 1$ , and  $\boldsymbol{w}_{L} \in \mathbb{R}^{m}$ . All  $L$  layers have the same width of  $m$ , and the biases of all layers are fixed to  $0$ .<sup>4</sup> We consider three activations: hyperbolic tangent  $\phi(t) = \tanh(t)$ , exponential linear unit  $\phi(t) = \mathrm{ELU}(t)$ , and linear  $\phi(t) = t$ .
Weight initialization. We perform gradient descent (GD) using five different randomly initialized sets of weights. The weights are initialized using Xavier initialization, and each layer is multiplied by a rescaling factor (gain) of  $\alpha$ . In the plots presented throughout this section, we mark the initialization points with an 'x' to distinguish them from other points on the trajectories.
Canonical reparameterization. We plot GD trajectories after applying the canonical reparameterization introduced in Definition 2.1:
$$
(p, q) := \left(f (\boldsymbol {x}; \boldsymbol {\Theta}) - y, \frac {2}{\eta \| \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x} ; \boldsymbol {\Theta}) \| _ {2} ^ {2}}\right),
$$
where  $\eta$  denotes the step size. For training on multiple data points, we employ the generalized canonical reparameterization as defined in Eq. (5):
$$
(p, q) := \left(P (f (\pmb {X}; \pmb {\Theta}) - \pmb {y}), \frac {2 n}{\eta \| \sum_ {i = 1} ^ {n} (\nabla_ {\pmb {\Theta}} f (\pmb {x} _ {i} ; \pmb {\Theta})) ^ {\otimes 2} \| _ {2}}\right),
$$
where  $P: \mathbb{R}^n \to \mathbb{R}$  can represent the mean value or vector norms. Specifically, we mainly focus on the mean value  $P(z) = \frac{1}{n}\sum_{i=1}^{n}z_i$ , but we also examine vector norms such as  $P(z) = \| z\|_1$ ,  $P(z) = \| z\|_2$ , and  $P(z) = \| z\|_\infty$ , where  $z = (z_1,z_2,\ldots,z_n) \in \mathbb{R}^n$ .
For large networks, explicitly calculating the  $\ell_2$  matrix norm  $\left\| \sum_{i=1}^{n} (\nabla_\Theta f(\pmb{x}_i; \Theta))^{\otimes 2} \right\|_2$  is infeasible. Therefore, we adopt a fast and efficient matrix-free method based on power iteration, as proposed by Yao et al. (2020). This method allows us to numerically compute the  $\ell_2$  norm of large-scale symmetric matrices.
4While we fix the biases of all layers to 0 to maintain consistency with our theory, the trajectory alignment phenomenon is consistently observed even for neural networks with bias. In Appendix A.4, we consider training networks with bias.
Figure 5: GD trajectories (1000 iterations) of 3-layer fully connected neural networks with tanh activation and a width of  $m = 256$ . The plots depict the trajectories for different initialization scales:  $\alpha = 0.5, 1.0, 2.0, 4.0$ . Smaller initialization scales ( $\alpha = 0.5, 1.0$ ) lead to trajectories in the gradient flow regime, while larger initialization scales ( $\alpha = 2.0, 4.0$ ) result in trajectories in the EoS regime.
(a)  $\alpha = 0.5$
(b)  $\alpha = 1.0$
(c)  $\alpha = 2.0$
(d)  $\alpha = 4.0$
A.2 Training on a single data point