SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Minhak Song
KAIST ISysE/Math haksong@kaist.ac.kr
Chulhee Yun
KAIST AI chulhee.yun@kaist.ac.kr
Cohen et al. (2021) empirically study the evolution of the largest eigenvalue of the loss Hessian, also known as sharpness, along the gradient descent (GD) trajectory and observe the Edge of Stability (EoS) phenomenon. The sharpness increases at the early phase of training (referred to as progressive sharpening), and eventually saturates close to the threshold of  $2/(step\ size)$ . In this paper, we start by demonstrating through empirical studies that when the EoS phenomenon occurs, different GD trajectories (after a proper reparameterization) align on a specific bifurcation diagram independent of initialization. We then rigorously prove this trajectory alignment phenomenon for a two-layer fully-connected linear network and a single-neuron nonlinear network trained with a single data point. Our trajectory alignment analysis establishes both progressive sharpening and EoS phenomena, encompassing and extending recent findings in the literature.
It is widely believed that implicit bias or regularization of gradient-based methods plays a key role in generalization of deep learning (Vardi, 2022). There is a growing literature (Gunasekar et al., 2017, 2018; Soudry et al., 2018; Arora et al., 2018, 2019; Ji and Telgarsky, 2019a; Woodworth et al., 2020; Chizat and Bach, 2020; Yun et al., 2021) studying how the choice of optimization methods induces an implicit bias towards specific solutions among the many global minima in overparameterized settings.
Cohen et al. (2021) identify a surprising implicit bias of gradient descent (GD) towards global minima with a certain sharpness<sup>1</sup> value depending on the step size  $\eta$ . Specifically, for reasonable choices of  $\eta$ , (a) the sharpness of the loss at the GD iterate gradually increases throughout training until it reaches the stability threshold<sup>2</sup> of  $2 / \eta$  (known as progressive sharpening), and then (b) the sharpness saturates close to or above the threshold for the remainder of training (known as Edge of Stability (EoS)). These findings have sparked a surge of research aimed at developing a theoretical understanding of the progressive sharpening and EoS phenomena (Arora et al., 2022; Lyu et al., 2022; Wang et al., 2022; Ahn et al., 2023; Chen and Bruna, 2023; Zhu et al., 2023). In this paper, we study these phenomena through the lens of bifurcation theory, both empirically and theoretically.
Motivating observations: Figure 1 illustrates the GD trajectories with different initializations and fixed step sizes trained on three types of two-dimensional functions: (a)  $\log (\cosh (xy))$ , (b)  $\frac{1}{2} (\tanh (x)y)^2$ , and (c)  $\frac{1}{2} (\mathrm{ELU}(x)y)^2$ , where  $x$  and  $y$  are scalars. The functions  $\mathcal{L}:\mathbb{R}^2\to \mathbb{R}$  have sharpness  $y^{2}$  at the global minimum  $(0,y)$  for all three models. These toy models can be viewed as examples of single-neuron models, where (a) represents a linear network with log-cosh loss, while (b) and (c) represent nonlinear networks with squared loss. These simple models can capture some interesting aspects of neural network training in the EoS regime, which are summarized below:
1Throughout this paper, the term "sharpness" means the maximum eigenvalue of the training loss Hessian.
For quadratic loss, GD becomes unstable if the sharpness is larger than a threshold of 2/(step size).
(a)  $\log (\cosh (xy))$
(b)  $\frac{1}{2} (\tanh (x)y)^2$
(c)  $\frac{1}{2} (\mathrm{ELU}(x)y)^2$
Figure 1: GD trajectories align on the bifurcation diagram at the Edge of Stability. We run GD on toy models with step size  $\eta = 2/25$  in (a), and  $\eta = 2/400$  in (b) and (c). Distinct colors indicate independent runs of GD with varying initializations. Top row: GD trajectories closely follow the bifurcation diagram of the map  $x \mapsto x - \eta \frac{\partial}{\partial x} \mathcal{L}(x, y)$  and asymptotically reaches the bifurcation point  $(0, \sqrt{2/\eta})$ . Bottom row: the sharpness reaches  $2/\eta$ , and  $x_{t}$  typically shows 2-period oscillating dynamics. The theoretical prediction  $\tilde{\lambda}(q_{t})$  (dashed lines, defined in Theorems 4.4 and C.4) in (a) and (b) approximates the sharpness along the GD trajectory and demonstrates progressive sharpening.
- EoS phenomenon: GD converges to a global minimum near the point  $(0, \sqrt{2 / \eta})$  with sharpness close to  $2 / \eta$ . During the convergence phase, the training dynamics exhibit period-2 oscillations.
- For different initializations, GD trajectories for a given step size align on the same curve. For example, Figure 1a shows that GD trajectories with different initializations closely follow a specific U-shaped curve until convergence. We call this phenomenon trajectory alignment.
- In Figures 1b and 1c, GD trajectories are aligned on a curve with a fractal structure that qualitatively resembles the bifurcation diagram of a typical polynomial map, such as the logistic map. Particularly, Figure 1c demonstrates a period-halving phase transition in the GD dynamics, shifting from period-4 oscillation to period-2 oscillation.
- Surprisingly, the curve that GD trajectories approach and follow coincides with the bifurcation diagram of a one-dimensional map  $x \mapsto x - \eta \frac{\partial}{\partial x} \mathcal{L}(x,y)$  with a fixed "control parameter"  $y$ . The stability of its fixed point  $x = 0$  changes at the bifurcation point  $(x,y) = (0,\sqrt{2 / \eta})$ , where period-doubling bifurcation occurs. Note that this point is a global minimum with sharpness  $2 / \eta$ .
Interestingly, such striking behaviors can also be observed in more complex models, up to a proper reparameterization, as we outline in the next subsection.
In this paper, we discover and study the trajectory alignment behavior of (reparameterized) GD dynamics in the EoS regime. To our best knowledge, we are the first to identify such an alignment with a specific bifurcation diagram independent of initialization. Our empirical findings are rigorously proven for both two-layer fully-connected networks and single-neuron nonlinear networks. Our main contributions are summarized below:
- In Section 2, we introduce a novel canonical reparameterization of training parameters, which incorporates the data, network, and GD step size. This reparameterization allows us to study the trajectory alignment phenomenon in a unified framework. Through empirical study, Section 3 demonstrates that the alignment property of GD trajectories is not limited to toy models but also occurs in wide and deep networks trained on real-world dataset. Furthermore, we find that the alignment trend becomes more pronounced as the network width increases.
- In Section 4, we use our canonical reparameterization to establish the trajectory alignment phenomenon for a two-layer fully-connected linear network trained with a single data point. Our theoretical analysis rigorously proves both progressive sharpening and the EoS phenomenon, extending the work of Ahn et al. (2023) to a much broader class of networks and also providing more accurate bounds on the limiting sharpness.
- Our empirical and theoretical analyses up to Section 4 are applicable to convex Lipschitz losses, hence missing the popular squared loss. In Section 5, we take a step towards handling the squared loss. Employing an alternative reparameterization, we prove the same set of theorems as Section 4 for a single-neuron nonlinear network trained with a single data point under squared loss.
The Edge of Stability (EoS) phenomenon has been extensively studied in recent years, with many works seeking to provide a deeper understanding of the evolution of sharpness and the oscillating dynamics of GD. Jastrzebski et al. (2019) and Jastrzebski et al. (2020) observe that step size affects the sharpness along the optimization trajectory. Cohen et al. (2021) first formalize EoS through empirical study, and subsequent works have built on their findings. Ahn et al. (2022) analyze EoS through experiments and identify the relations between the behavior of loss, iterates, and sharpness. Ma et al. (2022) suggest that subquadratic growth of the loss landscape is the key factor of oscillating dynamics. Arora et al. (2022) show that (normalized) GD enters the EoS regime, by verifying the convergence to some limiting flow on the manifold of global minimizers. Wang et al. (2022) divide GD trajectory into four phases and explain progressive sharpening and EoS by using the norm of output layer weight as an indicator of sharpness. Lyu et al. (2022) prove that normalization layers encourage GD to reduce sharpness. Damian et al. (2023) use the third-order Taylor approximation of the loss to theoretically analyze EoS, assuming the existence of progressive sharpening. Lee and Jang (2023) propose a new sharpness measure using batch gradient distribution and characterize EoS for SGD. Concurrent to our work, Wu et al. (2023) study the logistic regression problem with separable dataset and establish that GD exhibits an implicit bias toward the max-margin solution in the EoS regime, extending prior findings in the small step size regime (Soudry et al., 2018; Ji and Telgarsky, 2019b).
Some recent works rigorously analyze the full GD dynamics for some toy cases and prove that the limiting sharpness is close to  $2 / \eta$ . Zhu et al. (2023) study the loss  $(x,y)\mapsto \frac{1}{4} (x^{2}y^{2} - 1)^{2}$  and prove that the sharpness converges close to  $2 / \eta$  with a local convergence guarantee. Notably, Ahn et al. (2023) study the function  $(x,y)\mapsto \ell (xy)$  where  $\ell$  is convex, even, and Lipschitz, and provide a global convergence guarantee. The authors prove that when  $\ell$  is log-cosh loss or square root loss, the limiting sharpness in the EoS regime is between  $2 / \eta -\mathcal{O}(\eta)$  and  $2 / \eta$ . Our theoretical results extend their results on a single-neuron linear network to a two-layer fully-connected linear network and provide an improved characterization on the limiting sharpness, tightening the gap between upper and lower bounds to only  $\mathcal{O}(\eta^3)$ .
The trajectory alignment phenomenon is closely related to Zhu et al. (2023) which shows empirical evidence of bifurcation-like oscillation in deep neural networks trained on real-world data. However, their empirical results do not show the alignment property of GD trajectory. In comparison, we observe that GD trajectories align on the same bifurcation diagram, independent of initialization.
Very recently, Kreisler et al. (2023) observe a similar trajectory alignment phenomenon for scalar linear networks, employing a reparameterization based on the sharpness of the gradient flow solution. However, their empirical findings on trajectory alignment are confined to scalar linear networks, and do not provide a theoretical explanation. In contrast, our work employs a novel canonical reparameterization and offers empirical evidence for the alignment phenomenon across a wide range of networks. Moreover, we provide theoretical proofs for two-layer linear networks and single-neuron nonlinear networks.
Notations. For vectors  $\pmb{u}$  and  $\pmb{v}$ , we denote the  $\ell_p$  norm of  $\pmb{u}$  by  $\| \pmb{u} \|_p$ , their tensor product as  $\pmb{u} \otimes \pmb{v}$  and  $\pmb{u} \otimes \pmb{u}$  by  $\pmb{u}^{\otimes 2}$ . For a matrix  $\pmb{A}$ , we denote the spectral norm by  $\| \pmb{A} \|_2$ . Given a function  $\mathcal{L}$  and a parameter  $\Theta$ , we use  $\lambda_{\max}(\Theta) \coloneqq \lambda_{\max}(\nabla_\Theta^2 \mathcal{L}(\Theta))$  to denote the sharpness (i.e., the maximum eigenvalue of the loss Hessian) at  $\Theta$ . We use asymptotic notations with subscripts (e.g.,  $\mathcal{O}_{\ell}(\cdot)$ ,  $\mathcal{O}_{\delta, \ell}(\cdot)$ ) in order to hide constants that depend on the parameters or functions written as subscripts.
2.1 Problem settings
We study the optimization of neural network  $f(\cdot ;\Theta):\mathbb{R}^d\to \mathbb{R}$  parameterized by  $\Theta$ . We focus on a simple over-parameterized setting trained on a single data point  $\{(x,y)\}$ , where  $x\in \mathbb{R}^{d}$  and  $y\in \mathbb{R}$ . We consider the problem of minimizing the empirical risk
$$
\mathcal {L} (\boldsymbol {\Theta}) = \ell (f (\boldsymbol {x}; \boldsymbol {\Theta}) - y),
$$
where  $\ell$  is convex, even, and twice-differentiable with  $\ell''(0) = 1$ . We minimize  $\mathcal{L}$  using GD with step size  $\eta$ :  $\Theta_{t+1} = \Theta_t - \eta \nabla_\Theta \mathcal{L}(\Theta_t)$ . The gradient and the Hessian of the function are given by
$$
\begin{array}{l} \nabla_ {\Theta} \mathcal {L} (\Theta) = \ell^ {\prime} (f (\boldsymbol {x}; \Theta) - y) \nabla_ {\Theta} f (\boldsymbol {x}; \Theta), \\ \nabla_ {\Theta} ^ {2} \mathcal {L} (\Theta) = \ell^ {\prime \prime} (f (\boldsymbol {x}; \Theta) - y) \left(\nabla_ {\Theta} f (\boldsymbol {x}; \Theta)\right) ^ {\otimes 2} + \ell^ {\prime} (f (\boldsymbol {x}; \Theta) - y) \nabla_ {\Theta} ^ {2} f (\boldsymbol {x}; \Theta). \\ \end{array}
$$
Suppose that  $\Theta^{*}$  be a global minimum of  $\mathcal{L}$ , i.e.,  $f(x; \Theta^{*}) = y$ . In this case, the loss Hessian and the sharpness at  $\Theta^{*}$  are simply characterized as
$$
\nabla_ {\boldsymbol {\Theta}} ^ {2} \mathcal {L} (\boldsymbol {\Theta} ^ {*}) = \left(\nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta} ^ {*})\right) ^ {\otimes 2}, \text {a n d} \lambda_ {\max } \left(\boldsymbol {\Theta} ^ {*}\right) = \| \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta} ^ {*}) \| _ {2} ^ {2}. \tag {1}
$$
2.2 Canonical reparameterization
Definition 2.1 (canonical reparameterization). For given step size  $\eta$ , the canonical reparameterization of  $\Theta$  is defined as
$$
(p, q) := \left(f (\boldsymbol {x}; \boldsymbol {\Theta}) - y, \frac {2}{\eta \| \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x} ; \boldsymbol {\Theta}) \| _ {2} ^ {2}}\right). \tag {2}
$$
Under the canonical reparameterization,  $p = 0$  represents global minima, and Eq. (1) implies that the point  $(p,q) = (0,1)$  is a global minimum with sharpness  $2 / \eta$ . Note that  $(p,q)$  alone does not, in general, uniquely determine the value of  $\Theta$ . Rather, the motivation for this reparameterization technique is to effectively analyze the complex GD dynamics in the high-dimensional parameter space by reducing it to a 2-dimensional representation. The update of  $p$  can be written as
$$
\begin{array}{l} p _ {t + 1} = f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t + 1}) - y = f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t} - \eta \ell^ {\prime} (f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) - y) \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t})) - y \\ \approx f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) - \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) ^ {\top} (\eta \ell^ {\prime} (f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) - y) \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t})) - y \\ = \left(f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) - y\right) - \eta \ell^ {\prime} \left(f (\boldsymbol {x}; \boldsymbol {\Theta} _ {t}) - y\right) \| \nabla_ {\boldsymbol {\Theta}} f (\boldsymbol {x}; \boldsymbol {\Theta}) \| _ {2} ^ {2} \\ = p _ {t} - \frac {2 \ell^ {\prime} \left(p _ {t}\right)}{q _ {t}}, \tag {3} \\ \end{array}
$$
which can be obtained by first-order Taylor approximation on  $f$  for small step size  $\eta$ .<sup>3</sup>
2.3 Bifurcation analysis
Motivated from the approximated 1-step update rule given by Eq. (3), we conduct the bifurcation analysis on this one-dimensional map, considering  $q_{t}$  as a control parameter. We first review some basic notions used in bifurcation theory (Strogatz, 1994).
Definition 2.2 (stability of fixed point). Let  $z_0$  be a fixed point of a differentiable map  $f: \mathbb{R} \to \mathbb{R}$ , i.e.,  $f(z_0) = z_0$ . We say  $z_0$  is a stable fixed point of  $f$  if  $|f'(z)| < 1$ , and we say  $z_0$  is an unstable fixed point of  $f$  if  $|f'(z)| > 1$ .
Definition 2.3 (stability of periodic orbit). A point  $z_0$  is called a period- $p$  point of a map  $f: \mathbb{R} \to \mathbb{R}$  if  $z_0$  is the fixed point of  $f^p$  and  $f^j(z_0) \neq z_0$  for any  $1 \leq j \leq p - 1$ . The orbit of  $z_0$ , given by  $\{z_j = f^j(z_0) \mid j = 0, 1, \ldots, p - 1\}$  is called the period- $p$  orbit of  $f$ . A period- $p$  orbit is stable (unstable) if its elements are stable (unstable) fixed points of  $f^p$ , i.e.,  $\prod_{j=0}^{p-1} |f'(z_j)| < 1 (>1)$ .
Now we analyze the bifurcation of the one-parameter family of mappings  $f_{q}:\mathbb{R}\to \mathbb{R}$  given by
$$
f _ {q} (p) := p \left(1 - \frac {2 r (p)}{q}\right), \tag {4}
$$
where  $q$  is a control parameter and  $r$  is a differentiable function satisfying Assumption 2.4 below.
<sup>3</sup>The approximation is used just to motivate Lemma 2.1; in our theorems, we analyze the exact dynamics.
Figure 2: (a), (b) GD trajectories of two-layer fully-connected linear networks trained with different initialization scale  $\alpha$ . Each color corresponds to a single run of GD. Smaller initialization scale falls into the gradient flow regime, whereas larger initialization falls into the EoS regime. (c) In the EoS regime,  $\frac{q_t}{r(p_t)}$  approaches 1 in the early phase of training, whereas  $p_t$  converges to 0 relatively slowly.
Assumption 2.4. A function  $r: \mathbb{R} \to \mathbb{R}$  is even, continuously differentiable,  $r(0) = 1$ ,  $r'(0) = 0$ ,  $r'(p) < 0$  for any  $p > 0$ , and  $\lim_{p \to \infty} r(p) = \lim_{p \to -\infty} r(p) = 0$ . In other words,  $r$  is a smooth symmetric bell-shaped function with the maximum value  $r(0) = 1$ .
We note that Eq. (3) can be rewritten as  $p_{t + 1} = f_{q_t}(p_t)$  if we define  $r$  by  $r(p) \coloneqq \frac{\ell'(p)}{p}$  for  $p \neq 0$  and  $r(0) \coloneqq 1$ . Below are examples of  $\ell$  for which the corresponding  $r$ 's satisfy Assumption 2.4. These loss functions were previously studied by Ahn et al. (2023) to explain EoS for  $(x,y) \mapsto \ell(xy)$ .
- log-cosh loss:  $\ell_{\mathrm{log - cosh}}(p)\coloneqq \log (\cosh (p))$  . Note  $\ell_{\mathrm{log - cosh}}^{\prime}(p) = \tanh (p)$
- square-root loss:  $\ell_{\mathrm{sqrt}}(p) \coloneqq \sqrt{1 + p^2}$ . Note  $\ell_{\mathrm{sqrt}}'(p) = \frac{p}{\sqrt{1 + p^2}}$ .
If  $r$  satisfies Assumption 2.4, then for any  $0 < q \leq 1$ , there exists a nonnegative number  $p$  such that  $r(p) = q$ , and the solution is unique which we denote by  $\hat{r}(q)$ . In particular,  $\hat{r}:(0,1] \to \mathbb{R}_{\geq 0}$  is a function satisfying  $r(\hat{r}(q)) = r(-\hat{r}(q)) = q$  for any  $q \in (0,1]$ .
Lemma 2.1 (period-doubling bifurcation of  $f_{q}$ ). Suppose that  $r$  is a function satisfying Assumption 2.4. Let  $p^{*} = \sup \{p \geq 0 \mid \frac{x r^{\prime}(x)}{r(x)} > -1 \text{ for any } |x| \leq p\}$  and  $c = r(p^{*})$ . If  $p^{*} = \infty$ , we choose  $c = 0$ . Then, the one-parameter family of mappings  $f_{q} : \mathbb{R} \rightarrow \mathbb{R}$  given by Eq. (4) satisfies
(i) If  $q > 1$ ,  $p = 0$  is the stable fixed point.
(ii) If  $q \in (c, 1)$ ,  $p = 0$  is the unstable fixed point and  $\{\pm \hat{r}(q)\}$  is the stable period-2 orbit.
Proof. The map  $f_{q}$  has the unique fixed point  $p = 0$  for any  $q > 0$ . Since  $|f_{q}'(0)| = |1 - \frac{2}{q}|$ ,  $p = 0$  is a stable fixed point if  $q > 1$  and  $p = 0$  is an unstable fixed point if  $0 < q < 1$ . Now suppose that  $q \in (c, 1)$ . Then, we have  $f_{q}(\hat{r}(q)) = -\hat{r}(q)$  and  $f_{q}(-\hat{r}(q)) = \hat{r}(q)$ , which implies that  $\{\pm \hat{r}(q)\}$  is a period-2 orbit of  $f_{q}$ . Then,  $|f_{q}'(\hat{r}(q))| = |f_{q}'(-\hat{r}(q))| = \left|1 + \frac{2\hat{r}(q)r'(\hat{r}(q))}{q}\right| < 1$  implies that  $\{\pm \hat{r}(q)\}$  is a stable period-2 orbit.
According to Lemma 2.1, the stability of the fixed point  $p = 0$  undergoes a change at  $q = 1$ , resulting in the emergence of a stable period-2 orbit. The point  $(p,q) = (0,1)$  is referred to as the bifurcation point, where a period-doubling bifurcation occurs. A bifurcation diagram illustrates the points asymptotically approached by a system as a function of a control parameter. In the case of the map  $f_{q}$ , the corresponding bifurcation diagram is represented by  $p = 0$  for  $q \geq 1$  and  $p = \pm \hat{r}(q)$  (or equivalently,  $q = r(p)$ ) for  $q \in (c,1)$ .
It is worth noting that the period-2 orbit  $\{\pm \hat{r}(p)\}$  becomes unstable for  $q \in (0, c)$ . If we choose  $r$  to be  $r(p) = \frac{\ell'(p)}{p}$  for  $p \neq 0$  and  $r(0) = 1$ , then  $1 + \frac{pr'(p)}{r(p)} = \frac{\ell''(p)}{r(p)} > 0$  for all  $p$ , assuming  $\ell$  is convex. Consequently, for log-cosh loss and square root loss we have  $c = 0$ , indicating that the period-2 orbit of  $f_q$  remains stable for all  $q \in (0, 1)$ . However, in Section 5, we will consider  $r$  with  $c > 0$ , which may lead to additional bifurcations.
(a)  $m = 64,L = 3$
(b)  $m = 256,L = 3$
Figure 3: GD trajectories of tanh-activated neural networks with varying width and depth. Each color corresponds to a single run of GD. We observe that the wider network ( $m = 256$ ) exhibits a stronger trajectory alignment phenomenon compared to the narrower network ( $m = 64$ ). Figure 3c depicts the trajectories for a deeper network ( $L = 10$ ), which also shows the trajectory alignment phenomenon.
(c)  $m = 256, L = 10$
3 Trajectory Alignment of GD: An Empirical Study
In this section, we conduct experimental studies on the trajectory alignment phenomenon in GD dynamics under the canonical reparameterization proposed in Section 2.
We consider a fully-connected  $L$ -layer neural network  $f(\cdot ;\Theta):\mathbb{R}^d\to \mathbb{R}$  written as
$$
f (\boldsymbol {x}; \boldsymbol {\Theta}) = \boldsymbol {w} _ {L} ^ {\top} \phi \left(\boldsymbol {W} _ {L - 1} \phi \left(\dots \phi \left(\boldsymbol {W} _ {2} \phi \left(\boldsymbol {W} _ {1} \boldsymbol {x}\right)\right) \dots\right)\right),
$$
where  $\phi$  is an activation function,  $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$ . We minimize the empirical risk  $\mathcal{L}(\Theta) = \ell(f(\boldsymbol{x}; \Theta) - y)$ . We visualize GD trajectories under the canonical parameterization, where each plot shows five different randomly initialized weights using Xavier initialization multiplied with a rescaling factor of  $\alpha$ . For this analysis, we fix the training data point and hyperparameters as  $\boldsymbol{x} = \boldsymbol{e}_{1} = (1, 0, \dots, 0)$ ,  $y = 1$ ,  $\eta = 0.01$ ,  $d = 10$ , and focus on the log-cosh loss for  $\ell$ , with either  $\phi(t) = t$  (linear) or  $\phi(t) = \tanh(t)$ . We note that the trajectory alignment phenomenon is consistently observed in other settings, including square root loss, different activations (e.g., ELU), and various hyperparameters, in particular for sufficiently wide networks (additional experimental results are provided in Appendix A).