Three Mechanisms of Feature Learning in a Linear Network

## Abstract

Abstract We identify and exactly solve the learning dynamics of a one-hidden-layer linear model at any finite width whose limits exhibit both the kernel phase and the feature learning phase. We analyze the phase diagram of this model in different limits of common hyperparameters including width, layer-wise learning rates, scale of output, and scale of initialization. Our solution identifies three novel prototype mechanisms of feature learning: (1) learning by alignment, (2) learning by disalignment, and (3) learning by rescaling. In sharp contrast, none of these mechanisms is present in the kernel regime of the model. We empirically demonstrate that these discoveries also appear in deep nonlinear networks in real tasks.

## 1 Introduction

It has been shown that for a neural network under certain types of scaling towards infinite width (or certain parameters), the learning dynamics can be precisely described by the neural tangent kernel (NTK) dynamics [], or the “kernel regime.” We say that a model is in the kernel regime if the NTK of the model remains unchanged throughout training, and the learning dynamics is linear in the model parameters. When the learning dynamics is not linear, we say that the model is in the feature learning regime. Since then, a lot of works have been devoted to the study of how the kernel evolves during training as it sheds light on nonlinear mechanisms of learning [, , , , , , , , , ]. Ideally, we want a theory that allows us to precisely characterize the dynamics of NTK and whether the current model under training is in the kernel phase or the feature learning regime. However, despite the progress in understanding these infinite-width models, there have been limited results in understanding how feature learning actually happens and the relationship between these two limiting regimes. Arguably, the main theoretical gap is that until now, we do not know a single example of a finite-width exactly solvable model, whose NTK dynamics can be analytically and precisely described and exhibits both the NTK and feature learning regimes. Having such a model would help us both understand how feature learning happens in an infinite-width model and how the insights we have extensively derived for infinite-width models carry over to finite-width models.

Our main contribution is to analytically solve the evolution dynamics of the NTK of a minimal finite-width model for arbitrary hyperparameters and initialization choices. The minimal model we address is a one-hidden-layer linear network. While the model is simple in nature, its loss landscape is nonconvex and its training involves strongly coupled dynamics, and the exact solution of its learning dynamics is previously unknown. Our results reveal three novel mechanisms of learning that are only existent in the feature learning phase of the network.

## 2 Related Work

Kernel and feature learning. Due to the importance of the NTK in understanding the training and generalization of neural networks, a lot of works are devoted to understanding its structure and when and how it changes as training proceeds. Under the NTK scaling, it is shown that NTK remains unchanged in the infinite-width limit [, , ], where the network is asymptotically equivalent to the kernel regression using NTK. Higher order feature learning corrections of the NTK have also been studied [, , , ]. An important alternative to the NTK parameterization is the mean-field (or $\mu P$) parameterization where features evolve at infinite width [, , ]. Within this literature, the works closest to ours are those computing finite width corrections [, , ]. However, these results are perturbative in nature and applicable when the width is large. Our study has the same goal of understanding the learning dynamics but with a different approach. We solve an analytically solvable model that admits analysis both when the model size is finite and infinite.

Deep linear networks. Our work is also related to the analysis of deep linear networks, which has provided significant insights into the loss landscape [, ], optimization [, , , ], generalization [, ] and learning dynamics[, , ] of neural networks. Closely related to ours are Refs. [, , ], which solve the learning dynamics of linear models under special initializations. Our main advancement in this respect is to exactly solve the learning dynamics from arbitrary initializations, which we then utilize to analyze the effect of the initialization scale for learning.

## 3 An Exactly Solvable Model

This section presents and analyzes the solution of the learning dynamics. Proofs are provided in Section A. Experimental details are given in Section B.

### 3.1 Solution

Let us consider a two-layer linear network
$f(\boldsymbol{x})=\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ij}x_{j}$,
where $d_{0}$ is the input size, $d$ is the network width, $\boldsymbol{u}$ and $\mathbf{w}$ are the weight vector/matrix of the first and the second layer, respectively, and $\gamma$ is a normalization factor, which is essential for the mean-field scaling [, , ]. We consider that the network is trained on the MSE loss:

$$ $\tilde{L}=\mathbb{E}\bigg{[}(\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ij} \tilde{x}_{j}-y(\tilde{x}))^{2}\bigg{]},$ (1) $$

where we treated the target $y$ as a function of $\tilde{x}$. Throughout this work, we write $\mathbb{E}:=N^{-1}\sum_{\tilde{x}}$ to denote the averaging over the training set (although it does not have to be a finite sum). The training proceeds with the gradient flow algorithm. A little more general than the conventional study of NTK, we allow the two layers to have different learning rates, $\eta_{u}$ and $\eta_{w}$:

$$ $\frac{du_{i}}{dt}=-\eta_{u}\frac{\partial\tilde{L}}{\partial u_{i}},\quad\frac {dw_{i}}{dt}=-\eta_{w}\frac{\partial\tilde{L}}{\partial w_{ij}}.$ (2) $$

We restrict to when the data lies on a 1d manifold, and the following proposition shows that the learning dynamics under Eq.(1) is equivalent to that under a simplified loss.

###### Proposition 1 .

Let $\tilde{x}=an$, where $a\in\mathbb{R}$ is a random variable and $n$ is a fixed unit vector. Let $x=\sqrt{\mathbb{E}[\tilde{x}^{2}]}n$ and $y=\frac{\mathbb{E}[\tilde{x}y(\tilde{x})]}{\sqrt{\mathbb{E}[\tilde{x}^{2}]}}$. Then, the gradient flow of Eq.(1) equals the gradient flow of

$$ $L=\bigg{[}\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ij}x_{j}-y\bigg{]}^{2},$ (3) $$

Despite this being the simplest type of linear networks, its learning dynamics has not been analytically found in previous works. For this problem, only in two special settings, the dynamics of gradient descent (GD) have been solved. One is the standard kernel regime, where the dynamics are exactly linear []; the second case is when the two layers are initialized to be perfectly aligned, where $u$ is a left eigenvector of $w$ [, ]. The following theorem gives a precise characterization of the dynamics of $u_{i}$ and $w_{i}$ for arbitrary initialization and hyperparameter choices.

###### Theorem 1 .

Let

$$ $\begin{cases}&p_{i}(t):=\frac{1}{2\rho}(\sqrt{\eta_{u}}\sum_{j=1}^{d_{0}}w_{ij }(t)x_{j}+\sqrt{\eta_{w}}\rho u_{i}(t)),\\ &q_{i}(t):=\frac{1}{2\rho}(\sqrt{\eta_{u}}\sum_{j=1}^{d_{0}}w_{ij}(t)x_{j}- \sqrt{\eta_{w}}\rho u_{i}(t)),\end{cases}$ (4) $$

where $\rho:=\sqrt{\frac{1}{d_{0}}\sum_{i=0}^{d_{0}}x_{i}^{2}}$, and $P:=\frac{1}{d}\sum_{j=1}^{d}p_{j}(0)^{2},\ Q:=\frac{1}{d}\sum_{j=1}^{d}q_{j}(0
)^{2}$. If $P\neq 0$, then

$$ $\begin{cases}p_{i}(t)=p_{i}(0)\left[\frac{\alpha_{+}+\xi(t)\alpha_{-}}{1-\xi(t )}\right]^{1/2},\\ q_{i}(t)=q_{i}(0)\left[\frac{\alpha_{+}+\xi(t)\alpha_{-}}{1-\xi(t)}\right]^{-1 /2},\end{cases}$ (5) $$

where(^1^11Because $\alpha_{+}>0$, $\xi(t)<1$. So, (5) is well defined. Also, See Appendix A for the case $P=0$.)

$$ $\xi(t):=\frac{1-\alpha_{+}}{1+\alpha_{-}}\exp\left(-4t/t_{c}\right),$ (6) $$

$$ $t_{c}:=1/\left(\sqrt{\eta_{u}\eta_{w}\gamma^{2}\rho^{2}y^{2}+4\rho^{4}(\gamma^ {2}d)^{2}PQ}\right),$ (7) $$

$$ $\alpha_{\pm}:=\frac{1}{2(\gamma^{2}d)\rho^{2}P}\left(\sqrt{\eta_{u}\eta_{w}} \gamma\rho y\pm t_{c}^{-1}\right).$ (8) $$

Let us begin by analyzing each term and clarifying their meanings. In the theorem, we have transformed $u_{i}$ and $w_{ij}$ into an alternative basis $p_{i}$ and $q_{i}$, and $\xi(t)$ is the only time-dependent term. Note that $\xi$ decays exponentially towards zero at the time scale $t_{c}$.

By definition, we have that

