SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mingze Wang
Beijing, 100081, P.R. China
mingzewang@stu.pku.edu.cn
chaoma@stanford.edu
The training process of ReLU neural networks often exhibits complicated nonlinear phenomena. The nonlinearity of models and non-convexity of loss pose significant challenges for theoretical analysis. Therefore, most previous theoretical works on the optimization dynamics of neural networks focus either on local analysis (like the end of training) or approximate linear models (like Neural Tangent Kernel). In this work, we conduct a complete theoretical characterization of the training process of a two-layer ReLU network trained by Gradient Flow on a linearly separable data. In this specific setting, our analysis captures the whole optimization process starting from random initialization to final convergence. Despite the relatively simple model and data that we studied, we reveal four different phases from the whole training process showing a general simplifying-to-complicating learning trend. Specific nonlinear behaviors can also be precisely identified and captured theoretically, such as initial condensation, saddle-to-plateau dynamics, plateau escape, changes of activation patterns, learning with increasing complexity, etc.
Deep learning shows its remarkable capabilities across various fields of applications. However, the theoretical understanding of its great success still has a long way to go. Among all theoretical topics, one of the most crucial aspect is the understanding of the optimization dynamics of deep neural network (NN), particularly the dynamics produced by Gradient Descent (GD) and its variants. This topic is highly challenging due to the highly non-convex loss landscape and existing works usually work with settings that do not align well with realistic practices. For instance, the extended studied Neural Tangent Kernel (NTK) theory (Jacot et al., 2018; Du et al., 2018; Zou et al., 2018; Allen-Zhu et al., 2019) proves the global convergence of Stochastic gradient descent (SGD) to zero training error for highly over-parameterized neural networks; however, the optimization behaviors are similar to kernel methods and do not exhibit nonlinear behaviors, because neurons remain close to their initialization throughout training.
In reality, however, the training of practical networks can exhibit plenty of nonlinear behaviors (Chizat and Bach, 2018; Mei et al., 2019; Woodworth et al., 2020). In the initial stage of the training, a prevalent nonlinear phenomenon induced by small initialization is initial condensation (Maennel et al., 2018; Luo et al., 2021), where neurons condense onto a few isolated orientations. At the end of training, NNs trained by GD can directionally converge to the KKT points of some constrained max-margin problem (Nacson et al., 2019; Lyu and Li, 2019; Ji and Telgarsky, 2020). However, KKT points are not generally unique, and determining which direction GD converges to can be challenging. Nonlinear training behaviors besides initial and terminating stages of optimization are also numerous. For example, for square loss, Jacot et al. (2021) investigates the saddle-to-saddle dynamics where GD traverses a sequence of saddles during training, but it is unclear whether similar behavior can
occur for classification tasks using exp-tailed loss. Moreover, while in lazy regime, most activation patterns do not change during training ReLU networks, it remains uncertain when and how activation patterns evolve beyond lazy regime. Additionally, while it is generally conjectured that GD learns functions of increasing complexity (Nakkiran et al., 2019), this perspective has yet to be proven.
As reviewed in Section 2, works have been done to analyze and explain the nonlinear training behaviors listed above. However, due to the complexity of the training dynamics, most existing works only focus on one phenomenon and conduct analysis on a certain stage of the training process. Few attempts have been done to derive a full characterization of the whole training dynamics from the initialization to convergence, and the settings adopted by these works are usually too simple to capture many nonlinear behaviors (Phuong and Lampert, 2021; Lyu et al., 2021; Wang and Ma, 2022; Boursier et al., 2022).
In this work, we make an attempt to theoretically describe the whole neural network training dynamics beyond the linear regime, in a setting that many nonlinear behaviors manifest. Specifically, We analyze the training process of a two-layer ReLU network trained by Gradient Flow (GF) on a linearly separable data. In this setting, our analysis captures the whole optimization process starting from random initialization to final convergence. Despite the relatively simple model and data that we studied, we reveal multiple phases in training process, and show a general simplifying-to-complicating learning trend by detailed analysis of each phase. Specifically, by our meticulous theoretical analysis of the whole training process, we precisely identify four different phases that exhibit numerous nonlinear behaviors. In Phase I, initial condensation and simplification occur as living neurons rapidly condense in two different directions. Meanwhile, GF escapes from the saddle around initialization. In Phase II, GF gets stuck into the plateau of training accuracy for a long time, then escapes. In Phase III, a significant number of neurons are deactivated, leading to self-simplification of the network, then GF tries to learn using the almost simplest network. In Phase IV, a considerable number of neurons are reactivated, causing self-complication of the network. Finally, GF converges towards an initialization-dependent direction, and this direction is not even a local max margin direction. Overall, the whole training process exhibits a remarkable simplifying-to-complicating behavior.
2 Other Related Works
Initial condensation phenomenon are studied in (Maennel et al., 2018; Luo et al., 2021; Zhou et al., 2022a;b; Abbe et al., 2022a;b; Chen et al., 2023). Theoretically, Lyu et al. (2021); Boursier et al. (2022) analyze the condensation directions under their settings, which are some types of data average. Additionally, Atanasov et al. (2022) demonstrates that NNs in the rich feature learning regime learn a kernel machine due to the silent alignment phenomenon, similar to the initial condensation.
The end of training is extensively studied for classification tasks. Specifically, for classification with exponentially-tailed loss functions, if all the training data can be classified correctly, NNs trained by GD converge to the KKT directions of some constrained max-margin problem (Nacson et al., 2019; Lyu and Li, 2019; Chizat and Bach, 2020; Ji and Telgarsky, 2020; Kunin et al., 2023). In (Phuong and Lampert, 2021; Lyu et al., 2021), they analyze entire training dynamics and derive specific convergent directions that only depend on the data. Furthermore, another famous phenomenon in the end of training is the neural collapse (Papyan et al., 2020; Fang et al., 2021; Zhu et al., 2021; Han et al., 2021), which says the features represented by over-parameterized neural networks for data in a same class will collapse to one point, and such points for all classes converge to a simplex equiangular tight frame.
Saddle-to-saddle dynamics are explored for square loss in (Jacot et al., 2021; Zhang et al., 2022; Boursier et al., 2022; Pesme and Flammarion, 2023; Abbe et al., 2023). Furthermore, learning of increasing complexity, also called simplifying-to-complicating or frequency-principle, is investigated in (Arpit et al., 2017; Nakkiran et al., 2019; Xu et al., 2019; Rahaman et al., 2019).
Beyond lazy regime and local analysis, Phuong and Lampert (2021); Lyu et al. (2021); Wang and Ma (2022); Boursier et al. (2022) also characterize the whole training dynamics and exhibit a few of nonlinear behaviors. Specifically, Lyu et al. (2021) studies the training dynamics of GF on Leaky ReLU networks, which differ from ReLU networks because Leaky ReLU is always activated on any data. In (Safran et al., 2022), they study the dynamics of GF on one dimensional dataset, and characterizes the effective number of linear regions. In (Brutzkus et al., 2017), they study the
dynamics of SGD on Leaky ReLU networks and linearly separable dataset. Moreover, Boursier et al. (2022) characterizes the dynamics on orthogonally data for square loss. The studies closest to our work are Phuong and Lampert (2021); Wang and Ma (2022), exploring the complete dynamics on classifying orthogonally separable data. However, this data is easy to learn, and all the features can be learned rapidly (accuracy=100%) in initial training, followed by lazy training (activation patterns do not change). Unfortunately, this simplicity does not hold true for actual tasks on much more complex data, and NNs can only learn some features in initial training, which complicates the overall learning process. Furthermore, we provide a detailed comparison between our results and these works in Section 5. Another related work (Saxe et al., 2022) introduces a novel bias of learning dynamics: toward shared representations. This idea and the view of gating networks are enlightening for extending our two-layer theory to deep ReLU neural networks.
Our work also investigates the max-margin implicit bias of ReLU neural networks, and related works have been listed above. Although in homogenized neural networks such as ReLU, GD implicitly converges to a KKT point of the max-margin problem, it is still unclear where it is an actual optimum. A recent work (Vardi et al., 2022) showed that in many cases, the converged KKT point is not even a local optimum of the max margin problem. Besides, there are many other attempts to explain the implicit bias of deep learning (Vardi, 2023). Another popular implicit bias is the flat minima bias (Hochreiter and Schmidhuber, 1997; Keskar et al., 2016). Recent studies (Wu et al., 2018; Blanc et al., 2020; Ma and Ying, 2021; Li et al., 2021; Mulayoff et al., 2021; Wu et al., 2022; Wu and Su, 2023) provided explanations for why SGD favors flat minima and flat minima generalize well.
Basic Notations. We use bold letters for vectors or matrices and lowercase letters for scalars, e.g.  $\pmb{x} = (x_{1},\dots ,x_{d})^{\top}\in \mathbb{R}^{d}$  and  $\pmb {P} = (P_{ij})_{m_1\times m_2}\in \mathbb{R}^{m_1\times m_2}$ . We use  $\langle \cdot ,\cdot \rangle$  for the standard Euclidean inner product between two vectors, and  $\| \cdot \|$  for the  $l_{2}$  norm of a vector or the spectral norm of a matrix. We use progressive representation  $\mathcal{O},\Omega ,\Theta$  to hide absolute positive constants. For any positive integer  $n$ , let  $[n] = \{1,\dots ,n\}$ . Denote by  $\mathcal{N}(\boldsymbol {\mu},\boldsymbol {\Sigma})$  the Gaussian distribution with mean  $\pmb{\mu}$  and covariance matrix  $\boldsymbol{\Sigma}$ ,  $\mathbb{U}(S)$  the uniform distribution on a set  $S$ . Denote by  $\mathbb{I}\{E\}$  the indicator function for an event  $E$ .
3.1 Binary Classification with Two-layer ReLU Networks
Binary classification. In this paper, we consider the binary classification problem. We are given  $n$  training data  $S = \{(x_{i},y_{i})\}_{i = 1}^{n}\subset \mathbb{R}^{d}\times \{\pm 1\}$ . Let  $f(\cdot ;\theta)$  be a neural network model parameterized by  $\theta$ , and aim to minimize the empirical risk given by:
$$
\mathcal {L} (\boldsymbol {\theta}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(y _ {i} f \left(\boldsymbol {x} _ {i}; \boldsymbol {\theta}\right)\right), \tag {1}
$$
where  $\ell(\cdot): \mathbb{R} \to \mathbb{R}$  is the exponential-type loss function (Soudry et al., 2018; Lyu and Li, 2019) for classification tasks, including the most popular classification losses: exponential loss, logistic loss, and cross-entropy loss. Our analysis focuses on the exponential loss  $\ell(z) = e^{-z}$ , while our method can be extended to logistic loss and cross-entropy loss.
Two-layer ReLU Network. Throughout the following sections, we consider two-layer ReLU neural networks comprising  $m$  neurons defined as
$$
f (\boldsymbol {x}; \boldsymbol {\theta}) = \sum_ {k = 1} ^ {m} a _ {k} \sigma (\boldsymbol {b} _ {k} ^ {\top} \boldsymbol {x}),
$$
where  $\sigma (z) = \max \{z,0\}$  is the ReLU activation function,  $b_{1}\dots ,b_{m}\in \mathbb{R}^{d}$  are the weights in the first layer,  $a_1,\dots ,a_m$  are the weights in the second layer. And we consider the case that the weights in the second layer are fixed, which is a common setting used in previous studies (Arora et al., 2019; Chatterji et al., 2021). We use  $\pmb {\theta} = (\pmb {b}_1^\top ,\dots ,\pmb {b}_m^\top)^\top \in \mathbb{R}^{md}$  to denote the concatenation of all trainable weights.
3.2 Gradient Flow Starting from Random Initialization
Gradient Flow. As the limiting dynamics of (Stochastic) Gradient Descent with infinitesimal learning rate (Li et al., 2017; 2019), we study the following Gradient Flow (GF) on the objective function (1):
$$
\frac {\mathrm {d} \boldsymbol {\theta} (t)}{\mathrm {d} t} \in - \partial^ {\circ} \mathcal {L} (\boldsymbol {\theta} (t)), \quad t \geq 0. \tag {2}
$$
Notice that the ReLU is not differentiable at 0, and therefore, the dynamics is defined as a subgradient inclusion flow (Bolte et al., 2010). Here,  $\partial^{\circ}$  denotes the Clarke subdifferential, which is a generalization of the derivative for non-differentiable functions. Additionally, to address the potential non-uniqueness of gradient flow trajectories, we adopt the definition of solutions for discontinuous systems (Filippov, 2013). For formal definitions, please refer to Appendix B, G, and H.
Random Initialization. We consider GF (2) starting from the following initialization:
$$
\boldsymbol {b} _ {k} (0) \stackrel {\text {i . i . d .}} {\sim} \frac {\kappa_ {1}}{\sqrt {m}} \mathbb {U} (\mathbb {S} ^ {d - 1}) \text {a n d} a _ {k} = \mathrm {s} _ {k} \frac {\kappa_ {2}}{\sqrt {m}} \text {f o r} k \in [ m ];
$$
$$
\mathrm {s} _ {k} = 1 \text {f o r} k \in [ m / 2 ]; \mathrm {s} _ {k} = - 1 \text {f o r} k \in [ m ] - [ m / 2 ].
$$
Here,  $0 < \kappa_{1} < \kappa_{2} \leq 1$  controls the initialization scale. It is worth noting that since the distribution  $\mathcal{N}(\mathbf{0},\mathbf{I}_d / d)$  is close to  $\mathbb{U}(\mathbb{S}^{d - 1})$  in high-dimensional settings, our result can be extended to the initialization  $\pmb {b}_k\stackrel {\mathrm{i.i.d.}}{\sim}\mathcal{N}(\mathbf{0},\kappa_1^2\pmb {I}_d / md)$  with high probability guarantees.
3.3 Linearly Separable Data beyond Orthogonally Separable
In previous works (Phuong and Lampert, 2021; Wang and Ma, 2022), a special case of the linearly separable dataset was investigated, namely "orthogonally separable". A training dataset is orthogonally separable when  $\langle \pmb{x}_i,\pmb{x}_j\rangle \geq 0$  for  $i,j$  in the same class, and  $\langle \pmb{x}_i,\pmb{x}_j\rangle \leq 0$  for  $i,j$  in different classes. As mentioned in Section 2, in this case, GF can learn all features and achieve  $100\%$  training accuracy quickly, followed by lazy training. In this work, we consider data that is more difficult to learn, which leads to more complicated optimization dynamics. Specifically, we consider the following data.
Assumption 3.1. Consider the linearly separable dataset  $S = \{(x_i, y_i)\}_{i \in [n]} \subset \mathbb{R}^d \times \mathbb{R}$  such that  $(x_i, y_i) = \left\{ \begin{array}{ll} (x_+, 1), & i \in [n_+] \\ (x_-, -1), & i \in [n] - [n_+] \end{array} \right.$ , where  $x_+, x_- \in \mathbb{S}^{d-1}$  are two data points with a small angle  $\Delta \in (0, \pi/2)$ , and  $n_+, n_-$  are the numbers of positive and negative samples, respectively, with  $n = n_+ + n_-$ . We also use  $p := n_+/n_-$  to denote the ratio of  $n_+$  and  $n_-$ , which measures the class imbalance. Furthermore, we assume  $p \cos \Delta > 1$ .
Remark 3.2. We focus on the training dataset satisfying Assumption 3.1 with a small  $\Delta \ll 1$ . The margin of the dataset is  $\sin (\Delta /2)$ , which implies that the separability of this data is much weaker than that of orthogonal separable data. Additionally, the condition  $p\cos \Delta >1$  merely requires a slight imbalance in the data. These two properties work together to produce rich nonlinear behaviors during training.
4 Characterization of Four-phase Optimization Dynamics
In this section, we study the whole optimization dynamics of GF (2) starting from random initialization when training the two-layer ReLU network on linearly separable dataset satisfying Assumption 3.1 and using the loss function (1). To begin with, we introduce some additional notations.
Additional Notations. First, we identify several crucial data-dependent directions under Assumption 3.1. These include two directions that are orthogonal to the data, defined as  $\pmb{x}_{+}^{\perp} := \frac{\pmb{x}_{-} - \langle\pmb{x}_{-},\pmb{x}_{+}\rangle\pmb{x}_{+}}{\|\pmb{x}_{-} - \langle\pmb{x}_{-},\pmb{x}_{+}\rangle\pmb{x}_{+}\|}$  and  $\pmb{x}_{-}^{\perp} := \frac{\pmb{x}_{+} - \langle\pmb{x}_{+},\pmb{x}_{-}\rangle\pmb{x}_{-}}{\|\pmb{x}_{+} - \langle\pmb{x}_{+},\pmb{x}_{-}\rangle\pmb{x}_{-}\|}$ , which satisfy  $\langle \pmb{x}_+, \pmb{x}_+^\perp \rangle = \langle \pmb{x}_-, \pmb{x}_-^\perp \rangle = 0$ . Additionally, we define the label-average data direction as  $\mu := \frac{z}{\|z\|}$  where  $z = \frac{1}{n}\sum_{i=1}^{n}y_{i}\pmb{x}_{i}$ . One can verify that  $\langle \pmb{\mu},\pmb{x}_{+}\rangle > 0$  and  $\langle \pmb{\mu},\pmb{x}_{-}\rangle > 0$  under the condition  $p\cos \Delta > 1$ . In Figure 3, we visualize these directions.
Second, we use the following notations to denote important quantities during the GF training process. We denote the prediction on  $\pmb{x}_{+}$  and  $\pmb{x}_{-}$  by  $f_{+}(t) \coloneqq f(\pmb{x}_{+};\pmb{\theta}(t))$ ,  $f_{-}(t) \coloneqq f(\pmb{x}_{-};\pmb{\theta}(t))$ . We use
$\operatorname{Acc}(t) \coloneqq \frac{1}{n} \sum_{i=1}^{n} \mathbb{I}\{y_i f(\boldsymbol{x}_i; \boldsymbol{\theta}(t)) > 0\}$  to denote the training accuracy at time  $t$ . For each neuron  $k \in [m]$ , we use  $\boldsymbol{w}_k(t) \coloneqq \boldsymbol{b}(t) / \| \boldsymbol{b}(t) \|$  and  $\rho_k(t) \coloneqq \| \boldsymbol{b}(t) \|$  to denote its direction and norm, respectively. To capture the activation dynamics of each neuron  $k \in [m]$  on each data, we use  $\operatorname{sgn}_k^+(t) \coloneqq \operatorname{sgn}(\langle \boldsymbol{b}_k(t), \boldsymbol{x}_+ \rangle)$  to record whether the  $k$ -th neuron is activated with respect to  $\boldsymbol{x}_+$ , and  $\operatorname{sgn}_k^-(t) \coloneqq \operatorname{sgn}(\langle \boldsymbol{b}_k(t), \boldsymbol{x}_- \rangle)$  defined similarly, which we call ReLU activation patterns.
4.1 A Brief Overview of four-phase Optimization Dynamics
We illustrate different phases in the training dynamics by a numerical example. Specifically, we train a network on the dataset that satisfies Assumption 3.1 with  $p = 4$  and  $\Delta = \pi / 15$ . The directions and magnitudes of the neurons at some important times are shown in Figure 1, reflecting four different phases on the training behavior and activation patterns. More experiment details and results can be found in Appendix A.1.
(b)  $t = 200$
(c)  $t = 50000$
Figure 1: These figures visualize (in polar coordinates) the projections of all neurons  $\{\pmb{b}_k(t)\}_{k\in [m]}$  onto the 2d subspace  $\mathrm{span}\{\pmb{x}_{+},\pmb{x}_{-}\}$  during training. Each purple star represents a positive neuron  $(k\in [m / 2])$ , while each brown star represents a negative neuron  $(k\in [m] - [m / 2])$ . Additionally, the directions of  $\pmb{x}_{+},\pmb{x}_{-},\pmb{x}_{+}^{\perp},\pmb{x}_{-}^{\perp},\pmb{\mu}$  are plotted in blue, orange, green, red and pink colors, respectively. The complete version of these figures is Figure 4 in Appendix A.1.
(d)  $t = 60000$
(e)  $t = 150000$
From Fig 1(a) to (b) is the Phase I of the dynamics, marked by a condensation of neurons. Although the initial directions are random, we see that all neurons are rapidly divided into three categories: living positive neurons  $(k\in \mathcal{K}_{+})$  and living negative neurons  $(k\in \mathcal{K}_{-})$  condense in one direction each  $(\pmb{\mu}$  and  $x_{+}^{\perp})$ , while other neurons  $(k\notin \mathcal{K}_{+}\cup \mathcal{K}_{-})$  are deactivated forever. From the perspective of loss landscape, GF rapidly escapes from the saddle near  $\mathbf{0}$  where the loss gradient vanishes.
From Fig 1(b) to (c) is the Phase II of the dynamics, in which GF gets stuck into a plateau with training accuracy  $\frac{p}{1 + p}$  for a long time  $T_{\mathrm{plat}}$  before escaping. Once the dynamics escapes from the plateau, the training accuracy rises to a perfect  $100\%$ . Moreover, activation patterns do not change in this phase.
From Fig 1(c) to (d) is the Phase III of the dynamics. The phase transition from phase II to phase III sees a rapid deactivation of all the living positive neurons  $k \in \mathcal{K}_{+}$  on  $x_{-}$  rapidly, while other activation patterns are unchanged. This leads to a simpler network in phase III, in which only living positive neurons (in  $\mathcal{K}_{+}$ ) predict  $x_{+}$ , and only living negative neurons (in  $\mathcal{K}_{-}$ ) predict  $x_{-}$ . Hence, in this phase the GF tries to learn the training data using almost the simplest network by only changing the norms of the neurons.
Finally, Fig 1(d) to (e) shows Phase IV, starting from another "phase transition" when all the living negative neurons  $(k\in \mathcal{K}_{-})$  reactivate simultaneously on  $x_{+}$ . This leads to a more complicated network. After the phase transition, the activation patterns no longer change, and the neurons eventually converges towards some specific directions dependent on both data and initialization. Additionally, this direction is not even the local optimal max margin direction.
Overall, the whole dynamics exhibit a simplifying-to-complicating learning trend.
In the following four subsections, we present a meticulously detailed and comprehensive depiction of the whole optimization dynamics and nonlinear behaviors. For clarity, in Figure 2, we first display the timeline of our dynamics and some nonlinear behaviors.
In Appendix A.2, we further validate our theoretical bounds on the key time points in Figure 2 numerically. Additionally, in Appendix A.3, we relax the data Assumption 3.1 by perturbing the
Figure 2: Timeline of the four-phase optimization dynamics, containing some key time points  $T_{\mathrm{I}}, T_{\mathrm{II}}, T_{\mathrm{III}}, T_{\mathrm{plat}}$  and their theoretical estimates, and some basic nonlinear behaviors: ① initial condensation, ② saddle escape, ③ getting stuck in plateau, ④ plateau escape, ⑤ neuron deactivation, ⑥ neuron reactivation, ⑦ initialization-dependent directional convergence. Notice ①~⑦ are only some basic nonlinear behaviors. Moreover, ②+③ is saddle-to-plateau, ①+⑤+⑥ is simplifying-to-complicating.
data with random noise, and our experimental results illustrate that similar four-phase dynamics and nonlinear behaviors persist.
4.2 Phase I. Initial Condensation and Saddle Escape
Let  $T_{\mathrm{I}} = 10\sqrt{\frac{\kappa_1}{\kappa_2}}$ , and we call  $t \in [0, T_{\mathrm{I}}]$  Phase I. The theorem below is our main result in Phase I.
Theorem 4.1 (Initial Condensation). Let the width  $m = \Omega (\log (1 / \delta))$ , the initialization  $\kappa_{1},\kappa_{2} = \mathcal{O}(1)$  and  $\kappa_{1} / \kappa_{2} = \mathcal{O}(\Delta^{8})$ . Then with probability at least  $1 - \delta$ , the following results hold at  $T_{1}$ :