Title: Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: feature learning

rich regime

lazy regime

exact solutions

conserved quantities

balanced initialization

neural tangent kernel

grokking
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

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.

1 Introduction

Deep learning has transformed machine learning, demonstrating remarkable capabilities in a myriad of tasks ranging from image recognition to 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 a complex interplay between datasets, network architectures, and optimization algorithms. Within this framework, two distinct regimes, determined at initialization, have emerged: the lazy and the rich.

Lazy regime. Various investigations have revealed a notable phenomenon in overparameterized neural networks, where throughout training the networks remain close to their linearization [1, 2, 3, 4, 5]. 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 a 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 a deterministic NTK [6, 7], minimal movement in parameter space [8], static hidden representations, exponential learning curves, and implicit biases aligned with a reproducing kernel Hilbert space (RKHS) norm [9]. However, Chizat et al. [8] challenged this understanding, asserting that the lazy regime isn't a

0.8

0.4

0.0

0.4

0.8

t

=100

t

=1000

t

=10000

t

=100000

10

1

10 0

0.8

0.4

0.0

0.4

0.8

10

1

10 0

10

1

10 0

10

1

10 0

0.0

0.4

0.8

1.2

Test Loss

10

2

0.0

0.2

0.4

0.6

Kernel Distance

(a) Overall and relative scale impact feature learning

(b) A complex phase portrait of feature learning