$$ $\begin{cases}u_{i}(t)=\frac{1}{\sqrt{\eta_{w}}}(p_{i}(t)-q_{i}(t)),\\ w_{ij}(t)=w_{ij}(0)+\left(p_{i}(t)-p_{i}(0)+q_{i}(t)-q_{i}(0)\right)\frac{x_{j }}{\sqrt{\eta_{u}}\rho}.\end{cases}$ $$

The constants $\alpha_{+}\alpha_{-}={Q}/{P}$ are two asymptotic scale factors. In the limit $t\to\infty$, we have that

$$ $p_{i}(\infty)=p_{i}(0)\sqrt{\alpha_{+}},\ q_{i}(\infty)=q_{i}(0)/\sqrt{\alpha_ {+}}.$ (9) $$

This directly gives us the mapping between the initialization to the converged solution. Unlike a strongly convex problem where the solution is independent of the initialization, we see that the converged solution for our model is strongly dependent on the initialization and on the choice of hyperparameters. Perhaps surprisingly, because $\alpha_{\pm}$ in (8) are functions of the learning rates, the converged solution (9) depends directly on the magnitudes of the learning rates. This directly tells us the implicit bias of gradient-descent training for this problem. Another special feature of the solution is that for any direction orthogonal to $x$, the model will remain unchanged during training. Let $m\perp x$, we have that $\sum_{j}w_{ij}(t)m_{j}=\sum_{j}w_{ij}(0)m_{j}$. Namely, the output of the model in the subspace where there is no data remains constant during training.

In the theorem, what is especially important is the characteristic time scale $t_{c}$, which is roughly the time it takes for learning to happen. Notably, the squared learning speed $t_{c}^{-2}$ depends on two competing factors:

$$ $t_{c}^{-2}=\underbrace{\eta_{u}\eta_{w}\gamma^{2}\rho^{2}y^{2}}_{\text{ contribution from feature learning}}+\underbrace{4\rho^{4}(\gamma^{2}d)^{2}PQ} _{\text{contribution from kernel learning}}$ $$

The first factor depends on the input-output correlation and learning rate, which we will see is indicative of feature learning. The second term depends only on the input data and on the model initialization. We will see that when this term is dominant, the model is in the kernel regime. In fact, this result already invites a strong interpretation: the learning of the kernel regime is driven by the initialization and the input feature, whereas the learning in the feature learning regime is driven by the target mapping and large learning rates.(^2^22Another pair of important quantities are $P$ and $Q$, which are essentially the initialization scales of the model. For $\eta_{u}=\eta_{w}$, a small $P$ implies that $w_{i}(0)\approx-u_{i}(0)$ for all $i$, which implies that the model is close to anti-parallel at the start of training. Likewise, a small $Q$ implies that $w_{i}(0)\approx u_{i}(0)$ for all $i$; namely, the two layers start training when they are approximately parallel. When both $P$ and $Q$ are small, the model is initialized close to the origin, which is a saddle point. We will also see below that in the kernel regime, we always have $Q=P$.)

Using this theorem, one can compute the evolution of the NTK. Note that when different learning rates are used for different layers, the NTK needs to be defined slightly differently from the conventional definition. For the MSE loss, the NTK is the quantity $K$ that enters the following dynamics: $\frac{df(x)}{dt}=2K(x,x^{\prime})(f(x^{\prime})-y)$. This implies that for our problem,

$$ $K(x,x^{\prime})=\gamma^{2}x^{T}(\eta_{w}W^{T}W+\eta_{u}||\boldsymbol{u}||^{2}I )x^{\prime},$ (10) $$

which follows from Eq. (2) and (3). When $\eta_{w}=\eta_{u}$, this definition agrees with the standard NTK.

Figure: Figure 1: The evolution of $\zeta$ of two-layer networks with different settings. Specifically, we test linear, ReLU, sigmoid, swish, and leaky ReLU activations for both alignment (upper) and disalignment (lower) cases. For the linear network, we show the theoretical predictions obtained from (11) as lines and experimental results as points. The results for nonlinear networks are qualitatively similar.
Refer to caption: extracted/5576765/plots/fig1.jpg

###### Proposition 1 .

###### Theorem 1 .

### 3.2 Learning by Alignment and Disalignment

Before we discuss the various phase diagrams implied by Theorem 1, we first focus on a crucial effect predicted by this theorem, which differentiates it from previous results on similar problems. An important quantity our theory enables us to study is the evolution of the layer alignment $\zeta(t):={\boldsymbol{u}^{T}\boldsymbol{w}}/{||\boldsymbol{u}||||\boldsymbol{
w}||}$, which represents the cosine similarity between $u$ and $w$. Here, we set $x$ to be 1d, because Theorem 1 suggests that the dynamics of GD training has only a rank-1 effect on the model. This quantity is especially interesting to study because it tells us how well-aligned the two layers are during training. Notably, this quantity vanishes as $d\to\infty$ if and only if the model is in the kernel regime, so it serves as a great metric for probing how feature learning happens.

Letting $x=1$ and denoting $\alpha(t)=\frac{\alpha_{+}+\xi\alpha_{-}}{1-\xi}$, we have by Theorem 1:

$$ $\zeta(t)=\frac{\alpha(t)P-Q/\alpha(t)}{\sqrt{(\alpha(t)P+Q/\alpha(t))^{2}-( \frac{2}{d}\sum p_{i}q_{i})^{2}}},$ (11) $$

where $4p_{i}q_{i}=u_{i}^{2}-w_{i}^{2}=const$ does not change during training. In general, the angle evolves by an $O(1)$ amount during training. In fact, the angle remains unchanged only in the orthogonal initialization case or
in the kernel phase, where $\alpha(t)=0/1$ throughout training (see Appendix B).

Let us first consider the kernel case. Here, the easiest way to see that $\zeta$ remains zero is to note that in the kernel regime, $\|u\|$ and $\|w\|$ are of order $\sqrt{d}$, whereas $u^{T}w$ is always of order $1/\gamma$. Therefore, $\zeta(t)=o(1)$ in the NTK phase (see Section 4) and vanishes in the limit $d\to\infty$. Alternatively, one can see this from Theorem 1, which implies that in the kernel regime, $\alpha(t)=1$ is a constant (see Section 4), which in turn implies that $\zeta(t)=0$. Therefore, in the kernel regime, the two layers are essentially orthogonal to each other throughout training. This suggests one mechanism for the failure of the kernel learning phase. For a data point $x$, the hidden representation is $wx$, but predominantly many information in $wx$ is ignored after the the layer $u$. This implies that the model will have a disproportionately larger norm than what is actually required to fit the data, which could in turn imply strong overfitting.

The second case is when the two layers are initialized in a parallel way. This setting is often called the “orthogonal initialization” []. In the orthogonal initialization, $u$ is parallel to $w$, and so $p_{i}=Cq_{i}$ for a constant $C$. In this case, it is easy to very that $\zeta(t)^{2}=1$,
meaning that $u$ and $w$ remain parallel or anti-parallel throughout training. Therefore, prior theory offers no clue regarding how $\zeta$ evolves in general.

Our solution implies a rather remarkable fact: $\zeta$ is always a monotonic function of $t$. To see this,
its derivative is

$$ $\frac{d\zeta}{d\alpha}=\frac{(P+\frac{Q}{\alpha^{2}})(4PQ-(\frac{2}{d}\sum p_{ i}q_{i})^{2})}{[(\alpha(t)P+Q/\alpha(t))^{2}-(\frac{2}{d}\sum p_{i}q_{i})^{2}] ^{3/2}}.$ (12) $$

When $u$ and $w$ are parallel, this quantity is zero, in agreement with our discussion about orthogonal initialization. When they are not parallel, we have that $4PQ-(\frac{2}{d}\sum p_{i}q_{i})^{2}>0$ by the Cauchy inequality, and thus ${d\zeta}/{d\alpha}>0$. Because $\alpha(t)$ monotonically evolves from $1$ to $\alpha_{+}>0$, the evolution of $\zeta$ is also simple: $\zeta(t)$ monotonically increases if $\alpha_{+}>1$ or, equivalently, if

$$ $\frac{\sqrt{\eta_{u}\eta_{w}}y}{2\gamma d\rho P}+\sqrt{\left(\frac{\sqrt{\eta_ {u}\eta_{w}}y}{2\gamma d\rho P}\right)^{2}+\frac{Q}{P}}>1$ (13) $$

and monotonically decreases if $\alpha_{+}<1$. $\zeta$ does not change if $\alpha_{+}=1$.

Figure: (a) alignment in a 4-layer FCN
Refer to caption: extracted/5576765/plots/MNIST4layer_1.png

