Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning
Authors: Daniel Kunin
Allan Raventos
Clémentine Carla Juliette Dominé
Feng Chen
David Klindt
Andrew M Saxe
Surya Ganguli
Venue: NeurIPS 2024 spotlight (2024)
Keywords: feature learning
rich regime
lazy regime
exact solutions
conserved quantities
balanced initialization
neural tangent kernel
grokking
Primary area: learning_theory
TL;DR: We derive exact solutions to a minimal model that transitions between lazy and rich learning, precisely elucidating how unbalanced initialization variances and learning rates determine the degree of feature learning in a finite-width network.

Abstract: While the impressive performance of modern neural networks is often attributed to their capacity to efficiently extract task-relevant features from data, the mechanisms underlying this *rich feature learning regime* remain elusive, with much of our theoretical understanding stemming from the opposing *lazy regime*. In this work, we derive exact solutions to a minimal model that transitions between lazy and rich learning, precisely elucidating how unbalanced *layer-specific* initialization variances and learning rates determine the degree of feature learning. Our analysis reveals that they conspire to influence the learning regime through a set of conserved quantities that constrain and modify the geometry of learning trajectories in parameter and function space. We extend our analysis to more complex linear models with multiple neurons, outputs, and layers and to shallow nonlinear networks with piecewise linear activation functions. In linear networks, rapid feature learning only occurs from balanced initializations, where all layers learn at similar speeds. While in nonlinear networks, unbalanced initializations that promote faster learning in earlier layers can accelerate rich learning. Through a series of experiments, we provide evidence that this unbalanced rich regime drives feature learning in deep finite-width networks, promotes interpretability of early layers in CNNs, reduces the sample complexity of learning hierarchical data, and decreases the time to grokking in modular arithmetic. Our theory motivates further exploration of unbalanced initializations to enhance efficient feature learning.

Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning
Daniel Kunin"!
Allan Raventós"'
Clémentine Dominé
Feng Chen'
David Klindt
Andrew Saxe
Surya Ganguli'
Stanforr University
University College London
Cold Spring Harbor Laboratory
Abstract
While the impressive performance of modern neural networks is often attributed to their capacity to efficiently extract task-relevant features from data, the mech- anisms underlying this rich feature learning regime remain elusive, with much of our theoretical understanding stemming from the opposing lazy regime. In this work, we derive exact solutions to minimal model that transitions between lazy and rich learning, precisely elucidating how unbalanced layer-specife ini- tialization variances and learning rates determine the degree feature learning. Our analysis reveals that they conspire to influence the learning regime through a set of conserved quantities that constrain and modify the geometry of learning trajectories in parameter and function space. We extend our analysis to more complex linear models with multiple neurons, outputs, and layers and to shallow nonlinear networks with piecewise linear activation functions. In linear networks, rapid feature learning only occurs from balanced initializations, where all layers learn at similar speeds. While in nonlinear networks, unbalanced initializations that promote faster learning in earlier layers can accelerate rich learning. Through a series of experiments, feature learning in deep finite-width networks, promotes interpretability of early layers in CNNs, reduces the sample complexity of learning hierarchical data, and decreases the time to grokking in modular arithmetic. Our theory motivates further exploration of unbalanced initializations to enhance efficient feature learning.
provide evidence that this unbalanced rich regime drives
Introduction
Deep learning has transformed machine learning, demonstrating remarkable capabilities in a myriad of tasks ranging from image recognition natural language processing. It's widely believed that the impressive performance of these models lies in their capacity to efficiently extract task-relevant features from data. However, understanding this feature acquisition requires unraveling complex in- terplay between datasets, network architectures, and optimization algorithms. Within this framework, two distinct regimes, determined initialization, have emerged: the lazy and the rich.
Lazy regime. Various investigations have revealed a notable phenomenon overparameterized neural networks, where throughout training the networks remain close their linearization 2 3 4 []
Seminal work by Jacot et al. (6], demonstrated that in the infinite-width limit, the Neural Tangent Kernel (NTK), which describes the evolution of the neural network through training, converges to deterministic limit. Consequently, the network learns a solution akin to kernel regression with the NTK matrix. Termed the lazy or kernel regime, this domain has been characterized by deterministic NTK [6] Z) minimal movement in parameter space [8] static hidden representations, exponential learning curves, and implicit biases aligned with reproducing kemnel Hilbert space (RKHS) norm 9 However, Chizat et al. g challenged this understanding, asserting that the lazy regime isn 't
Equal contribution. Correspondence to kunin@stanford, edu and aravento@stant ord, edu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Rich [T- 5-0
Rich IT- 5-
t= 100
t 1000
t- 10000
t 109000
USB Sot----a
0.
04-
5501 150g
0.06
0.38
0.4
-0.3
0.00
0:20
- P5000 P M
0.4
-0.16
-8.16
-0.4
-0.3
10
109
10
10t4 105
10
109
(a) Overall and relative scale impact feature learning
(b) A complex phase portrait of feature learning
Figure I: Unbalanced initializations lead to rapid rich learning and generalization. We follow the experimental setup used in Fig. 1 of Chizat al. [ a wide two-layer student ReLU network f(:; Σi= a, max(O, w wf²) trained on a dataset generated from narrow two-layer teacher ReLU network. The student parameters are initialized as w, Unif(Sd-1 (2)) and a; Hat, such that 7 > 0 controls the overall scale of the function, while a > 0 controls the relative scale of the first and second layers through the conserved quantity 8 (2
²). (a) Shows the training trajectories of |a.] Re; (color denotes agm(a;)) when d = 2 for four different settings of T, 8. The left plot confirms that small overall scale leads to rich and large overall scale to lazy. The right plot shows that even at small overall scale, the relative scale can move the network between rich and lazy as well. Here an upstream initialization a downstream initialization å < 0 shows no alignment. (b) Shows the test loss and kernel distance from initialization computed through training over a sweep of r and 6 when d' = 100. Lazy learning happens when r is large, rich learning happens when is small, and rapid rich learning happens when both 7 is small and is large - an upstream initialization. This initialization also leads to the smallest test loss. See Fig. 10 in Appendix D.I for supporting figures.
shows striking alignment to the teacher (dotted lines), while
product of the infinite-width architecture, but is contingent on the overall scale of the network initialization. They demonstrated that given any finite-width model f(x; 0) whose output is zero at initialization, scaled version of the model ff () will enter the lazy regime as the scale T diverges. However, they also noted that these scaled models often perform worse in test error. While the lazy regime offers insights into the network's 's convergence to global minimum, it does not fully capture the generalization capabilities of neural networks trained with standard initializations. It is thus widely believed that a different regime, driven by small or vanishing initializations, underlies the many successes of neural networks.
Rich regime. In contrast to the lazy regime, the rich or feature-learning or active regime is distin- guished by a learned NTK that evolves through training, non-convex dynamics traversing between saddle points [10 [2] sigmoidal learning curves, and simplicity biases such as low-rankness [I3] or sparsity [14). Yet, the exact characterization of rich learning and the features it learns frequently depends on the specific problem at hand, with its definition commonly simplified as what it is not: lazy. Recent analyses have shown that beyond overall scale, other aspects of the initialization can substantially impact the extent of feature learning, such as the effective rank [5], layer-specific initialization variances [16]
17 18] and large learning rates [19] 20 21 22 Azulay et al. [9] demonstrated that in two-layer linear networks, the relative difference in weight magnitudes between the first and second layer, termed the relative scale in our work, can impact feature learning, with balanced initializations yielding rich learning dynamics, while unbalanced ones tend to induce lazy dynamics. However, shown in Fig- D for nonlinear networks unbalanced initializations can induce both rich and lazy dynamics, creating a complex phase portrait of learning regimes influenced by both overall and relative scale. Building on these observations, our study aims to precisely understand how layer-specifi, initialization variances and learning rates determine the transition between lazy and rich learning in knite-width networks. Moreover, we endeavor to gain insights into the inductive biases of both regimes, and the transition between them, during training and at interpolation, with the ultimate goal of elucidating how the rich regime acquires features that facilitate generalization.
Our contributions. Our work begins with an exploration of the two-layer single- linear network proposed by Azulay et al. [9] as a minimal model displaying both lazy and rich learning. In Section 3 we derive exact solutions for the gradient flow dynamics with layer-specific learning rates of this model by employing a combination of hyperbolic and spherical coordinate transformations.

---

Alongside recent work by Xu and Ziyin 23 our analysis stands out as one of the few analytically tractable models for the transition between lazy and rich learning in a finite-width network, marking a notable contribution to the field. Our analysis reveals that the layer-specific initialization variances and learning rates conspire to influence the learning regime through simple set of conserved quantities that constrain the geometry of learning trajectories. Additionally, it reveals that crucial aspect of the relative scale overlooked in prior analysis is its directionality, While a balanced initialization results in all layers learning at similar rates, an unbalanced initialization can cause faster learning in either earlier layers, referred to as an upstream initialization, or later layers, referred to as a downstream initialization. Due to the depth-dependent expressivity of layers in a network, upstream and downstream initializations often exhibit fundamentally distinct learning trajectories. In Section we extend our analysis of the relative scale developed in the single-neuron model to more complex linear models with multiple neurons, outputs, and layers and in Section 5 to two-layer nonlinear networks with piecewise linear activation functions. We find that in linear networks, rapid rich learning can only occur from balanced initializations, while in nonlinear networks, upstream initializations can actually accelerate rich learning. Finally, through series of experiments, we provide evidence that upstream initializations drive feature learning in deep finite-width networks, promote interpretability of early layers in CNNs, reduce the sample complexity of learning hierarchical data, and decrease the time to grokking in modular arithmetic.
Notation. In this work, we consider feedforward network f(; Rd - R parameterized by B E RRm. Unless otherwise specified, 1. The network is trained by gradient flow Ó = -Mje- Vol(O), with an initialization 0o and layer-specific learning rate 179 E R to minimize the mean squared error [()) = = (f 0) )2 computed over a dataset (121,81)...
(Fn, Bn)} of size We denote the input matrix as X E Ruxdd with rows z; E Rd and the label vector as E IR" The network's 's output f(x; 0) evolves according to the differential equation. O, (x; = Σi= xx; (y: - f(z;; 8)), where 0(z,,';8)
Rd x Rd + R is the Neural Tangent Kernel (NTK), defined as O(z,z'; = Σx= 170 00, 0)0,, ((): The NTK quantifies how one gradient step with data point x' affects the evolution of the networks's 's output evaluated at another data point x. When nop is shared by all parameters, the NTK is the kernel associated with the feature map (af 0) RTH We also define the NTK matrix K E RTx", which is computed across the training data such that Ki = - 0). The NTK matrix evolves from its initialization Ko to convergence Koc through training. Lazy and rich learning exist on spectrum, with the extent of this evolution serving as the distinguishing factor. Various studies have proposed different metrics to track the evolution of the NTK matrix [24 25 (6). We kernel distance [27], defined as S(t1, t2) 1 1- (Ke, : / K,, || Ke, Il. which is a scale invariant measure of similarity between the NTK at two times. In the lazy regime S(,, 1) R O, while in the rich regime 0 < (S(,,) 1.
2 Related Work
Linear networks. Significant progress in studying the rich regime has been achieved in the context of linear networks. In this setting, (2; 0) Tz is linear in its input but can exhibit highly nonlinear dynamics in parameter 0 and function 3(0) space. Foundational work by Saxe et al. [10] provided exact solutions to gradient flow dynamics in linear networks with task-aligned initializations. They achieved this by solving a system of Bernoulli differential equations that prioritize learning the most salient features first, which can be beneficial for generalization
(28). This analysis has been extended to wide [29 30 and deep [3 32 33 lincar networks with more flexible initialization schemes [34 B5| 36). It has also been applied to study the evolution of the NTK (37) and the influence of the scale on the transition between lazy and rich learning 12 23 In this work, we present novel
exact solutions for minimal model utilizing mix of Bernoulli and Riccati equations to showcase complex phase portrait of lazy and rich learning with separate aligmment and fitting phases.
Implicit bias. An effective analysis approach to understanding the rich regime studies how the initialization influences the inductive bias at interpolation. The aim is to identify a function Q() such that the network converges to first-order KKT point minimizing Q(0) among all possible interpolating solutions. Foundational work by Soudry et al. (38) pioneered this approach for a
linear classifier trained with gradient descent, revealing max margin bias. These findings have
been extended to deep linear networks [39 40 41), homogeneous networks 42 43 44, and quasi-
homogeneous networks 45
A similar line of research expresses the learning dynamics of networks
"u and Ziyin
23 presented exact NTK dynamics for linear model trained with one- -e-dimensional data.
3

---

trained with mean squared error as a mirror flow for some potential (33), such that the inductive bias can be expressed as a Bregman divergence [46]. This approach has been applied to diagonal linear networks, revealing an inductive bias that interpolates between El and l norms in the rich and lazy regimes respectively [14. However, finding the potential ((3) is problem-specific and requires solving a second-order differential equation, which may not be solvable even in simple settings 48].
Azulay et al. [9] extended this analysis to a time-warped mirror flow, enabling the study of a broader class of architectures. In this work we derive exact expressions for the inductive bias of our minimal model and extend the results in Azulay et al. 9 to wide and deep linear networks.
47
Two-layer networks. Two-layer, or single-hidden hidden layer, piecewise linear networks have emerged as a key setting for advancing our understanding of the rich regime. Maennel et al. [49 observed that in training two-layer ReLU networks from small initializations, the first-layer weights concentrate along fixed directions determined by the training data, irrespective of network width. This phenomenon, termed quantization, has been proposed as simplicity bias inherent to the rich regime, driving the network towards low-rank solutions when feasible. Subsequent studies have aimed to precisely
elucidate this effect by introducing structural constraints on the training data [50] [1 52 53 54 55
Across these analyses, a consistent observation is that the learning dynamics involve distinct phases: an initial alignment phase characterized by quantization, followed by fitting phases where the task is learned. All of these studies assumed a balanced (or nearly balanced) initialization between the first and second layer. In this study, we explore how unbalanced initializations influence the phases of learning, demonstrating that it can eliminate or augment the quantization effect.
infinite-width networks. Many recent advancements in understanding the rich regime have come from studying how the initialization variance and layer-wise learning rates should scale in the infinite- width limit to ensure constant movement in the activations, gradients, and outputs. In this limit, analyzing dynamics becomes simpler in several respects: random variables concentrate and quantities will either vanish to zero, remain constant, or diverge to infinity [17]. A set of works used tools from statistical mechanics to provide analytic solutions for the rich population dynamics of lay-layer nonlinear neural networks initialized according to the mean field parameterization [56] 57 58 59)
These ideas were extended to deeper networks through a tensor program framework, leading to the derivation of maximal update parametrization (μP) [16] 18). The μP parameterization has also been derived through a ceff-consistent dynamical mean field theory [60] and a spectral scaling analysis
61.. In this study, focus fnith neural networks, but discuss the connection between our work and these width-dependent parameterizations in Section [
A Minimal Model of Lazy and Rich Learning with Exact Solutions
Here we explore an illustrative setting simple enough to admit exact gradient flow dynamics, yet complex enough to showcase lazy and rich learning regimes. We study a two-layer linear network with single hidden neuron defined by the map f((x;) 0)) = Tæ
4
1
T. where a E R, we E Rd are the parameters. We examine how the parameter initializations do, wo and the layer-wise learning rates 7a The influence the training trajectory in pa- rameter space, function space (defined by the product B aw), and the evolu- tion of the the NTK matrix,
(a) 6 -2
(b) 8 = 0
(c) 6 = 2
Figure 2: Balance determines geometry of trajectory.
The quantity 8 = n= 1/ [2| is conserved through gradient flow, which constrains the trajectory to: (a) a one- sheeted hyperboloid for downstream initializations, (b) double cone for balanced initializations, and (c) a two- sheeted hyperboloid for upstream initializations. Gradient flow dynamics for three different initializations dg, wo with the same product By ag are shown. The minima manifold is shown in red and the manifold of equivalent So initializations in gray. The surface colored according to training loss, with blue representing higher loss and red representing lower loss.
K = X (n + naww)) XT. (1)
Except for a measure zero set of ini- tializations which converge to saddle point& all gradient flow trajectories will converge to a global minimum, de- termined by the normal equations X XTy. However, even when XTX X is invertible such that the global minimum B. is unique, the rescaling symmetry between a and w results in manifold
"The set of saddle points {(a, w)} is the d 1 dimensional subspace satisfying 0 and wT XTy 0.
4

---

of minima in parameter space. The minima manifold is one mne-dimensional hyperbola where w x B. and has two distinct branches for positive and negative a. The symmetry also imposes constraint on the network's 's trajectory, maintaining the difference ö = 77..a 1/ | || E R throughout training (see Appendix for details). This confines the parameter dynamies to the surface of hyperboloid where the magnitude and sign of the conserved quantity determines the geometry, as shown in Fig. 2
An upstream initialization occurs when > O, a balanced initialization when 6 = 0, and a downstream initialization when å < 0.
Deriving exact solutions in parameter space. We initially assume whitened input XTX = La such that the ordinary least squares solution is B_ = yyy, and the gradient flow dynamics simplify to à == 7e (wT B. all w| I1). 2 Tw (aB, - a² a²w). Notice that w(t) E span( (wwg, B+)). and through training, w aligns in direction to =3, depending on the basin of attraction the parameters are initialized in. Therefore, we can monitor the dynamics by tracking the hyperbolic geometry between a and (t)ll and the spherical angle between w(t) and B... We study the variables μ = al|wo[], an T invariant under the rescale symmetry, and Φ ||| 1 the cosine of the spherical angle. From these two scalar quantities plt), Φ(!) and the initialization ao, wn. we can determine the trajectory a(l) and w(() in parameter space. The dynamics for jr, @ are given by the coupled nonlinear ODEs,
Na7 2/ 13.11
jr V 82 Aw (0113a1) p),
S
(1 d²)
82 + 4ma
Amazingly, this system can be solved exactly, as discussed in Appendix A.2 and shown in Fig. 3 Without delving into the specifics, we can develop an intuitive understanding of the solutions by examining the influence of the relative scale 8.
Upstream. When 8 >> 0, the up- dates for both μ and & diverge, but updates much more rapidly. We
Sioned Signed Magnitude p1
Karnel Distance SID,
Wogmentt
00.75
9.6
can decouple the dynamics of
0.4
and Φ by separation of their time scales and assume & has reached its steady-state of +-1 before / has updated. Then, the dynamics of a linear and proceeds exponen- tially to + =|. This regime ex- hibits minimal kernel movement (see Fig. (c)) because the kernel dominated by the 7/ I term, whereas it is mainly that updates.
0.4
0.2
0.00
15
15
Time
(a) p(t)
(b) ((t)
(c) t),,t)
Figure 3: Exact solutions for the single hidden neuron model. Our theoretical predictions (black dashed lines) agree with gradi- ent flow simulations (solid lines, color-coded based on values), shown here for three key metries: μ (left), & (middle), and (1 t) (right). Each metric starts at the same value for all 8, but varying 5 has a pronounced effect on the metric's dynamics. For upstream initializations (ō >> 0), f changes only slightly, ó exponentially aligns, and S remains near zero, indicative of the lazy regime. For balanced initializations (ō 0), both f and Φ change significantly and S quickly moves away from zero indicative of the rich regime. For downstream initializations (ō < 0), f quickly drops zero, then μ and slowly climb back to one. Similarly, S remains small before a sudden tran- sition towards one, indicative of delayed rich regime. See Appendix for further details.
Balanced. When lows a Bernoulli differential equa- tion driven by a time-dependent signal @ |]. and o follows a Ric- cati equation evolving from an ini- tial value to 11 depending on the basin of attraction. For vanishing initialization || + 0, the tempo- ral dynamics of μ and ạ decouple such that there are two phases of learning: an initial alignment phase where Φ + +1, followed by a fitting phase where f - +1 In the first phase, aligns to 3. resulting in rank-one update to the NTK, identical to the silent alignment effect described in Atanasov et al. [37]. In the second phase, the dynamics of μu simplify to the Bernoulli equation studied in Saxe et al. [10) and the kernel evolves solely in overall scale.
/ fol-
Downstream. When < O, the updates for μ diverge, while the updates for 6 vanishes. In this regime the dynamics proceed by an initial fast phase where μ converges exponentially to its steady state of ø ق||,3_ I|. Plugging this steady state into the dynamics of 6 gives a Bemoulli differential equation
We relax this assumption when considering the dynamics of ( function space and their implicit bias. 'The basin is given by sgn(eo) for ê ' or sgn(w] + 44 (8 /8 + A11..13). for ò 0. See A.2.5
5

---

e = 11. 0(1 -() Due to the coefficient 18|_1, the second alignment phase proceeds very slowly as @ approaches =1, assuming é. p & O, which saddle point. In this regime, the dynamics proceed by an initial lazy fitting phase, followed by rich alignment phase, where the delay determined by the magnitude of 6.
Identifying regimes of learning in func- tion space. Here we take an alternative route towards understanding the influence of the relative scale by directly examining the dynamics in function space, an analy- broader sis strategy we will generalize
1.25
0.9
0.75
0.6
6,5
0.3
setups in Sections 4 and [ The network's s function is determined by the product / aw and governed by the ODE,
0.00
-12
-0.5
0.9
0.39
0.25
1.00
1.25
/ (nw La + Ta ww)) XTP, P,
(3)
81
8
(a) Whitened X
(b) Low-Rank X
where p 3 y is the residual. These dynamics can be interpreted as precondi- tioned gradient flow on the loss in function space where the preconditioning matrix M depends on time through its dependence on a² and T Whenever ||.3|| * O, we can express M directly in terms of and as
Figure 4: Balance modulates 3 dynamics and implicit bias. Here we show the dynamics of [3 = aw with dif- ferent values of ó, but the same initial 30. When X is whitened (left), we can solve for the dynamics ex- actly using our expressions for p. (black dashed lines). Upstream initializations follow the trajectory of gradi- ent flow on B, downstream initializations first move in the direction of 3o before sweeping around towards B., and balanced initializations take an intermediate trajectory between these two. When X X is low-rank (right), then we can only predict the trajectorics in the limit of 8 too. If the interpolating manifold is one- dimensional, then we can solve for the solution in terms of 5 exactly (black dots). See Appendix AA4 for details.
K + 6 La
R - 8 38T
M
(4)
2
2 1812
where A V 82 + 4n | (see Ap-
pendix A.3 for a derivation). This estab-
lishes a self-consistent equation for the dy- namics of / regulated by 8. Additionally, notice that M characterizes the NTK matrix Eq. [T] Thus, understanding the evolution of M along the trajectory 8o to B_ offers a method to discern between lazy and rich learning. Upstream. When >> O, M x 8Ig, and the dynamics of 33 converge to the trajectory of linear regression trained by gradient flow. Along this trajectory the NTK matrix remains constant, confirming the dynamics are lazy. Balanced. When 8 = MM- Now ||
[aT ). Here the dynamics balance between following the lazy trajectory and attempting to fit the task by only changing in norm. As a result the NTK changes in both magnitude and direction through training, confirming the dynamics are rich. Downstream. When 8 < 0, M |ö]
||[ +
1997
B I| and B3 follows projected gradient descent trajectory, attempting to reach 3, in the direction of S. Along this trajectory the NTK matrix doesn' 't evolve. However, if Bu is not aligned to 5.,, then at some point the dynamics of 3 will slowly align. In this second alignment phase the NTK matrix will change, confirming the dynamics are initially lazy followed by delayed rich phase. See Appendix 1.3.1 for a derivation of the NTK dynamics K.
Determining the implicit bias via mirror flow, So far we have considered whitened or full rank XTX, ensuring the existence of a unique least squares solution B.. In this setting, ô influences the trajectory the model takes from 3o to 3., as shown in Fig. 4 (a). Now we consider low-rank XTX, such that there exist infinitely many interpolating solutions in function space. By studying the structure of M, we can characterize how 5 determines the interpolating solution the dynamics converge to. Extending time- wareed mirror flow analysis strategy pioneered by Azulay al. (9 to allow 8 < 0 (see Appendix [A.4 for details), we prove the following theorem, which shows a tradeoff between reaching the minimum norm solution and preserving the direction of the initialization 80- Theorem 3.1 (Extending Theorem 2 in Azulay et al. 1). For a single hidden neuron linear network, for any ó IR, and initialization 8o such that (() + 0 for all 2 0, if the gradient flow solution 3(00) satisfies X (xx) = y. then,
B(oc) = arg min V. v5()
tha
Dg |
B
s.t.
XB=
BERU
where vz( (v82+4133P) = 28)
82
1 + and 4/ v
V82 + 11

---

We observe that for vanishing initializations there is functionally no difference between the inductive bias of the upstream (5 & O), it is the second term preserving the direction of the initialization that dominates the inductive bias. This tradeoff in inductive bias as a function of 8 presented in Fig. 4 (b), where the null
> 0) and balanced (6 = 0) settings. However, in the downstream setting
space of X X is one-dimensional, we can solve for ((o) in closed form (see Appendix A.4).
Wide and Deep Linear Networks
We now show how the analysis techniques used study the influence of relative scale in the single-neuron setting can be applied to linear networks with multiple neurons, outputs, and layers.
Wide linear networks. We consider the dynamics of a two-layer linear network with h hidden neurons and c outputs, fx;; 0) = ATWx, x. where W E xd and A E RRA xe, We assume 2 min(d, c), such that this parameterization can represent all linear maps from R + R The rescaling symmetry between A and W implies the h x h matrix A Au-AA - Th WoTT is conserved throughout gradient flow 62). Drawing insights from our analysis of the single- scenario (h = c = 1), we consider the dynamics of [ WT E Rdxe.
VCC = (nw ATA @ Na WTW vec(XTX3 - -
M
where vec(.) denotes the vectorization operator and denotes the Kronecker sum As in the single-neuron setting, we find that the dynamics of 3 are preconditioned by matrix M that depends on quadratics of A and W and characterizes the NTK matrix K = (Ic @ X) M (L. XT). We now show how M can be expressed in terms of the rank-1 1 matrices fAx wg g E Rdxc, which represent the contribution to 33 of a single neuron with parameters Wkk ak and conserved quantity 6g = Aak-
Theorem 4.1. Whenever ||Bx]|r # 0 for all lo E [M], the matrix M can be expressed as the sum M- = Eh= Ms over hidden neurons where My is defined as,
M₁ = Vo + 41/a Jm0 3R
B. + Oz
3 3R
(V8 + 4M
3x 1T
Ax 1
@
(7)
2
1.3
2
By studying the dependence of M on the conserved quantities 5g and the dimensions d, h and c, can determine the influence of the relative scale on the learning regime. When min(d, c) S h < max(d, c), and assuming independent initializations for all 3k, then networks which narrow from input to output (d > c) enter the lazy regime when all 6% >> O, whereas networks which expand from input to output (d c) do so when all du < 0. However, with opposite signs for ōR, and assuming all & (0) & B., these networks enter a delayed rich regime. As elaborated in Appendix 1.5 this occurs because in these regimes a solution 3_ does not exist within the space spanned by M at initialization. When h 2 max(d, d, c) all networks enter the lazy regime when all 6g >>>>>>>>>>> 0 or all dr << O. Conversely, as all 8% - O, all networks transition into the rich regime regardless of dimensions. While Theorem 4. offers valuable insight into the learning regimes in the limits of &g. understanding the transition between regimes remains challenging. To achieve this, we aim to express M in terms of B, rather than 3r, by introducing structure on the conserved quantities A.
Theorem 4.2. When A = oI, and h = d 8 < 0 or = if > 0, then the matrix M can be
expressed as M = V 3T L 2 I. +1 @ V 33T + 82 T La-
From Theorem 4.2 the resulting dynamics of 3 simplify to a self-consistent equation regulated by 8,
82 = - 3T 3T + 7 I. V 3,3 + 4 LaXTP,
82
(8)
where P = XB - Y is the residual. Under our isotropic assumption on the conserved quanitities A SI. these dynamics are exact. Concurrent to our work, Tu et al. (63) finds that 8 appraxi- mately follows these dynamics in the overparameterized setting h >>>>" max(d, e) under a Gaussian initialization (00 o²) of the parameters where o²h is analogous 8.
'The Kronecker sum defined for square matrices G E ReXe and E Rdxd as C D = C +I D.
When / 1 we recover Eq. 4 presented in the single-neuron setting directly from Eq. 7
7

---

Equipped with a self consistnntt equation for the dynamics of 3 we now aim to interpret these dynamics as a mirror flow with 6-dependent potential. As presented in Theorem [B.6] the dynamies of the singular values of 3 can be described as a mirror flow with a hyperbolic entropy potential, which smoothly interpolates between an e and (2 penalty on the singular values for the rich (6 - 0) and lazy (6 + +ooo regimes respectively. This potential was first identified as the inductive bias for diagonal linear networks by Woodworth et al. [14] and the same mirror flow on the singular values is derived from different initialization choice in prior work by Varre et al. [64].
Deep linear networks. As presented in Theorem B.10 we generalize the inductive bias derived for rich two-layee linear networks by Azulay et al.
[9 to deep linear networks. For a depth- (L + 1) linear network, f=:; AT AT Wix, where 3 I =1 WT A, we find that the inductive bias of the rich || . 8... This inductive bias strikes a depth-dependent balance between attaining the minimum norm solution and preserving the initialization direction.
regime is ((3) T,2 (1+1))1
Piecewise Linear Networks
We now take a first step towards extending our analysis from linear networks to piecewise linear networks with activation functions of the form (=))
7=)) The input-output map of piecewise linear network with L hidden layers and h hidden neurons per layer is comprised of potentially (l*")) convex activation regions [65]. Each region is defined by unique activation
pattern of the hidden neurons. The input-output map is linear within each region and continuous at the boundary between regions. Collectively, the activation regions form 2-colorable" convex convexx
partition of input space, as shown in Fig. 5 We investigate how the relative scale influences the evolution of this partition and the linear maps within each region.
Two-layer network. We consider the dynamics of a two-layer piecewise linear network without biases, (a:; 0) = aTo(Wx), where W E Rexd and E RA, Following the approach in Section 4 we consider the contribution to the input-output map from a single hidden neuron k E [h) with
parameters wy E Rd, ag E R and conserved quantity dg nn 7a I 1 62]. However, unlike the linear setting, the neuron's 's contribution to [ 0) is regulated by whether the input xy is in
the neuron's 's active halfspaace Let C' E Rhxa be the matrix with elements Cti = o'(w[xx;), which determines the activation of the th neuron for the şth data point. The dynamics of Ba = Og Wy are,
3x (Th La + Ma (.ww)
Ei= Sha 0) vi).
The matrix M& E Raxd is preconditioning matrix on the dynamics, and when 84 * O. it can be expressed in terms of 3& and 6z. Unlike the linear setting, & E Rd driving the dynamics is not shared for all neurons because of its dependence on Cki. Additionally, the NTK matrix in this setting depends on M₁ and C, with elements Ki, Σk=1 C+a2. Myjr,kk. To examine the evolution of K, we consider a signed spherical coordinate transformation separating the dynamics of 3, into its directional 3 = (gn(as) Be and radial f4 = sgn(as (ax)))) 3x&|| components, such that 3x = - Ax ||
determines the direction and orientation of the halfspace where the 1.th neuron active, while ji determines the slope of the contribution in this halfspace. These coordinates evolve according to,
j4 V 62 4aa
B =
0g + 4njahuu + du (Ia- (a-BB Ex-
(10)
2ph
Downstream. When de < E 0, Rs 1 and the dynamics are approximately 0 = 0 and | 1 Irrespective of Ex. (t) (0), which implies the overall partition map
0.g4: - doesn't change (Fig. 5 bottom), nor the activation patterns C, nor M4. Only μ4 changes to fit the data, while the NTK remains constant. If the number of hidden neurons is insufficient to fit the data, there is a delayed rich alignment phase where the kernel will change, with 16g determining the delay.
Balanced. When & 0, My Tha lu (Ia + BL and the dynamics simplify to, 0, -Vnahhssnn((s)) (I. 3.. and 0,44: -2 -2/na Thel / I Here both the direction and magnitude of 3r evolve, resulting in changes to the activation regions, patterns C. and NTK K. For vanishing initializations where || (0)l + 0 for all & E [b]. we can decouple the dynamics into two
To our knowledge, this property has not been recognized before. See Appendix C. for formal statement.
8

---

100
11000
<= 10000
100000
Phase Es Partition Aiganent
Phase &t Data Firiing
Hamming Disaarc Paromatt
04
Rich
404
Food Partiion Data
0.
Lazy
10 0.0
10 30 100
10
10
Te
(a) Evolution of ReLU network's 's input-output map (h) Hamming and parameter distance over T-ō sweep
Figure 5: Rapid feature learning is caused by large activation changes with minimal parameter movement. (a) We show the surface of two-layer ReLU network trained on an XOR-like task, starting with a near-zero input-output map. [(;; 0b) R 0. The surface consists of convex conic regions, each with a distinet activation pattern, colored by the parity of actíve neurons. A lazy initialization (bottom) maintains a fixed activation partition throughout training, reweighting the hidden neurons to fit the data. In contrast, rich balanced or upstream initialization (top) features an initial alignment phase where the partition map changes rapidly while the input-output map remains close to zero, followed by a data-fitting phase. (b) We show the evolution of Hamming distance in activation patterns and parameter distance, relative to = 0, as function of overall and relative scales (same experiments as in Fig. [b)). Rapid feature learning occurs from a small-r upstream initialization that promotes faster learning in early layers, driving large change in Hamming distance, but a small change parameter space. contrast, small-r downstream initializations require large parameter movement to fit the data in the delayed rich regime.
distinct phases of training (Fig. 5 top), analogous to the rich regime discussed in Section Phase I: Purtition alignment. At vanishing scale, the output (a; 0o) R for all input x, such that the vector driving the dynamics time, the radial dynamics slow down relative to the directional dynamics, and the function's 's output will remain small as each neuron aligns to certain data-dependent fixed points, decoupled from the rest. Prior works have introduced structural constraints on the training data, such as orthogonally separable [50 533 554) pair- wise orthonormal [52]. linearly separable and symmetric [51] or small angle (55), to analytically determine the fixed points of this alignment phase. Phase I1: Data fitting. After enough time, the magnitudes of 3x have grown such that can no longer assume fx; x 0 and thus the residual will depend on all 3x. In this phase, the radial dynamics dominate the learning driving the network to fit the data. However, it is possible for the directions to continue change, and therefore some prior works have further decomposed this phase into multiple stages.
~ - -E -i=1= Cki2,z/ is independent of the other hidden neurons. At the same
Upstream. When oz > 0, Me ss 6g and the dynamics are approximately 0,,² = -6 (Ia (e and 0,f4, = -5, Again, both the direction and magnitude of 3r change. However, unlike the balanced setting, in this setting Mr is independent of B8 and stays constant through training. Yet, as 3r change in direction, so can C, and thus the NTK. This setting is unique, because is rich due to changing activation pattern, but the dynamics do not move far in parameter space. Furthermore, unlike in the balanced scenario where scale adjusts the speed of radial dynamics, here it regulates the speed of directional dynamics, with vanishing initializations prompting an extremely fast alignment phase, as observed in Fig. and in Fig,
Connections to infinite-width. Our study of learning regimes in finite witth two-layer ReLU networks as a function of the overall and relative scale is consistent with existing infinite- width analysis of feature learning. For example, Luo et al. [7] they consider a network f(x) =
1 Ek_1 (w.[x)a) with weights initialized as ag ~ NN, .22 and Wi ~ (o La as width h + 0.. They obtain a phase diagram at infinite width capturing the dependence of learning regime on the overall function scale B.BSw/e /a and the relative initialization scale Ba/ /Bw, each suitably normalized
function of width. The resulting phase portrait is analogous to ours in Fig. (b). where we use the conserved quantity ö rather than the relative scale Ba/Bw. Specifically, there is a lazy regime that includes the NTK parameterization, which always achieved at large scale (as in the large-7 regions of Fig. (b)), but is also achieved at small scale if the first layer variance is sufficiently larger than the second (as in the downstream initializations at small 7 in Fig. (b)) On the other side of the phase boundary is the infinite-width analog of rapid rich learning, where all neurons condense to few directions. This is induced either at small function scale, or at larger scales if Ba/ Sw is sufficiently large, such that W learns fast enough relative to a. The phase boundary, in

---

104
Accu Lur
8.75
Distance
0.25
1071
Pacameter
Hamming
1o
Epachs
PrPras
Steo
10
(a) Kernel Distance
(b) Convolutional Filters (c) Random Hicrarchy (d) Transformer Grokking
Figure 6: Impact of upstream initializations in practice. Here we provide evidence that an upstream initialization (a) drives feature learning through changing activation patterns, (b) promotes interpretability of early layers in CNNs, (c) reduces the sample complexity of learning hierarchical data, and (d) decreases the time to grokking in modular arithmetic. In these experiments, we regulate the first layer's learning speed relative to the rest of the network by dividing its initialization by a. For models without normalization layers, we also scale the last layer's ' initialization by a to preserve the input-output map. a represents standard parameterization, while a >> and a < correspond to upstream and downstream initializations, respectively. See Appendix D.3 for details.
turn, which exists only at infinite width, contains a range of parametrizations, including the mean-field parametrization. More broadly, across width-dependent parametrizations, the random initialization of weights induces a distribution over per-neuron conserved quantities. While the distinction between the NTK and the mean-fielle parametrizations has been extensively studied, both lead the same distribution of per-neuron conserved quantities, which is zero in expectation with a non-vanishing variance. more thorough study of what role the distribution of per-neuron conserved quantities plays in feature learning at finite-widths is left to future work.
Unbalanced initializations in practice. Our analysis shows that upstream initializations can drive rapid rich learning in nonlinear networks. Further experiments in Fig. show that upstream initializa- tions are relevant across various domains of deep learning: (a) Standard initializations see significant NTK evolution early in training [27]. We show the movement is linked to changes in activation patterns rather than large parameter shifts. Adjusting the initialization variance of the first and last layers can amplify or diminish this movement. (b) Filters in CNNs trained on image classification tasks often align with edge detectors [66]. We show that adjusting the learning speed of the first layer can enhance or degrade this alignment. (c) Deep learning models are believed to avoid the curse of dimensionality and learn with limited data by exploiting hierarchical structures in real-world tasks. Using the Random Hierarchy Model, introduced by Petrini et al. 66 as framework for synthetic hierarchical tasks, we show that modifying the relative scale can decrease or increase the sample complexity of learning. (d) Networks trained on simple modular arithmetic tasks will suddenly generalize long after memorizing their training data (68) This behavior, termed grokking, is thought to result from a transition from lazy to rich learning (69 70 71] and believed to be important towards understanding emergent phenomena (72). We show that decreasing the variance of the embedding in single-layer transformer (< 6% of all parameters) significantly reduces the time grokking.
6 Conclusion
In this work, we derived exact solutions to a minimal model that can transition between lazy and rich learning to precisely elucidate how unbalanced layer-specific initialization variances and learning rates determine the degree of feature learning. We further extended our analysis wide and deep linear networks and shallow piecewise linear networks. We find through theory and empirics that unbalanced initializations, which promote faster learning at earlier layers, can actually accelerate rich learning. Limitations. The primary limitation lies in the difficulty to extend our theory to deeper nonlinear networks. In contrast to linear networks, where additional symmetries simplify dynamics, nonlinear networks require consideration of the activation pattern's 's impact on subsequent layers. One potential solution involves leveraging the path framework used in Saxe et al. (73]. Another limitation is our omission of discretization and stochastic effects of SGD, which disrupt the conservation laws central to our study and introduce additional simplicity biases 77,, 75 76 77). Future work. Our theory encourages further investigation into unbalanced initializations to optimize efficient feature learning. Understanding how the learning speed profile across layers impacts feature learning, inductive biases, and generalization is important direction for future work.
10

---

Acknowledgments and Disclosure of Funding
We thank Francisco Acosta, Alex Atanasov, Yasaman Bahri, Abby Bertics, Blake Bordelon, Nan Cheng, Alex Infanger, Mason Kamb, Guillaume Lajoie, Nina Miolane, Cengiz Pehlevan, Ben Sorscher, Javan Tahir, Atsushi Yamamura for helpful discussions. D.K. thanks the Open Philanthropy AI Fellowship for support. S.G. thanks the James S. McDonnell and Simons Foundations, NTT Research, and an NSF CAREER Award for support. This research was supported part by grant NSF PHY-1748955 the Kavli Institute for Theoretical Physics (KITP).
Author Contributions
This project originated from conversations between Daniel and Allan the Kavli Institute for Theoretical Physics. Daniel, Allan, and Feng are primarily responsible for the single neuron analysis in Section 3 Daniel, Clem, Allan, and Feng primarily responsible for the wide and deep linear analysis in Section Daniel primarily responsible for the nonlinear analysis in Section Allan, Feng, and David are primarily responsible for the empiries in Fig. and Fig. 6 Daniel is primarily responsible for writing the main sections. All authors contributed to the writing of the appendix and the polishing of the manuscript.
References
[1] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054 2018.
[2] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685, PMLR, 2019.
[3] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
[4] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameetriizatoo.. In International conference on machine learning, pages 242-252. PMLR, 2019.
[5] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine learning, 109:467-492, 2020.
[6] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[7] Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14 14548, 2020.
[8] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in neural information processing systems, 32, 2019.
[9] Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference Machine Learning, pages 468-477. PMLR, 2021.
[10] Andrew M Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 2013.
[11] Andrew M Saxe, James L. McClelland, and Surya Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116 -133 2019.
11

---

[12] Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to- saddle dynamics in deep linear networks: Small initialization training. symmetry, and sparsity. arXiv preprini arXiv: /06. 15933, 2021.
[13] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. arXiv preprint vxXiv:202.00983) 2020.
[14] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673. PMLR, 2020.
[15] Yuhan Helena Liu, Aristide Baratin, Jonathan Cornford, Stefan Mihalas, Eric Shea-Brown, and Guillaume Lajoie. How connectivity structure shapes rich and lazy learning in neural circuits. ArXiv, 2023.
[16] Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011. /4522, 2020.
[17] Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at widdh limit. Journal of Machine Learning Research, 22(71): 2021.
[18] Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.cm 03466, 2022.
[19] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218, 2020.
[20] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
[21] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Catapults in sgd: spikes in the training loss and their impact on generalization through feature learning. arXiv preprint 230.::2306.04815 2023.
[22] Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue M Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402.044980, 2024.
[23] Yizhou Xu and Liu Ziyin. When does feature learning happen? perspective from an analytically solvable model. arXiv preprint urXiv:2401.0777 100055, 2024.
[24] Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research, 13(1):795 82008899-88888 2012.
[25] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment. 2020(11):1 1 13301, 2020.
[26] Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien. Implicit regularization via neural feature alignment. In International Conference on Artificial Intelligence and Statistics, pages 2269-2277. PMLR, 2021.
[27] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances Neural Information Processing Systems, 33:5850-5861, 2020.
12

---

[28] Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning deep linear networks. arXiv preprint arXiv:1809.10374. 2018.
[29] Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 1004): 1998.
[30] Lukas Braun, Clémentine Carla Juliette Dominé, James E Fitzgerald, and Andrew M Saxe. Exact learning dynamies of deep linear networks with prior knowledge. In Advances Neural Information Processing Systems, 2022.
[31] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparaeeterization. In International conference machine learning, pages 244-253, PMLR, 2018.
[32] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[33] Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. Advances in Neural Information Processing Systems, 35:24446-24455 2022.
[34] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[35] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In International Conference on Machine Learning, pages 10153-10161. PMLR, 2021.
[36] Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv: 1909, 12051, 2019.
[37] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2021.
[38] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-2878. 2018
[39] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv: 181000003 2018
[40] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
[41] Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale training accuracy. Advances in neural information processing systems,
2020.
[42] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint (eXXiv1906.058900 2019.
[43] Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexico- graphic and sensitive margins in homogeneous and non- non-homogencous deep models. In International Conference on Machine Learning. pages 4683 4692. PMLR, 2019.
[44] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on learning theory, pages 1305-1338. PMLR, 2020.
[45] Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. arXiv preprint 22x.i::22000388.00 2022.
13

---

[46] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages I 1832-1841. PMIR, 2018.
[47] Suriya Gunasekar, Blake Woodworth, and Nathan Srebro. Mirrorless mirror descent: A natural derivation of mirror descent. In International Conference on Artificial Intelligence and Statistics, PMLR, 2021.
pages
[48] Zhiyuan Li, Tianhao Wang, Jason D Lee, and Sanjeev Arora. Implicit bias of gradient descent on reparametrized models: On equivalence mirror descent. Advances in Neural Information Processing Systems, 35:34626- -4644, 2022.
[49] Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint arXiv:1803. 08367, 2018.
[50] Mary Phuong and Christoph H Lampert. The inductive bias of relu networks orthogonally separable data. In International Conference on Learning Representations, 2020.
[51] Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34, 2021.
[52] Etienne Boursier, Loucas Pillaud-Vivis and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances in Neural Information Processing Systems, 35:20105- 18, 2022.
[53] Mingze Wang and Chao Ma. Early stage convergence and global convergence of training mildly parameterized neural networks. Advances in Neural Information Processing Systems, 35: 743-756, 2022.
[54] Hancheng Min, René Vidal, and Enrique Mallada. Early neuron alignment in two-layer relu networks with small initialization. arXiv preprint arXiv:2307.12851, 2023.
(55] Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of relu networks. Advances Neural Information Processing Systems, 36, 2024.
[56] Song Mei, Andrea Montanari, and Phan-Minh Nguyen, A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018.
[57] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[58] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):725-752, 2020.
[59] Grant Rotskoff and Eric Vanden-Fijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 15(9):1889-1955, 2022.
[60] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022.
[61] Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.178/11 2023.
[62] Simon Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homoge- neous models: Layers are automatically balanced. Advances in Neural Information Processing Systems, 31, 2018.
[63] Zhenfeng Tu, Santiago Aranguri, and Arthur Jacot. Mixed dynamics in linear networks: Unifying the lazy and active regimes. arXiv preprint arXiv:2405.17580 2024.
14

---

[64] Aditya Vardhan Varr,, Maria- Vladarean, Loucas Pillaud- Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. Advances in Neural Information Processing Systems, 36, 2024.
[65] Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jascha Sohl- Sin--Dickttein On the expressive power of deep neural networks. In international conference on machine learning. pages 2847-2854. PMLR, 2017.
[66] Alex Krizhevsky, Ily Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):84-90, 2017.
[67] Leonardo Petrini, Francesco Cagnetta, Umberto M Tomasini, Alessandro Favero, and Matthieu Wyart. How deep neural networks learn compositional data: The random hierarchy model. arXiv preprint arXiv:2307.02 022199 2023.
[68] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmic datasets, arXiv preprini arXiv:2201.02117 2022.
[69] Tanishq Kumar, Blake Bordelon, Samuel J Gershman, and Cengiz Pehlevan. Grokking the transition from lazy to rich training dynamics. arXiv preprint arXiv:2310.061 10, 2023.
[70] Kaifeng Lyu, Jikai Jin, Zhiyuan Li, Simon Shaolei Du, Jason D Lee, and Wei Hu. Dichotomy of early and late phase implicit biases can provably induce grokking. In The Twelfth International Conference on Learning Representations, 2023.
[71] Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as first order phase transition two layer networks, In The Twelfth International Conference on Learning Representations, 2024. URL httppttt//peeereriiiii
[72] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.0522 2023.
[73] Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The neural race reduction: Dynamics of abstraction in gated networks. In International Conference on Machine Learning, pages 19287-19309. PMLR, 2022.
[74] Daniel Kunin, Javier Sagastuy- Bggastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanies: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:20/2.047288 2020.
[75] Hidenori Tanaka and Daniel Kunin. Noether's learning dynamics: Role of symmetry breaking in neural networks. Advances in Neural Information Processing Systems, 34:25646-25666 2021.
[76] Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: How gra- dient noise attracts sgd dynamics towards simpler subnetworks. Advances in Neural Information Processing Systems, 36, 2024.
[77] Liu Ziyin, Mingze Wang, and Lei Wu. The implicit bias of gradient noise: A symmetry perspective. arXiv preprint arXiv:2402.07193, 2024.
[78] David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters, 74(21):4337, 1995.
[79] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Generalisation dynamics of online learning in over-parameterised neural networks. arXiv preprint 1rXiv: 00000,,..00,....1 2019.
[80] Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference Learning Theory, pages 4224-4258. PMLR, 2021.
[81] Razvan Pascanu, Guido Montufar, and Yoshua Bengio. On the number of response regions of deep feed forward networks with piece-wise linear activations, arXiv preprint arXiv: 1,02,60 2013.
15

---

[82] Guido F Montufar, Razvan Pascanu, Kyunghyun Cho, and Yoshua Bengio. On the number of linear regions of deep neural networks. Advances in neural information processing systems, 27, 2014.
[83] Matus Telgarsky. Representation benefits of deep feedforward networks. arXiv preprint 50xiv:1500.08/01, 2015.
[84] Raman Arora, Amitabh Basu, Poorya Mianjy, and Anirbit Mukherjee. Understanding deep neural networks with rectified linear units. arXiv preprint arXiv: 2016.
[85] Thiago Serra, Christian Tjandraatmadja, and Srikumar Ramalingam. Bounding and counting linear regions of deep neural networks. In International Conference on Machine Learning, pages 4558-4566 PMLR, 2018.
[86] Boris Hanin and David Rolnick. Complexity of linear regions deep networks. International Conference on Machine Learning, pages 2596-2604. PMLR, 2019.
[87] Boris Hanin and David Rolnick. Deep relu networks have surprisingly few activation patterns. Advances neural information processing systems, 32, 2019.
[88] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 1):2278-232, 1998.
[89] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages
2015.
16

---

A Single-Neuron Linear Network
In this section, we provide detailed analysis of the two-layer linear network with a single hidden neuron discussed in Section 3 The network is defined by the function f(z; 0) = awTx, x, where E R and w E IR are the parameters, We aim to understand the impact of the initializations ao, wo and the layer-wise learning rates Vla: Thu on the training trajectory in parameter space, function space (defined by the product / = aw), and the evolution of the Neural Tangent Kernel (NTK) matrix K:
K = X (nw Ia + naww)) X
(11)
The gradient flow dynamics are governed by the following coupled ODEs:
T/ T (XTXaw aw xx A = - a (XTX
a(O) ao
(12)
X aw - x)y)
w(0) = wi-
(13)
The global minima of this problem are determined by the normal equations XTXaw XTy, Even when XTX is invertible, yielding a unique global minimum in function space 3_ (XTX)- the symmetry between a and w. permitting scaling transformations, a - aa and w + w/a a for any * 0 without changing the product aw, results in manifold of minima in parameter space. This minima manifold is a onc-dimensional hyperbola where aw = 3.. with two distinct branches for positive and negative satisfying a = O and T X = O. Except for a measure zero set of initializations that converge to the saddle points, all gradient flow trajectories will converge to global minimum. In Appendix (A.2.5 we detail the basin of attraction for each branch of the minima manifold and the d-dimensional surface of initializations that converge saddle points, separating the two basins.
The set of saddle points {(a, w]} forms (d -)ddieensional subspace
A.1 Conserved quantity
The scaling symmetry between and to results in a conserved quantity 5 E R throughout training, as noted in many prior works [10 62 74 where
= mm Ta |/]
(14)
This can be casily verified by explicitly writing out the dynamics of 6. Define for succinct notation, such that
(XTXaw XTy)
ò am TT;
a (-naw² p) 2na T ap)
0.
The conserved quantity confines the parameter dynamics to the surface of a hyperboloid where the magnitude and sign of the conserved quantity determines the geometry, shown in Fig. 2 A hyperboloid of the form Σ x? Ei=k4 x = a, with a 2 0, exhibits varied topology and geometry based on & and a. It has two sheets when k and one sheet otherwise. Its geometry
is primarily dictated by at as a tends to infinity, curvature decreases, while at a = 0, a singularity occurs at the origin.
A.2 Exact solutions
To derive exact dynamics we assume the input data is whitened such that X Ia and B, = X such that B. * 0. The dynamics of a and w can then be simplified as
à = (w) 3_ al wl[
a()) =
(15)
t 1/o (a.3. - a² d²ww)
(0) wn-
(16)
A.2.1 Deriving the dynamics for μ and Φ
As discussion in Section we study the variables μ = al[w|] | an invariant under the rescale symmetry, and Φ |e]|||
- 8, the cosine of the angle between w and 8. This change of variables can also be understood as signed spherical decomposition of B: a is the signed magnitude of / and o is the
17

---

cosine angle between 3 and B.. Through chain rule, we obtain the dynamics for and o which can be expressed as
00 + no (0113-11 pt)
(( au Wool
(17)
1a 20 3.11
w 3.
(1 (1-6)) (1-0²),
<(0)
(18)
8 + 4m
6
wo |.
We leave the derivation to the reader, but emphasize that a key simplification used to express the sum a + Na [|/l in terms of 8,
a + No |0/l2 = v 82 + 4oa
(19)
Additionally, notice that 7a and The only appear in the dynamics for μ and Φ as the product Ta or in the expression for 6. If we were to define p' nMalt and B = /uu//u then it is not hard to show that the product Th is absorbed into the dynamics. Thus, without loss of generality we can assume the product nh
resulting in the following coupled system of nonlinear ODEs,
i 0 82 4n² (6113_11 -1),
(0)) aol wo
(20)
3.113.11
w 3 ((0) = w 3+
82 + 41,2 (1 (1-6²)-
(21)
We will now show how to solve this system equations for μ and Φ. We will solve this system when 5 0, 5 > 0, and 8 < separately. We will then in Appendix A.2.6 show general treatment how to obtain the individual coordinates of and w from the solutions for fa and Φ.
A.2.2 Balanced 8 =
When d 0, the dynamics for μr, Φ are,
A sgn(n)2p(( 3.11|| m))
n(0) ao W ||
(22)
+ (gn( )²),
((0) w 9 Two ||...
(23)
First, we show that the sign of ju cannot change through training and sgn(n) = sgn(a). Because double cone with singularity the origin
5 0, the dynamics of a and w are constrained (a = 0, w = 0). This point is saddle point of the dynamics, so the trajectory cannot pass through this point to move from one cone to the other. In other words, the cone where the dynamics are initialized on is the cone they remain on. Without loss of generality, we assume op > 0, and solve the dynamics. The dynamics of μ is Bernoulli differential equation driven by dime-dependent signal @|1.3.1- ||. The dynamics of ó is decoupled from and is in the form of a Riccati equation evolving from an initial value de to 1, as we have assumed an initialization with positive ag. This ODE separable with the solution.
((t) = tanh (cs + [13.1(t,
(24)
where c tanh -2 (do). Plugging this solution into the dynamics for μ gives Bernoulli differential equation,
A = 3| tanh (co + ||33, (t) μ 2n1².
(25)
with the solution,
2 cosh² (ca + [13+
m(() = 2 (co + (1 sinh (2(cg + [.3..(t) + cp'
(26)
where C, = 2p1² cosh (cØ) - (2cs +
Note, Φn = - 1, then @ = 0, and the dynamics of a will be driven O, which saddle point.
A.2.3 Upstream 5 0
When 5 >> O, the dynamics are constrained to hyperboloid composed of two identical sheets determined by the sign of op (as shown in Fig. (c)). Without loss of generality we assume ag > which ensures a(t) 0 for all t > 0. However, unlike in the balanced setting, the dynamics of
18

---

and á do not decouple, making it difficult to solve, Instead, we consider v wTAe_ which evolves according to the Riccati equation.
v 1|3.112 du -v²
v(0)
"o..O. dg
(27)
The solution is given by,
vt) 2 Ruo cosh (RI) + (211 12 5vb) sinh (Rt) 2Rcoohh (Ri) + (2 d) sinh (Ri)
(28)
where = 20 3. 11. The trajectory of a(t) is given by the Bernoulli equation,
((((() + a²);
a(O) = do,
(29)
which can be solved analytically using v(t). For au > O, we have that
a(l) = 26.55/ 1| sech² (Y(1)) 4e" 1e. (6² + 4 3. 113) (113.112 aZ) + bb)
1/ a 1.1...12 anbró
-1/2
-6 (ōcohh (2Y V82 + +|1 sinh (2Y (
2hg
where by w. and (()) = 2 / 82 313.111 atanh
From the solutions for
+01+1100.11
M, we can easily obtain dynamics for μr d.
A.2.4 Downstream á <
When 0, the dynamics are constrained to hyperboloid composed of a single sheet (as shown in Fig. (a)). However, unlike in the upstream setting, may change sign. A zero-crossing in a leads to a finite time blowup in v. Consequently, applying the approach used solve for the dynamics the upstream setting becomes more intricate. First we show the following lemma:
Lemma A.1. af + 0 or w.B. B. + 0, then (()) 3, = 0 has at most one solution for 2
Proof. Let (t) = w())T B.. The two-dimensional dynamics of a(t) and w(t) are given by,
w - a(a² - 8),
(30)
W = ==]11._12 a²w.
(31)
Consider the orthant O+ ((a, (]) > 0, w 0}. The boundary 80+ formed by orthogonal
subspaces. On {(a (aa = 0, 2 0), à 2 0. On ((a,w))l 2 0, = 0), à 2 0. Therefore, O+ is positively invariant set. Similarly, O
((a,w)la < 0,w 0) is positively invariant set. On the boundary 80+ U 000 = ((aww) law = 0}. the flow is contained only at the origin = 0, = 0, which represents all saddle points of the dynamics of (a, w). By assumption, (a, w) is not initialized at a saddle point, and thus the origin is not reachable for t 2 0. As a result, the trajectory (a(t), (t)) will at most intersect the boundary 80+ U UO0 once.
From Lemma A.1 we conclude that either a crosses zero, u 8, crosses zero, or neither crosses zero. When doesn't 't cross zero, then v is well-defined for 2 and our argument from Appendix [A.2.3
still holds, leading to solutions for μ., é. When does cross zero, instead of v, we consider v: = w
the inverse of In this case, we know from Lemma A. that wT B. does not cross zero and thus is well-defined for 2 0 and evolves according to the Riccatti equation,
ύ= 1 80 - u2.
v(0)
4g
(32)
w D.
These dynamics have solution similar to Eq. 28). which we leave to the reader. With v(f), we can then solve for the dynamics wT TT. Let T B., then w evolves according the Bernoulli equation,
= w -
w = wO)
(33)
which can be solved analytically using (1), analogous the solution for a(t) in Appendix A.2.3 From the solutions for can easily obtain dynamics for μ, ó.
19

---

A.2.5 Basins of attraction
From Lemma A. we know that can cross zero no more than once during its trajectory. Consequently. identify the basin of attraction by determining the conditions under which changes sign. This analysis is crucial because initial conditions leading to a sign change a correspond to scenarios where initial positive and negative values of do drawn towards the negative and positive branches of the minima manifold, respectively. From Eq. (28) we can immediately see that will change sign when the denominator vanishes. This can happen v + 3, <
-2% - For < 0, this is satisfied if Po < (-ō 8 4 3.1))) which gives the hyperplane w. + 2 (5 0 + = 0 that separates between initializations for which changes sign and initializations for which it does not (Fig.
7). Consequently, letting S be the set of initializations
attracted the minimum manifold with 0, we have that:
st st=((tos (wn, ao)
4n
if 2 0 if i <0 0
(34)
w 3.
Ô +
82 3.12
where the bottom inequality means that 3o sufficiently aligned to 8, in the case of ag 2 or
sufficiently misaligned in the of da O. We can similarly define the analogous S An
initialization on the separating hyperplane will converge to a saddle point where w 3. = a = 0.
5.0
2.5
0.0
-2.5
-5.0
5.0
0.0
-2.5
0.0
2.5
-5.0
Figure Two basins of attraction. For this model, parameter space partitioned into two basins of attraction, for the positive and negative branch of the minima manifold. The surface separating = 0. For given 6, this equation describes hyperplane through the origin. However, given can only be achieved on the surface of some hyperboloid. Thus, the separating surface is the union of the intersections of hyperplane and a hyperboloid, both parameterized by d. This intersection empty if 8 > 0. Initializations exactly on the separating surface will travel along the surface to saddle point where w 3, = a
the basins of attraction is determined by the equation w.B.3 + 2 (ō -
A.2.6 Recovering parameters (a, w) from (H,
We now discuss how recover the dynamics of the parameters (a, w) from solutions for (ni, Φ)) We can recover a and ||w||| from p.. Using Eq. (19h discussed previously, we can show
88 + 4p² +
V82 + 4112 8
a = ==sg((n))
|||||
(35)
2
2
20

---

We now discuss how to obtain the vector w from Ø. The key observation, discussed in Section that w only moves in the span of wh and B.. This means we can express w(t) as
is
w(t) = ct())
(
Bx [|3.]]
+ c2(t)
///moo (
(/
(36)
1|2
:II 115-11
where (t) is the coefficient in the direction of B. and (2 is the coefficient in the direction orthogonal to 3. on the two-dimensional plane defined by wp. From the definition of á we can easily obtain the coefficients
||| and C2 = / c.. We always choose the positive square root for c2, (t) 2 0 for all See Appendix D.2 for experimental details of how we ran our simulations and a notebook generating these exact solutions.
v(t)
a(t)
W1(t)
W2(t)
6=.22
0.5
6=
0.0
6-0,02
5=02
-0.5
5-2
1.5
1.5
8=
1.0
5=0022
1.0
5-02
6=1
0.5
5=2
0.5
1.5
6=-2
-1
0.50 0.75 1.00
00
6= 002
1.0
6=02
0.50 0.75
5-1
5
10
0
5
10
0
5
10
5
10
Time
Time
Time
Time
Figure 8: Exact temporal dynamics of relevant variables in single-hidden neuron model. Our theory recovers the time evolution under gradient flow of the quantities considered in this section, specifically v, ip, and Ç. as well the resulting dynamics of the model parameters {a, W1, W2). The true [* is unit vector pointing in T/ / direction; :30) is a unit vector pointing towards 3z /2, -xx -//4, and π/ / directions, respectively, for each of the three rows. 8 then defines how a(0) and [[) are chosen for particular 3(()) where by convention we choose a(0) > 0.
A.3 Function space dynamics of 3
The network's function is determined by the product 3 aw and governed by the ODE,
aud + dw = a + na www (XTX3 - -XTy)
(37)
XTP
21

---

Notice, that the vector XTp driving the dynamics of [8 is the gradient of the loss with respect to 3, p = VgL. Thus, these dynamics can be interpreted as preconditioned gradient flow on the loss in space where the preconditioning matrix M depends on time through its dependence on a² and ww The matrix M also characterizes the NTK matrix, K = X As discussed in Section 3 our goal understand the evolution of M along a trajectory { Rd t t 0} solving Eq. (37).
First, notice that by expanding || = a²1 ||000 12 in terms of the conservation law, we can show
a² 82 4m |||2 +
(38)
which is the unique positive solution of the quadratic expression n 5a² na 11331|2 = When a² > 0 we can use this solution and the outer product 3 = TwwwT to solve for wwT in terms of B,
082 mo ||3]|2 8 3,3T
ww
(39)
21/a
11
Plugging these expressions into M gives
82 + Mm Tue
| 2
+
8
8 +
1812 - 5 3,97
M
la
(40)
2
2
||
Thus, given any initialization do: we such that a(t)² * 0 for all f 2 0, we can express the dynamics
of 3 entirely in terms of B. This is true for all initializations with 5 2 0, except if initialized on the saddle point the origin. It is also true for all initializations with 5
where the sign of does
not switch signs. In the next section we will show how to interpret these trajectories as time-warped mirror flows for potential that depends on 8. As a means of keeping the analysis entirely in space, we will make the slightly more restrictive assumption to only study trajectories given any initialization Bu such that [|3(()]l > 0 for all 2 0.
Notice, that 7a and yu only appear in the dynamics for 3 as the product nhe or in the expression for 5. By defining 8 Bu and y' = yyahyyy and studying the dynamics of B', can absorb Mae into the / terms in M and the additional factor /Ma7h into the / and terms in p. This transformation of 33 and y merely rescales & space without changing the loss landscape or location of critical points. As a result, from here on we will, without loss of generality, study the dynamics of B assuming The 1.
A.3.1 Kernel dynamics
The dynamics of the NTK matrix K = XMX is determined by M. From Eq, (4 which is derived in this section, we can write M 21188 (Ia + -33)). + (3,3) where B = THT From this expression we see that the change in M driven by two terms, one that depends on the change the magnitude of 3 and another that depends on the change in the direction of 3. As done in the main text, we consider 8 >> 0, 8 < O, and 8 = 0 to identify different regimes of learning. For >> 0, the coefficients in front of both terms vanish, and thus, irrespective of the trajectory taken from 3(0) to 3., the change in the NTK is vanishing, indicative of a lazy regime. For < O, the coefficient for the first term vanishes, while the coefficient on the second term diverges. Here, the change in the NTK is driven solely by the change in the direction of B. This is why for large negative delta we observe a delayed rich regime, where the eventual alignment of 33 to B. leads to dramatic change in the kernel. When = 0, the coefficients for both terms are roughly of the same order, and thus changes in both the magnitude and direction of 3 contribute to change in the kernel, indicative of a rich regime.
B
A.4 Deriving the inductive bias
Until now, we have primarily considered that XTX is either whitened or full rank, ensuring the existence of unique least squares solution 3.. In this setting, 5 inffuences the trajectory the model takes from initialization to convergence, but all models eventually converge to the same point, as shown in Fig. 4 Now we consider the peeapearaeteri setting where we have more features d than observations n such that X low-rank and there exists infinitely many interpolating solutions in function space. By studying the structure of M we can characterize or even predict how 8 determines which interpolating solution the dynamies converge to among all possible interpolating solutions. To do this we will extend time-warped mirror flow analysis strategy pioncered by Azulay et al. [9]
22

---

A.4.1 Overview of time-warped mirror flow analysis
Here we recap the standard analysis for determining the implicit bias of linear network through mirror flow. As first introduced in Gunasekar et al. [46]. the learning dynamics of the predictor can be expressed as mirror flow for some strictly convex potential & (3),
= (= XTT
ptp,
(41)
where p (X - y) is the residual, then the limiting solution of the dynamics is determined by the constrained optimization problem.
B(0c) arg min D4. ((,)
s.t.
X3 = y.
(42)
BER
where De. = 2.) (V is the Bregman divergence defined with o... To understand the relationship between mirror flow Eq. (41 and the optimization problem Eq. (42).
consider an equivalent constrained optimization problem
B(00) arg min Q(3) s.t. XB BER
(43)
where Q(3) = (3) V (3 (3))) which often referred to as the implicit bias. Q(3) is strictly convex, and thus is sufficient show that 3(0o) is a first order KKT point of the constrained
optimization (43). This true iff there exists v E R" such that VQ( = XTv. The goal is to derive from the mirror flow Eq. (1). Notice, we can rewrite Eq. (41) as, (v) = -XTp, which integrated over time gives
V (3(00)) (B, -Vx f p(t)dt.
(44)
The LHS is VQ(3(x)). Thus, by defining v = - p(l)dl, which assumes the residual decays fast enough such that this is well defined, then have shown the desired KKT condition. Crucial to this analysis is that there exists solution to the second- -oddrr differential equation
(3) (VoBV
(45)
which even for extremely simple Jacobian maps may not be true [47) Azulay et al. [ showed that if there exists a smooth positive function g33) Rd (0, oo) such that the ODE,
V (3) . (V (VDDV
(46)
has solution, then the previous interpretation holds for $.)3) with v= J ((((()) As before, is crucial that this integral exists and finite. Azulay et al. (9) further explained that this scalar function g(3) can be considered as warping time 7(t) = 1 ((((t)) on the trajectory taken
in predictor space ((r(t).. So long as this warped time doesn' 't "stall out", that is we require that 7(20)- = co. then this will not change the interpolating solution.
A.4.2 Applying time-warped mirror flow analysis
Here show how to apply the warped mirror flow analysis to the dynamics of 3 derived in Appendix A.3 where Vo M. We will only consider initializations f0 such that [|(((t) for all 2 U, such that M can be expressed as
82
4
3
+
82 +
1
4
8
3,3T
M
la
(47)
2
1|
Computing -1 Whenever ||||| 0, then M is positive definite matrix with a unique inverse
that can be derived using the Sherman-Morrist formula, (A + uoT)-1 = A-1 A 1 1-1
Here
14
we can define A, u, and
80 + 11 +
0 1 5
A
Id. u =
1112 U-U=8
(48)
2
23

---

- After some algebra, M is
First notice the following simplification, uTA-'v
482+411811+55
+1111111
M
33
(49)
882 + 4113112 + +8)
|| 82 + 3|
To make notation simpler we will define the following two scalar functions, + 4r
2
fo(x)
ho(2) zV8 Ax (vo² 4x 60)
(50)
v02 + 4r
such that we can express M- (113112) Ia h8 (13112) 3,3..
Proving M is not Hessian map. If . is the Hessian of some potential, then show that the dynamics of /3 are mirror flow. However, from our expression for MM-1 we can actually prove that not a Hessian map. As discussed Gunasekar et al. [47], symmetric matrix (() is the Hessian of some potential (33) if and only if it satisfies the condition,
(H2
(l1,(33
vB E IR", Vi,j, E [m]]
(51)
03
033
We will use this property to show M is not Hessian map. First, notice this condition trivially true when i = j = k. Second, notice that for all + + k,
aM 05
aM 03,
sk
2 (13113) 8,3,3
(52)
Thus, M-1 is a Hessian map if and only if for all i i M- the LHS is
BM. oB;
aM, 1 03.
Using our expression for
OM,1
f4 (11313) ) B3 2Vhh (11813) B508
(53)
03,
while the RHS is
0M, -1
05, ha (118 3 2V/8 (131 3,5
(54)
Thus, M-1 a Hessian map if and only if 2V [a hk(x) 0. Plugging our definitions of fz(z) and hz(x) we find
[x)
(55)
V8a 1z( V8 + 4r 8)2
which does not equal zero and thus M- is not a Hessian map.
Finding scalar function (5(x) such that gs(()) (11 is a Hessian map. While we have shown that -- is not Hessian map, it is very close Hessian map. Here we will show that there exists scalar function ga(x) such that 96 (I[3)13 M- is Hessian map. For any ga(x)) can define gs (113 -1 1 in terms of two new functions (xx) and ha(x) evaluated at æ ||311
g8 (118119) M-1 g8 (113112) Ss (11812) g8 (113112) ha (118112) 1,3T. fa(1)0113)
(56)
he([))1P)
Thus, as derived in the previous section, we get the analogous condition on fa(x) and hg(x)) for gs (113112) M- to be Hessian map,
2 (o(()))((a) + g(x) + (a(x))h(())
(57)
ha(m)
Rearranging terms we find that gz( must solve the ODE
Vgs(2)= -(fff(=)) (2V + hs(x)) 95(x).
(58)

---

Using our previous expressions (Eq. 50 and Eq. 55
we find
1
- -(ff(()) (2V + hs hz())) V82 4z(v82 82 + 4x +4x 8)
(59)
which implies (xr solves the differential equation, Vgs(a)) = (8 The solution is go (z) = cv V52 + 4x + 6, where e E is a constant. Let 1. Plugging in our expressions for go (11311). (113 has (11331)), we get that
(a(m)
8214
2
V8 1811+45
98 (118113) M M-1
Ia
333
(60)
|| |13/880 + 4113112
82 4 3
a Hessian map for some unknown potential (3)
Solving for the potential (3)).. Take the ansatz that there exists some function scalar q(x) such that (B) qo (I]@11) + cs where cæ constant such that (3(3)) for all and 4b((0) 0.
The Hessian this ansatz takes the form,
b.) = V V (1311)
Vg((l31) |
((131111 |
la
3
(61)
|
Equating terms from expression for 96 (118113) M (equation 60 ()(()))
get the expression for
2|/1
Vg(1131)) = V + 2||31|2 +
(62)
which plugged into the second term gives the expression for (31(1111..
182441181366
2
+1133115
8 + + 3
+
V ((1
(63)
82 V
113 +
88 + 11.
80 + 111311
We now look for function (ax such that both these conditions (Eq. 62 and Eq. 63 are true. Consider the following function and its derivatives,
((x)
1
1/02 -20)
28) v V82 82 + 4x2 +
(64)
3
2x
Vg(x)
(65)
Vo + 2 + 5
08 4x + 8
V2((2)
(66)
8 + 4x
Letting x = ||ل]| notice ( and (²(([1B1]) satisfies the previous conditions. Furthermore.
V2((x) > 0 for all 5 long # 0 and thus ((x) is convex function which achieves its
minimum at x = 0. Thus, the constant Cs = -q(0) is
0
if < = max 0, (gg))()) V
(67)
V266
if >
3
and the potential (B) is
(B) 1 (88++13- 28
VV B| +8 max
0, sgn(6) v2|0/3
(68)
Finally, putting it all together, express the inductive bias in Theorem 3.1
25

---

A.4.3 Connection to Theorem 2 in Azulay et al. [9]
We discuss how Theorem 3.1 connects to Theorem 2 in Azulay et al. 19, which we rewrite:
Theorem A.2 (Theorem 2 from Azulay et al. (9). For a depth 2 fully connected network with
single hidden neuron (h- = 1), any 6 2 O, and initialization So such that Bo F O. if the gradient flow solution B(00) satisfies X (00) = y, then,
3(xo) = arg min 98 (1131) + T s.t. X3= DER
(69)
x-- ( ( / +2
where (z))
and 2 === ViBooll
6 Ap
2 1
The most striking difference is in the expressions for the inductive bias. Azulay et al. [ take an alternative route towards deriving the inductive bias by inverting M in terms of the original parameters and w and then simplifying M in terms of B, which results in quite a different expression for their inductive bias. However, they are actually functionally equivalent. It requires a bit of algebra, but one can show that
( = 2V2
28 ((1 + cg.
(70)
3
Another important distinction between our two theorems lies in the assumptions we make. Azulay al. (9] consider only initializations such that 2 and 90 * 0. We make less restrictive assumption by considering initializations Bo such that [|](()]] > 0 for all t 2 O. which allows for both positive and negative 6. Except for measure zero set of initializations, all initializations considered by Azulay et al. [9] also satisfy our assumptions. In both cases, our assumptions ensure that M is invertible for the entire trajectory from initialization to interpolating solution. However, it is worth considering whether the theorems would hold even when there exists a point on the trajectory where M is low-rank. As discussed in Appendix A..3 this can only happen for an initialization with 8 < 0 and where the sign of changes. Only at the point where a(t) = 0 does M become low-rank. A similar challenge arose in this setting when deriving the exact solutions presented in Appendix [A.2.4] We were able to circumvent the issue in part by introducing Lemma A.I proving that this sign change could only happen at most once given any initialization. This lemma was based on the setting with whitened input, but similar statement likely holds for the general setting. If this were the case, could define M at this unique point on the trajectory in terms of the limit of M as it approached this point. This could potentially allow to extend the time-warped mirror flow analysis to all initializations such that ||&ol| > O.
A.4.4 Exact solution when interpolating manifold is one-dimensional
When the null space of X is one-dimensional, the constrained optimization problems in The- orem 3.T and Theorem A.2 have an exact analytic solution. In this case we can parameterize all interpolating solutions 3 with a single scalar E R such that B = 3, + ov where X = 0 and ||u||| = 1. Using this description of 8, we can then differentiate the inductive bias with respect to a, set to zero, and solve for a. We will use the following expressions,
82
8
V Vag(2)=
= sign(T)) V
12
Va
Vo= T3 = ZTv.
(71)
2
4
2'
|
We will also use the expression, ||3112 = ||3,1 1 following equation for a,
+ a²
Pulling these expressions together we get the
|13 + a²
8² 8
a
22T
(72)
4
2
/
/1.
a
3
If we let k = 2 3
the solution for a is
12 +
122 +
2
a= a=
+
+
+ 12
(73)
2
2
This solution always works for the initializations we considered in Theorem 3.1 Interestingly, it appears that 9 3. - ar also works for initializations not previously considered. This includes trajectories that pass through the origin, resulting in change in the sign of a.
26

---

B Wide and Deep Linear Networks
Here discuss how our analysis techniques, developed in the previous section for single-neuron linear network, can be extended to linear networks with multiple neurons, outputs, and layers.
B.1 Wide linear networks
We consider the dynamics of two-layer linear network with h hidden neurons and outputs, [(;; 0) = ArWx, where W E Rexd and A E RRxxe_ We assume that h 2 min(d, c), such that this parameterization can represent all linear maps from RR + R.. As the single-neuron setting, the rescaling symmetry in this model between the first and second layer implies the h x h matrix = A Wo determined at initialization remains conserved throughout gradient flow 62]
This can be easily shown from the temporal dynamics of A and W,
À - W (X3 - Y): W = - Y)AT.
(74)
(75)
Extending derivations in [30, the NTK matrix can be expressed as
K = (Le @ X) (T A @ 1) WTW) (I. 2 XT).
(76)
where @ and @ denote the Kronecker product and sum respectively. The Kronecker sum is defined for square matrices C Rexe and D E Raxd C' e D C @ La +I D.
B.1.1 Parameter space dynamics
Inspired by our analysis of the single-neures setting, we introduce two coordinate transformations to study the parameter space dynamics of wide two-layer linear network. In both analyses assume whitened input X X = Ia and let 7/ = The 1. However, we will find that the analysis of the dynamics in function space, for general unwhitened data, is more tractable.
Parameter dynamics when e = 1. Drawing insights from our analysis of the single-neuron scenario (h = 1), we might consider a combination of hyperbolic and spherical coordinate transformations to study the parameter space dynamics of a wide two-layer linear network. We consider the following two quantities for each hidden neuron & [h]:
B |w |
ΦR
(77)
We will also consider matrix quantity Q Rhxh with elements QRK resulting dynamics for μ and & can be entirely written in terms p, Φ A:
we
The
wprll
j = DDagg( 4)2 + 4D)ag(()22 (Ø Qu),
(78)
à MDiag(u) (1118.112 -T#) + Diag())Quu
(79)
where M 2 (iaa(A)²= 4 (Diag(u) Diag(A)) (4)) Using the conserved structure of A can express Q as a function of μ and M,
Q = Mp M - 1/2 1/2
(80)
This approach yields a coupled nonlinear dynamical system with 2h variables. Imposing additional assumptions the initialization, such permutation invariance between hidden neurons, can simplify the system of differential equations. A similar approach was used by Saad and Solla (78) to derive set of differential equations for soft committee machine model, capturing its online learning dynamics a reacher-student setup, which Goldt et al. [79 extended to its generalization error dynamics.
Parameter dynamics when = h. In this analysis we assume an initialization such that the conserved quantities A öI, an assumption we will discuss further in Appendix [B.66) and that A invertible throughout training. Let 3, XTY, which for whitened input, is the unique minimum of the dynamics in function space. We consider the variable A W 8, E Rexe Using the identity
27

---

that A = - A'AA and our assumption on A, we find that the matrix // evolves according to the matrix Riccati ODE,
V 31 8v v².
(81)
Additionally, consider the variable C' CA=ATA, A, which evolves according to the matrix Bernoulli ODE
Č= C(v + oI,) + (v+ C -0
(82)
Taken together we have found a change of variables, analogous the one introduced in Ap- pendix[A.2] for the single-neuron setting, that evolves according to a matrix Riccati and Bernoulli equation,
à == 3 dv - v²: Č= C(v oI,) + (v + TC 20
(0) = A WoB., C(0) Aooo
(83)
(84)
However, solving this system exactly as did in the single-neuron setting is challenging. Unless we assume that y/ and BI B. share the same eigenspace allowing us to decouple the dynamics of v into of scalar Riccati equations the system cannot be easily solved. Instead, we will find that the dynamics of the product WT in function space is more tractable and requires fewer assumptions.
B.1.2 Function space dynamics
We consider the dynamics of 3 = WTA E Raxe in function space, which is governed by the ODE,
= WT A X Y) 7) W -
(85)
Vectorizing using the identity vec(ABBC) = (CT @ A)) equation 85
becomes
vec 3 vec (1 - T Y)L),,
(86)
= (Te La + noLe
W
(87)
(nw na WTW) vec XTY).
(88)
As in the single-neuron setting, we find that the dynamics of B can be expressed as gradient flow preconditioned by matrix M that depends on quadratics of A and W.
B.1.3 Proving Theorem 4.1
We first prove Theorem 4.1
Consider single hidden neuron k [h] of the multi-output model defined by the parameters wy Rd and ag E R Let By wya be the Rdxe matrix representing the contribution of this hidden neuron to the input-output map of the network = the two gram matrices 3 E Reee and Sx3T e Rdxd
Σk=n 8k- Consider
3 3
BR ||ay (laxl2wyee Pww..
(89)
Notice that we can express || 3 as
| 11 TT(333) - ||au||||00,111
(90)
At each hidden neuron we have the conserved quantity
8 7m | 11?
na 12 6g where d5. E R.
Using this quantity we can invert the expression for ||
f to get
V
18
F dz
|
(91)
217wr
V5 +
44/a
1
| 12
(92)
21%a
As long 1, then the surface of this d + hyperboloid is always connected, however its topology will depend on the relationship between d and
28

---

When ||5 > 0, we can use these expressions to solve for the outer products ax and fmsm in terms of 2g and òg.
Vog + Mna
13x F
8%
BT
(93)
2/h
11
08 + 4/a
11xx 1/
8/
Bk
(94)
2na
11
By substituting these expressions into the decompositions ATA Ek Wx W. we derive the representation for M presented in Theorem
Σ
ax and WTW
-
4.1
M
Σk=
Mr where
V88 + Th 3& P + 5k
ST SK 1
Vog + 4n
The
5
11
BKST.
Mx =
+
(95)
2
1
2
| ||
B.1.4
Understanding M when there is a single-neuron /1
When there is single-hidden neuron h min(d, c) 1. the expression for M presented in Theorem 4.] simplifies allowing us to precisely understand the influence of 6 on the learning regime.
When = c 1, then
AT = Therefore, Eq. ( simplifies to |
v 82 + no Ma 11 + ó Ia + 82 + 1a 4 1
8 1,9T
M
(96)
2
2
||
and we recover Eq. simplifies to,
presented in Section
3
When h d 1. then
BBT
1 and thus Eq.
M V 82 + nm 4 3 +
BTB |1
88 + na B1
+
I
(97)
2
2
In both settings, M is the weighted sum of the identity matrix and a rank-one projection matrix. While these equations are strikingly similar there is an interesting distinction that arises in the limits of 5. As 8 expression for M becomes proportional the rank-1 projection S7S T
then the first expression for M becomes proportional to La, while the second Conversely, as 6 -00,
then the first expression for M becomes proportional to the rank- projection [| while the second
A8T7
expression for M becomes proportional to I.. When h = d and thus in both limits of - too, becomes a constant independent of 3. all settings, when 5 M depends 8. In other words, the influence of ó on whether the dynamics are lazy, rich, or delayed rich, crucially depends on the relative sizes of dimensions d, h, and c.
1, then M = =V82 Ma ||
B.1.5 Interpreting M in different limits and architectures
We now seek to more generally understand the influence of the conserved quantities 6, and the relative sizes of dimensions d, and on the learning regime. For a matrix Rdxe, let Row(A) C R and Col( () S Rd denote the row and column space of A respectively.
Theorem B.1. The dynamics are the lazy regime, for all 2 0, i + oo for all k E [h] and there exists least squares solution 3,, E Iddxe such that
(Ú ((RO((B2(0))) k=
Row (B.) C Span
(98)
or 6% + -000 all E [h] and there exists a solution such that
Ü
f Col (Bx(0))
Col(3.) S Span
(99)
k=1
29

---

Proof. As 62 + 00, o0, - 16xl 8.0
&Ig implying 34 10x DC 2/
P Notice that 15+11
1.00c 11 11
11
is the unique orthogonal projection matrix onto the one-dimensional row space of 3a. Thus, the dynamics of each 3g follow projected gradient descent in their row space. As a result, Mr will not change and thus the NTK will be static. By assumption there exists least squares solution B, such that the rows of B. are in the span of the rows of 3g. Thus, a solution will be reached as t + oc, while the Mr remain static.
As 0z -o for all & [M], M₁ Ie @6 |ōk|
BAAT
and analogous argument can be made. Pell
Note that the assumptions in Theorem B. can be more intuitively expressed in terms of the parameter space (W, Except in highly degenerate cases, the assumption (ow(3.) S Span (Uk=1 Row (8x is equivalent to the existence of 8, whose rows lie in the span of {ax k=)) or, equivalently, to the existence of matrix W such that 9. WT A(0). Similarly, the condition (ol( _ Span (Ui_ Col (3x(()) is in most cases equivalent to the existence of matrix A such that 3. WO
A direct consequence of Theorem B. is that networks which narrow from input output (d > c) must enter the lazy regime with probability 1 as all 5g - oo whenever 2 e and assuming independent initializations for all Bx. In this case, the rows of {3],..., 3r, span all of R and thus the condition on the least squares solution is trivially true. By the same logic, networks which expand from input to output (d < c) do so all 8% + -0 whenever h d and assuming independent initializations for all 5g. Additionally, when h 2 max(d, c) and assuming independent initializations for all Se. then all networks enter the lazy regime as either all 8x + oo or all 8% + -00.
Another interesting implication of Theorem B.I is that if there does not exist a least squares solution B. with rows in the span of the rows of (31,.... 8,,, then the network will enter delayed rich regime as all 6 -> oc, where the magnitude of the 6% will determine the delay. In this setting, the network is initially lazy, attempting to fit the solution within the row space of the 8g, but eventually the direction of the rows must change in order to fit the problem, leading to a rich phase. A similar statement involving the columns of B, true all 8R -00.
B.1.6 Simplifying M through assumptions on A
We now consider how introducing structures on A can lead to simpler expressions for M. A natural assumption to consider is the following:
Assumption B.2 (Isotropic initialization). Let A E RRix and W E R be initialized such that ōI,
A = Tu A() na Paa(()) (O))(())
In square networks, where the dimensions the input, hidden, and output layers coincide (d = c). and the weights are initialized as A ~ (0, o²/ /c)c and W ~ (0,, /d/d). this assumption is naturally satisfied with o2 og as the dimension h -> oc. However, a limitation of this assumption is that for general ö it requires h < min(d, e). Specifically, when d > 0, the isotropic initialization requires that A(O)) > O, which implies h S c. Similarly, when 8 < O, the isotropic initialization requires that W (O)W (O)T > O. which implies h < d. Now we prove two important implications of the isotropic initialization assumption.
Lemma B.3. Let A ål, If either 8 > or 8 < and h 2 d, we have that
1
8
82
WTW
la + Tw 38T + - Is
(100)
7o
2
Proof. The quantity Th AAT - 7a WWT = SI, is conserved in gradient flow. Multiplying on the left by WT and the right by W we have that
na (WT + SWTW 33T.
(101)
Completing the square by adding
8² 4no
La to both sides and dividing by na we get the equality,
WTW +
8
La
2
8²
Ia +
The 33T
(102)
21aa
4
aa
30

---

For 8 2 0, WTW +
5
Ia = 0. For 5 < 0, then we know from the conserved quantity that
24
2n I, > O, which implies when 2 that WTW + 2% result, we can take the principal square root of each side,
WWT+ + 2n,. In Th AAT
Ia - As
8 La
8²
WTW +
Ia
BBT,
(103)
2M
V
4ma
which rearranged gives the final result.
Lemma B.4. Ler A òIn. If either 8 < or 6 and h 2 C, we have that
8 Ie
82
A A
+
H noT Le
(104)
Th
2
Proof. The proof is analogous to the proof of Lemma B.3
From Lemma B.3 and Lemma B.4 we can prove Theorem 4.2 as shown below.
Proof.
We start from
vec
(Tw AT
Ne WTW)
-
(105)
M
Plugging in expressions for WTW from Lemma write,
B.3
and ATA from Lemma B.4
directly
82
8 La + V
82
M
I I.
1 no +
Ic
+
Nall:: +
In
(106)
4
82
82
ToThe BT
&
2 V noT
(107)
4
From this expression for (3) we can easily consider how it simplifies in limiting settings of ō:
5
M >
Mo BT B
La I. @ P3T
5
(108)
5 - 30.
As -> +000, M >> dlde. and the dynamics are lazy. this limit, the dynamics of 3 converge to the trajectory of linear regression trained by gradient flow and along this trajectory the NTK matrix remains constant. When 0, /Tunuu, @ I + I, & Vijaauw and the dynamics are rich. Here the NTK changes in both magnitude and direction through training. In the section will attempt to better understand these dynamics for intermediate values of 8 through the lens of mirror flow.
B.1.7 Deriving a mirror flow for the singular values of
For matrix 8, the dynamics of one of its singular values given by ô (8, where u and the corresponding left and right singular vectors. This equality can be derived from chain rule and the fact that ||u|| = ||ll| 1:
Be + uT Br + Bi úuuo + u Bu + ouT = uT uT...
(109)
In the last equality used that fact that for any vector with fixed norm, || 2 0. Letting diag Raxe + pmmind,e) be the operator that, given rectangular matrix, returns vector of the elements on the main diagonal, we can then write.
diag (UT (V)
(110)
31

---

where À E Rmii(dde) is the vector of singular values of 3. In the following lemma, we use the shared singular vector structure between 3 and and W rewrite these dynamics as
À MV
(111)
where M is diagonal matrix and V Cs is the gradient of the loss with respect the singular values of 3. Without loss of generality we consider
Y/u 1.
Lemma B.S. Let A ōI, We then have that X = MV xh, where M Ruin(de))xmi((eee) is diagonal matrix with
(V82+4X) +
S min(d,
Ma
(112)
0
otherwise
Proof. First note that
À diag(UT (U)
(113)
= diag V
(114)
= diag (U ( ( U -Y))
(115)
where we let WTW = UE UT and AT - VT3 using the fact that, under I. the
eigenvectors of ATA are the right singular vectors of 3 and the eigenvectors of WTW are the left singular vectors of B. This expression rewrites as
Á -Mdiag (UTXT(XB Y)V)
(116)
where M E [miindd))xmi((da is diagonal matrix with Ma (E3)se + (Ew) )-) For i & min(d, h, c), one can show that Mü 0 + 4X2. This is because for min(d, h, c).
(E = (Ew) 8 from the conservation law and (Ely (E))) = X from the definition of X2. which is a quadratic equation in (E )a- If h min(d, c) then M, = 0 for i > min(d, c) accounting for rank deficiency of both A and W in this case. Additionally, in our setting of MSE loss, it is straightforward to show that
A. Together this implies (Ely )a (5 + (E))))
OC
= XU - (Y))
(117)
We then have that VxC = diag (U (U (X3 which, combined with our expression for M, completes the proof.
Leveraging Lemma B.5 we can show that the singular values of 3 evolve under mirror flow in the following theorem.
Theorem B.6. == SI, and assume 2 min(d, c) and o We then have that the dynamics of X, the singular values of B, are given by the mirror flow
A (==(v VXC.
(118)
where (A( min((d))) 98 (Ai) and qs the hyperbolic entropy potential
(1(22 2x
(5(x)= = 2x sinh
4x + 82 + +181
(119)
|8]
Proof. When A öI, then by Lemma B.5 the dynamics of the singular values of B can be expressed as À = - -V C. Furthermore, when h 2 min(d, c) and 8 * 0, we have that M = V + 4 min(d,e)) where X2 is element-wise, which is always invertible. Observe, this expression for M is the inverse Hessian of the potential (A)) = E: go (A,) for qs specified in the theorem statement. Thus, the dynamics for the singular values are the mirror flow = - =-(v2t,(>)) -1 V.C.
32

---

Theorem B.6 implies that the dynamics for the singular values of 3 can be described as mirror flow with d-dependent potential. This potential was first identified as the inductive bias for diagonal linear networks by Woodworth et al. [14]. Termed hyperbolic entropy, this potential smoothly interpolates between an e and e2 penalty on the singular values for the rich (ō -> 0) and lazy (6 + +00) regimes respectively. Unfortunately, in our setting we cannot adapt our mirror flow interpretation into a statement on the inductive bias at interpolation because the singular vectors evolve through training. If we introduce additional assumptions task-aligned initialization such that the singular vectors of 30 are aligned with those of B. - we can ensure that the singular vectors remain constant and thus derive an inductive bias the singular values. However, in this setting the dynamics decouple completely, implying there is no difference between applying an l' or l penalty on the singular values. Consequently, even though the dynamics will depend on 6, the final interpolating solution will be independent of ó, making a statement on the inductive bias insignificant.
specifically, whitened input data (XTX La) and
B.2 Deep linear networks
We now consider the influence of depth by studying a depth-(L + 1) linear network, f(x;;)
a II',, Wiz, where W1 E Rhxd, W1 E Rhxh for 1 < 1 S. L, and a E RR., We assume that the dimensions d = h and that all parameters share the same learning rate 17 1. For this model the predictor coefficients are computed by the product 8 = I WTa Rd, Similar our analysis of a two-layer setting, we assume an isotropic initializations of the parameters.
Definition B.7. There exists a d E R such that aa - WI oI, and for all I E [Z. 1]
WI,, W4+1 = W.W.T.
This assumption can easily be achieved by setting a = 0 and W a0; for all I E [Z]. where O E Rdxd is an random orthogonal matrix and a > 0. In this case N aa Further, notice this parameterization is naturally achieved in the high-dimensional limit as d + oo under standard Gaussian initialization with a variance inversely proportional with width. As in the two-layer setting, this structure of the initialization will remain conserved throughout gradient flow. We now show how two natural quantities of B, its squared norm || and its outer product BBT, can always be expressed as polynomials of ||a][ and WTW; respectively.
Lemma B.8. For a depth- linear network with square width (d = h) and isotropic initialization. then for all t 2 0,
|| |all2 (Ilal1- -0112-8) 6)" 3 (WIW1) 1+1 + +W(WW))
(120)
(121)
Proof. The norm of the regression coefficients is the product ||,||| at (IIL wi) (I' wi) T Using the conservation of the initial conditions between consecutive weight matrices, WT,, W W.W,, we can express this telescoped product as || aT WW> When plugging in the conservation between last two layers, this implies || = T (aat 8I,) the desired result.
)a a, which expanded gives
The outer product of the regression coefficients is 0,3T = (I' w;) aaT (II' Wi). Using the conserved initial conditions of the last weights we can factor the outer product as the sum, 3,3T (T' wi)" WLWZ (II'_ wi) + & ( wi) (II' Both these telescoping products factor using the conservation of the initial conditions between consecutive weight matrices giving the desired result.
We now demonstrate how the quadratic terms |a|2 and WTW, significantly influence the dynamics of 8, similar to our analysis in the two-layer setting.
Lemma B.9. The dynamics of 3 given by a differenntial equation á - -XX where M is positive semi- deninte matrix that solely depends on ||a|2. WTW,, and 6,
T--1
(1aal12 5Y (WTWa) (-1))
M = (WTW,) + ||a|2
Σ
(122)
33

---

Proof. Using a similar telescoping strategy used in the previous proof we obtain the form of M.
Finally, we consider how the expression for M simplifies in the limit as 5 0 allowing us to be precise about the inductive bias in this setting.
Theorem B.10. For depth-(L (L + 1) linear network with square width (d = h) and isotropic the gradient flow
initialization 3o such that |[t(t)] > 0 for all t Z O, then in the limit as
solution ((0o) satisfies X (00) = y, then,
3(x0) = arg min BERd
L + 1 L+2
|||||
L+2
3(0)
B s.t. XB y-
(123)
[13(0))
Proof. Whenever ||||| > 0 and in the limit as 6 + O, then can find a unique expression for |||]
and WT terms of ||@|| and 3,3,
|a|l2 ||||| पि
WTW; || 2 3BT.
(124)
Plugged into the previous expression for M results in positive definite rank-one perturbation the identity,
M = ||||| - La + || पे. BBg
(125)
Using the Sherman-Morrison formula we find that M- is
M = ||||| ||:
|||||
42
38T
(126)
+++1
We can now apply time-warped mirror flow analysis similar to the analysis presented in Ap-
pendix A.4 Consider the time-warping function ga(11]1) = ||||| पीं and the potential ((3)
3
122
then its not hard to show M-1 98 ( (1 v This gives the desired re-
2-2
sult.
This theorem is a generalization of Proposition 1 derived in 9 for two-layer linear networks in the rich limit to deep linear networks the rich limit. We find that the inductive bias, Q(3)
(+ + || , 3JB, strikes depth-dependent balance between attaining the minimum norm solution and preserving the initialization direction.
34

---

C Piecewise Linear Networks
Our goal to extend the tools developed in our analysis of linear networks to piecewise linear networks and understand their limitations. We focus on the dynamics of the input-output map, rather than on the inductive bias of the interpolating solutions. As discussed in Azulay et al. [9, Vardi and Shamir [80), extending mirror flow style analysis directly to non-trivial piecewise linear networks is very difficult or provably impossible. In this section, we first describe the properties of the input-output map of a piecewise linear function, then describe the dynamics of a two-layer network, and finally discuss the challenges in extending this analysis to deeper networks and potential directions for future work.
Here, we elaborate on the theoretical results presented in Section 5
C.1 Surface of a piecewise linear network
The input-output map of a piecewise linear network f(; f), with I hidden layers and h hidden neu- rons per layer, is comprised of potentially O(ha)) connected linear regions, each with their own vec- tor of predictor coefficients [65) The exploration of this complex surface has been the focus of nu- merous prior works, the vast majority of them fo- cused on counting and bounding the number of lin- regions
Lase
(a) With Biases
(b) Without Biases
Figure 9: Surface of a ReLU network. Here we depict the surface of three-layer ReLU network f(x; 0) R² - R with twenty hid- den units per layer at initialization, comparing configurations with biases (left) and without biases (right). The network with biases parti- tions input space into convex polytopes that tile input space. The network without biases partitions input space into convex conic sec- tions emanating from the origin. Each region exhibits distinct activation pattern, allowing the partition to be colored with two colors based on the parity of active neurons. The network operates linearly within each region
function of the width and depth 85 86 87). The central object in all of these studies is the activation region.
88
82
83
84
65
Definition C.I. For a piecewise linear network (a;;0,, comprising N hidden neurons with pre- activation 2, (2; 0) for i E [ let the activation pat- tern A represent an assignment of signs a, E { {-1.1} to each hidden neuron. The activation region R(A, 0)
is the subset of input space that generates A,
R) = {x E Rd Vi (1; > 0}. (127)
The input-output map linear within each non- con-empty activation region and continuous at the boundary be- and maintains continuity across boundaries. tween regions. Linearity implies that every non-empty?
Compttyaativvatiion activation region is associated with a linear predicior vector Ax E Rd such that for all x E R(A; 0), 3. [=r(x; 0). Continuity implies that the boundary between regions is formed by hyperplane determined by where the pre-activation for neuron is exactly zero, {z : 2, (z; 0) 0}. When the neighboring regions have different linear predictor: [10 then this hyperplane is orthogonal to their difference, which is a vector in the span of the first-layer weights. Taken together, this implies that the union of all activation regions forms convex partition of input space, as shown in Fig. 9 We now present a surprisingly simple, yet to the best of our knowledge not previously understood property of this partition:
Proposition C.2 (2-colorable). If [xx; 0) lacks redundant neurons, implying that every neuron influences an activation region, then the partition of input space can be colored with two distinct colors such that neighboring regions do not share the same color.
The justification for this proposition is straightforward. There is one color for regions with an even number of active neurons and another for regions with an odd number of active neurons. Because I (a; 0) lacks redundant neurons, there does not exist a boundary between activation regions where two neurons activations change simultaneously. In this work, we solely utilize this proposition for visualization purposes, as shown in Fig. 9 Nonetheless, we believe it may be of independent interest as it strengthens the connection between the surface of piecewise linear networks and the mathematics of paper folding, connection previously alluded to in the literature [82]
WWhile it is trivial to see that for network f(z; 8) with N hidden neurons there are 2 distinct activation patterns, not all activation patterns are attainable. See Raghu al. (65 for a discussion.
iI is possible for neighboring regions to have the same linear predictor. Some works define linear regions as maximally connected component of input space with the same linear predictor [87].
35

---

C.2 Dynamics of a two-layer piecewise linear network
We consider the dynamics of two-layer piecewise linear network without biases, f(z; a aTa (W)), where W E RAxd and E R* The activation function is a(z) m2x for E [0, 1), which includes ReL.U = 0 and Leaky ReLU y E (0, We permit h > d, which in the limit h - co, ensures the network possesses the functional expressivity represent any continuous nonlinear function from R to I passing through the origin. Following a similar strategy used in Section we consider the contribution the input-output map from a single hidden neuron E [A] with parameters Wk E Rd and dg E RR As in the linear setting, each hidden neuron is associated with conserved quantity, 5% Mu - T/a ||w Il?. Unlike in the linear setting, this neuron's 's contribution to the output fx,; 8) is regulated by whether the input is in the neuron's s active halfspace, {ar E IRd : w > 0}. Let C E Rhxx be the matrix with elements Cji = a o'([[rr).
which determines the activation of the hth neuron for the jth training data point. The subgradient and d'(z) = y if z O. These activation functions exhibit positive homogeneity, implying (z) o' (z)z. Thus, we can express a = Ck x..
d'(z) - if > o' (z) E [y, 1] if =
allowing us express the gradient flow dynamics for wy and aj
àg --7a Σ CLETFPE)
vr Uk Σ Exittaph
(128)
where Pj = fx,; 0) Yi is the residual associated with the 5th training data point. If we let 3g = 0g which determines the contribution of each hidden neuron the output (1n; 0), then its not hard to see that the gradient flow dynamics of 3R are
BR U La + Ta w (E Dis CReesps)
(129)
My
Ea
As in the linear setting, the matrix M& Raxdd appears as preconditioning matrix on the dynamics Using the exact same derivation presented in Appendix A.3 entirely in terms of Bg and os.
whenever oB K O, we can express M₂.
M / 1 + 41 To 3xl 2 + òx Ia
0 + 4n B 86
3r 1.3g 1
(130)
2
2
However, unlike in the linear setting, the vector &a Rd driving the dynamics is not shared for all neurons because of its dependence on Cks. Additionally, the NTK matrix in this setting depends on M& and C, with elements Kij = > k=1 c (nh La + Th uy w.ww) zyC4j. Thus, in order to assess the temporal dynamics of the NTK matrix, we must understand the dynamics of M2 and C. We consider signed spherical coordinate transformation separating the dynamics of 83 into its 6t and radial j4 = sgn(a))) II 8 components, such that B ma. directional Ax sgn(ax) [
Here, 3r determines the orientation and direction of the halfspace where the th neuron is active, while मı determines the slope of the linear region in this halfspace. These coordinates evolve according to.
Vog + 4na + 6y (Ia- &
i4 = -0
02 41
(131)
2p
These equations can be derived directly from Eq. 128 through chain rule similar Appendix A.2. In fact its worth noting that the this change of coordinates is similar to the change of coordinates used in the single-neuron analysis. Expressed in terms of the parameters, 3x
and μ4 = a4 a w T
wi
36

---

D Experimental Details
We used Google Cloud Platform (GCP) nodes to run all experiments. Figure experiments were run node with 360 AMD Genoa CPU cores with runtime totaling approximately 90 minutes including averaging over seeds as described below. Neural network training and NTK calculation for Figure was performed on single A100 GPU nodes. Runtime was approximately 20 hours for Figure 5(a), four hours for 5(b), 12 hours for 5(c) (with individual runs ranging from five to 30 minutes depending on the number of datapoints), and 12 hours for 5(d). Figures 2, 3, and are not compute-hes heaey, and these experiments were run on personal computer. Overall, we estimate approximately 200 hours of single A100 runtime well as 100 hours of the 360-core node accounting for failed runs and exploratory experiments.
D.1 Figure 1: Teacher-Student with Two-layer ReLU Networks
For Fig. | we consider a student-teacher setup similar to that in { with one-hidden layer ReLU networks of the form (x; 0) Ei (, x). where f: Rd R and o is the ReLU activation function. The teacher model, freacher has
k hidden neurons initialized w.eaacee iid.
Unif(S and a; We use a symmetrized initialization, as considered in [8] where for i S h/2, we sample wy ~ Sd-1 we symmetrize by setting wl = w;-h/2 and a, -a--/22 This ensures that fstudent predicts on any input at initialization.
lid ~ Unif((#1}) for S k. The f in turn, has h hidden neurons. iid
Lid. Unif({+1))) and then for Z and 0;
Note that the base student initialization described thus far is perfectly balanced each neuron, that 5, = for i E [m]: we also define this to be our setting where the scale 7 is 1. In order to transform the base initialization into particular setting of and 6, first solve for the relative layer scaling a in 52 = r2(22 - a 2- and then scale each w; by T/a and each Os by TO. We obtain a training dataset (2x(0, (9) = by sampling (+) iid. 5d-1 and computing noiseless labels as (d) = ftencher (x)( gieacher). The student is then trained with full-batch gradient descent on a mean square loss objective.
Figure 1 (a).
Here the setting is: d 2, h 50, = and n = 20. We sample a single teacher and then train four students with the same base initialization but different configurations of and ó: (r 0.1, 8 0) and (r = 2, 8 = 0) for the left subfigure, and (r = 0.1, 1) and (r = 0.1, 5 = -1) for the right subfigure. Training is for million steps learning rate of le-4.
Figure 1 (b).
Here the setting is: d = 100, m = 50, A: = 3, and n = 1000, as in Fig. Ic of 8 Training is performed with learning rate of 5e-3/>". Test error is computed mean square error over held-out set of 10,000 datapoints. We sweep over 7 over a logarithmic scale in the range (0.1,2) and 8 over linear scale in the [ sampling of: the teacher weights greacher the base initialization of gstudent and the training data (z(0) yx In this way, each random seed is used for sweep over all combinations of 7 and 6 in the sweep; we simply apply the scaling described above to get to each point on the (r, grid. The kernel distance computed is defined in (27). where here compute it at time t relative to the kernel at initialization, i.e. S(t) = (Ko. K,)/ [|| II l | p). In Fig. 10 we additionally plot Hamming and parameter distances relative initialization, well as training loss, while training for ten times longer than in Fig. (b).
We average over 16 random seeds, where the seed controls the
Notebooks generating all two-layer experiment figures provided here
37

---

100
t= 1000
10000
100000
1000000
0.8
0.6
Kernel Distance
0.4
0,4
0.0
0.4
0.2
0.8
0.0
21
Hamming Distance Parameter Distance
0.8
18
0.4
15
12
0.0
9
-0.4
-0.8
3.2
0.8
2.8
0.4
2.4
2.0
0.0
1.6
1.2
0.4
0.8
-0.8
0.4
0.0
0.8
0.4
10
Train Loss
0.0
100
0.4
0.8
10
0.8
1.2
0.4
Test Loss
0.8
0.0
-0.4
-0.8
0.0
10-1
10
10-1
100
10-1
10
10
10a
10-1
10
T
T
T
Figure 10: Supporting figures for Fig. (b). We plot Hamming distance, parameter distance, and training loss, on top of the test loss and kernel distance considered Fig. (b), and train for ten times longer than in Fig. (b). We observe that although training loss still drops between 10s and 10t steps, the test loss and other distances considered remain largely unchanged. Training loss saturated at le-10.
D.2 Figures 2, 3, 4: Single-Neuron Linear Network
Figures 2, 3, and were generated by simulating gradient flow using scipy. integrate solve_ _ivp function with the RK45 method for solving the ODEs, with relative tolerance of 1 x 10 and
time span of (0, 20). In the experiments with full-rank data, used X L, 8. -14] and
3. For the experiment with low-rank data, used XTX 05.25 0,3 8, -00.11 and
3. 0.05 J. See the discussion in Appendix A.2 for details how we determined our thcoretical predictions. A notebook generating the figures is provided here
D.3 Figure 5:
Kernel Distance
We trained LeNet-5 88
We use He initialization [89] and divide the first layer weights by a and multiply the last layer weights by a at initialization, which keeps the network functionally the same at initialization. We trained the model
(with ReLU nonlinearity and Max Pooling) on MNIST 88
38

---

for 500 epochs with a learning rate of le-4 and a batch size of 512. The parameter distance is defined as the L2 distance between all the parameters. To quantify the distance between the activations, we binarize the hidden activation with representing active neuron. We evaluate Hamming distance over all the binarized hidden activations normalized by the the total number of the activations. We use kernel distance [27]. defined as S(1,; (2) invariant measure of similarity between the NTK at two points in time. We subsample 10% of MNIST to evaluate the Hamming distance and kernel distance. All curves in the figure are averaged over 8 runs.
(Kr. /2,)/ K,, || K2 Il),, which is scale
Gabor Filters
We are training a small ResNet based the CIFARIO script provided in the DAWN benchmark (code available here). The only modifications to the provided code hase are we increase the convolution kemel size from x 3 15 15, hetter observe the learned spatial patterns, and set the weight decay parameter to 0 to avoid confounding variables. Moreover, we are dividing the convolutional filters weights by a parameter a (after standard initialization) which controls the balancedness of the network. To quantify the smoothness of the filters, we compute the normalized Laplacian of each filter Wij E [R15x15, over input (1, 2, 3) and output j
(1, ..., 64) channels
wij
smoothness(wet) =
(132)
|waa l|
where the Laplacian kernel is defined as
0.25
-0.5
-0.25
=
0.5
-0.5
(133)
0.25
- 0.5
-0.25
Validation Accuracy
Train Accuracy
0.90
0.90
0.0
Smoothness
0.75
0.75
0.60
10.0
3.3
-3.0
1.0
1-0
0.45
8
0.45
0m0,3
-4.5
0.30
0.30
0
10
20
0
10
20
0.10
0.33
1.00
3.33
10.00
Epochs
Epochs
&
(a) Accuracy
(b) Gabor Smoothness
Figure 11: Interpreting convolutional filters. CNN experiments on CIFARIO. We can see A) that all networks achieve comparable training and test accuracy, despite the modification in initialization. However, B) we see that networks with small initialization (a < 1) learn much smoother filters, giving quantiative support to results in Fig. 6 The smoothness is defined as the normalized Laplacian of the filters (see text, eq. 132
Random Hierarchy Model
We refer to [67]. who originally proposed the random hierarchy model (RHM) as tool for studying how deep networks learn compositional data, for in-depth treatment. Here briefly recap the setup following the notation used in [67].
An RHM essentially lets us build a random classification task with clear hierarchical structure. The top level of the RHM specifies m equivalent high-level features for each class label in {I....,n.]] ne.). where each feature has length and n. is the number of classes. For example, suppose the vocabulary at the top level is V = {a. b, c). n. = 2, m = 3, and s = 2. Then particular instantiation of this RHM, we might have that Class 1 has ab, oa, and precisely the example used in Fig. of (67). Class will then have three random high-level features, with the constraint that they are not features for Class for example, bb, be, ac.
equivalent high-level features (this is
Each successive level specifies m equivalent lower-level features for each "token" the vocabulary at the previous level. For example, if V {d, f) might have that can be equivalently
39

---

represented as de, of, or ff; ff; b and e will each have equivalent representations of their own. We assume that the vocabulary size, v, is the same all levels. Therefore, sampling an RHM with hyperparameters Re, m, 8, requires sampling mn, + (T. 1)me rules.
In order sample datapoint from RHM, first sample class label (e.g. Class 1), then uniformly sample one of the highest level features, (c.g, ab), then for each "token" in this feature we sample lower level features (e.g. - de, b b ee), and so on recursively. The generated sample will therefore have length sh and a class label. For training a neural network to perform this classification task, each input is converted into one-hot representation, which will be of shape (s", v), and is then flattened.
We the code released by [67] to train MLP of width 64 with three hidden layers to learn RHM with = 3, nh = 8, m 4, = 2, = 8. The main change we make allowing for scaling the initialization of the first layer by 1/ax and the initialization the readout layer by a. We then sweep over a E {0.03, 0.1, 1.1.0...... 0.3, 1, 3, 10} and over the number of datapoints in the training set, which is specified fraction of the total number of datapoints the RHM can generate. We average test accuracy, which is by default computed on held-out set of 20,000 samples, over six random seed configurations, where each configuration seeds the RHM, the neural network, and the data generation.
We train with the default settings used in [67, that is stochastic gradient descent with momentum of 0.9, run for 250 epochs with a learning rate initialized at 6.4 (0.1 times width) and decayed with cosine schedule down to 80% of epochs. The batch size of 128; we do not use biases or weight decay.
Grokking
We training one layer transformer model the modular arithmetic task in Power al. 68 Our experimental code is based on an existing Pytorch implementation (code available here). The only modifications the provided code base is that the default 2-layer model). Prior analysis in Nanda et al. 72 has shown that this model can learn minimal nnttention-based circuit that solves the task.
single transformer layer (instead of
We study the effects on grokking time (defined as 2 0.99 accuracy on the validation data) of two manipulations. Firstly, we divide the embedding weights of the positional and token embeddings by the same balancedness parameter a as in the CNN gabor experiments. Secondly, like in Kumar et al. 669. multiply the output of the model (i.e., the logits) by factor and divide the learning rate by
10a
10a
2.4e+
7.5e+01
Steps until Trained
Steps until Grokked
1.8e+ 04
6.0c+01
10
8
109
1.2e+ 04
4.5e+01
6.0e+ 03
3.0e+01
10
0.6
1.0
1.2
1.4
10-1
0.8
1.2
1.4
0.0e+00
(a) Steps until training
(b) Steps until grokked
Figure 12: Transformer Grokking in Modular Arithmetic Task. A) Shows the number of training steps required until the training accuracy passes a predefined threshold of 99%; we sample scaling (0.5, 0.75, 1.0, 1.25, 1.5} [69] and balance n = random initializations with maximal computational budget of m = 30, 000 training steps. B) Same A), but reporting the number of training steps required until the test performance passes the predefined threshold of 99%. We clearly see the fastest grokking in an unbalanced rich setting.
1. 1.0, 3.0, 10)
regular grid over
40

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
The abstract clearly lays out the four main contributions sections of our paper (I) novel minimal model of the transition between rich and lazy learning with exact solutions, (2) extension of this model to more complex linear networks, (3) application this analysis to shallow nonlinear networks, (4) demonstration of the relevance of the unbalanced rich regime identified in practical deep learning settings. The claims made accurately match the theoretical and experimental results in the paper. These four main contributions (also discussed in our contributions section in Section [ reflect Section Section 4 and Section 5
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: In the conclusion we address some of our primary limitations, which include the difficulty in extending our theory to deeper nonlinear network settings, as well as discretization and stochastic effects of SGD. On top of that, our theoretical and empirical assumptions are documented and stated throughout the paper and appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
41

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
For all main theoretical lemmas and theorems we either clearly state assump- tions in the theorem statement
Justification:
introduce the assumptions carlier in the paper. We focus on high-level intuition in the main body of the text and include complete proofs in the appendix; where each section of the document has corresponding appendix section.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We disclose all details needed to reproduce experiments the paper, partieu- larly in the appendix. We plan to release all code used for reproducing empirical results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., s., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make clear how reproduce that algorithm.
42

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c)
If the contribution either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
new model (e.g.. large language model), then there should
(d) We recognize that reprodueibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Some of our experimental results use existing open-souree source codebases; for those we provide specific instructions on how to reproduce results. For the rest of our experiments we plan to release code with instructions for faithful reproduction outlined the appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
public//uiiees/CoeSttmii ottp///ddeauumiissssoooooooo/o// for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https SroodeSubmissiooPPooe
//nips.ccl httppsscc/ppblligguuddeees
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We make sure to outline all necessary details for faithful reproduction of our deep learning experiments in the appendix. Sufficient detail to appreciate and make sense of the empirical results is presented in the main paper with specific details placed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
43

---

7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes].
Justification: For our main empirical results, shown in Figure 5, we make sure to run each experiment with multiple seeds and provide standard deviations on all plots. Details on how we average over seeds are outlined carefully in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide estimates of the total runtime and compute type required to reproduce our results. We also provide estimates of our total runtime including failed experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't ' make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpps:/neeurips. httph/ewwwrspssscseeeereeeeeeeee http:/pp....pppleoeacccccceeeeea
Answer: [Yes]
Justification: We make sure to conform with the NeurIPS Code of Ethics. Items in the code particularly relevant to our work pertain disclosing details necessary for reproducibility.
Guidelines:
44

---

The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This paper primarily a theoretical paper attempting to better understand one of the great mysteries of deep learning, how neural networks acquire task-relevant features. We do not see direct path negative applications of our work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not introduce new models or datasets that are not already open sourced or published.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
45

---

12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We make sure cite the papers whose open-source code or datasets we used for our experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi wwwode....d..... com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
46

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not involve human study experiments and as such IRB approvals are not relevant to our work.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
47


Reviews:
[{"review_id": "TnZvbErZSD", "reviewer": "Reviewer_hh3v", "summary": "The authors derive exact solutions to a minimal model that transitions between lazy and rich learning, elucidating how unbalanced layer-specific initialization variances and learning rates determine the degree of feature learning in a finite-width network. \nThey provide evidence that this unbalanced rich regime drives feature learning in deep finite-width networks, promotes interpretability of early \nlayers in CNNs, reduces the sample complexity of learning hierarchical data, and decreases the time to grokking in modular arithmetic.", "questions": "1. In Ref. [1] which you mention, the formula for the NTK with unbalanced LRs is (using your notation) $K = X \u0000(\\eta_a a^2 I_d + \\eta_w ww^T) X^T$  (eq. therein), i.e. the LRs have switched places. I'm not sure where this discrepancy comes from or if it is an error. \n\n2. Line 145 - can you explain why in the infinite width limit non-linearities act linearly? It is true that the model becomes linear in the weights, but do you mean that each layer acts linearly on its pre-activation? \n\n\n[1] https://arxiv.org/abs/2401.07085", "limitations": "The authors adequately addressed the limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is very well written: \n- clear visuals and exposition, going from a minimal model and gradually adding complexity and realism. \n- The quantities of interest (e.g. $\\delta$) give an intuitive picture. \n- The notations are very clear and easy to follow. One small comment is that it might be better to use a different index for $\\theta$ and $x$. \n- The contributions are distinguished from previous works and relevant connections to them are clearly stated. \n\n2. Addressing the rich regime is arguably more interesting than the lazy regime and is much less studied, making the questions addressed here timely. \n\n3. The experiments presented in Fig 5 show that the effects studied analytically in previous sections also appear in realistic settings, making this work relevant to a broad audience.", "weaknesses": "1. To yield the exact solutions, the authors evoke the assumption of whitened input, which is rather unrealistic. The low-rank case is also addressed but then one needs to additionally assume that the interpolating manifold is one-dimensional to find the solution in terms of $\\delta$ exactly. \n\nmissing citations: \n- In line 398, when mentioning grokking and the transition from lazy to rich learning - it is worth noting Ref. [1]. \n\n[1] https://arxiv.org/abs/2310.03789", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your comprehensive review and the time you have taken to suggest improvements for our study. We are also happy to hear you think our work is a timely and important contribution to the field. We will respond individually to the weaknesses and questions you raised regarding our paper:\n\n**Notation for indices.** We agree with you that using a different index for parameters, data, and layers can help with readability. This is a more comprehensive change, but something we will implement in our updated draft.\n\n**Whitened input assumption.** We agree that the assumption of whitened input for our minimal model is quite strong, although a common assumption used in many prior works (see Saxe et al. 2014 for example). That said, we actually do relax this assumption throughout the analysis. As you noticed, a strength of our work is that we start with a simplified setting where we can derive exact solutions, then “gradually adding complexity and realism”, we show that the key-takeaways remain. To be specific, section 3 is composed of three parts:\n\n1. *Deriving exact solutions in parameter space.* In order to solve analytically the coupled system of ODEs we need whitened input.\n2. *Interpreting the dynamics as preconditioned gradient flow in function space.* This analysis does not require whitened input. We only require that $X^\\intercal X$ is full-rank such that there exists a unique OLS solution. As shown in equation 2 we did not replace $X^\\intercal X$ with $\\mathbf{I}_d$. \n3. *Identifying the implicit bias when $X^\\intercal X$ is low-rank.* Actually the main theorem (Theorem A.2) in this analysis does not make any assumptions on the dimension of the null space. This theorem expresses the interpolating solution as a solution to a constrained optimization problem (which in general can only be solved numerically). We can then interpret the objective function being minimized in different limits of $\\delta$. If we additionally assume the null space is one-dimensional then we can analytically express the solution to this constrained optimization problem.\n\nIn summary, we do need the whitened input assumption for the first part of section 3, however the remaining subsections relax this assumption to full-rank and then low-rank. We have added comments in this section to make this more clear and we have brought theorem A.2 up from the appendix into the main to make it very clear that the assumption on a one-dimensional null space is only needed to solve analytically the constrained optimization problem in this theorem.\n\n**Missing citation.** Thank you for this additional reference we were unaware of. We have added this citation to line 398 as suggested.\n\n**Formula for the NTK with unbalanced LRs.** The formula (equation 10) in the paper you referenced seems incorrect to us. Here is an explanation for the form of the NTK in our paper (we also added this derivation to the top of appendix A):\n\nAs discussed in our notation section, $K_{ij}  = \\Theta(x_i, x_j; \\theta) = \\sum_{k = 1}^p \\eta_{\\theta_k} \\partial_{\\theta_k} f(x_i;\\theta)\\partial_{\\theta_k} f(x_j;\\theta)$. For the minimal model we study $f(x;\\theta) = aw^\\intercal x$ implies $\\partial_a f = w^\\intercal x$ and $\\partial_w f = ax$. Thus, we see that the term associated with the gradient of $a$, and thus the learning rate $\\eta_a$, depends on $w$, while the term associated with the gradient of $w$, and thus the learning rate $\\eta_w$, depends on $a$. That is why the NTK is defined by the matrix $M = \\eta_wa^2 \\mathbf{I}_d + \\eta_aww^\\intercal$ in our expression. It seems to us that the equation from the paper you referenced has a notational mistake.\n\n**Clarifying line 145.** We agree this sentence is misleading. We were referencing a common simplification used in many infinite width analyses where the features before and after the nonlinearity are assumed to be of the same scale (see “A Spectral Condition for Feature Learning” by Yang et al. for example). These works will often build intuition for their analysis by replacing the nonlinearities with linear activations, without affecting their results. However, as line 145 is currently written, it sounds like the nonlinearities are always linear on the preactivation, which is not true. Thus, we have modified this sentence to read, “In this limit, analyzing dynamics becomes simpler in several respects: random variables concentrate and quantities will either vanish to zero, remain constant, or diverge to infinity.” Also, it is worth noting we have added a longer discussion on how our theory connects to infinite width limit analyses in our updated manuscript (see responses to Reviewers JJEZ and K8a2 for details).\n\nPlease let us know if you have any other questions regarding our work. Overall we are happy to hear that you enjoyed our paper and we hope you will consider our paper an important contribution to the NeurIPS community. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HyHtC9pyiR", "reviewer": "Reviewer_KNvb", "summary": "The paper studies the learning dynamics of a deep network by leveraging the conserved quantities due to symmetries", "questions": "See the weakness section", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The theoretical finding that unbalancedness in the layers drives feature learning is a novel and interesting insight\n\nThe technical tool of using symmetries and conserved quantities to characterize the dynamics is also novel\n\nFinding exact solutions can help us understand better the causal aspects of the phenomena in deep learning, and should be encouraged\n\nOverall, I am positive towards this work", "weaknesses": "I think the paper could improve by explaining its results in more detail. For example, I find the following points rather confusing\n\n1. Does theorem 4.1 hold after training? Or it holds at any time in training？\n\n2. The paper argues that it finds an exact solution, but I am sure what this is referring to. Is theorem 4.1 the \"exact solution\" the paper finds? If so, what is it a solution to?\n\n3. The title claims \"rapid feature learning\", but almost no place in paper discuss what it means to be rapid, nor what it means to be doing \"feature learning.\" The only part in the paper that I find to be peripherally related to this claim is the discussion in lines 380-383. This result holds true without needing any result from sections 3-4. Then, what is the point of 3-4? If the title is the main claim of the paper, can the authors write it out in a more explicit manner, and with much greater mathematical detail?\n\n4. In figure 5, a crucial quantity is $\\alpha$, but it is not explained in the context. What is $\\alpha$? It seems to me that this quantity is the scaling factor introduced in lines 171-172? Is this case? Even if so, the authors should have explained it much better in the immediate context\n\n5. To be fair, I think the authors are putting too much content both in the appendix and in the main text. It feels to me much better to move one of section 3 or 4 to the appendix and expand what is in the main text to make it much more readable\n\n6. A constructive suggestion is that it would be quite insightful to compare the results in section 4 to the results derived for SGD in deep linear models in https://arxiv.org/abs/2402.07193. SGD and GD have very different behaviors for these models and this discussion will be very helpful to the audience", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to thoroughly review our paper and highlight areas needing further clarification. We appreciate your positive feedback on our work. We will address each of the weaknesses you mentioned individually, hoping this will enhance your confidence in the significance of our research.\n\n**Theorem 4.1.** This theorem holds throughout training. We will make this more clear in the theorem statement by adding a time dependence to $\\beta_i(t)$ and write for all $t \\ge 0$.\n\n**Exact solution.** The exact solutions we refer to are for the gradient flow dynamics of $a$ and $w$ in the minimal model we present in section 3. The gradient flow dynamics for this model boil down to a complex coupled system of nonlinear ODEs shown in equation 1. As stated on line 191, we solve this system of ODEs exactly, derived in appendix A in detail. See figures 2, 3, and 4; in all these figures the dashed black lines represent theoretical predictions (i.e exact solutions) while the colored lines are empirical (the results of training on a computer). We will make this more clear by referencing section 3 when we discuss exact solutions in the paragraph titled “our contributions” in section 1.\n\n**Rapid feature learning.** As discussed in the related work section (lines 45-62), feature learning in the context of this work is synonymous with rich learning. We mathematically define rich learning on lines 96-101, as a change in the NTK through training measured by the kernel distance metric proposed in Fort et al. 2020. The main claim of our paper is that an upstream initialization leads to rapid feature learning *relative* to a balanced or downstream initialization. By \"rapid\" in our title, we refer to this specific claim. We demonstrate this empirically in figure 1. Sections 3 and 4 provide the necessary analysis to prove this claim in section 5. This explains why our title states “exact solutions reveal how unbalanced initializations promote rapid feature learning.” To clarify this further, we will revise the writing on lines 63-84 under \"Our contributions.\" This section outlines our contributions and sets up the paper’s layout. We aim to make this section more concise to clearly explain why we use the approach in sections 3 and 4 to prove our main claim and its connection to the title.\n\n**Alpha in figure 5.** Thank you for catching this. Yes you are correct, $\\alpha$ in this context is equivalent to the scaling factor used in lines 171-172.  It is also described in the caption for figure 1 and in the experimental details in appendix D, however, we forgot to thoroughly describe it in the caption for figure 5. We will revise the caption to make this more clear.\n\n**Content distribution between the main and appendix.** We appreciate this comment and we agree that the paper is dense. We have tried to make the work as readable as possible, but we believe that section 3 and 4 are necessary steps to build the analysis used in section 5. With the additional page allotted for the final version of this paper we have moved up some aspects of the appendix into section 3 and 4 which should make it more readable, added clarifying transitions between the sections, and expanded the content in section 5 to clarify its connection to the previous sections and the main claims of our paper. We think these changes will significantly improve the readability.\n\n**Related work on SGD.** Thank you for highlighting this reference. We will include it in our discussion on the limitations of our work (lines 405-411). In this section, we explain how SGD disrupts the conservation laws that are central to our study.  As you pointed out, one of our key findings is how conserved quantities arising from symmetry affect the degree of feature learning. Other works have focused on the influence of stochasticity in SGD on these conserved quantities. We agree integrating these analyses to understand the influence of stochasticity on the degree of feature learning would be an insightful direction for future work. However, this is beyond the scope of our current paper, which is already densely packed with content.\n\nThank you again for your constructive feedback. We hope we addressed your points regarding our work. If we have, we would appreciate it if you could raise your score to reflect this. Thank you!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gBpzEhWjSP", "reviewer": "Reviewer_JJEZ", "summary": "The paper studies the dynamics of training in neural networks with the scope of identifying how the variance of weights' initialization together with layer-wise learning rates determines different learning regimes, encompassing lazy, rich, and the transition between them. The paper identifies a conserved quantity $delta$ that is preserved throughout training (i.e. $\\dot{\\delta} = 0$) and depends on both the weight's magnitude and the learning rates. Thus, the sign and magnitude of $\\delta$ at initialization affect learning and the geometry of the landscape that is traversed through gradient flow. The paper finds three different regimes, starting from a solvable model of two-layer linear network. \n\n1. $\\delta > 0$ (named *upstream initialization*) induces lazy dynamics and corresponds to the case where $\\eta_w a^2 >> \\eta_a ||w||^2$, where $w$ and $a$ are the first and last layer weights, resp. \n\n2. $\\delta = 0$ (*Balanced*): corresponds to rich dynamics.\n\n3. $\\delta < 0$: initial lazy fitting regime and second rich phase.\n\nThe authors then extend these findings to wider and deeper networks and piece-wise linear activations.", "questions": "See weaknesses", "limitations": "Addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper finds a simplified model that can be thoroughly analyzed and gives very precise statements on what causes rich and lazy regimes based on the single conserved quantity $\\delta$, which makes the results very clean. This quantity is then naturally connected to the NTK, which makes intuitive sense given that the NTK is the ultimate factor that determines the learning regime (rich, lazy, or somewhere in between).\n\n2. The formula for $\\delta=\\eta_w a^2 - \\eta_a ||w||^2$ also is nicely interpretable, elucidating the interplay between learning rates and weights magnitude. \n\n3. The authors test their theory in various settings and observe a close alignment between the experimental verification and the theoretical predictions, which makes me confident of the correctness of their theory.", "weaknesses": "1. I would have appreciated more on the NTK dynamics.  There is a clear connection between $\\delta$ and (part of) the NTK $K$. I wonder whether this connection can be made more explicit by studying the NTK dynamics $\\dot{K}$. This would make the connection between the conserved quantity and the feature learning more explicit. \n\n2. Some equations are not included. The authors justify this in Line 1022: \"We omit the solution due to its complexity, but provide a notebook used to generate our figures encoding the solution\". I would just provide the formulas for the sake of completeness. Also, I do not see the equations describing the kernel distance S, which is related to the NTK movement, thus to my previous point. \n\n3. In section 4 the paper includes various extensions, including wide networks. It would have been nice to recover known results by relating the conserved quantity $\\Delta$ and specific parameterizations of the network as a function of the width. This would have connected existing results in the NTK and $\\mu$P literature to the conserved quantities identified in this paper and the learning regime.\n\n4. The connection between the scale of the initialization and the learning rate is not entirely new. This is quite obvious at grokking problems, where to exhibit grokking is often necessary to play with the initialization parameters of different layers. In this context, Kumar et al [61] already clarified that the delayed generalization was caused by the NTK dynamics transitioning from lazy to rich regimes. \n\nOverall, I am in favor of the paper's acceptance because it devised a model where the three different learning regimes are clean and crystallized.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and constructive suggestions. We appreciate your positive feedback and address the weaknesses you highlighted individually. We hope this will increase your confidence in the importance of our work.\n\n**NTK dynamics.**  This is a good point; we can in fact study the dynamics of the NTK matrix directly, which leads to an argument similar to the one provided in lines 239 - 264. Here we outline this analysis, which we will add to appendix A.\n\nThe dynamics of the NTK matrix $K = X M X^\\intercal$ is determined by $\\dot{M}$. From Equation 3 in the main text, we can write $\\dot{M} = \\frac{2 \\|\\beta\\|}{\\kappa} (I_d + \\hat{\\beta}\\hat{\\beta}^\\intercal) \\partial_t \\|\\beta\\| + \\frac{\\kappa - \\delta}{2} \\partial_t (\\hat{\\beta}\\hat{\\beta}^\\intercal)$. From this expression we see that the change in $M$ is driven by two terms, one that depends on the change in the magnitude of $\\beta$ and another that depends on the change in the direction of $\\beta$. As done in the main text, we consider the limits as $\\delta\\to\\pm\\infty$ and when $\\delta = 0$ to identify different regimes of learning. For $\\delta\\to\\infty$, the coefficients in front of both terms vanish, and thus, irrespective of the trajectory taken from $\\beta(0)$ to $\\beta_*$, the change in the NTK is vanishing, indicative of a lazy regime. For $\\delta \\to -\\infty$, the coefficient for the first term vanishes, while the coefficient on the second term diverges. Here, the change in the NTK is driven solely by the change in the direction of $\\beta$. This is why for large negative delta we observe a delayed rich regime, where the eventual alignment of $\\beta$ to $\\beta_*$ leads to a dramatic change in the kernel. When $\\delta = 0$, the coefficients for both terms are roughly of the same order, and thus changes in both the magnitude and direction of $\\beta$ contribute to a change in the kernel, indicative of a rich regime.\n\n**Adding equations for completeness.** We agree that it is important to include expressions for completeness. We will include explicit solutions to the one hidden neuron model dynamics in the appendix. Specifically, we will write out the expression for $a(t)$, which is currently omitted because of its verbosity. We will also write out the expression for the kernel distance presented in figure 3c and include this in the added subsection of the appendix describing the kernel dynamics (as we discussed above).\n\n**Connections to width-dependent neural network parameterizations.** We agree, discussing how our analysis of finite-width networks connects to existing analyses of infinite-width networks is important. We outline a discussion we will add to the main:\n\nIn a width-dependent parametrization, such as mean field or NTK, the random initialization of weights leads to a distribution over conserved quantities. For example, in the two-layer setting, the mean-field parametrization leads to the per-neuron conserved quantity being 0 in expectation, but with a non-vanishing variance. The NTK parametrization has the same distribution over conserved quantities, but at a larger function scale, leading to lazy dynamics. Thus, between NTK and mean-field, a change in function scale modulates the degree of feature learning. In our work, we show that relative scale between layers (i.e. $\\delta$) is a separate knob that can also be tuned to influence the degree of feature learning.\n\nWhen considering these parametrizations in the infinite-width limit we can recover a phase diagram analogous to our results in the finite-width setting. Reference [14] previously considered $f(x) = \\frac{1}{\\alpha}\\sum_{i\\leq k}a_i\\sigma(w_i^\\intercal x)$ with weights initialized as $a_i \\sim \\mathcal{N}(0, \\beta_a^2)$ and $w_i \\sim \\mathcal{N}(0, \\beta_W^2I_d)$ as width $k\\to\\infty$. They obtain a phase diagram at infinite width capturing the dependence of learning regime on the overall function scale $\\beta_a\\beta_W/\\alpha$ and the relative scale $\\beta_a/\\beta_W$. The resulting phase portrait is analogous to ours in Figure 1 (b), which considers the conserved quantity $\\delta$ rather than the relative scale $\\beta_a/\\beta_W$. In particular, there is a lazy regime, which is always achieved at large scale (just as in the large-$\\tau$ regions of Figure 1 (b)), but is also achieved at smaller scale if the first layer variance is sufficiently larger than the second (as in the downstream initializations at small $\\tau$ in Figure 1 (b)). On the other side of the phase boundary is the infinite width analog of rapid rich learning, with all neurons condensing to a few directions. This is induced either at small function scale, or at larger function scale if the relative scale is sufficiently large, such that $W$ learns faster than $a$.\n\n**Connections to prior work in grokking.** Kumar et al. [61] explored how grokking arises as the result of the transition from lazy to rich dynamics, studying this transition as a function of overall function scale, initial NTK alignment with the test labels, and train set size. However, we add more nuance to this picture in that we show how grokking can be induced not just by modulating the overall function scale, but by simply changing the initialization geometry, that is by scaling the embedding weights for positional and token embeddings without affecting overall scale (as visualized in Figure 5 (d)). This leads to the complex phase diagram of time to grokking as a function of both scale and geometry, as presented in Figure 11. We believe this adds to the picture presented in Kumar et al. by showcasing the unique role of layer-wise initialization scales.\n\n\nWe hope we addressed your points regarding our work. If we have, we would appreciate it if you would consider raising your score to reflect this. Thank you again for your constructive feedback.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c3v29bGb6A", "reviewer": "Reviewer_K8a2", "summary": "This paper analyses the lazy versus rich learning dynamics of minimal but finite neural networks by deriving exact solutions under arbitrary layerwise initialization and learning rates. The theoretical insights are successively extended to networks of increasing complexity and corroborated by several experiments.", "questions": "- Why do you call the layerwise initialization variances and learning rates *initialization geometry* when learning rate scalings clearly only influence the training dynamics? Similarly, the term geometry is unspecific and does not highlight the interplay between the layers. I encourage the authors to find a more suitable terminology for the choice of hyperparameters.\n- Why is the kernel distance in Figure 1 eventually maximized by a downstream initialization $\\delta<0$, and not by an upstream initialization $\\delta>0$ where feature learning occurs rapidly both directionally and radially?", "limitations": "The main limitations are acknowledged in the Conclusion section. It is acceptable that an exact analysis of learning dynamics does not include deep nonlinear networks trained with large learning rates.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and provides an exact understanding of the learning dynamics of very simple neural networks, in increasing complexity over the course of the paper. The figures complement the derivations very comprehensively. In this way, the paper makes a significant contribution to the understanding of neural network training dynamics. I am not able to judge the claimed novelty, as I am not entirely familiar with the related work.", "weaknesses": "Overall the paper leaves little room for criticism. The results being limited to shallow, often linear neural networks trained with small learning rates is justified by the detailed understanding of the training dynamics, and mitigated by experiments. It could have been acknowledged more clearly how the maximal update parameterization by Yang and Hu (2021) already shows that upstream initializations are necessary for non-vanishing feature learning in large models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback and constructive questions. We are grateful for your positive feedback on our work and are committed to improving our manuscript by addressing each of the weaknesses you identified.\n\n**Connections to infinite width parametrizations.** We agree with the review and are adding a subsection to Section 5 carefully outlining how our results can be extended to infinite width settings and acknowledging connections to existing parametrizations. In the two-layer setting, muP corresponds to the mean-field parameterization, which, for input dimension constant in width, can be written as $f(x) = 1/k a^T\\sigma(Wx)$ with $a_i \\sim \\mathcal{N}(0, 1)$ and $W_{ij} \\sim \\mathcal{N}(0, 1/d)$. We note that this actually leads to the per-neuron conserved quantity being 0, or balanced, in expectation, with a non-vanishing variance. Please see our response to reviewer JJEZ for a more detailed discussion on the connection between our analysis of finite-width networks and existing works on infinite-width networks.\n\n**Initialization geometry.** We chose the term initialization geometry because the per-layer learning rates and layer magnitudes collectively determine the geometry of the surface that the dynamics are constrained to. This geometry is leveraged consistently in our theory through the use of conserved quantities and is clearly visualized for the one hidden neuron model in Figure 2. \n\n**Kernel movement in downstream initialization.** In Figure 1(b), the small-scale downstream initialization ($\\delta<0$) starts off lazy, attempting to fit the training data by only changing the small readout weights $a$ with minimal movement in the large first-layer weights $\\{w_i\\}$. Up to this stage, the downstream initialization exhibits smaller kernel movement than the upstream initialization, which undergoes a rapid change in the kernel. However, in order to interpolate the training data, the network with downstream initialization needs to align its hidden neurons by a non-trivial amount, thus undergoing a change from lazy to rich learning. This alignment of the hidden neurons, which are large in scale, results in a dramatic change in the kernel. In the case of the upstream initialization, the network is able to interpolate the training data by simply aligning its hidden neurons (which are small in scale) directionally and radially as needed, leading to a smaller movement in the kernel overall. See our response to reviewer JJEZ for a detailed discussion on the dynamics of the NTK matrix.\n\nWe hope we addressed your points regarding our work. Thank you again for your constructive feedback!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all the reviewers for their careful and detailed comments. We greatly appreciate the time and effort you put into reviewing our paper, which we believe has significantly improved our work. We have addressed each reviewer’s questions individually and provided linked responses for similar questions. We hope you will consider our paper a valuable contribution to the NeurIPS community. We look forward to your continued feedback!

Meta-review: This paper studies minimal models of lazy and rich learning with exact solutions and uses their derivations to gain insights into the learning dynamics of neural networks.
The latter is controlled by the quantity delta that combines layer-specific initialization variances and learning rates. It is argued that a layer-wise initialization out of balance can promote (delayed) feature learning (also in a finite-width network), which is related to Grokking. In contrast to previous work, the authors also consider layer-wise scale imbalance that does not affect the overall function scale.

All reviewers agree that the presented work has merit for the NeurIPS community and recommend acceptance.
They particularly highlight positively that the paper is well written (even though dense) and appreciate the fact that (small) tractable models are introduced that provide insights into the rich learning regime and also include finite-width results.
Decision: Accept (spotlight)