product of the infinite-width architecture, but is contingent on the overall scale of the network at initialization. They demonstrated that given any finite-width model f ( x ; θ ) whose output is zero at initialization, a scaled version of the model τf ( x ; θ ) will enter the lazy regime as the scale τ 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 convergence to a 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 distinguished by a learned NTK that evolves through training, non-convex dynamics traversing between saddle points [10, 11, 12], sigmoidal learning curves, and simplicity biases such as low-rankness [13] 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 [15], 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, as shown in Fig. 1, 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-specific initialization variances and learning rates determine the transition between lazy and rich learning in finite-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-neuron 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] 1 , 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 a simple set of conserved quantities that constrain the geometry of learning trajectories. Additionally, it reveals that a 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 4 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 a 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 a feedforward network f ( x ; θ ) : R d → R c parameterized by θ ∈ R m . Unless otherwise specified, c = 1 . The network is trained by gradient flow ˙ θ = -η θ · ∇ θ L ( θ ) , with an initialization θ 0 and layer-specific learning rate η θ ∈ R m + , to minimize the mean squared error L ( θ ) = 1 2 ∑ n i =1 ( f ( x i ; θ ) -y i ) 2 computed over a dataset { ( x 1 , y 1 ) , . . . , ( x n , y n ) } of size n . We denote the input matrix as X ∈ R n × d with rows x i ∈ R d and the label vector as y ∈ R n . The network's output f ( x ; θ ) evolves according to the differential equation, ∂ t f ( x ; θ ) = ∑ n i =1 Θ( x, x i ; θ )( y i -f ( x i ; θ )) , where Θ( x, x ′ ; θ ) : R d × R d → R is the Neural Tangent Kernel (NTK) , defined as Θ( x, x ′ ; θ ) = ∑ m p =1 η θp ∂ θ p f ( x ; θ ) ∂ θ p f ( x ′ ; θ ) . The NTK quantifies how one gradient step with data point x ′ affects the evolution of the networks's output evaluated at another data point x . When η θp is shared by all parameters, the NTK is the kernel associated with the feature map ∇ θ f ( x ; θ ) ∈ R m . We also define the NTK matrix K ∈ R n × n , which is computed across the training data such that K ij = Θ( x i , x j ; θ ) . The NTK matrix evolves from its initialization K 0 to convergence K ∞ through training. Lazy and rich learning exist on a 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, 26]. We use kernel distance [27], defined as S ( t 1 , t 2 ) = 1 -⟨ K t 1 , K t 2 ⟩ / ( ∥ K t 1 ∥ F ∥ K t 2 ∥ F ) , which is a scale invariant measure of similarity between the NTK at two times. In the lazy regime S (0 , t ) ≈ 0 , while in the rich regime 0 ≪ S (0 , t ) ≤ 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, f ( x ; θ ) = β ( θ ) ⊺ x is linear in its input x , but can exhibit highly nonlinear dynamics in parameter θ and function β ( θ ) 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 [31, 32, 33] linear networks with more flexible initialization schemes [34, 35, 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 a minimal model utilizing a mix of Bernoulli and Riccati equations to showcase a complex phase portrait of lazy and rich learning with separate alignment 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 a first-order KKT point minimizing Q ( θ ) among all possible interpolating solutions. Foundational work by Soudry et al. [38] pioneered this approach for a linear classifier trained with gradient descent, revealing a max margin bias. These findings have been extended to deep linear networks [39, 40, 41], homogeneous networks [42, 43, 44], and quasihomogeneous networks [45]. A similar line of research expresses the learning dynamics of networks

trained with mean squared error as a mirror flow for some potential Φ( β ) , 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 ℓ 1 and ℓ 2 norms in the rich and lazy regimes respectively [14]. However, finding the potential Φ( β ) is problem-specific and requires solving a second-order differential equation, which may not be solvable even in simple settings [47, 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.

Two-layer networks. Two-layer, or single-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 a 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, 51, 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 infinitewidth 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 two-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 self-consistent dynamical mean field theory [60] and a spectral scaling analysis [61]. In this study, we focus on finite-width neural networks, but discuss the connection between our work and these width-dependent parameterizations in Section 5.

3 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 a single hidden neuron defined by the map f ( x ; θ ) = aw ⊺ x where a ∈ R , w ∈ R d are the parameters. We examine how the parameter initializations a 0 , w 0 and the layer-wise learning rates η a , η w influence the training trajectory in parameter space, function space (defined by the product β = aw ), and the evolution of the the NTK matrix,

Except for a measure zero set of initializations which converge to saddle points 2 , all gradient flow trajectories will converge to a global minimum, de-

3.0

1.5

0.0

1.5

3.0

w

1

3.0

1.5

0.0

1.5

3.0

w

2

3.0

1.5

0.0

1.5

3.0

a

(a)

δ

=

-

2

3.0

1.5

0.0

1.5

3.0

w

1

3.0

1.5

0.0

1.5

3.0

w

2

3.0

1.5

0.0

1.5

3.0

a

(b)

δ

= 0

3.0

1.5

0.0

1.5

3.0

w

1

3.0

1.5

0.0

1.5

3.0

w

2

3.0

1.5

0.0

1.5

3.0

a

(c)

δ

= 2

Figure 2: Balance determines geometry of trajectory. The quantity δ = η w a 2 -η a ∥ w ∥ 2 is conserved through gradient flow, which constrains the trajectory to: (a) a onesheeted hyperboloid for downstream initializations, (b) a double cone for balanced initializations, and (c) a twosheeted hyperboloid for upstream initializations. Gradient flow dynamics for three different initializations a 0 , w 0 with the same product β 0 = a 0 w 0 are shown. The minima manifold is shown in red and the manifold of equivalent β 0 initializations in gray. The surface is colored according to training loss, with blue representing higher loss and red representing lower loss.

termined by the normal equations X ⊺ Xaw = X ⊺ y . However, even when X ⊺ X is invertible such that the global minimum β ∗ is unique, the rescaling symmetry between a and w results in a manifold

of minima in parameter space. The minima manifold is a one-dimensional hyperbola where w ∝ β ∗ and has two distinct branches for positive and negative a . The symmetry also imposes a constraint on the network's trajectory, maintaining the difference δ = η w a 2 -η a ∥ w ∥ 2 ∈ R throughout training (see Appendix A.1 for details). This confines the parameter dynamics to the surface of a hyperboloid where the magnitude and sign of the conserved quantity determines the geometry, as shown in Fig. 2. An upstream initialization occurs when δ > 0 , a balanced initialization when δ = 0 , and a downstream initialization when δ < 0 .

Deriving exact solutions in parameter space. We initially assume 3 whitened input X ⊺ X = I d such that the ordinary least squares solution is β ∗ = X ⊺ y , and the gradient flow dynamics simplify to ˙ a = η a ( w ⊺ β ∗ -a ∥ w ∥ 2 ) , ˙ w = η w ( aβ ∗ -a 2 w ) . Notice that w ( t ) ∈ span( { w 0 , β ∗ } ) , and through training, w aligns in direction to ± β ∗ depending on the basin of attraction 4 the parameters are initialized in. Therefore, we can monitor the dynamics by tracking the hyperbolic geometry between a and ∥ w ( t ) ∥ and the spherical angle between w ( t ) and β ∗ . We study the variables µ = a ∥ w ∥ , an invariant under the rescale symmetry, and ϕ = w ⊺ β ∗ ∥ w ∥∥ β ∗ ∥ , the cosine of the spherical angle. From these two scalar quantities µ ( t ) , ϕ ( t ) and the initialization a 0 , w 0 , we can determine the trajectory a ( t ) and w ( t ) in parameter space. The dynamics for µ, ϕ are given by the coupled nonlinear ODEs,

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 δ .

Upstream. When δ ≫ 0 , the updates for both µ and ϕ diverge, but ϕ updates much more rapidly. We can decouple the dynamics of µ and ϕ by separation of their time scales and assume ϕ has reached its steady-state of ± 1 before µ has updated. Then, the dynamics of µ is linear and proceeds exponentially to ±∥ β ∗ ∥ . This regime exhibits minimal kernel movement (see Fig. 3 (c)) because the kernel is dominated by the η w a 2 I d term, whereas it is mainly w that updates.

Balanced. When δ = 0 , µ follows a Bernoulli differential equation driven by a time-dependent signal ϕ ∥ β ∗ ∥ , and ϕ follows a Riccati equation evolving from an initial value to ± 1 depending on the basin of attraction. For vanishing initialization ∥ β 0 ∥ → 0 , the temporal dynamics of µ and ϕ decouple such that there are two phases of

0

5

10

15

20

Time

t

0.25

0.50

0.75

1.00

(a)

µ

(

t

)

0

5

10

15

20

Time

t

0.00

0.25

0.50

0.75

1.00

Alignment

=4

=2

=1

=0

=

1

=

2

=

4

(b)

ϕ

(

t

)

0

5

10

15

20

Time

t

0.0

0.2

0.4

0.6

0.8

Kernel Distance

S

(0,

t

)

(c)

S

(0

, t

)

learning: an initial alignment phase where ϕ →± 1 , followed by a fitting phase where µ →±∥ β ∗ ∥ . In the first phase, w aligns to β ∗ resulting in a 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 µ simplify to the Bernoulli equation studied in Saxe et al. [10] and the kernel evolves solely in overall scale.

Downstream. When δ ≪ 0 , the updates for µ diverge, while the updates for ϕ vanishes. In this regime the dynamics proceed by an initial fast phase where µ converges exponentially to its steady state of ϕ ∥ β ∗ ∥ . Plugging this steady state into the dynamics of ϕ gives a Bernoulli differential equation

̸

˙ ϕ = η a η w ∥ β ∗ ∥ 2 | δ | -1 ϕ (1 -ϕ 2 ) . Due to the coefficient | δ | -1 , the second alignment phase proceeds very slowly as ϕ approaches ± 1 , assuming ϕ, µ = 0 , which is a saddle point. In this regime, the dynamics proceed by an initial lazy fitting phase, followed by a rich alignment phase, where the delay is determined by the magnitude of δ .

Identifying regimes of learning in function space. Here we take an alternative route towards understanding the influence of the relative scale by directly examining the dynamics in function space, an analysis strategy we will generalize to broader setups in Sections 4 and 5. The network's function is determined by the product β = aw and governed by the ODE,

where ρ = Xβ -y is the residual. These dynamics can be interpreted as preconditioned gradient flow on the loss in function space where the preconditioning matrix M depends on time through its dependence on a 2 and ww ⊺ . Whenever ∥ β ∥ ̸ = 0 , we can express M directly in terms of β and δ as

where κ = √ δ 2 +4 η a η w ∥ β ∥ 2 (see Appendix A.3 for a derivation). This establishes a self-consistent equation for the dynamics of β regulated by δ . Additionally,

1.2

0.9

0.6

0.3

0.0

1

0.0

0.3

0.6

0.9

1.2

2

=4

=2

=1

=0

=

1

=

2

=

4

(a) Whitened

X

⊺

X

0.50

0.75

1.00

1.25

1.50

1.75

1

0.00

0.25

0.50

0.75

1.00

1.25

2

=

=

(b) Low-Rank

X

⊺

X

Figure 4: Balance modulates β dynamics and implicit bias. Here we show the dynamics of β = aw with different values of δ , but the same initial β 0 . When X ⊺ X is whitened (left), we can solve for the dynamics exactly using our expressions for µ, ϕ (black dashed lines). Upstream initializations follow the trajectory of gradient flow on β , downstream initializations first move in the direction of β 0 before sweeping around towards β ∗ , and balanced initializations take an intermediate trajectory between these two. When X ⊺ X is low-rank (right), then we can only predict the trajectories in the limit of δ = ±∞ . If the interpolating manifold is onedimensional, then we can solve for the solution in terms of δ exactly (black dots). See Appendix A.4 for details.

notice that M characterizes the NTK matrix Eq. (1). Thus, understanding the evolution of M along the trajectory β 0 to β ∗ offers a method to discern between lazy and rich learning. Upstream. When δ ≫ 0 , M ≈ δ I d , and the dynamics of β 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 δ = 0 , M = √ η a η w ∥ β ∥ ( I d + ββ ⊺ ∥ β ∥ 2 ) . 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 δ ≪ 0 , M ≈ | δ | ββ ⊺ ∥ β ∥ 2 , and β follows a projected gradient descent trajectory, attempting to reach β ∗ in the direction of β 0 . Along this trajectory the NTK matrix doesn't evolve. However, if β 0 is not aligned to β ∗ , then at some point the dynamics of β will slowly align. In this second alignment phase the NTK matrix will change, confirming the dynamics are initially lazy followed by a delayed rich phase. See Appendix A.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 X ⊺ X , ensuring the existence of a unique least squares solution β ∗ . In this setting, δ influences the trajectory the model takes from β 0 to β ∗ , as shown in Fig. 4 (a). Now we consider low-rank X ⊺ X , such that there exist infinitely many interpolating solutions in function space. By studying the structure of M , we can characterize how δ determines the interpolating solution the dynamics converge to. Extending a time-warped mirror flow analysis strategy pioneered by Azulay et al. [9] to allow δ < 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 β 0 .

̸

Theorem 3.1 (Extending Theorem 2 in Azulay et al. [9]) . For a single hidden neuron linear network, for any δ ∈ R , and initialization β 0 such that β ( t ) = 0 for all t ≥ 0 , if the gradient flow solution β ( ∞ ) satisfies Xβ ( ∞ ) = y , then,

We observe that for vanishing initializations there is functionally no difference between the inductive bias of the upstream ( δ ≫ 0 ) and balanced ( δ = 0 ) settings. However, in the downstream setting ( δ ≪ 0 ), 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 δ is presented in Fig. 4 (b), where if the null space of X ⊺ X is one-dimensional, we can solve for β ( ∞ ) in closed form (see Appendix A.4).

4 Wide and Deep Linear Networks

We now show how the analysis techniques used to 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. Weconsider the dynamics of a two-layer linear network with h hidden neurons and c outputs, f ( x ; θ ) = A ⊺ Wx , where W ∈ R h × d and A ∈ R h × c . We assume h ≥ min( d, c ) , such that this parameterization can represent all linear maps from R d → R c . The rescaling symmetry between A and W implies the h × h matrix ∆ = η w A 0 A ⊺ 0 -η a W 0 W ⊺ 0 is conserved throughout gradient flow [62]. Drawing insights from our analysis of the single-neuron scenario ( h = c = 1 ), we consider the dynamics of β = W ⊺ A ∈ R d × c ,

where vec( · ) denotes the vectorization operator and ⊕ denotes the Kronecker sum 5 . As in the single-neuron setting, we find that the dynamics of β are preconditioned by a matrix M that depends on quadratics of A and W and characterizes the NTK matrix K = ( I c ⊗ X ) M ( I c ⊗ X ⊺ ) . We now show how M can be expressed 6 in terms of the rank-1 matrices β k = w k a ⊺ k ∈ R d × c , which represent the contribution to β of a single neuron with parameters w k , a k and conserved quantity δ k = ∆ kk .

̸

Theorem 4.1. Whenever ∥ β k ∥ F = 0 for all k ∈ [ h ] , the matrix M can be expressed as the sum M = ∑ h k =1 M k over hidden neurons where M k is defined as,

By studying the dependence of M on the conserved quantities δ k and the dimensions d , h and c , we can determine the influence of the relative scale on the learning regime. When min( d, c ) ≤ h < max( d, c ) , and assuming independent initializations for all β k , then networks which narrow from input to output ( d > c ) enter the lazy regime when all δ k ≫ 0 , whereas networks which expand from input to output ( d < c ) do so when all δ k ≪ 0 . However, with opposite signs for δ k , and assuming all β k (0) ̸∝ β ∗ , these networks enter a delayed rich regime . As elaborated in Appendix B.1.5, this occurs because in these regimes a solution β ∗ does not exist within the space spanned by M at initialization. When h ≥ max( d, c ) all networks enter the lazy regime when all δ k ≫ 0 or all δ k ≪ 0 . Conversely, as all δ k → 0 , all networks transition into the rich regime regardless of dimensions. While Theorem 4.1 offers valuable insight into the learning regimes in the limits of δ k , understanding the transition between regimes remains challenging. To achieve this, we aim to express M in terms of β , rather than β k , by introducing structure on the conserved quantities ∆ .

Theorem 4.2. When ∆ = δ I h and h = d if δ < 0 or h = c if δ > 0 , then the matrix M can be expressed as M = √ η a η w β ⊺ β + δ 2 4 I c ⊗ I d + I c ⊗ √ η a η w ββ ⊺ + δ 2 4 I d .

From Theorem 4.2 the resulting dynamics of β simplify to a self-consistent equation regulated by δ ,

where P = Xβ -Y is the residual. Under our isotropic assumption on the conserved quanitities ∆ = δ I h , these dynamics are exact. Concurrent to our work, Tu et al. [63] finds that β approximately follows these dynamics in the overparameterized setting h ≫ max( d, c ) under a Gaussian initialization N (0 , σ 2 ) of the parameters where σ 2 h is analogous to δ .

Equipped with a self-consistent equation for the dynamics of β we now aim to interpret these dynamics as a mirror flow with a δ -dependent potential. As presented in Theorem B.6, the dynamics of the singular values of β can be described as a mirror flow with a hyperbolic entropy potential, which smoothly interpolates between an ℓ 1 and ℓ 2 penalty on the singular values for the rich ( δ → 0 ) and lazy ( δ →±∞ ) 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 a 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-layer linear networks by Azulay et al. [9] to deep linear networks. For a depth-( L +1) linear network, f ( x ; θ ) = A ⊺ ∏ L l =1 W l x , where β = ∏ L l =1 W ⊺ l A , we find that the inductive bias of the rich regime is Q ( β ) = ( L +1 L +2 ) ∥ β ∥ L +2 L +1 -∥ β 0 ∥ -L L +1 β ⊺ 0 β . This inductive bias strikes a depth-dependent balance between attaining the minimum norm solution and preserving the initialization direction.

5 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 σ ( z ) = max( z, γz ) . The input-output map of a piecewise linear network with L hidden layers and h hidden neurons per layer is comprised of potentially O ( h dL ) convex activation regions [65]. Each region is defined by a 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 a 2-colorable 7 convex 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, f ( x ; θ ) = a ⊺ σ ( Wx ) , where W ∈ R h × d and a ∈ R h . Following the approach in Section 4, we consider the contribution to the input-output map from a single hidden neuron k ∈ [ h ] with parameters w k ∈ R d , a k ∈ R and conserved quantity δ k = η w a 2 k -η a ∥ w k ∥ 2 [62]. However, unlike the linear setting, the neuron's contribution to f ( x i ; θ ) is regulated by whether the input x i is in the neuron's active halfspace . Let C ∈ R h × n be the matrix with elements c ki = σ ′ ( w ⊺ k x i ) , which determines the activation of the k th neuron for the i th data point. The dynamics of β k = a k w k are,

̸

The matrix M k ∈ R d × d is a preconditioning matrix on the dynamics, and when β k = 0 , it can be expressed in terms of β k and δ k . Unlike the linear setting, ξ k ∈ R d driving the dynamics is not shared for all neurons because of its dependence on c ki . Additionally, the NTK matrix in this setting depends on M k and C , with elements K ij = ∑ h k =1 c ki x ⊺ i M k x j c kj . To examine the evolution of K , we consider a signed spherical coordinate transformation separating the dynamics of β k into its directional ˆ β k = sgn( a k ) β k ∥ β k ∥ and radial µ k = sgn( a k ) ∥ β k ∥ components, such that β k = µ k ˆ β k . ˆ β k determines the direction and orientation of the halfspace where the k th neuron is active, while µ k determines the slope of the contribution in this halfspace. These coordinates evolve according to,

Downstream. When δ k ≪ 0 , M k ≈ | δ k | ˆ β k ˆ β ⊺ k , and the dynamics are approximately ∂ t ˆ β k = 0 and ∂ t µ k = -| δ k | ˆ β ⊺ k ξ k . Irrespective of ξ k , ˆ β k ( t ) = ˆ β k (0) , which implies the overall partition map doesn't change (Fig. 5, bottom), nor the activation patterns C , nor M k . Only µ k 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 | δ k | determining the delay.

Balanced. When δ k = 0 , M k = √ η a η w | µ k | ( I d + ˆ β k ˆ β ⊺ k ) , and the dynamics simplify to, ∂ t ˆ β k = - √ η a η w sgn( µ k )( I d -ˆ β k ˆ β ⊺ k ) ξ k and ∂ t µ k = -2 √ η a η w | µ k | ˆ β ⊺ k ξ k . Here both the direction and magnitude of β k evolve, resulting in changes to the activation regions, patterns C , and NTK K . For vanishing initializations where ∥ β k (0) ∥ → 0 for all k ∈ [ h ] , we can decouple the dynamics into two

Rich

Lazy

Phase II: Data Fitting

Phase I: Partition Alignment

Fixed Partition Data Fitting

Time

Figure 5: Rapid feature learning is caused by large activation changes with minimal parameter movement. (a) We show the surface of a two-layer ReLU network trained on an XOR-like task, starting with a near-zero input-output map, f ( x ; θ 0 ) ≈ 0 . The surface consists of convex conic regions, each with a distinct activation pattern, colored by the parity of active neurons. A lazy initialization ( bottom ) maintains a fixed activation partition throughout training, reweighting the hidden neurons to fit the data. In contrast, a 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 t = 0 , as a function of overall and relative scales (same experiments as in Fig. 1(b)). Rapid feature learning occurs from a smallτ upstream initialization that promotes faster learning in early layers, driving a large change in Hamming distance, but a small change in parameter space. In contrast, smallτ 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 3. Phase I: Partition alignment. At vanishing scale, the output f ( x ; θ 0 ) ≈ 0 for all input x , such that the vector driving the dynamics ξ k ≈ -∑ n i =1 c ki x i y i is independent of the other hidden neurons. At the same time, the radial dynamics slow down relative to the directional dynamics, and the function'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, 53, 54], pair-wise orthonormal [52], linearly separable and symmetric [51] or small angle [55], to analytically determine the fixed points of this alignment phase. Phase II: Data fitting. After enough time, the magnitudes of β k have grown such that we can no longer assume f ( x ; θ ) ≈ 0 and thus the residual will depend on all β k . 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 to change, and therefore some prior works have further decomposed this phase into multiple stages.

Upstream. When δ k ≫ 0 , M k ≈ δ k I d , and the dynamics are approximately ∂ t ˆ β k = -δ k µ -1 k ( I d -ˆ β k ˆ β ⊺ k ) ξ k and ∂ t µ k = -δ k ˆ β ⊺ k ξ k . Again, both the direction and magnitude of β k change. However, unlike the balanced setting, in this setting M k is independent of β k and stays constant through training. Yet, as β k change in direction, so can C , and thus the NTK. This setting is unique, because it is rich due to a 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. 1 and in Fig. 5.

Connections to infinite-width. Our study of learning regimes in finite-width 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, in Luo et al. [17] they consider a network f ( x ) = 1 α ∑ h k =1 a k σ ( w ⊺ k x ) with weights initialized as a k ∼ N (0 , β 2 a ) and w k ∼ N (0 , β 2 W I d ) as width h →∞ . They obtain a phase diagram at infinite width capturing the dependence of learning regime on the overall function scale β a β W /α and the relative initialization scale β a /β W , each suitably normalized as a function of width. The resulting phase portrait is analogous to ours in Fig. 1 (b), where we use the conserved quantity δ rather than the relative scale β a /β W . Specifically, there is a lazy regime that includes the NTK parameterization, which is always achieved at large scale (as in the largeτ regions of Fig. 1 (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 τ in Fig. 1 (b)). On the other side of the phase boundary is the infinite-width analog of rapid rich learning, where all neurons condense to a few directions. This is induced either at small function scale, or at larger scales if β a /β W is sufficiently large, such that W learns fast enough relative to a . The phase boundary, in

=10

(a) Kernel Distance

10

3

10

2

10

1

10 0

P

/

P

10

5

10

4

10

3

10

2

10

1

10 0

=33

=10

=3

=1.0

=0.3

=0.1

=3

=1

=0.3

=0.1

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-field parametrizations has been extensively studied, both lead to the same distribution of per-neuron conserved quantities, which is zero in expectation with a non-vanishing variance. A 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. 6 show that upstream initializations 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. [67] as a 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 a single-layer transformer ( < 6% of all parameters) significantly reduces the time to 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 to 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 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 [74, 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 an important direction for future work.

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 in part by grant NSF PHY-1748958 to the Kavli Institute for Theoretical Physics (KITP).

Author Contributions

This project originated from conversations between Daniel and Allan at 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 are primarily responsible for the wide and deep linear analysis in Section 4. Daniel is primarily responsible for the nonlinear analysis in Section 5. Allan, Feng, and David are primarily responsible for the empirics in Fig. 1 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

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054 , 2018.

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.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International conference on machine learning , pages 242-252. PMLR, 2019.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep relu networks. Machine learning , 109:467-492, 2020.

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.

Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14548 , 2020.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems , 32, 2019.

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 on Machine Learning , pages 468-477. PMLR, 2021.

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.

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 (23):11537-11546, 2019.

Arthur Jacot, François Ged, Berfin ¸ Sim¸ sek, Clément Hongler, and Franck Gabriel. Saddle-tosaddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106.15933 , 2021.

Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. arXiv preprint arXiv:2012.09839 , 2020.

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.

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.

Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522 , 2020.

Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at infinite-width limit. Journal of Machine Learning Research , 22(71):1-47, 2021.

Greg Yang, Edward J 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.03466 , 2022.

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.

Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems , 35:37932-37946, 2022.

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 arXiv:2306.04815 , 2023.

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.04980 , 2024.

Yizhou Xu and Liu Ziyin. When does feature learning happen? perspective from an analytically solvable model. arXiv preprint arXiv:2401.07085 , 2024.

Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research , 13(1):795-828, 2012.

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):113301, 2020.

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.

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 in Neural Information Processing Systems , 33:5850-5861, 2020.

Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809.10374 , 2018.

Kenji Fukumizu. Effect of batch learning in multilayer neural networks. Gen , 1(04):1E-03, 1998.

Lukas Braun, Clémentine Carla Juliette Dominé, James E Fitzgerald, and Andrew M Saxe. Exact learning dynamics of deep linear networks with prior knowledge. In Advances in Neural Information Processing Systems , 2022.

Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In International conference on machine learning , pages 244-253. PMLR, 2018.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems , 32, 2019.

Liu Ziyin, Botao Li, and Xiangming Meng. Exact solutions of a deep linear network. Advances in Neural Information Processing Systems , 35:24446-24458, 2022.

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.

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.

Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv:1909.12051 , 2019.

Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2021.

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research , 19(1):2822-2878, 2018.

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv:1810.02032 , 2018.

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.

Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Advances in neural information processing systems , 33:22182-22193, 2020.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890 , 2019.

Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning , pages 4683-4692. PMLR, 2019.

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.

Daniel Kunin, Atsushi Yamamura, Chao Ma, and Surya Ganguli. The asymmetric maximum margin bias of quasi-homogeneous neural networks. arXiv preprint arXiv:2210.03820 , 2022.

Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning , pages 1832-1841. PMLR, 2018.

Suriya Gunasekar, Blake Woodworth, and Nathan Srebro. Mirrorless mirror descent: A natural derivation of mirror descent. In International Conference on Artificial Intelligence and Statistics , pages 2305-2313. PMLR, 2021.

Zhiyuan Li, Tianhao Wang, Jason D Lee, and Sanjeev Arora. Implicit bias of gradient descent on reparametrized models: On equivalence to mirror descent. Advances in Neural Information Processing Systems , 35:34626-34640, 2022.

Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint arXiv:1803.08367 , 2018.

Mary Phuong and Christoph H Lampert. The inductive bias of relu networks on orthogonally separable data. In International Conference on Learning Representations , 2020.

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.

Etienne Boursier, Loucas Pillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances in Neural Information Processing Systems , 35:20105-20118, 2022.

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.

Hancheng Min, René Vidal, and Enrique Mallada. Early neuron alignment in two-layer relu networks with small initialization. arXiv preprint arXiv:2307.12851 , 2023.

Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of relu networks. Advances in Neural Information Processing Systems , 36, 2024.

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.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

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.

Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics , 75(9):1889-1935, 2022.

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.

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813 , 2023.

Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. Advances in Neural Information Processing Systems , 31, 2018.

Zhenfeng Tu, Santiago Aranguri, and Arthur Jacot. Mixed dynamics in linear networks: Unifying the lazy and active regimes. arXiv preprint arXiv:2405.17580 , 2024.

Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas Pillaud-Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. Advances in Neural Information Processing Systems , 36, 2024.

Maithra Raghu, Ben Poole, Jon Kleinberg, Surya Ganguli, and Jascha Sohl-Dickstein. On the expressive power of deep neural networks. In international conference on machine learning , pages 2847-2854. PMLR, 2017.

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM , 60(6):84-90, 2017.

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.02129 , 2023.

Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02177 , 2022.

Tanishq Kumar, Blake Bordelon, Samuel J Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics. arXiv preprint arXiv:2310.06110 , 2023.

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.

Noa Rubin, Inbar Seroussi, and Zohar Ringel. Grokking as a first order phase transition in two layer networks. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=3ROGsTX3IR .

Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217 , 2023.

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.

Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728 , 2020.

Hidenori Tanaka and Daniel Kunin. Noether's learning dynamics: Role of symmetry breaking in neural networks. Advances in Neural Information Processing Systems , 34:25646-25660, 2021.

Feng Chen, Daniel Kunin, Atsushi Yamamura, and Surya Ganguli. Stochastic collapse: How gradient noise attracts sgd dynamics towards simpler subnetworks. Advances in Neural Information Processing Systems , 36, 2024.

Liu Ziyin, Mingze Wang, and Lei Wu. The implicit bias of gradient noise: A symmetry perspective. arXiv preprint arXiv:2402.07193 , 2024.

David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters , 74(21):4337, 1995.

Sebastian Goldt, Madhu S Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Generalisation dynamics of online learning in over-parameterised neural networks. arXiv preprint arXiv:1901.09085 , 2019.

Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference on Learning Theory , pages 4224-4258. PMLR, 2021.

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:1312.6098 , 2013.

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.

Matus Telgarsky. Representation benefits of deep feedforward networks. arXiv preprint arXiv:1509.08101 , 2015.

Raman Arora, Amitabh Basu, Poorya Mianjy, and Anirbit Mukherjee. Understanding deep neural networks with rectified linear units. arXiv preprint arXiv:1611.01491 , 2016.

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.

Boris Hanin and David Rolnick. Complexity of linear regions in deep networks. In International Conference on Machine Learning , pages 2596-2604. PMLR, 2019.

Boris Hanin and David Rolnick. Deep relu networks have surprisingly few activation patterns. Advances in neural information processing systems , 32, 2019.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

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 1026-1034, 2015.

A Single-Neuron Linear Network

In this section, we provide a 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 ( x ; θ ) = aw ⊺ x , where a ∈ R and w ∈ R d are the parameters. We aim to understand the impact of the initializations a 0 , w 0 and the layer-wise learning rates η a , η w 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 :

The gradient flow dynamics are governed by the following coupled ODEs:

̸

The global minima of this problem are determined by the normal equations X ⊺ Xaw = X ⊺ y . Even when X ⊺ X is invertible, yielding a unique global minimum in function space β ∗ = ( X ⊺ X ) -1 X ⊺ y , the symmetry between a and w , permitting scaling transformations, a → aα and w → w/α for any α = 0 without changing the product aw , results in a manifold of minima in parameter space. This minima manifold is a one-dimensional hyperbola where aw = β ∗ , with two distinct branches for positive and negative a . The set of saddle points { ( a, w ) } forms a ( d -1) -dimensional subspace satisfying a = 0 and w ⊺ X ⊺ y = 0 . Except for a measure zero set of initializations that converge to the saddle points, all gradient flow trajectories will converge to a 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 to saddle points, separating the two basins.

A.1 Conserved quantity

The scaling symmetry between a and w results in a conserved quantity δ ∈ R throughout training, as noted in many prior works [10, 62, 74], where

This can be easily verified by explicitly writing out the dynamics of δ . Define ρ = ( X ⊺ Xaw -X ⊺ y ) for succinct notation, such that

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, as shown in Fig. 2. A hyperboloid of the form ∑ k i =1 x 2 i -∑ n i = k +1 x 2 i = α , with α ≥ 0 , exhibits varied topology and geometry based on k and α . It has two sheets when k = 1 and one sheet otherwise. Its geometry is primarily dictated by α : as α tends to infinity, curvature decreases, while at α = 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 ⊺ X = I d and β ∗ = X ⊺ y such that β ∗ = 0 . The dynamics of a and w can then be simplified as

A.2.1 Deriving the dynamics for µ and ϕ

As discussion in Section 3 we study the variables µ = a ∥ w ∥ , an invariant under the rescale symmetry, and ϕ = w ⊺ β ∗ ∥ w ∥∥ β ∗ ∥ , the cosine of the angle between w and β ∗ . This change of variables can also be understood as a signed spherical decomposition of β : µ is the signed magnitude of β and ϕ is the

cosine angle between β and β ∗ . Through chain rule, we obtain the dynamics for µ and ϕ , which can be expressed as

We leave the derivation to the reader, but emphasize that a key simplification used is to express the sum η w a 2 + η a ∥ w ∥ 2 in terms of δ ,

Additionally, notice that η a and η w only appear in the dynamics for µ and ϕ as the product η a η w or in the expression for δ . If we were to define µ ′ = √ η a η w µ and β ′ ∗ = √ η a η w β ∗ , then it is not hard to show that the product η a η w is absorbed into the dynamics. Thus, without loss of generality we can assume the product η a η w = 1 , resulting in the following coupled system of nonlinear ODEs,

We will now show how to solve this system of equations for µ and ϕ . We will solve this system when δ = 0 , δ > 0 , and δ < 0 separately. We will then in Appendix A.2.6 show a general treatment on how to obtain the individual coordinates of a and w from the solutions for µ and ϕ .

A.2.2 Balanced δ = 0

When δ = 0 , the dynamics for µ, ϕ are,

First, we show that the sign of µ cannot change through training and sgn( µ ) = sgn( a ) . Because δ = 0 , the dynamics of a and w are constrained to a double cone with a singularity at the origin ( a = 0 , w = 0 ). This point is a 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 a 0 > 0 , and solve the dynamics. The dynamics of µ is a Bernoulli differential equation driven by a time-dependent signal ϕ ∥ β ∗ ∥ . The dynamics of ϕ is decoupled from µ and is in the form of a Riccati equation evolving from an initial value ϕ 0 to 1 , as we have assumed an initialization with positive a 0 . This ODE is separable with the solution,

where c ϕ = tanh -1 ( ϕ 0 ) . Plugging this solution into the dynamics for µ gives a Bernoulli differential equation,

with the solution,

where c µ = 2 µ -1 0 cosh 2 ( c ϕ ) -(2 c ϕ +sinh(2 c ϕ )) . Note, if ϕ 0 = -1 , then ˙ ϕ = 0 , and the dynamics of µ will be driven to 0 , which is a saddle point.

A.2.3 Upstream δ > 0

When δ > 0 , the dynamics are constrained to a hyperboloid composed of two identical sheets determined by the sign of a 0 (as shown in Fig. 2 (c)). Without loss of generality we assume a 0 > 0 , which ensures a ( t ) > 0 for all t ≥ 0 . However, unlike in the balanced setting, the dynamics of µ

and ϕ do not decouple, making it difficult to solve. Instead, we consider ν = w ⊺ β ∗ a , which evolves according to the Riccati equation,

The solution is given by,

where R = 1 2 √ δ 2 +4 ∥ β ∗ ∥ 2 . The trajectory of a ( t ) is given by the Bernoulli equation,

which can be solved analytically using ν ( t ) . For a 0 > 0 , we have that

where b 0 = w ⊺ 0 β ∗ , and Y ( t ) = 1 2 √ δ 2 +4 ∥ β ∗ ∥ 2 t +atanh ( 2 b 0 a 0 + δ √ δ 2 +4 ∥ β ∗ ∥ 2 ) . From the solutions for ν, a , we can easily obtain dynamics for µ, ϕ .

A.2.4 Downstream δ < 0

When δ < 0 , the dynamics are constrained to a hyperboloid composed of a single sheet (as shown in Fig. 2 (a)). However, unlike in the upstream setting, a may change sign. A zero-crossing in a leads to a finite time blowup in ν . Consequently, applying the approach used to solve for the dynamics in the upstream setting becomes more intricate. First we show the following lemma:

̸

Lemma A.1. If a 0 = 0 or w ⊺ 0 β ∗ = 0 , then a ( t ) w ( t ) ⊺ β ∗ = 0 has at most one solution for t ≥ 0 .

̸

Proof. Let ω ( t ) = w ( t ) ⊺ β ∗ . The two-dimensional dynamics of a ( t ) and ω ( t ) are given by,

Consider the orthant O + = { ( a, ω ) | a > 0 , ω > 0 } . The boundary ∂O + is formed by two orthogonal subspaces. On { ( a, ω ) | a = 0 , ω ≥ 0 } , ˙ a ≥ 0 . On { ( a, ω ) | a ≥ 0 , ω = 0 } , ˙ ω ≥ 0 . Therefore, O + is a positively invariant set. Similarly, O -= { ( a, ω ) | a < 0 , ω < 0 } is a positively invariant set. On the boundary ∂O + ∪ ∂O -= { ( a, ω ) | aω = 0 } , the flow is contained only at the origin a = 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 ≥ 0 . As a result, the trajectory ( a ( t ) , ω ( t )) will at most intersect the boundary ∂O + ∪ ∂O -once.

From Lemma A.1, we conclude that either a crosses zero, w ⊺ β ∗ crosses zero, or neither crosses zero. When a doesn't cross zero, then ν is well-defined for t ≥ 0 , and our argument from Appendix A.2.3 still holds, leading to solutions for µ, ϕ . When a does cross zero, instead of ν , we consider υ = a w ⊺ β ∗ , the inverse of ν . In this case, we know from Lemma A.1 that w ⊺ β ∗ does not cross zero and thus υ is well-defined for t ≥ 0 and evolves according to the Riccatti equation,

These dynamics have a solution similar to Eq. (28), which we leave to the reader. With υ ( t ) , we can then solve for the dynamics of w ⊺ β ∗ . Let ω = w ⊺ β ∗ , then ω evolves according to the Bernoulli equation,

which can be solved analytically using υ ( t ) , analogous to the solution for a ( t ) in Appendix A.2.3. From the solutions for υ, ω , we can easily obtain dynamics for µ, ϕ .

A.2.5 Basins of attraction

From Lemma A.1 we know that a can cross zero no more than once during its trajectory. Consequently, we can identify the basin of attraction by determining the conditions under which a changes sign. This analysis is crucial because initial conditions leading to a sign change in a correspond to scenarios where initial positive and negative values of a 0 are drawn towards the negative and positive branches of the minima manifold, respectively. From Eq. (28) we can immediately see that a will change sign when the denominator vanishes. This can happen if √ δ 2 +4 ∥ β ∗ ∥ 2 <

-2 ν 0 -δ . For δ < 0 , this is satisfied if ν 0 < 1 2 ( -δ -√ δ 2 +4 ∥ β ∗ ∥ 2 ) , which gives the hyperplane w ⊺ 0 β ∗ + a 0 2 ( δ + √ δ 2 +4 ∥ β ∗ ∥ 2 ) = 0 that separates between initializations for which a changes sign and initializations for which it does not (Fig. 7). Consequently, letting S + be the set of initializations attracted to the minimum manifold with a > 0 , we have that:

where the bottom inequality means that β 0 is sufficiently aligned to β ∗ in the case of a 0 ≥ 0 or sufficiently misaligned in the case of a 0 ≤ 0 . We can similarly define the analogous S -. An initialization on the separating hyperplane will converge to a saddle point where w ⊺ β ∗ = a = 0 .

5.0

2.5

0.0

2.5

5.0

w

1

5.0

2.5

0.0

2.5

5.0

w

2

5.0

2.5

0.0

2.5

5.0

a

A.2.6 Recovering parameters ( a, w ) from ( µ, ϕ )

We now discuss how to recover the dynamics of the parameters ( a, w ) from our solutions for ( µ, ϕ ) . We can recover a and ∥ w ∥ from µ . Using Eq. (19) discussed previously, we can show

We now discuss how to obtain the vector w from ϕ . The key observation, as discussed in Section 3, is that w only moves in the span of w 0 and β ∗ . This means we can express w ( t ) as

where c 1 ( t ) is the coefficient in the direction of β ∗ and c 2 ( t ) is the coefficient in the direction orthogonal to β ∗ on the two-dimensional plane defined by w 0 . From the definition of ϕ we can easily obtain the coefficients c 1 = ∥ w ∥ ϕ and c 2 = √ ∥ w ∥ 2 -c 2 1 . We always choose the positive square root for c 2 , as c 2 ( t ) ≥ 0 for all t . See Appendix D.2 for experimental details of how we ran our simulations and a notebook generating these exact solutions.

2

0

2

(

t

)

0

1

a

(

t

)

=

2

=

1

=

0.2

=0.02

=0.2

=1

=2

0.5

0.0

0.5

w

1(

t

)

1

0

w

2(

t

)

0

1

2

0.5

1.0

1.5

=

2

=

1

=

0.2

=0.02

=0.2

=1

=2

0.5

1.0

1.5

1

0

0

5

10

Time

1

2

0

5

10

Time

0.5

1.0

1.5

=

2

=

1

=

0.2

=0.02

=0.2

=1

=2

0

5

10

Time

0.50

0.75

1.00

0

5

10

Time

0.50

0.75

1.00

A.3 Function space dynamics of β

The network's function is determined by the product β = aw and governed by the ODE,

Notice, that the vector X ⊺ ρ driving the dynamics of β is the gradient of the loss with respect to β , X ⊺ ρ = ∇ β L . 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 2 and ww ⊺ . The matrix M also characterizes the NTK matrix, K = XMX ⊺ . As discussed in Section 3, our goal is to understand the evolution of M along a trajectory { β ( t ) ∈ R d : t ≥ 0 } solving Eq. (37).

First, notice that by expanding ∥ β ∥ 2 = a 2 ∥ w ∥ 2 in terms of the conservation law, we can show

which is the unique positive solution of the quadratic expression η w a 4 -δa 2 -η a ∥ β ∥ 2 = 0 . When a 2 > 0 we can use this solution and the outer product ββ ⊺ = a 2 ww ⊺ to solve for ww ⊺ in terms of β ,

Plugging these expressions into M gives

̸

Thus, given any initialization a 0 , w 0 such that a ( t ) 2 = 0 for all t ≥ 0 , we can express the dynamics of β entirely in terms of β . This is true for all initializations with δ ≥ 0 , except if initialized on the saddle point at the origin. It is also true for all initializations with δ < 0 where the sign of a does not switch signs. In the next section we will show how to interpret these trajectories as time-warped mirror flows for a potential that depends on δ . 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 β 0 such that ∥ β ( t ) ∥ > 0 for all t ≥ 0 .

Notice, that η a and η w only appear in the dynamics for β as the product η a η w or in the expression for δ . By defining β ′ = √ η a η w β and y ′ = √ η a η w y and studying the dynamics of β ′ , we can absorb η a η w into the β terms in M and the additional factor √ η a η w into the β and y terms in ρ . This transformation of β 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 β assuming η a η w = 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 = 2 ∥ β ∥ κ ( I d + ˆ β ˆ β ⊺ ) ∂ t ∥ β ∥ + κ -δ 2 ∂ t ( ˆ β ˆ β ⊺ ) where ˆ β = β ∥ β ∥ . 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 β and another that depends on the change in the direction of β . As done in the main text, we consider δ ≫ 0 , δ ≪ 0 , and δ = 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 β (0) to β ∗ , the change in the NTK is vanishing, indicative of a lazy regime. For δ ≪ 0 , 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 β . This is why for large negative delta we observe a delayed rich regime, where the eventual alignment of β to β ∗ leads to a 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 β contribute to a change in the kernel, indicative of a rich regime.

A.4 Deriving the inductive bias

Until now, we have primarily considered that X ⊺ X is either whitened or full rank, ensuring the existence of a unique least squares solution β ∗ . In this setting, δ influences 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 over-parameterized setting where we have more features d than observations n such that X ⊺ X is 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 δ determines which interpolating solution the dynamics converge to among all possible interpolating solutions. To do this we will extend a time-warped mirror flow analysis strategy pioneered by Azulay et al. [9].

A.4.1 Overview of time-warped mirror flow analysis

Here we recap the standard analysis for determining the implicit bias of a linear network through mirror flow. As first introduced in Gunasekar et al. [46], if the learning dynamics of the predictor β can be expressed as a mirror flow for some strictly convex potential Φ α ( β ) ,

where ρ = ( Xβ -y ) is the residual, then the limiting solution of the dynamics is determined by the constrained optimization problem,

where D Φ α ( p, q ) = Φ α ( p ) -Φ α ( q ) -⟨∇ Φ α ( q ) , p -q ⟩ is the Bregman divergence defined with Φ α . To understand the relationship between mirror flow Eq. (41) and the optimization problem Eq. (42), we consider an equivalent constrained optimization problem

where Q ( β ) = Φ α ( β ) - ∇ Φ α ( β (0)) ⊺ β , which is often referred to as the implicit bias . Q ( β ) is strictly convex, and thus it is sufficient to show that β ( ∞ ) is a first order KKT point of the constrained optimization (43). This is true iff there exists ν ∈ R n such that ∇ Q ( β ( ∞ )) = X ⊺ ν . The goal is to derive ν from the mirror flow Eq. (41). Notice, we can rewrite Eq. (41) as, ˙ ( ∇ Φ α ( β )) = -X ⊺ ρ , which integrated over time gives

The LHS is ∇ Q ( β ( ∞ )) . Thus, by defining ν = ∫ ∞ 0 ρ ( t ) dt , which assumes the residual decays fast enough such that this is well defined, then we have shown the desired KKT condition. Crucial to this analysis is that there exists a solution to the second-order differential equation

which even for extremely simple Jacobian maps may not be true [47]. Azulay et al. [9] showed that if there exists a smooth positive function g ( β ) : R d → (0 , ∞ ) such that the ODE,

has a solution, then the previous interpretation holds for Φ α ( β ) with ν = ∫ ∞ 0 g ( β ( t ′ )) ρ ( t ′ ) dt . As before, it is crucial that this integral exists and is finite. Azulay et al. [9] further explained that this scalar function g ( β ) can be considered as warping time τ ( t ) = ∫ t 0 g ( β ( t ′ )) dt ′ on the trajectory taken in predictor space β ( τ ( t )) . So long as this warped time doesn't 'stall out', that is we require that τ ( ∞ ) = ∞ , then this will not change the interpolating solution.

A.4.2 Applying time-warped mirror flow analysis

Here show how to apply the time-warped mirror flow analysis to the dynamics of β derived in Appendix A.3 where ∇ θ β ∇ θ β ⊺ = M . We will only consider initializations β 0 such that ∥ β ( t ) ∥ > 0 for all t ≥ 0 , such that M can be expressed as

Computing M -1 . Whenever ∥ β ∥ > 0 , then M is a positive definite matrix with a unique inverse that can be derived using the Sherman-Morrison formula, ( A + uv ⊺ ) -1 = A -1 -A -1 uv ⊺ A -1 1+ u ⊺ A -1 v . Here we can define A , u , and v as

First notice the following simplification, u ⊺ A -1 v = √ δ 2 +4 ∥ β ∥ 2 -δ √ δ 2 +4 ∥ β ∥ 2 + δ . After some algebra, M -1 is

To make notation simpler we will define the following two scalar functions,

such that we can express M -1 = f δ ( ∥ β ∥ 2 ) I d -h δ ( ∥ β ∥ 2 ) ββ ⊺ .

Proving M -1 is not a Hessian map. If M -1 is the Hessian of some potential, then we can show that the dynamics of β are a mirror flow. However, from our expression for M -1 we can actually prove that it is not a Hessian map. As discussed in Gunasekar et al. [47], a symmetric matrix H ( β ) is the Hessian of some potential Φ( β ) if and only if it satisfies the condition,

̸

We will use this property to show M -1 is not a Hessian map. First, notice this condition is trivially true when i = j = k . Second, notice that for all i = j = k ,

̸

̸

Thus, M -1 is a Hessian map if and only if for all i = j , ∂M -1 ii ∂β j = ∂M -1 ij ∂β i . Using our expression for M -1 , the LHS is

while the RHS is

Thus, M -1 is a Hessian map if and only if 2 ∇ f δ ( x ) + h δ ( x ) = 0 . Plugging in our definitions of f δ ( x ) and h δ ( x ) we find

which does not equal zero and thus M -1 is not a Hessian map.

Finding a scalar function g δ ( x ) such that g δ ( ∥ β ∥ 2 ) M -1 is a Hessian map. While we have shown that M -1 is not a Hessian map, it is very close to a Hessian map. Here we will show that there exists a scalar function g δ ( x ) such that g δ ( ∥ β ∥ 2 ) M -1 is a Hessian map. For any g δ ( x ) can define g δ ( ∥ β ∥ 2 ) M -1 in terms of two new functions ˜ f δ ( x ) and ˜ h δ ( x ) evaluated at x = ∥ β ∥ 2 ,

Thus, as derived in the previous section, we get the analogous condition on ˜ f δ ( x ) and ˜ h δ ( x ) for g δ ( ∥ β ∥ 2 ) M -1 to be a Hessian map,

Rearranging terms we find that g δ ( x ) must solve the ODE

Using our previous expressions (Eq. (50) and Eq. (55)) we find

which implies g δ ( x ) solves the differential equation, ∇ g δ ( x ) = g δ ( x ) √ δ 2 +4 x ( √ δ 2 +4 x + δ ) . The solution is g δ ( x ) = c √ √ δ 2 +4 x + δ , where c ∈ R is a constant. Let c = 1 . Plugging in our expressions for g δ ( ∥ β ∥ 2 ) , f δ ( ∥ β ∥ 2 ) , h δ ( ∥ β ∥ 2 ) , we get that

is a Hessian map for some unknown potential Φ δ ( β ) .

Solving for the potential Φ δ ( β ) . Take the ansatz that there exists some function scalar q ( x ) such that Φ δ ( β ) = q δ ( ∥ β ∥ ) + c δ where c δ is a constant such that Φ δ ( β ) > 0 for all β = 0 and Φ δ (0) = 0 . The Hessian of this ansatz takes the form,

̸

Equating terms from our expression for g δ ( ∥ β ∥ 2 ) M -1 (equation 60) we get the expression for ∇ q ( ∥ β ∥ )

which plugged into the second term gives the expression for ∇ 2 q ( ∥ β ∥ ) ,

We now look for a function q ( x ) such that both these conditions (Eq. (62) and Eq. (63)) are true. Consider the following function and its derivatives,

̸

Letting x = ∥ β ∥ notice ∇ q ( ∥ β ∥ ) and ∇ 2 q ( ∥ β ∥ ) satisfies the previous conditions. Furthermore, ∇ 2 q ( x ) > 0 for all δ as long as x = 0 and thus q ( x ) is a convex function which achieves its minimum at x = 0 . Thus, the constant c δ = -q (0) is

and the potential Φ δ ( β ) is

Finally, putting it all together, we can express the inductive bias as in Theorem 3.1.

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. [9], which we rewrite:

̸

Theorem A.2 (Theorem 2 from Azulay et al. [9]) . For a depth 2 fully connected network with a single hidden neuron ( h = 1 ), any δ ≥ 0 , and initialization β 0 such that β 0 = 0 , if the gradient flow solution β ( ∞ ) satisfies Xβ ( ∞ ) = y , then,

The most striking difference is in the expressions for the inductive bias. Azulay et al. [9] take an alternative route towards deriving the inductive bias by inverting M in terms of the original parameters a and w and then simplifying M -1 in terms of β , 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 √

̸

Another important distinction between our two theorems lies in the assumptions we make. Azulay et al. [9] consider only initializations such that δ ≥ 0 and β 0 = 0 . We make a less restrictive assumption by considering initializations β 0 such that ∥ β ( t ) ∥ > 0 for all t ≥ 0 , which allows for both positive and negative δ . Except for a 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 δ < 0 and where the sign of a 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.1 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 a similar statement likely holds for the general setting. If this were the case, we 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 us to extend the time-warped mirror flow analysis to all initializations such that ∥ β 0 ∥ > 0 .

A.4.4 Exact solution when interpolating manifold is one-dimensional

When the null space of X ⊺ X is one-dimensional, the constrained optimization problems in Theorem 3.1 and Theorem A.2 have an exact analytic solution. In this case we can parameterize all interpolating solutions β with a single scalar α ∈ R such that β = β ∗ + αv where X ⊺ Xv = 0 and ∥ v ∥ = 1 . Using this description of β , we can then differentiate the inductive bias with respect to α , set to zero, and solve for α . We will use the following expressions,

We will also use the expression, ∥ β ∥ 2 = ∥ β ∗ ∥ 2 + α 2 . Pulling these expressions together we get the following equation for α ,

If we let k = -2 z ⊺ v 3 , the solution for α is

This solution always works for the initializations we considered in Theorem 3.1. Interestingly, it appears that β = β ∗ -αv also works for initializations not previously considered. This includes trajectories that pass through the origin, resulting in a change in the sign of a .

B Wide and Deep Linear Networks

Here we discuss how our analysis techniques, developed in the previous section for a 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 a two-layer linear network with h hidden neurons and c outputs, f ( x ; θ ) = A ⊺ Wx , where W ∈ R h × d and A ∈ R h × c . We assume that h ≥ min( d, c ) , such that this parameterization can represent all linear maps from R d → R c . As in the single-neuron setting, the rescaling symmetry in this model between the first and second layer implies the h × h matrix ∆ = A 0 A ⊺ 0 -W 0 W ⊺ 0 determined at initialization remains conserved throughout gradient flow [62]. This can be easily shown from the temporal dynamics of A and W ,

Extending derivations in [30], the NTK matrix can be expressed as

where ⊗ and ⊕ denote the Kronecker product and sum respectively. The Kronecker sum is defined for square matrices C ∈ R c × c and D ∈ R d × d as C ⊕ D = C ⊗ I d + I c ⊗ D .

B.1.1 Parameter space dynamics

Inspired by our analysis of the single-neuron setting, we introduce two coordinate transformations to study the parameter space dynamics of a wide two-layer linear network. In both analyses we assume whitened input X ⊺ X = I d and let η a = η w = 1 . However, we will find that the analysis of the dynamics in function space, for general unwhitened data, is more tractable.

Parameter dynamics when c = 1 . Drawing insights from our analysis of the single-neuron scenario ( h = c = 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 k ∈ [ h ] :

We will also consider a new matrix quantity Q ∈ R h × h with elements Q kk ′ = w ⊺ k w k ′ ∥ w k ∥∥ w k ′ ∥ . The resulting dynamics for µ and ϕ can be entirely written in terms µ, ϕ, ∆ :

where M = 2 ( √ Diag(∆) 2 +4Diag( µ ) 2 -Diag(∆) ) -1 . Using the conserved structure of ∆ we can express Q as a function of µ and M ,

This approach yields a coupled nonlinear dynamical system with 2 h variables. Imposing additional assumptions on the initialization, such as permutation invariance between hidden neurons, can simplify the system of differential equations. A similar approach was used by Saad and Solla [78] to derive a set of differential equations for a soft committee machine model, capturing its online learning dynamics in a teacher-student setup, which Goldt et al. [79] extended to its generalization error dynamics.

Parameter dynamics when c = h . In this analysis we assume an initialization such that the conserved quantities ∆ = δ I h , an assumption we will discuss further in Appendix B.1.6, and that A is invertible throughout training. Let β ∗ = X ⊺ Y , which for whitened input, is the unique minimum of the dynamics in function space. We consider the variable ν = A -1 Wβ ∗ ∈ R c × c . Using the identity

that ˙ A -1 = -A -1 ˙ AA -1 and our assumption on ∆ , we find that the matrix ν evolves according to the matrix Riccati ODE,

Additionally, consider the variable C = A ⊺ A , which evolves according to the matrix Bernoulli ODE

Taken together we have found a change of variables, analogous to the one introduced in Appendix A.2.3 for the single-neuron setting, that evolves according to a matrix Riccati and Bernoulli equation,

However, solving this system exactly as we did in the single-neuron setting is challenging. Unless we assume that ν and β ⊺ ∗ β ∗ share the same eigenspace - allowing us to decouple the dynamics of ν into a set of scalar Riccati equations - the system cannot be easily solved. Instead, we will find that the dynamics of the product W ⊺ A in function space is more tractable and requires fewer assumptions.

B.1.2 Function space dynamics

We consider the dynamics of β = W ⊺ A ∈ R d × c in function space, which is governed by the ODE,

Vectorizing using the identity vec( ABC ) = ( C ⊺ ⊗ A )vec( B ) equation 85 becomes

As in the single-neuron setting, we find that the dynamics of β can be expressed as gradient flow preconditioned by a matrix M that depends on quadratics of A and W .

B.1.3 Proving Theorem 4.1

We first prove Theorem 4.1. Consider a single hidden neuron k ∈ [ h ] of the multi-output model defined by the parameters w k ∈ R d and a k ∈ R c . Let β k = w k a ⊺ k be the R d × c matrix representing the contribution of this hidden neuron to the input-output map of the network β = ∑ h k =1 β k . Consider the two gram matrices β ⊺ k β k ∈ R c × c and β k β ⊺ k ∈ R d × d ,

Notice that we can express ∥ β k ∥ 2 F as

At each hidden neuron we have the conserved quantity 8 η w ∥ a k ∥ 2 -η a ∥ w k ∥ 2 = δ k where δ k ∈ R . Using this quantity we can invert the expression for ∥ β k ∥ 2 F to get

When ∥ β k ∥ 2 F > 0 , we can use these expressions to solve for the outer products a k a ⊺ k and w k w ⊺ k in terms of β k and δ k ,

By substituting these expressions into the decompositions A ⊺ A = ∑ h k =1 a k a ⊺ k and W ⊺ W = ∑ h k =1 w k w ⊺ k , we derive the representation for M presented in Theorem 4.1: M = ∑ h k =1 M k where

B.1.4 Understanding M when there is a single-neuron h = 1

When there is a single-hidden neuron h = min( d, c ) = 1 , the expression for M presented in Theorem 4.1 simplifies allowing us to precisely understand the influence of δ on the learning regime. When h = c = 1 , then β ⊺ β ∥ β ∥ 2 F = 1 . Therefore, Eq. (7) simplifies to

and we recover Eq. (4) presented in Section 3. When h = d = 1 , then ββ ⊺ ∥ β ∥ 2 F = 1 and thus Eq. (7) simplifies to,

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 δ . As δ → ∞ , then the first expression for M becomes proportional to I d , while the second expression for M becomes proportional to the rank-1 projection β ⊺ β ∥ β ∥ 2 . Conversely, as δ → -∞ , then the first expression for M becomes proportional to the rank-1 projection ββ ⊺ ∥ β ∥ 2 , while the second expression for M becomes proportional to I c . When h = d = c = 1 , then M = √ δ 2 + η a η w 4 ∥ β ∥ 2 and thus in both limits of δ →±∞ , M becomes a constant independent of β . In all settings, when δ = 0 , M depends on β . 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 .

B.1.5 Interpreting M in different limits and architectures

We now seek to more generally understand the influence of the conserved quantities δ i and the relative sizes of dimensions d , h and c on the learning regime. For a matrix A ∈ R d × c , let Row( A ) ⊆ R c and Col( A ) ⊆ R d denote the row and column space of A respectively.

Theorem B.1. The dynamics are in the lazy regime, for all t ≥ 0 , if δ k →∞ for all k ∈ [ h ] and there exists a least squares solution β ∗ ∈ R d × c such that

or δ k →-∞ for all k ∈ [ h ] and there exists a solution such that

Proof. As δ k →∞ , M k →| δ k | β ⊺ k β k ∥ β k ∥ 2 F ⊗ I d , implying ˙ β k = -| δ k | ∂ L ∂β ( β ⊺ k β k ∥ β k ∥ 2 F ) . Notice that ( β ⊺ k β k ∥ β k ∥ 2 F ) is the unique orthogonal projection matrix onto the one-dimensional row space of β k . Thus, the dynamics of each β k follow a projected gradient descent in their row space. As a result, M k will not change and thus the NTK will be static. By assumption there exists a least squares solution β ∗ such that the rows of β ∗ are in the span of the rows of β k . Thus, a solution will be reached as t → ∞ , while the M k remain static.

Note that the assumptions in Theorem B.1 can be more intuitively expressed in terms of the parameter space ( W,A ) . Except in highly degenerate cases, the assumption Row( β ∗ ) ⊆ Span ( ⋃ h k =1 Row( β k (0)) ) is equivalent to the existence of a β ∗ whose rows lie in the span of { a k (0) } h k =1 , or, equivalently, to the existence of a matrix W such that β ∗ = W ⊺ A (0) . Similarly, the condition Col( β ∗ ) ⊆ Span ( ⋃ h k =1 Col ( β k (0)) ) is in most cases equivalent to the existence of a matrix A such that β ∗ = W (0) ⊺ A .

Adirect consequence of Theorem B.1 is that networks which narrow from input to output ( d > c ) must enter the lazy regime with probability 1 as all δ k →∞ whenever h ≥ c and assuming independent initializations for all β k . In this case, the rows of { β 1 , . . . , β h } span all of R c 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 as all δ k →-∞ whenever h ≥ d and assuming independent initializations for all β k . Additionally, when h ≥ max( d, c ) and assuming independent initializations for all β k , then all networks enter the lazy regime as either all δ k →∞ or all δ k →-∞ .

Another interesting implication of Theorem B.1, is that if there does not exist a least squares solution β ∗ with rows in the span of the rows of { β 1 , . . . , β h } , then the network will enter a delayed rich regime as all δ k →∞ , where the magnitude of the δ k will determine the delay. In this setting, the network is initially lazy, attempting to fit the solution within the row space of the β k , 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 β ∗ is true as all δ k →-∞ .

B.1.6 Simplifying M through assumptions on ∆

We now consider how introducing structures on ∆ can lead to simpler expressions for M . A natural assumption to consider is the following:

Assumption B.2 (Isotropic initialization) . Let A ∈ R h × c and W ∈ R h × d be initialized such that ∆ = η w A (0) A (0) ⊺ -η a W (0) W (0) ⊺ = δ I h .

In square networks , where the dimensions of the input, hidden, and output layers coincide ( d = h = c ), and the weights are initialized as A ∼ N (0 , σ 2 a /c ) and W ∼ N (0 , σ 2 w /d ) , this assumption is naturally satisfied with δ = σ 2 a -σ 2 w as the dimension h →∞ . However, a limitation of this assumption is that for general δ it requires h ≤ min( d, c ) . Specifically, when δ > 0 , the isotropic initialization requires that A (0) A (0) ⊺ ≻ 0 , which implies h ≤ c . Similarly, when δ < 0 , the isotropic initialization requires that W (0) W (0) ⊺ ≻ 0 , which implies h ≤ d . Now we prove two important implications of the isotropic initialization assumption.

Lemma B.3. Let ∆ = δ I h . If either δ ≥ 0 or δ < 0 and h ≥ d , we have that

Proof. The quantity η w AA ⊺ -η a WW ⊺ = δ I h is conserved in gradient flow. Multiplying on the left by W ⊺ and on the right by W we have that

Completing the square by adding δ 2 4 η a I d to both sides and dividing by η a we get the equality,

For δ ≥ 0 , W ⊺ W + δ 2 η a I d ⪰ 0 . For δ < 0 , then we know from the conserved quantity that WW ⊺ + δ 2 η a I h = η w η a AA ⊺ -δ 2 η a I h ≻ 0 , which implies when h ≥ d that W ⊺ W + δ 2 η a I d ≻ 0 . As a result, we can take the principal square root of each side,

which rearranged gives the final result.

Lemma B.4. Let ∆ = δ I h . If either δ ≤ 0 or δ > 0 and h ≥ c , we have that

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

Plugging in expressions for W ⊺ W from Lemma B.3 and A ⊺ A from Lemma B.4 we can directly write,

From this expression for M ( β ) we can easily consider how it simplifies in limiting settings of δ :

As δ →±∞ , M → δ I dc , and the dynamics are lazy. In this limit, the dynamics of β converge to the trajectory of linear regression trained by gradient flow and along this trajectory the NTK matrix remains constant. When δ = 0 , M = √ η a η w β ⊺ β ⊗ I d + I c ⊗ √ η a η w ββ ⊺ , and the dynamics are rich. Here the NTK changes in both magnitude and direction through training. In the next section we will attempt to better understand these dynamics for intermediate values of δ through the lens of a mirror flow.

B.1.7 Deriving a mirror flow for the singular values of β

For a matrix β , the dynamics of one of its singular values are given by ˙ σ = u ⊺ ˙ βv , where u and v are the corresponding left and right singular vectors. This equality can be derived from chain rule and the fact that ∥ u ∥ = ∥ v ∥ = 1 :

In the last equality we used that fact that for any vector z with a fixed norm, ˙ ∥ z ∥ 2 = 2˙ z ⊺ z = 0 . Letting diag : R d × c → R min( d,c ) be the operator that, given a rectangular matrix, returns a vector of the elements on the main diagonal, we can then write,

where λ ∈ R min( d,c ) is the vector of singular values of β . In the following lemma, we use the shared singular vector structure between β and A and W to rewrite these dynamics as

where M is a diagonal matrix and ∇ λ L is the gradient of the loss with respect to the singular values of β . Without loss of generality we consider η a = η w = 1 .

Lemma B.5. Let ∆ = δ I h . We then have that ˙ λ = -M ∇ λ L , where M ∈ R min( d,c ) × min( d,c ) is a diagonal matrix with

Proof. First note that

where we let W ⊺ W = U Σ 2 W U ⊺ and A ⊺ A = V Σ 2 A V ⊺ , using the fact that, under ∆ = I h , the eigenvectors of A ⊺ A are the right singular vectors of β and the eigenvectors of W ⊺ W are the left singular vectors of β . This expression rewrites as

where M ∈ R min( d,c ) × min( d,c ) is a diagonal matrix with M ii = (Σ 2 A ) ii + (Σ 2 W ) ii . For i ≤ min( d, h, c ) , one can show that M ii = √ δ 2 +4 λ 2 i . This is because for i ≤ min( d, h, c ) , (Σ 2 A ) ii = (Σ 2 W ) ii + δ from the conservation law and (Σ 2 W ) ii (Σ 2 A ) ii = λ 2 i from the definition of λ . Together this implies (Σ 2 W ) ii ( δ +(Σ 2 W ) ii ) = λ 2 i , which is a quadratic equation in (Σ 2 W ) ii . If h < min( d, c ) then M ii = 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

We then have that ∇ λ L = diag ( U ⊺ X ⊺ ( Xβ -Y ) V ) , which, combined with our expression for M , completes the proof.

Leveraging Lemma B.5, we can show that the singular values of β evolve under a mirror flow in the following theorem.

̸

Theorem B.6. Let ∆ = δ I h and assume h ≥ min( d, c ) and δ = 0 . We then have that the dynamics of λ , the singular values of β , are given by the mirror flow

where Φ δ ( λ ) = ∑ min( d,c ) i =1 q δ ( λ i ) and q δ is the hyperbolic entropy potential

̸

Proof. When ∆ = δ I h , then by Lemma B.5 the dynamics of the singular values of β can be expressed as ˙ λ = -M ∇ λ L . Furthermore, when h ≥ min( d, c ) and δ = 0 , we have that M = √ δ 2 +4 λ 2 I min( d,c ) , where λ 2 is element-wise, which is always invertible. Observe, this expression for M is the inverse Hessian of the potential Φ δ ( λ ) = ∑ i q δ ( λ i ) for q δ specified in the theorem statement. Thus, the dynamics for the singular values are the mirror flow ˙ λ = -( ∇ 2 Φ δ ( λ ) ) -1 ∇ λ L .

Theorem B.6 implies that the dynamics for the singular values of β can be described as a mirror flow with a δ -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 ℓ 1 and ℓ 2 penalty on the singular values for the rich ( δ → 0 ) and lazy ( δ →±∞ ) 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 - specifically, whitened input data ( X ⊺ X = I d ) and a task-aligned initialization such that the singular vectors of β 0 are aligned with those of β ∗ -wecan ensure that the singular vectors remain constant and thus derive an inductive bias on the singular values. However, in this setting the dynamics decouple completely, implying there is no difference between applying an ℓ 1 or ℓ 2 penalty on the singular values. Consequently, even though the dynamics will depend on δ , the final interpolating solution will be independent of δ , making a statement on the inductive bias insignificant.

B.2 Deep linear networks

We now consider the influence of depth by studying a depth-( L + 1) linear network, f ( x ; θ ) = a ⊺ ∏ L l =1 W l x , where W 1 ∈ R h × d , W l ∈ R h × h for 1 < l ≤ L , and a ∈ R h . We assume that the dimensions d = h and that all parameters share the same learning rate η = 1 . For this model the predictor coefficients are computed by the product β = ∏ L l =1 W ⊺ l a ∈ R d . Similar to our analysis of a two-layer setting, we assume an isotropic initializations of the parameters.

Definition B.7. There exists a δ ∈ R such that aa ⊺ -W L W ⊺ L = δ I h and for all l ∈ [ L -1] W ⊺ l +1 W l +1 = W l W ⊺ l .

This assumption can easily be achieved by setting a = 0 and W l = αO l for all l ∈ [ L ] , where O l ∈ R d × d is an random orthogonal matrix and α ≥ 0 . In this case δ = -α 2 . Further, notice this parameterization is naturally achieved in the high-dimensional limit as d → ∞ under a 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 β , its squared norm ∥ β ∥ 2 and its outer product ββ ⊺ , can always be expressed as polynomials of ∥ a ∥ 2 and W ⊺ 1 W 1 respectively.

LemmaB.8. For a depth-( L +1) linear network with square width ( d = h ) and isotropic initialization, then for all t ≥ 0 ,

Proof. The norm of the regression coefficients is the product ∥ β ∥ 2 = a ⊺ ( ∏ L l =1 W l )( ∏ L l =1 W l ) ⊺ a . Using the conservation of the initial conditions between consecutive weight matrices, W ⊺ l +1 W l +1 = W l W ⊺ l , we can express this telescoped product as ∥ β ∥ 2 = a ⊺ ( W L W ⊺ L ) d a . When plugging in the conservation between last two layers, this implies ∥ β ∥ 2 = a ⊺ ( aa ⊺ -δ I h ) d a , which expanded gives the desired result.

The outer product of the regression coefficients is ββ ⊺ = ( ∏ L l =1 W l ) ⊺ aa ⊺ ( ∏ L l =1 W l ) . Using the conserved initial conditions of the last weights we can factor the outer product as the sum, ββ ⊺ = ( ∏ L l =1 W l ) ⊺ W L W ⊺ L ( ∏ L l =1 W l ) + δ ( ∏ L l =1 W l ) ⊺ ( ∏ L l =1 W l ) . 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 W ⊺ 1 W 1 significantly influence the dynamics of β , similar to our analysis in the two-layer setting.

Lemma B.9. The dynamics of β are given by a differential equation ˙ β = -MX ⊺ ρ where M is a positive semi-definite matrix that solely depends on ∥ a ∥ 2 , W ⊺ 1 W 1 , and δ ,

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 δ → 0 allowing us to be precise about the inductive bias in this setting.

Theorem B.10. For a depth-( L + 1) linear network with square width ( d = h ) and isotropic initialization β 0 such that ∥ β ( t ) ∥ > 0 for all t ≥ 0 , then in the limit as δ → 0 , if the gradient flow solution β ( ∞ ) satisfies Xβ ( ∞ ) = y , then,

Proof. Whenever ∥ β ∥ > 0 and in the limit as δ → 0 , then we can find a unique expression for ∥ a ∥ 2 and W ⊺ 1 W 1 in terms of ∥ β ∥ 2 and ββ ⊺ ,

Plugged into the previous expression for M results in a positive definite rank-one perturbation to the identity,

Using the Sherman-Morrison formula we find that M -1 is

We can now apply a time-warped mirror flow analysis similar to the analysis presented in Appendix A.4. Consider the time-warping function g δ ( ∥ β ∥ ) = ∥ β ∥ -L L +1 and the potential Φ( β ) = ( L +1 L +2 ) ∥ β ∥ L +2 L +1 , then its not hard to show M -1 = g δ ( ∥ β ∥ ) ∇ 2 Φ( β ) . This gives the desired result.

This theorem is a generalization of Proposition 1 derived in [9] for two-layer linear networks in the rich limit to deep linear networks in the rich limit. We find that the inductive bias, Q ( β ) = ( L +1 L +2 ) ∥ β ∥ L +2 L +1 -∥ β 0 ∥ -L L +1 β ⊺ 0 β , strikes a depth-dependent balance between attaining the minimum norm solution and preserving the initialization direction.

C Piecewise Linear Networks

Here, we elaborate on the theoretical results presented in Section 5. Our goal is 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 a 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.

C.1 Surface of a piecewise linear network

The input-output map of a piecewise linear network f ( x ; θ ) , with l hidden layers and h hidden neurons per layer, is comprised of potentially O ( h dl ) connected linear regions, each with their own vector of predictor coefficients [65]. The exploration of this complex surface has been the focus of numerous prior works, the vast majority of them focused on counting and bounding the number of linear regions as a function of the width and depth [81, 82, 83, 84, 65, 85, 86, 87]. The central object in all of these studies is the activation region ,

Definition C.1. For a piecewise linear network f ( x ; θ ) , comprising N hidden neurons with preactivation z i ( x ; θ ) for i ∈ [ N ] , let the activation pattern A represent an assignment of signs a i ∈ {-1 , 1 } to each hidden neuron. The activation region R ( A ; θ ) is the subset of input space that generates A ,

The input-output map is linear within each non-empty activation region and continuous at the boundary be-

(a) With Biases

tween regions. Linearity implies that every non-empty 9 activation region is associated with a linear predictor vector β R ∈ R d such that for all x ∈ R ( A ; θ ) , β R = ∇ x f ( x ; θ ) . Continuity implies that the boundary between regions is formed by a hyperplane determined by where the pre-activation for a neuron is exactly zero, { x : z i ( x ; θ ) = 0 } . When the neighboring regions have different linear predictors 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 a 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 f ( x ; θ ) 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 f ( x ; θ ) 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 , a connection previously alluded to in the literature [82].

C.2 Dynamics of a two-layer piecewise linear network

We consider the dynamics of a two-layer piecewise linear network without biases, f ( x ; θ ) = a ⊺ σ ( Wx ) , where W ∈ R h × d and a ∈ R h . The activation function is σ ( z ) = max( z, γz ) for γ ∈ [0 , 1) , which includes ReLU γ = 0 and Leaky ReLU γ ∈ (0 , 1) . We permit h > d , which in the limit as h →∞ , ensures the network possesses the functional expressivity to represent any continuous nonlinear function from R d to R passing through the origin. Following a similar strategy used in Section 4, we consider the contribution to the input-output map from a single hidden neuron k ∈ [ h ] with parameters w k ∈ R d and a k ∈ R . As in the linear setting, each hidden neuron is associated with a conserved quantity, δ k = η w a 2 k -η a ∥ w k ∥ 2 . Unlike in the linear setting, this neuron's contribution to the output f ( x i ; θ ) is regulated by whether the input x i is in the neuron's active halfspace , { x ∈ R d : w ⊺ k x > 0 } . Let C ∈ R h × n be the matrix with elements c ki = σ ′ ( w ⊺ k x i ) , which determines the activation of the k th neuron for the i th training data point. The subgradient σ ′ ( z ) = 1 if z > 0 , σ ′ ( z ) ∈ [ γ, 1] if z = 0 , and σ ′ ( z ) = γ if z < 0 . These activation functions exhibit positive homogeneity, implying σ ( z ) = σ ′ ( z ) z . Thus, we can express σ ( w ⊺ k x i ) = c ki w ⊺ k x i , allowing us to express the gradient flow dynamics for w k and a k as

where ρ i = f ( x i ; θ ) -y i is the residual associated with the i th training data point. If we let β k = a k w k , which determines the contribution of each hidden neuron to the output f ( x i ; θ ) , then its not hard to see that the gradient flow dynamics of β k are

̸

As in the linear setting, the matrix M k ∈ R d × d appears as a preconditioning matrix on the dynamics Using the exact same derivation presented in Appendix A.3, whenever a 2 k = 0 , we can express M k entirely in terms of β k and δ k ,

However, unlike in the linear setting, the vector ξ k ∈ R d driving the dynamics is not shared for all neurons because of its dependence on c ki . Additionally, the NTK matrix in this setting depends on M k and C , with elements K ij = ∑ h k =1 c ki x ⊺ i ( η w a 2 k I d + η a w k w ⊺ k ) x j c kj . Thus, in order to assess the temporal dynamics of the NTK matrix, we must understand the dynamics of M k and C . We consider a signed spherical coordinate transformation separating the dynamics of β k into its directional ˆ β k = sgn( a k ) β k ∥ β k ∥ and radial µ k = sgn( a k ) ∥ β k ∥ components, such that β k = µ k ˆ β k . Here, ˆ β k determines the orientation and direction of the halfspace where the k th neuron is active, while µ k determines the slope of the linear region in this halfspace. These coordinates evolve according to,

These equations can be derived directly from Eq. (128) through chain rule similar to Appendix A.2.1. 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, ˆ β k = w k ∥ w k ∥ and µ k = a k ∥ w k ∥ .

D Experimental Details

We used Google Cloud Platform (GCP) nodes to run all experiments. Figure 1 experiments were run on a 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 5 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 4 are not compute-heavy, and these experiments were run on a personal computer. Overall, we estimate approximately 200 hours of single A100 runtime as 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. 1 we consider a student-teacher setup similar to that in [8], with one-hidden layer ReLU networks of the form f ( x ; θ ) = ∑ m i =1 a i σ ( w ⊺ i x ) , where f : R d → R and σ is the ReLU activation function. The teacher model, f teacher , has m = k hidden neurons initialized as w teacher i i.i.d. ∼ Unif( S d -1 ) and a i i.i.d. ∼ Unif( {± 1 } ) for i ≤ k . The student, f student , in turn, has h hidden neurons. We use a symmetrized initialization, as considered in [8], where for i ≤ h/ 2 , we sample w i i.i.d. ∼ S d -1 and a i i.i.d. ∼ Unif( {± 1 } ) , and then for i ≥ h 2 + 1 we symmetrize by setting w i = w i -h/ 2 and a i = -a i -h/ 2 . This ensures that f student predicts 0 on any input at initialization.

Note that the base student initialization described thus far is perfectly balanced at each neuron, that is δ i = 0 for i ∈ [ m ] ; we also define this to be our setting where the scale τ is 1. In order to transform the base initialization into a particular setting of τ and δ , we first solve for the relative layer scaling α in δ 2 = τ 2 ( α 2 -α -2 ) and then scale each w i by τ/α and each a i by τα . We obtain a training dataset { x ( i ) , y ( i ) } n i =1 by sampling x ( i ) i.i.d. ∼ S d -1 and computing noiseless labels as y ( i ) = f teacher ( x ( i ) ; θ teacher ) . 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 , k = 3 , and n = 20 . We sample a single teacher and then train four students with the same base initialization but different configurations of τ and δ : ( τ = 0 . 1 , δ = 0) and ( τ = 2 , δ = 0) for the left subfigure, and ( τ = 0 . 1 , δ = 1) and ( τ = 0 . 1 , δ = -1) for the right subfigure. Training is for 1 million steps at a learning rate of 1e-4.

Figure 1 (b).

Here the setting is: d = 100 , m = 50 , k = 3 , and n = 1000 , as in Fig. 1c of [8]. Training is performed with learning rate of 5e-3 /τ 2 . Test error is computed as mean square error over a held-out set of 10,000 datapoints. We sweep over τ over a logarithmic scale in the range [0 . 1 , 2] and δ over a linear scale in the range [ -1 , 1] . We average over 16 random seeds, where the seed controls the sampling of: the teacher weights θ teacher , the base initialization of θ student , and the training data { x ( i ) } n i =1 . In this way, each random seed is used for a sweep over all combinations of τ and δ in the sweep; we simply apply the scaling described above to get to each point on the ( τ, δ ) grid. The kernel distance computed is as defined in [27], where here we compute it at time t relative to the kernel at initialization, i.e. S ( t ) = 1 -⟨ K 0 , K t ⟩ / ( ∥ K 0 ∥ F ∥ K t ∥ F ) . In Fig. 10, we additionally plot Hamming and parameter distances relative to initialization, as well as training loss, while training for ten times longer than in Fig. 1 (b).

Notebooks generating all two-layer experiment figures are provided here.

0.8

0.4

0.0

0.4

0.8

t

=100

t

=1000

t

=10000

t

=100000

t

=1000000

0.8

0.4

0.0

0.4

0.8

0.8

0.4

0.0

0.4

0.8

0.8

0.4

0.0

0.4

0.8

10

1

10 0

0.8

0.4

0.0

0.4

0.8

10

1

10 0

10

1

10 0

10

1

10 0

10

1

10 0

0.0

0.2

0.4

0.6

Kernel Distance

0

3

6

9

12

15

18

21

Hamming Distance

0.0

0.4

0.8

1.2

1.6

2.0

2.4

2.8

3.2

Parameter Distance

10

10

10

7

10

4

Train Loss

0.0

0.4

0.8

1.2

Test Loss

10

2

D.2 Figures 2, 3, 4: Single-Neuron Linear Network

Figures 2, 3, and 4 were generated by simulating gradient flow using scipy.integrate.solve_ivp function with the RK45 method for solving the ODEs, with a relative tolerance of 1 × 10 -6 and time span of (0 , 20) . In the experiments with full-rank data, we used X ⊺ X = I 2 , β ∗ = [ 0 1 ] , and β 0 = [ -1 0 ] . For the experiment with low-rank data, we used X ⊺ X = [ 0 . 25 0 . 5 0 . 5 1 ] , β ∗ = [ 0 . 44 0 . 88 ] , and β 0 = [ 0 . 4 0 . 05 ] . See the discussion in Appendix A.2 for details on how we determined our theoretical predictions. A notebook generating all the figures is provided here.

D.3 Figure 5:

Kernel Distance

We trained LeNet-5 [88] (with ReLU nonlinearity and Max Pooling) on MNIST [88]. We use He initialization [89] and divide the first layer weights by α and multiply the last layer weights by α at initialization, which keeps the network functionally the same at initialization. We trained the model

for 500 epochs with a learning rate of 1e-4 and a batch size of 512. The parameter distance is defined as the L 2 distance between all the parameters. To quantify the distance between the activations, we binarize the hidden activation with 1 representing an 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 ( t 1 , t 2 ) = 1 -⟨ K t 1 , K t 2 ⟩ / ( ∥ K t 1 ∥ F ∥ K t 2 ∥ F ) , which is a scale 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.

Gabor Filters

We are training a small ResNet based on the CIFAR10 script provided in the DAWN benchmark (code available here). The only modifications to the provided code base are we increase the convolution kernel size from 3 × 3 to 15 × 15 , to better observe the learned spatial patterns, and we set the weight decay parameter to 0 to avoid confounding variables. Moreover, we are dividing the convolutional filters weights by a parameter α (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 w ij ∈ R 15 × 15 , over input i = (1 , 2 , 3) and output j = (1 , ..., 64) channels

where the Laplacian kernel is defined as

0

10

20

Epochs

0.30

0.45

0.60

0.75

0.90

Train Accuracy

=10.0

=3.3

=1.0

=0.3

=0.1

0

10

20

Epochs

0.30

0.45

0.60

0.75

0.90

Validation Accuracy

=10.0

=3.3

=1.0

=0.3

=0.1

(a) Accuracy

0.10

0.33

1.00

3.33

10.00

4.5

3.0

1.5

0.0

Smoothness

(b) Gabor Smoothness

Random Hierarchy Model

We refer to [67], who originally proposed the random hierarchy model (RHM) as a tool for studying how deep networks learn compositional data, for a more in-depth treatment. Here we briefly recap the setup following the notation used in [67].

An RHM essentially lets us build a random classification task with a clear hierarchical structure. The top level of the RHM specifies m equivalent high-level features for each class label in { 1 , . . . , n c } , where each feature has length s and n c is the number of classes. For example, suppose the vocabulary at the top level is V L = { a, b, c } , n c = 2 , m = 3 , and s = 2 . Then in a particular instantiation of this RHM, we might have that Class 1 has ab , aa , and ca as equivalent high-level features (this is precisely the example used in Fig.1 of [67]). Class 2 will then have three random high-level features, with the constraint that they are not features for Class 1, for example, bb, bc, ac .

Each successive level specifies m equivalent lower-level features for each 'token" in the vocabulary at the previous level. For example, if V L -1 = { d, e, f } , we might have that a can be equivalently

represented as de , df , or ff ; b and c will each have m equivalent representations of their own. We assume that the vocabulary size, v , is the same at all levels. Therefore, sampling an RHM with hyperparameters n c , m, s, v requires sampling mn c +( L -1) mv rules.

In order to sample a datapoint from an RHM, we first sample a class label (e.g. Class 1), then uniformly sample one of the highest level features, (e.g. ab ), then for each 'token" in this feature we sample lower level features (e.g. a → de , b → ee ), and so on recursively. The generated sample will therefore have length s L and a class label. For training a neural network to perform this classification task, each input is converted into a one-hot representation, which will be of shape ( s L , v ) , and is then flattened.

We use the code released by [67] to train an MLP of width 64 with three hidden layers to learn an RHMwith L = 3 , n c = 8 , m = 4 , s = 2 , v = 8 . The main change we make is allowing for scaling the initialization of the first layer by 1 /α and the initialization the readout layer by α . We then sweep over α ∈ { 0 . 03 , 0 . 1 , 0 . 3 , 1 , 3 , 10 } and over the number of datapoints in the training set, which is specified as a fraction of the total number of datapoints the RHM can generate. We average test accuracy, which is by default computed on a 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 a cosine schedule down to 80% of epochs. The batch size of 128; we do not use biases or weight decay.

Grokking

We are training a one layer transformer model on the modular arithmetic task in Power et al. [68]. Our experimental code is based on an existing Pytorch implementation (code available here). The only modifications to the provided code base is that we use a single transformer layer (instead of the default 2-layer model). Prior analysis in Nanda et al. [72] has shown that this model can learn a minimal (attention-based) circuit that solves the task.

We study the effects on grokking time (defined as ≥ 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 α as in the CNN gabor experiments. Secondly, like in Kumar et al. [69], we multiply the output of the model (i.e., the logits) by a factor τ and divide the learning rate by τ 2 .

0.6

0.8

1.0

1.2

1.4

10

1

10 0

10 1

3.0e+01

4.5e+01

6.0e+01

7.5e+01

Steps until Trained

(a) Steps until training

0.6

0.8

1.0

1.2

1.4

10

1

10 0

10 1

0.0e+00

6.0e+03

1.2e+04

1.8e+04

2.4e+04

Steps until Grokked

(b) Steps until grokked

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]

Justification: The abstract clearly lays out the four main contributions sections of our paper (1) a 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 of 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 1) reflect Section 3, Section 4, and Section 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in 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.

It is fine to include aspirational goals as motivation as long as it 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 out 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 important 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]

Justification: For all main theoretical lemmas and theorems we either clearly state assumptions in the theorem statement or we introduce the assumptions earlier 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 a 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 crossreferenced.

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 to provide a 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We disclose all details needed to reproduce experiments in the paper, particularly 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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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., 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Some of our experimental results use existing open-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 in the appendix.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) 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 run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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 in supplemental material (appended to the 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, hyperparameters, 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 to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes] .

Justification: For our main empirical results, shown in Figure 5, we make sure to run each experiment with multiple seeds and we 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 not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It 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 one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 of compute required for each of the individual experimental runs as well as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We make sure to conform with the NeurIPS Code of Ethics. Items in the code particularly relevant to our work pertain to disclosing details necessary for reproducibility. Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 is 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 a direct path to 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 or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: 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)?

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 no 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 a best faith effort.

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 are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We make sure to 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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 not 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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.

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 an equivalent approval/review based on the requirements of your country or 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 crowdsourcing 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nOverall 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee weaknesses\n\nLIMITATIONS:\nAddressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the learning dynamics of a deep network by leveraging the conserved quantities due to symmetries\n\nSTRENGTHS:\nThe 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\n\nWEAKNESSES:\nI 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\n\nQUESTIONS:\nSee the weakness section\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}