When does condition  (13) hold? Let us focus on the case $y>0$ because the theory is symmetric in the sign of $y$. The first observation is that it holds whenever $Q\geq P$, which is equivalent to $u^{T}(0)w(0)<0$. Namely, if the model is making wrong predictions from the beginning, it will learn by aligning different layers. Moreover, this quantity also depends on the balance of the two learning rates. Notably, when the learning rates for the two different layers are the same, the change in $\zeta$ is independent of the learning rate. This implies that under the standard GD, the angle can be quite independent of the learning rates.
However, the dependence on the learning rate becomes significant once we use different learning rates on the two layers. For example, when $\eta_{u}\gg\eta_{w}$ (or vice versa), this condition depends monotonically and (essentially) linearly on $\eta_{w}$, and making $\eta_{w}$ close to $\eta_{u}$ has the effect of making the two layers more aligned.

Why does the alignment effect depend on the ratio $Q/P$? Because when $Q$ is small, the model layers are initialized to be aligned and are likely to make predictions that are too large, and the learning process necessarily involves decreasing the model output on the data points, which can be achieved in one of the two ways: (1) decrease the scale $\|u\|\|w\|$, or (2) decrease the alignment $\zeta$. When condition (13) is not satisfied, GD employs both mechanisms for learning. Lastly, it is also worth noting that for this problem, SGD has been shown to converge to a perfectly aligned distribution of solutions []. This comparison thus shows a qualitative difference between GD and SGD – using GD, alignment is a strong function of the initialization, whereas in SGD, the alignment is quite independent of the initialization. The difference between SGD and GD is of order $1$ in this problem, even if the noise is very small. This could imply that when the model is poorly initialized, the performance difference between SGD and GD will be the most significant – we leave the empirical validation of this prediction to future work.

See Figure 1, where we show that the evolution of $\zeta$ between $u$ and $v$ of two-layer networks with $d=10000$. It is trained on a regression task. Similar experimental results are observed for a classification task trained with the cross-entropy loss (Appendix B). We choose $\gamma=1/\sqrt{d}$ for the kernel phase and $\gamma=1/d$ for the feature learning phase. The initial weights are sampled from i.i.d. Gaussian distribution $\mathcal{N}(0,1)$. Therefore, we have $P\approx Q$ and the initial $\zeta(0)\approx 0$. From (8) we have
$\alpha_{+}>1$ if $y>0$, so $\zeta(t)$ monotonically increases, but the increase is negligible in the kernel phase. Therefore, in this case, the model learns features by alignment. Meanwhile, the orthogonal initialization refers to the initialization scheme in [], that is, $\zeta$ is initialized to be $1$, so it remains $1$ as predicted. In the near orthogonal case of Figure 1, we set $\zeta(0)\approx 1$ and the initial model output to be large. As predicted, $\zeta(t)$ monotonically decreases. In this case, the model learns by disalignment. From Figure 1, we also see that this phenomenon holds for all non-linear activation functions.

The layer alignment and disalignment effects can be easily generalized to higher dimension and deeper networks. Here, we define $\zeta:=\frac{||UW||}{||U||||W||}$, where $U,W$ are the weight matrices of two consecutive layers and the L2 norm for matrices is used. See Figure 2 for a four-layer fully connected network (FCN) with ReLU activation and different initialization scales trained on MNIST datasets. Figure 2 (a) shows that the alignment between consecutive layers increases during training, and Figure 2 (b) demonstrates that layers stop being aligned for large initialization. These results are consistent with simpler settings, verifying that our analysis applies to deeper networks with non-linear activation.

### 3.3 Learning by Rescaling

Learning can also happen by rescaling the output. The evolution of $||\boldsymbol{u}||$ and $||\boldsymbol{w}||$ are given by

$$ $||\boldsymbol{u}||^{2}=d(\alpha P+\frac{Q}{\alpha})-2\sum_{i=1}^{d}p_{i}q_{i}, \ ||w||^{2}=d(\alpha P+\frac{Q}{\alpha})+2\sum_{i=1}^{d}p_{i}q_{i},$ $$

and, thus, $\frac{d||\boldsymbol{u}||^{2}}{d\alpha}=\frac{d||\boldsymbol{w}||^{2}}{d\alpha
}=d(P-\frac{Q}{\alpha^{2}})$, which is positive when $\zeta>0$, and negative when $\zeta<0$. Thus, the rescaling coincides with the alignment, namely, $||u||$ and $||w||$ become larger when they are being aligned ($|\zeta|$ gets larger), and become smaller when they are being disaligned ($|\zeta|$ gets smaller). More explicitly, (1) $P>Q$ and $\alpha_{+}>1$, or $P<Q$ and $\alpha_{+}<1$, or $P=Q$: $||u||$ and $||w||$ monotonically increase. (2) $P>Q$ and $1>\alpha_{+}\geq\sqrt{Q/P}$, or $P<Q$ and $1<\alpha_{+}\leq\sqrt{Q/P}$: $||u||$ and $||w||$ and monotonically decrease. (3) $P>Q$ and $\alpha_{+}<\sqrt{Q/P}$, or $P<Q$ and $\alpha_{+}>\sqrt{Q/P}$: $||u||$ and $||w||$ first decrease, and then increase. (4) $\alpha_{+}=1$: everything keeps unchanged. Again, in the kernel phase, the scale change of the model vanishes. In the orthogonal initialization, however, this quantity changes by an $O(1)$ amount. Therefore, the orthogonal initialization essentially learns by rescaling the magnitude of the output.

Figure: Figure 3: The initialization scale $\sigma$ correlates negatively with the performance of Resnet-18 on the CIFAR-10 dataset. Left: test accuracy. Here, $\sigma$ is a constant multiplier we apply to the initialized weights of the model under the Kaiming init. Right: the norm of all weights. While all models achieve a 100% training accuracy, models initialized with a larger scale converge to solutions with higher weight norms, which is a sign that the layers are misaligned.
Refer to caption: extracted/5576765/plots/resnet1.png

### 3.4 How does feature learning happen?

Our analysis thus suggests three mechanisms for feature learning, both of which are absent in the kernel phase. The first two mechanisms are the alignment and disalignment in the hidden layer, which is driven by the initialization balancing between the two layers. The second mechanism is the rescaling output, which is a simple operation and is unlikely to be related to learning actual features. This argument also agrees with the common technique that even if we normalize the layer output, the performance of the network does not deteriorate [].

The second question is whether we want alignment or disalignment. From our theory, the intuitive answer seems to be that alignment should be preferred over disalignment. Because aligned layers require a smaller model norm to make the same prediction, whereas a disaligned model requires a very large model norm to make the prediction. Our theory thus offers an explanation of why small initialization is always preferred and leads to better performance in deep learning – when the model has a large initialization, it will learn by disalignment, whereas a small initialization prefers alignment. This is in agreement with the common observation that a larger initialization variance correlates strongly with a worse performance [, , ]. A numerical verification is presented in Figure 3, where a larger initialization leads to worse performance. Note that this example can only be explained through the disalignment effect because (1) the model achieves 100% train accuracy in all settings, yet (2) a larger initialization leads to a larger norm at the end of the training, which also correlates with worse performance. Another piece of evidence is the commonly observed underperformance of kernel models. In the kernel phase, the model norm diverges and the model alignment is always zero, which could be a hint of strong overfitting. Therefore, our theory suggests that it would be a great idea for future works to develop algorithms that maximize layer alignment while minimizing the change in the output scale.

**Table 1: Phases of learning in different scaling limits. For brevity, the learning rates of the two layers are set to be equal. The first block shows that the models can be frozen or unstable if we do not scale $\eta$ accordingly. The second block shows that one can always choose $\eta$ such that the model training is stable and does not freeze. The third and fourth blocks show that one can always choose a pair of $\eta$ and $\gamma$ such that the model is either in the feature learning phase or the kernel phase.**
| scaling | NTK | Mean Field [] | Xavier init. | Kaiming init. | lazy [] |
| --- | --- | --- | --- | --- | --- |
| $c_{d}$ | 1 | 1 | 1 | 1 | 0 |
| $c_{\gamma}$ | -1/2 | -1 | 0 | 0 | 1 |
| $c_{u}$ | 0 | 0 | -1 | 0 | 0 |
| $c_{w}$ | 0 | 0 | -1 | -1 | 0 |
| $c_{\eta}$ | 0 | 0 | 0 | 0 | 0 |
| phase | kernel | frozen | learning | unstable | unstable |
| $c_{\eta}^{*}$ | 0 | 1 | 0 | -1 | -2 |
| phase | kernel | learning | learning | kernel | kernel |
| $c_{\eta}^{+}$ | 1 | 1 | 0 | 1 | 0 |
| $c_{\gamma}^{+}$ | -1 | -1 | 0 | -1 | 0 |
| phase | learning | learning | learning | learning | learning |
| $c_{\eta}^{-}$ | 0 | 0 | -2 | -1 | -2 |
| $c_{\gamma}^{-}$ | -1/2 | -1/2 | 1 | 0 | 1 |
| phase | kernel | kernel | kernel | kernel | kernel |

## 4 Phase Diagrams

Our theory can be applied to study the learning of different scaling limits, where we scale the hyperparameters with a scaling parameter $\kappa$ towards infinity. Here, $\kappa$ is an abstract quantity that increases linearly, and all the hyperparameters including the width are a power-law function of $\kappa$. Conventionally, the choice of $\kappa$ is the model width; however, this excludes the discussion of the lazy training regime in the theory, where the model width is kept fixed and the scaling parameter is the model output scale $\gamma$.

We first establish the necessary and sufficient condition for learning to happen: the learning time $t_{c}$ needs to be of order $\Theta(1)$. When it diverges, learning is stationary and frozen at initialization. When it vanishes to zero, the discrete-time SGD algorithm will be unstable, a point that is first pointed out by []. Therefore, we first study the condition for $t_{c}$ to be of order $1$, which is equivalent to the condition that (assuming $x$, $y$ are order $1$)

$$ $\eta_{u}\eta_{w}\gamma^{2}+(\gamma^{2}d)^{2}PQ=\Theta(1).$ (14) $$

For Gaussian initialization $u_{i0}\sim\mathcal{N}(0,\sigma_{u}^{2})$ and $w_{i0}\sim\mathcal{N}(0,\sigma_{w}^{2})$, $P$ and $Q$ are random variables with expectation $(\eta_{w}\sigma_{u}^{2}+\eta_{u}\sigma_{w}^{2})/{4}$ and variance ${(\eta_{w}\sigma_{u}^{2}+\eta_{u}\sigma_{w}^{2})^{2}}/{8d}$. Generally, all hyperparameters are powers of $\kappa$: $d\propto\kappa^{c_{d}}$, $\gamma\propto\kappa^{c_{\gamma}}$, $\sigma_{w}^{2}\propto\kappa^{c_{w}}$, $\sigma_{u}^{2}\propto\kappa^{c_{u}}$, $\eta_{w}\propto\kappa^{c_{\eta_{w}}}$ and $\eta_{u}\propto\kappa^{c_{\eta_{u}}}$. For simplicity, we set the input dimension $d_{0}$ to be a constant.

Equation (14) implies

$$ $\max\left\{2c_{\gamma}+c_{\eta_{u}}+c_{\eta_{w}},\ 2c_{\gamma}+c_{d}+\max\{c_{ \eta_{w}}+c_{u},c_{\eta_{u}}+c_{w}\}\right\}=0.$ (15) $$

Whatever choice of the exponents that solves the above equation is a valid learning limit for a neural network. The phase of the network depends on the relative order of the above two terms.

###### Definition 1 .

A model is in the kernel phase if (1) Eq. (10) is independent of $t$ as $\kappa\to\infty$ (2) $NTK=\Theta(1)$.

When $t_{c}=\Theta(1)$, a model is said to be in the feature learning phase if it is not in the kernel phase.

###### Theorem 2 .

When Eq. (15) holds, a model is in the kernel phase if and only if $\lim_{\kappa\to\infty}{P}/{Q}=1$, a.s.. and

$$ $c_{d}+\max\{c_{\eta_{w}}+c_{u},c_{\eta_{u}}+c_{w}\}>c_{\eta_{u}}+c_{\eta_{w}}.$ (16) $$

The necessary condition $P/Q=1$ for the model being in the kernel phase is interesting and highlights the important role of initialization in deep learning. There are three common cases when this holds:

- 1.
$d\to\infty$ and $u_{0}$ and $w_{0}$ are independent (standard NTK);
- 2.
$d$ is finite and the initial model output is zero: $\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ij}x_{j}=0$ (lazy training)
- 3.
$d$ is finite, $\kappa\to\infty$ and $c_{u}+c_{\eta_{w}}\neq c_{w}+c_{\eta_{u}}$;

The first case is the standard way of initialization, from which one can derive the classic analysis of the kernel phase by invoking the law of large numbers.
The second case is the assumption used in the lazy training regime []. ([] assumes $c_{\gamma}=1$, $c_{\eta_{u}}=c_{\eta_{w}}=-2$ and $c_{u}=c_{w}=0$, satisfying the conditions of the exponents (15) and (16).) This case, however, relies on a special initialization, and thus our results better illustrate the occurrence of the kernel phase for advanced initialization methods where different weights can be correlated. The third case happens when the learning rate and the initialization are not balanced. This suggests that to achieve feature learning, one should make sure that the learning rate and the initialization are well balanced: $c_{u}=c_{w}$.

In conclusion, the overall phase is (1) kernel phase, if the first term in (15) is strictly smaller than the second term: $0=c_{\gamma}+c_{d}+\max\{c_{\eta_{w}}+c_{u},c_{\eta_{u}}+c_{w}\}>c_{\eta_{u}}+
c_{\eta_{w}}$ and $\lim_{\kappa\to\infty}{P}/{Q}=1$, (2) feature learning phase if otherwise. A key difference between these two phases is whether the evolution of the NTK is $O(1)$, or equivalently whether the model learns features.

An intriguing fact is that layers tend to align in the feature learning phase. Suppose $P\approx Q$ and $\sum_{i=1}^{d}p_{i}q_{i}\approx 0$, which holds for Gaussian initialization and $d$ sufficiently large, (11) leads to
$\zeta(t)\approx\frac{\alpha(t)^{2}-1}{\alpha(t)^{2}+1},$
which monotonously changes from 0 0 to $\frac{\alpha_{+}^{2}-1}{\alpha_{+}^{2}+1}$. In the feature learning phase, two terms in (7) are of the same order, so we can suppose $\sqrt{\eta_{u}\eta_{w}}\gamma\rho y\geq 2K\rho^{2}\gamma^{2}dP$ without loss of generality, where $K$ is a certain constant. This further leads to a non-zero upper bound of the alignment $\zeta(\infty)\geq\frac{(K+1)^{2}-1}{(K+1)^{2}+1},$
in clear contrast to the kernel phase, where the alignment remains asymptotically zero for Gaussian initialization.

###### Definition 1 .

###### Theorem 2 .

### 4.1 Phases Diagram of Infinite-Width Models

Now, let us focus on the case when $\kappa=d\to\infty$ (thus $c_{d}=1$), corresponding to the infinite width limit considered in the NTK and feature learning literature [, , ]. In this limit, $\lim_{\kappa\to\infty}{P}/{Q}=1$ naturally holds by law of large numbers. Therefore, a sufficient and necessary condition for the kernel phase is $0=c_{\gamma}+c_{d}+\max\{c_{\eta_{w}}+c_{u},c_{\eta_{u}}+c_{w}\}>c_{\eta_{u}}+
c_{\eta_{w}}$.

Figure: Figure 4: A two-layer fully connected ReLU net with $d$ neurons trained on the CIFAR-10 dataset for $10000$ epochs with batch size $128$. The kernel phase is shown in solid lines and the feature learning phase is shown in dashed lines. As the theory predicts, both types of initialization can be turned into either the feature learning or the kernel phase by choosing different combinations of $\gamma$ and $\eta$. Left: the best test accuracy during training. Right: relative distance from the initialization.
Refer to caption: extracted/5576765/plots/Kaiming1.png

The following corollaries are direct consequences of Eq. (15).

###### Corollary 1 .

For any $c_{\gamma}$, $c_{u}$ and $c_{w}$, choosing $c_{\eta_{u}}=c_{\eta_{w}}=\min\{-c_{\gamma},-2c_{\gamma}-c_{d}-\max\{c_{u},c_{
w}\}\}$ ensures that the model is stable.

###### Corollary 2 .

For any $c_{u}$ and $c_{w}$, choosing $c_{\eta_{u}}=c_{\eta_{w}}=c_{\eta}$ and $c_{\gamma}=-c_{\eta}$ with $c_{\eta}\geq c_{d}+\max\{c_{u},c_{w}\}$ leads to a feature learning phase.

###### Corollary 3 .

For any $c_{u}$ and $c_{w}$, choosing $c_{\gamma}=-\frac{1}{2}(c_{d}+\max\{c_{u},c_{w}\}+c_{\eta})$ and $c_{\eta}<c_{d}+\max\{c_{u},c_{w}\}$ leads to a kernel phase.

They imply two important messages: for every initialization scheme, (1) one can choose an optimal learning rate such that the learning is stable; (2) one can choose an optimal pair of learning rate and output scale $\gamma$ such that the model is in the feature learning phase. Point (1) agrees with the analysis in [], whereas point (2) is a new insight we offer. See Table 1 for the classification of different common scalings. We choose scalings according to Corollary 2 and 3, to turn each model into the feature learning or the kernel phase.

See Figure 4. We implement a two-layer FCN on the CIFAR-10 dataset with ReLU activation. We run experiments with the scalings of the standard NTK, standard mean-field, Kaiming model, and Xavier model. $c_{\gamma}$ and $c_{\eta}$ are chosen according to Table 1. Here, we use the superscript $+$ to denote the type of scaling that leads to a feature learning phase, and $-$ denotes the kernel phase. For the Kaiming and Xavier model, we choose both $c_{\eta}^{\pm}$ and $c_{\gamma}^{\pm}$, and refer them as Kaiming^± and Xavier^±, respectively. The left figure shows that turning the Kaiming model into the feature learning phase improves the test accuracy by approximately $5\%$, similar to the gap between the standard NTK model and the mean-field model. Meanwhile, turning the Xavier model into the kernel phase decreases the test accuracy by approximately $10\%$. This is because the fixed kernel restricts the generalization ability in the kernel phase, and the difference between these models in the kernel phase might be attributed to their different kernels. Thus, in agreement with the theory, choosing different combinations of the output scale $\gamma$ and $\eta$ can turn any initialization into the feature learning phase. This insight could be very useful in practice, as the Kaiming init. is predominantly used in deep learning practice and is often observed to have better performance at common widths of the network. Our result thus suggests that it is possible to keep its advantage even if we scale up the network. Further, our results also imply that any valid learning regimes transfer well when the model gets larger, while in the feature learning phase, a larger model generally leads to better performance. This is consistent with earlier work [].

###### Corollary 1 .

###### Corollary 2 .

###### Corollary 3 .

### 4.2 Phase Diagram for Initializations

Now, we study the case when $d$ is kept fixed, while other variables scale with $\kappa\to\infty$. In this case, the phase diagram is also given by Theorem 2. See Figure 5 for an experiment. We set $c_{u}=\max\{0,c_{w}\}$ and $c_{\gamma}=\min\{-c_{w}/2,0\}$. This choice ensures that the initial model output is $O(1)$. In this case, (15) is satisfied. By Theorem 2, the network is in the kernel phase if and only if $c_{w}>0$. One important example for this section is the lazy training regime, where $c_{u}=c_{w}=c_{d}=0$, and we can choose $c_{\gamma}=1$ and $c_{\eta}=-2$ according to Corollary 3, leading to a kernel phase in finite width.

Another example is to consider large initialization, i.e., $c_{u}=c_{w}=c>0$. In this case, we can choose $c_{\eta}=c$ and $c_{\gamma}=-c$ according to Corollary 2, leading to a feature learning phase. Actually, this choice of the normalization factor $\gamma$ cancels out the scaling of the initialization. On the other hand, if we choose $\gamma=1$ as commonly done, we have to choose $c_{\eta}=-c$ according to Corollary 3, leading to a kernel space. This might be another possible explanation that larger initialization often leads to worse performance empirically. Like before, we implement a two-layer fully connected ReLU network on the CIFAR-10 dataset with $d=2000$. We choose $\kappa=10$ for illustration purposes. A clear distinction is observed between the feature learning phase and the kernel phase. (1) In Figure 5(a), the training accuracy can reach $1.0$ in the feature learning phase but not the kernel phase, because the NTK in the kernel phase is fixed, and thus the best training accuracy is limited by the fixed kernel. (2) As discussed in the previous section, the test accuracy In Figure 5(a) is about $5\%$ higher in the feature learning phase due to its trainable kernel. (3) In Figure 5(b), the weight matrices evolve significantly in the feature learning phase but not the kernel phase.

Figure: Figure 5: A two-layer FCN with different initialization scales trained on the CIFAR-10 dataset. We see that finite-width models can also exhibit qualitative differences between the feature learning and the kernel phases when other hyperparameters are scaled toward infinity. Notably, this scaling is different from the lazy training scaling, implying that there are numerous (actually infinitely many) ways for the model to enter the kernel phase, even at a finite width.
Refer to caption: extracted/5576765/plots/Init1.png

## 5 Conclusion

Solving minimal models has been a primary approach in natural sciences to understand how controllable parameters are causally related to phenomena. In this work, we have solved the learning dynamics of a minimal finite-width model of a two-layer linear network. Through a comprehensive analysis of its learning dynamics and phase diagrams, we have uncovered valuable insights into how feature learning happens and the impact of various scalings on the training dynamics of non-linear neural networks. Our theory is obviously limited: the analytical results only hold for inputs lying in a one-dimensional subspace. This limitation arises from the inadequacy of conservation laws in more general cases [], which potentially implies the impossibility of solving a more general model than ours. Lastly, it is worth applying our results to study the differences between the GD and SGD algorithms.

## Appendix A Theoretical Concerns

### A.1 Proof of Proposition 1

###### Proof.

To make the analysis more concrete, we consider the standard loss function $\tilde{L}(u,w)=\frac{1}{N}\sum_{k=1}^{N}(\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}
}u_{i}w_{ij}\tilde{x}_{jk}-\tilde{y}_{k})^{2}$, where $N$ is the size of the training set.

Data points lie in a 1d-subspace, meaning that $\tilde{x}_{jk}=a_{k}n_{j}$ for a constant unit vector $\boldsymbol{n}$. Because of the 1d nature of the data, the training dynamics on this loss function is completely identical to training on the following loss $L(u,w)=(\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ij}x_{j}-y)$, where $x_{j}=\sqrt{\sum_{k=1}^{N}a_{k}^{2}}n_{j}$ and $y=\frac{\sum_{k=1}^{N}a_{k}\tilde{y}_{k}}{\sqrt{\sum_{k=1}^{N}a_{k}^{2}}}$. This is because

$$ $\displaystyle\tilde{L}(u,w)$ $\displaystyle=(\sum_{k=1}^{N}a_{k}^{2})(\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}} u_{i}w_{ij}n_{j})^{2}-2(\sum_{k=1}^{N}a_{k}\tilde{y}_{k})(\gamma\sum_{i=1}^{d} \sum_{j=1}^{d_{0}}u_{i}w_{ij}n_{j})+\sum_{k=1}^{N}y_{k}^{2}$ (17) $\displaystyle=L(u,w)+\sum_{k=1}^{N}y_{k}^{2}-y^{2}.$ $$

Therefore, without loss of generality, the training on the standard loss $\tilde{L}$ is identical to the training on $L$ because the difference is only by a constant that does not affect gradient descent training. This setting is thus equivalent to the case when the dataset contains only a single data point $(x,y)$.(^3^33Essentially, this is because we only need two points to specify a line. Also, it is trivial to extend to the case when $y$ is a vector that spans only a one-dimensional subspace.) As is clearly shown from this example, using the notation in terms of $x$ and $y$ is much simpler to understand than using $\tilde{x}_{jk}$ and $y_{k}$. We believe that this notation is necessary and greatly facilitates the later discussions once the readers accept it.

Finally, all these notations can also be written in terms of $\mathbb{E}_{x}:=\frac{1}{N}\sum_{k=1}^{N}$, which is the notation we chose for introducing the lemma.
∎

###### Proof.

### A.2 Proof of Theorem 1

###### Proof.

By the definition of the gradient flow algorithm,

$$ $\displaystyle\frac{du_{i}}{dt}=-\eta_{u}\frac{\partial L}{\partial u_{i}}=-2 \eta_{u}\gamma\sum_{j=1}^{d_{0}}w_{ij}x_{j}\left(\gamma\sum_{i=1}^{d}\sum_{j=1 }^{d_{0}}u_{i}w_{ij}x_{j}-y\right),$ (18) $\displaystyle\frac{dw_{ij}}{dt}=-\eta_{w}\frac{\partial L}{\partial w_{ij}}=-2 \eta_{w}\gamma u_{i}x_{j}\left(\gamma\sum_{i=1}^{d}\sum_{j=1}^{d_{0}}u_{i}w_{ ij}x_{j}-y\right),$ $$

which implies the following two conservation laws:

$$ $\frac{d}{dt}(\eta_{u}\sum_{j=1}^{d_{0}}w_{ij}^{2}-\eta_{w}u_{i}^{2})=0,$ (19) $$

$$ $\frac{d}{dt}\left(\frac{w_{ij}}{x_{j}}-\frac{w_{ij^{\prime}}}{x_{j^{\prime}}} \right)=\frac{1}{x_{j}}\frac{dw_{ij}}{dt}-\frac{1}{x_{j^{\prime}}}\frac{dw_{ij ^{\prime}}}{dt}=0.$ (20) $$

From Eq. (18), we can denote $\frac{dw_{ij}}{dt}=u_{i}x_{j}A$, which implies

$$ $\displaystyle\frac{d}{dt}\left(\sum_{j=1}^{d_{0}}w_{ij}^{2}-\frac{1}{\sum_{j=1 }^{d_{0}}x_{j}^{2}}(\sum_{j=1}^{N}w_{ij}x_{j})^{2}\right)$ (21) $\displaystyle=2\sum_{j=1}^{d_{0}}u_{i}w_{ij}x_{j}A-2\frac{\sum_{j=1}^{d_{0}}u_ {i}w_{ij}x_{j}}{\sum_{j=1}^{d_{0}}x_{j}^{2}}\sum_{j=1}^{d_{0}}x_{j}^{2}A=0.$ $$

Now we denote $p_{i}(t):=\frac{1}{2\rho}(\sqrt{\eta_{u}}\sum_{j=1}^{d_{0}}w_{ij}(t)x_{j}+
\sqrt{\eta_{w}}\rho u_{i}(t))$ and $q_{i}(t)\equiv:=\frac{1}{2\rho}(\sqrt{\eta_{u}}\sum_{j=1}^{d_{0}}w_{ij}(t)x_{j
}-\sqrt{\eta_{w}}\rho u_{i}(t))$, and thus

$$ $p_{i}(t)q_{i}(t)=\frac{1}{4\rho^{2}}\left(\eta_{u}(\sum_{j=1}^{N}w_{ij}(t)x_{j })^{2}-\eta_{w}\rho^{2}u_{i}(t)^{2}\right).$ (22) $$

Take derivatives on both sides and use (19) and (21). Then we have

$$ $\begin{aligned} \frac{d}{dt}\left(p_{i}(t)q_{i}(t)\right)&=\frac{1}{4}\frac{d} {dt}\left(\frac{\eta_{u}}{\sum_{j=1}^{d_{0}}x_{j}^{2}}(\sum_{j=1}^{N}w_{ij}x_{ j})^{2}-\eta_{w}u_{i}^{2}\right)\\ &=\frac{1}{4}\frac{d}{dt}\left(\eta_{u}\sum_{j=1}^{d_{0}}w_{ij}^{2}-\eta_{w}u_ {i}^{2}\right)=0\end{aligned}.$ (23) $$

Further, substituting (18) into the definition of $p_{i}$ and $q_{i}$, we have

$$ $\frac{dp_{i}}{dt}=-2\gamma\sqrt{\eta_{u}\eta_{w}}p_{i}\rho\left(\sum_{j=1}^{d} (p_{j}^{2}-q_{j}^{2})\frac{\gamma\rho}{\sqrt{\eta_{u}\eta_{w}}}-y\right).$ (24) $$

(24) implies $\frac{1}{p_{i}}\frac{dp_{i}}{dt}=\frac{1}{p_{i^{\prime}}}\frac{dp_{i^{\prime}}
}{dt}$, further leading to another conservation law

$$ $\frac{d}{dt}\frac{p_{i}(t)}{p_{i^{\prime}}(t)}=0.$ (25) $$

for all $i,i^{\prime}=1,2,\cdots,d$. Then according to (23) and (25), we have $p_{i^{\prime}}(t)=p_{i^{\prime}}(0)\frac{p_{i}(t)}{p_{i}(0)}$ and $q_{i^{\prime}}(t)=q_{i^{\prime}}(0)\frac{p_{i^{\prime}}(0)}{p_{i^{\prime}}(t)}$. Substituting them into (24), and we obtain a differential equation with only one variable $p_{i}$

$$ $\frac{dp_{i}}{dt}=-2p_{i}\left(\frac{(\gamma^{2}d)\rho^{2}P}{p_{i}(0)^{2}}p_{i }^{2}-\frac{(\gamma^{2}d)\rho^{2}Qp_{i}(0)^{2}}{p_{i}^{2}}-\gamma\rho y\sqrt{ \eta_{u}\eta_{w}}\right),$ (26) $$

where

$$ $P=\frac{1}{d}\sum_{i=1}^{d}p_{i}(0)^{2},\ Q=\frac{1}{d}\sum_{i=1}^{d}q_{i}(0)^ {2}.$ (27) $$

This differential equation is analytically solvable by integration

$$ $t=-\int_{p_{i}(0)^{2}}^{p_{i}^{2}}\frac{d\zeta}{4\left(\frac{(\gamma^{2}d)\rho ^{2}P}{p_{i}(0)^{2}}\zeta^{2}-\gamma xy\sqrt{\eta_{u}\eta_{w}}\zeta-(\gamma^{2 }d)\rho^{2}Qp_{i}(0)^{2}\right)}$ (28) $$

Because the denominator as a quadratic polynomial has two different roots $\alpha_{\pm}$, the result of the integration is

$$ $t=-\frac{t_{c}}{4}\log\frac{p_{i}(t)^{2}/p_{i}(0)^{2}-\alpha_{+}}{p_{i}(t)^{2} /p_{i}(0)^{2}-\alpha_{-}}+const,$ (29) $$

leading to

$$ $\frac{p_{i}(t)^{2}/p_{i}(0)^{2}-\alpha_{+}}{p_{i}(t)^{2}/p_{i}(0)^{2}-\alpha_{ -}}=\frac{1-\alpha_{+}}{1-\alpha_{-}}\exp\left(-4t/t_{c}\right),$ (30) $$

which gives (5).
∎

###### Proposition 2 .

Under the condition in Theorem 1, if $P=0$ and $Q\neq 0$, the result becomes

$$ $p_{i}(t)=0$ (31) $$

$$ $q_{i}(t)=q_{i}(0)\sqrt{\frac{\alpha^{\prime}\xi^{\prime}(t)}{1-\xi^{\prime}(t)}}$ (32) $$

where

$$ $\xi^{\prime}(t):=\frac{1}{1+\alpha^{\prime}}\exp\left(-4\sqrt{\eta_{u}\eta_{w} }\gamma\rho yt\right),$ (33) $$

and

$$ $\alpha^{\prime}:=\frac{\sqrt{\eta_{u}\eta_{w}}\gamma\rho y}{(\gamma^{2}d)\rho^ {2}Q}.$ (34) $$

Specially, if $P=Q=0$, we have $p_{i}(t)=q_{(}t)=0$, so the gradient flow will be stuck at the trivial saddle point.

The proof is similar to the proof of Theorem 1, because we can similarly obtain

$$ $\frac{dq_{i}}{dt}=-2q_{i}\left(\frac{(\gamma^{2}d)\rho^{2}Q}{q_{i}(0)^{2}}q_{i }^{2}+\gamma\rho y\sqrt{\eta_{u}\eta_{w}}\right).$ (35) $$

Its solution gives Proposition 2.

We note that the behavior of the solution is quite different from $P\neq 0$: when $y\leq 0$, we can obtain a solution with zero loss in the end, but when $y>0$, the gradient flow will converge to the trivial saddle point $p_{i}=q_{i}=0$.

###### Proof.

###### Proposition 2 .

### A.3 Proof of Theorem 2

###### Proof.

Using $P/Q=1$ and (16), we have

$$ $\lim_{\kappa\to\infty}\alpha_{+}=\lim_{\kappa\to\infty}\frac{2\rho^{2}(\gamma^ {2}d)\sqrt{PQ}}{2\rho^{2}(\gamma^{2}d)P}=1.\ a.s.$ (36) $$

By definition, $\xi(t)$ is a monotonic function. As $\frac{\alpha_{+}-\xi\alpha_{-}}{1-\xi}=\frac{\alpha_{+}-\alpha_{-}}{1-\xi}+
\alpha_{-}$ is monotonous to $\xi$, it evolves from $1$ to $\alpha_{+}$ monotonously. Then according to Equation (10), $\lim_{\kappa\to\infty}K(x,x^{\prime})(t)=\lim_{\kappa\to\infty}K(x,x^{\prime})
(0)$ if and only if $\lim_{\kappa\to\infty}\alpha_{+}=1$.

From Equation (15), Equation (16) also implies

$$ $2c_{\gamma}+c_{d}+\max\{c_{\eta_{w}}+c_{u},c_{\eta_{u}}+c_{w}\}=0.$ (37) $$

Therefore, we can see that the NTK remains $\Theta(1)$ because

$$ $\gamma^{2}d\alpha_{+}P=\Theta(\kappa^{2c_{\gamma}+c_{d}+\max\{c_{\eta_{w}}+c_{ u},c_{\eta_{u}}+c_{w}\}})=\Theta(1).$ (38) $$

The proof is complete.
∎

Figure: (a) alignment in a 6-layer FCN
Refer to caption: extracted/5576765/plots/MNIST6layer_1.png

###### Proof.

## Appendix B Additional Experimental Concerns

### B.1 Additional Experiments in Section 3.2

In Figure 1, we choose $x=1,y=2$, and for others we randomly sample 100 points from $\mathcal{N}(0,1)$ as data points $x$, and set $y=2x+\mathcal{N}(0,1)/10$ as the target. The learning rates are chosen such that the model converges well within given iterations. For the orthogonal initialization, we initialize the model as $u\sim\mathcal{N}(0,10I_{d})$ and $w\sim u+\mathcal{N}(0,0.1I_{d})$.

In Figure 2, to avoid the implicit bias of SGD to make layers aligned [], we consider full-batch GD with batch size 2000 and constant learning rate. The learning rates are chosen separately for each model such that the model converges well in 1000 iterations, with training accuracy above 95%. All models use the standard Kaiming initialization, but we scale each layer by $\sigma$. The results in Figure 2 also extend to deeper networks, although the training dynamics of deeper FCNs are less stable, as shown in Figure 6.

Moreover, we observe qualitatively the same phenomenon for all kinds of activation functions in the classification task in Figure 7, where the task is to classify training samples from $\mathcal{N}(0,1)$ and $\mathcal{N}(4,1)$. Initialization is the same as in Figure 1, but the binary cross-entropy loss is used. From Figure 7 we can also see that layers tend to align in the feature learning regime when they are initialized to be disaligned, and vice versa. Note that because of the binary cross-entropy loss, $\zeta$ keeps decreasing even after the loss converges. Further, because of the binary cross-entropy loss, $\zeta$ deviates from one for non-linear activation functions other than ReLU.

Figure: Figure 7: The evolution of the alignment angle $\zeta$ between $u$ and $v$ across two-layer ReLU, sigmoid, swish, and leaky ReLU networks with $d=10000$. The task is to classify two Gaussian distributions.
Refer to caption: extracted/5576765/plots/fig2.jpg

### B.2 Experiments in Section 3.3

In Figure 3, we train a Resnet18 network on the CIFAR-10 dataset with hyperparameters borrowed from https://github.com/kuangliu/pytorch-cifar. The only difference is that we scale each layer by $\sigma$ and record the test accuracy together with the sum of the norm of all layers.

### B.3 Experiments in Section 4.1

In Section 4.1, we utilize a two-layer FCN with the ReLU activation and $d$ hidden units. The input is vectorized and normalized, so the input dimension is $d_{0}=3072$. The cross-entropy loss and the stochastic gradient descent without moment or weight decay are used during training. We use a batch size of $128$ and report the best training and test accuracy among all epochs.

We choose $\gamma=\frac{1}{\sqrt{d}}$ and $\eta=0.05$ for the standard NTK model, $\gamma=\frac{10}{d}$ and learning rate $\eta=0.05d/100$ for the standard mean-field model, $\gamma=1$ and $\eta=0.05d/100$ for the Kaiming^- model, $\gamma=\frac{100}{d}$ and $\eta=0.05d/100$ for the Kaiming^+ model, $\gamma=1$ and $\eta=0.05$ for the Xavier^+ model, $\gamma=0.01d$ and $\eta=0.05(100/d)^{2}$ for the Xavier^- model. The choice of hyperparameters guarantees that the standard NTK model and the standard mean-field model, the Kaiming^+ and Kaiming^- model, and the Xavier^+ and Xavier^- model are the same for $d=100$, respectively.

### B.4 Experiments in Section 4.2

The experiment in Section 4.2 is similar to that in 4.1. The only difference is that we fix $d=2000$ and change the initialization scale. More specifically, we set $\kappa=10$, $\sigma_{u}^{2}=\kappa^{c}$, $\sigma_{w}^{2}=\kappa^{\max\{c,0\}}$ and $\gamma=\kappa^{-\min\{0,-c/2\}}$. We also fix $\eta=0.005$.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. The authors claim larger initialization leads to worse performance. They then provide an experiment with residual skip connections (Figure 3) to try verifying this claim. I agree that with residual skip connections, networks can be initialized with zero weight variance with little consequence to training speed or performance. However, a non-residual feedforward network (such as the linear networks the authors actually analyze) the network is stuck at a saddle point as the weight variance shrinks to zero, leading to slower optimization. This is the reason for the long escape times in Saxe et al 2013. 
2. Are the results/scalings of Table 1 covered by the $a,b,c$ symmetries in the Yang et al [paper](https://arxiv.org/abs/2011.14522) or in this [note](https://mlschool.princeton.edu/sites/g/files/toruqf5946/files/documents/Princeton___Lecture_Notes_0.pdf)? If so, it would be good to comment on. 
3. The authors claim that they characterize the role of width. How does width enter their formulas, is it primarily through the initial alignment $\zeta$? Or is width to be interpreted as a possible $\kappa$ scaling quantity?
4. Could Figure 4 be plotted on a log-log scale? I suspect that $|W-W_0| / |W_0 |$ should scale as a power law in $d$ with exponent that is different for different parameterizations. 
5. In Figure 5, it appears that the best performance is not at *very small* weights, but rather there is some optimal intermediate weight scale. How do the authors reconcile this 
6. Generally throughout the paper some measured metrics are implicitly related to performance/generalization, without this connection being made clear. For example "In the kernel phase, the model norm diverges and the model alignment is always zero, which could be a hint of strong overfitting."
soundness: 4
strengths: This paper studies an important problem of how model initialization impacts learning dynamics in neural networks. They provide exact solutions for the training in a specific setting (linear networks on one dimensional inputs). The solutions indicate a crucial dependence on output scale, learning rates, initialization variance, etc. The authors connect their findings to kernel limits, mean field limits, and $\mu$P scalings.
summary: This paper studies the role of weight initialization on feature learning dynamics in linear networks. The authors provide exact solutions for  two layer linear networks where the input data is one dimensional. In this case, the relevant initial condition is the alignment $\zeta$ which measures the cosine similarity between the readin and readout vectors. Depending on the scale of the initial weights, the scale of the output multiplier, and the initial alignment between these vectors, the model may or may not exit the kernel regime when fitting the data. The authors show that large weights can lead to worse alignment compared to starting with small weights. They provide a list of possible weight, learning rate and output multiplier scalings that can lead to kernel or feature learning.
weaknesses: The assumptions of the model make the scope of the theoretical results fairly limited since it relies on one-dimensional inputs and two layer linear networks (though the authors also have some results for other activations). That said, the authors do make an effort to test their ideas in more realistic settings, which is appreciated. There are also some questions which should be addressed (see below).

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: - Am I correct in understanding that $p_i, q_i$ are something like "order parameters" in the parlance of some statistical physics analyses of network dynamics (DMFT, etc)? ie. you are studying the dynamics of some "sufficient statistics" which are informally some rotated/transformed version of our true weights? 
- Why don't either $p_i, q_i$ or initialization variables show up in the feature learning component -- what is the intuition for this? Why is there dependence on the input norm $\rho$ here?
soundness: 3
strengths: Overall, I liked this paper. The final rating would be even higher if I was confident in my expertise in this topic (linear networks/optimization trajectories). In the end, I am not such an expert, so my rating reflects the fact that I cannot be sure other related works do not present very similar results. But in my opinion, the strong point of this paper was providing 1) an exact set of solutions, in contrast to the usual perturbative approach from infinite-width dynamics (which is usually extremely complicated and does not yield much insight) and 2) spending several pages interpreting the equations that result -- this is very important: the notion that the feature learning component can come from three distinct "types of feature learning" was new and interesting to me.
summary: The paper studies a linear network setup with single-index inputs and derives exact learning dynamics under gradient flow, and interprets the resulting solutions. In doing so, the paper finds that the time to learn the target function depends on a kernel and feature learning component and that the feature learning component in particular can come from several mechanisms, which are each interpreted and discussed. This results in insights about the role of unbalanced initializations, and output rescaling, the latter being a quantity central to understanding feature learning from past work.
weaknesses: - You never write down the dynamics for the NTK in Equation (13), even though as I understand it you claim the theory implies NTK evolution. Maybe this is immediate from your solution on the weight dynamics, but some interpretation for what the evolving kernel means would be useful. 
- The word "disalignment" was confusing to me. As I understand it, it just refers to learning due to asymmetry between either the initialization or learning rate for the two layers, where feature learning doesn't take place under symmetry/orthogonal initialization. I'm not sure what "disalignment" means in this context, whereas for instance the meaning of "alignment" is intuitively clear. 
- The idea that at some intermediate initialization scale we are not quite in the kernel regime but also not quite aligning layers to the task in some intuitive sense was interesting, and I wish there was more discussion on this topic (ie. Section 3.4). 
- The point made in passing at the end about "inadequacy of conservation laws" and the impossibility of stronger analytical results is provocative and insufficiently explained. If you are making such a comment, you must explain it in more detail. 
- Minor spelling/grammar, eg. Theorem 1.1 statement "Dynamics" typo
- I think a set of bullet points of contributions in simple (nontechnical) prose would be useful in the Intro or Section 3. There are many claims made throughout and results derived, and it's sometimes not clear what the overall aim of the paper is. Having one bullet point for each section to illustrate 1) what exactly is the new result and 2) how it all relates to each other would be very useful.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: If I’m not mistaken, at line 891, if we take $\beta=1$, and plug it into equation (30), then $F_i(x) = ln(x) + c$, not $e^x$. I believe this should be the inverse of $F_i$? In which case I do not understand how (33) reduces to (35).

Shouldn’t there be $x_{ij}$ with a tilde in Proposition 1? And similarly the loss rewriting in eqn. 3 shouldn’t there be an expectation?
I don’t understand what the authors mean in line 218 i.e. “the dynamics of GD training only has a rank-1 effect”. 

If I am not misunderstanding, in Line 209, I believe now there is a vector $w$ instead of the previous matrix $W$ because $x$ is taken to be 1d, but this is only mentioned later on in line 217-218, which is confusing. 

Could the authors elaborate what they mean in line 239, more specifically why would a larger norm (assuming of the weights?) imply strong overfitting?

Would it be possible as future work to analyze, in the same framework, the solutions for arbitrary nonlinearities? Maybe by using Hermite polynomials?

Is there any intuition for why the alignment increases much faster for Sigmoid/Swish/LeakyReLU in Figure 1?
soundness: 3
strengths: The main strength of the paper lies in Theorem 1, which completely characterizes the dynamics of the layer weights across time, under the assumption that the data lies on a 1D subspace. The authors obtain solutions for general power activations, but focus on giving closed form solutions for the case of a linear network $\beta=1$ due to the problem of computing $F_i^{-1}$. By studying the alignment between the layer weights (in the case of $x$ being 1D), the authors characterize the 3 mechanisms of feature learning - which are governed by the angle $\zeta$. Finally, the authors use the theoretical framework developed to study scaling limits in Section 4.

In general I think the paper is quite thorough and does a great job covering the mechanisms of feature learning in this toy setting. To my knowledge, the theoretical analysis of these dynamics in the finite width case is novel and interesting, bringing a great addition to the community.
summary: The paper studies three mechanisms of feature learning (alignment, disalignment, rescaling) in a one hidden layer (linear) neural network. Under a data assumption, the authors characterize the solution of the network fully in terms of width, initialization scale and learning rate across different regimes. The authors also confirm their findings with empirical evidence.
weaknesses: In my opinion the main weakness of this paper is the exposition. There are quite a bit of variables involved in the main body which makes it hard to track which quantity is which. Given that the whole dynamics description in the paper revolves around Theorem 1, it might be worthwhile to restate it in simpler terms. 

Another (minor in my opinion, given the analytical intractability of network that is more complicated) weakness is the very simplified setting studied, in particular the restriction over the data lying in a 1D subspace.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: ### **Difference from uv model**
How does the proposed model differ from the well-known catapult (uv) model? Suppose that we rotate the 1-dim data space $n$ to the direction of a unit vector $(1,0,\dots,0)$ with a rotation matrix $V$. Then, defining the new variable $W(t)= W^*(t) V$, we have:
$\frac{dW}{dt} = -\frac{\partial L}{\partial W}$
equivalent to
$\frac{dW^*}{dt} = -\frac{\partial L}{\partial W^*}$
since the rotational transformation is orthogonal (that is, the Jacobian $J=\partial W^*/\partial W$ satisfies $J^T J=I$). In other words, the proposed model in this work reduces to

$L(u, W^*) = \left[\gamma \left(\sum_{i} u_i w^{*\beta}_i  x  - y \right)\right]^2$

with $x:=\sqrt{E[\tilde{x}^2]}^\beta$.

This suggests that the proposed model is merely a nonlinear parameterization of the uv model, with $v_i = w_i^\beta$. For $\beta=1$, which is the focus of the current work, the model is identical to the uv model. Since the authors references Lewkowycz et al. (2020), they  likely know about the solvability of the uv model, yet there is no explanation of this relationship or reference to further developments such as 

*Kalra et al., "Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos", arXiv:2311.02076.

*Kalra & Barkeshli, "Phase diagram of early training dynamics in deep networks: effect of the learning rate, depth, and width", NeurIPS 2023.

Since  the model is indeed equivalent to the uv model for $\beta=1$, it is crucial to mention such prior studies and to carefully explain in what aspects the problem setting differs.

### **Difference from muP**
- It will be better to provide a clearer explanation in Corollaries 1-3 about how these findings differ from the muP parameterization [Yang & Hu (2020)]. The authors state
>one can choose an optimal pair of learning rate and output scale $\gamma$ such that the model is in the feature learning phase. Point (1) agrees with the analysis in Yang & Hu (2020), whereas point (2) is a new insight we offer

 in lines 471-473. My understanding is that point (2) would be a main difference (in the inifinite width). However, $\gamma$ also appears in muP as $1/(width)^a$ in the final layer. I recommend that the authors more explicitly state how their "optimal pair of learning rate and output scale γ" differs from or extends the muP parameterization.

If the claim of this work is that the authors have identified a previously unknown feature-learning (or kernel) regime in the infinite-width case, this would imply that certain conditions in prior studies were restrictive, leading to oversight. It seems better to clearly explain which specific conditions from prior work the current work has relaxed or generalized.

- The authors also state
>A key difference mentioned is that the results apply to finite-width networks with arbitrary initialization. 

This sentence mentioning the finite width appears in Section 4.1, which consider the case of $c_d=1$, that is, the *infinite* width limit. This is confusing. If there's a distinction between their finite-width and infinite-width results, clearly delineate which parts of the paper apply to each case.

### **Typo**
- Equation 1: $x \to \tilde{x}$

###After the rebuttal########

I raised the score because the author(s) revised unclear points and also promised to explain related research in essentially the same model (but different focus) in the final version.
soundness: 3
strengths: - The analysis is technically solid and concrete.
- This work addresses the challenging yet crucial problem of identifying the feature-learning and kernel regimes and characterizing the dynamics in feature learning. It provides concrete insights from a solvable model perspective.
- The obtained parameterizations are quite general which cover various width, learning rate, and initialization scales.
summary: This work proposes a solvable shallow (linear) network to gain insight into feature/kernel-regime classification and its parameterization. For a one-dimensional data space, the continuous-time learning dynamics can be reduced to certain variables $(p_i, i=1, ..., \text{width})$, allowing for explicit solutions in a specific case $(\beta=1)$. From this solution, particularly from the learning speed $(t_c^{-2})$, the authors identify the parameterization in as general a form as possible, in which the learning dynamics achieve either the feature-learning or kernel regime.
weaknesses: - The analyzed model seems essentially equivalent to the catapult model (Lewkowycz et al. 2020) or uv model (Kalra et al. 2023). In this sense, the solvability of this model itself seems not so surprising. In addition, It is unclear how the continuous-time analysis of the current work contributes to the original literature of the catapult (uv) model, which claims that a non-trivial dyanamics toword feature learning appears for *discrete* update with a large learning rate.

- Certainly, discrete updates have not been explicitly solved and achieving analytical solvability even in the continuous limit would be novel. However, for alignment, the dynamics is essentially equal to the uv product dynamics, which can already be tracked with any width (Lewkowycz et al. 2020). Therefore, the analytical solvability for the finite width might be a bit straightforward.

- Although this would be a presentation matter, it is not easy to follow how the parameterization derived from the current analysis in Section 4.1 compares to the muP parameterization by Yang & Hu (2020). 

Please see the **Questions** section for more detailed points.

Meta-review: ### Summary

The paper proposes a theoretical framework to investigate different regimes of learning.
The authors proposed a novel model that allows a solvable solution. 
The authors used this framework to interpolate between feature and lazy learning regimes
and characterised three mechanisms of learning within their framework.

### Strengths

The paper is very clear. The authors made a great job in presenting the theory and explaining the insights of their equations.
Focusing on the technical contributions: there are very few theoretical frameworks that allow us to analyse the transitions from feature learning to lazy learning (spanning several of the limits considered in other works) and even fewer are amenable to an explicit solution. 
On top of this, the results from the analyses are very interpretable and 
the authors managed to identify 3 mechanisms emerging from the equations.
The theory is based on a simplified model of learning that, however, shows some degree of universality. 
Indeed the authors provided numerical experiments confirming some of the predictions reported in their study.

### Weaknesses

The main weakness is the simplicity of the theory. 
The authors effectively study a 1-dimensional input which is far from a realistic setting.
The theory should allow explicit solutions for polynomial activations with exponents \beta = 1, 2, and 3, 
however the paper analyses only the linear case (\beta = 1) because the other exponents presented solutions hard to interpret.
In the words of the authors: "it is difficult to write the results in a comprehensible form".
Effectively this is a linear theory for 1-dimensional inputs. 

### Reasons for acceptance

Despite the simplicity of the model, the theory is solid and shows some degrees of universality in more complex applications.
The Reviewers (and the AC) agree on the paper's relevance and consider it an important contribution that should be presented at the conference.

Decision: Accept (Poster)