Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning

Authors: daniel  kunin
allan  raventós
clémentine  dominé
feng  chen
david  klindt
andrew  saxe
surya  ganguli
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 04:14:20

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.

Main paper: 
 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 Kernel Distance (b) A complex phase portrait of feature learning Figure 1 : Unbalanced initializations lead to rapid rich learning and generalization. We follow the experimental setup used in Fig. 1 of Chizat et al. [8] -a wide two-layer student ReLU network f (x; θ) = h i=1 a i max(0, w ⊺ i x) trained on a dataset generated from a narrow two-layer teacher ReLU network. The student parameters are initialized as w i ∼ Unif(S d-1 ( τ α )) and a i = ±ατ , such that τ > 0 controls the overall scale of the function, while α > 0 controls the relative scale of the first and second layers through the conserved quantity δ = τ 2 (α 2 -α -2 ). (a) Shows the training trajectories of |a i |w i (color denotes sgn(a i )) when d = 2 for four different settings of τ, δ. The left plot confirms that small overall scale leads to rich and large overall scale to lazy. The right plot shows that even at small overall scale, the relative scale can move the network between rich and lazy as well.Here an upstream initialization δ > 0 shows striking alignment to the teacher (dotted lines), while a downstream initialization δ < 0 shows no alignment. (b) Shows the test loss and kernel distance from initialization computed through training over a sweep of τ and δ when d = 100. Lazy learning happens when τ is large, rich learning happens when τ is small, and rapid rich learning happens when both τ is small and δ is large -an upstream initialization. This initialization also leads to the smallest test loss. See Fig. 10 in Appendix D.1 for supporting figures. 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 t1 , K t2 ⟩/ (∥K t1 ∥ F ∥K t2 ∥ 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. 
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. 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. 
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,K = X η w a 2 I d + η a ww ⊺ X ⊺ . (1)Except for a measure zero set of initializations which converge to saddle points 2 , all gradient flow trajectories will converge to a global minimum, determined 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 assume3 whitened input X ⊺ X = I d such that the ordinary least squares solution is β * = X ⊺ y, and the gradient flow dynamics simplify to ȧ = η a w ⊺ β * -a∥w∥ 2 , ẇ = η 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,μ = δ 2 + 4η a η w µ 2 (ϕ∥β * ∥ -µ) , φ = η a η w 2µ∥β * ∥ δ 2 + 4η a η w µ 2 -δ 1 -ϕ 2 .(2) 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 δ. Our theoretical predictions (black dashed lines) agree with gradient flow simulations (solid lines, color-coded based on δ values), shown here for three key metrics: µ (left), ϕ (middle), and S(0, t) (right). Each metric starts at the same value for all δ, but varying δ has a pronounced effect on the metric's dynamics. For upstream initializations (δ ≫ 0), µ changes only slightly, ϕ exponentially aligns, and S remains near zero, indicative of the lazy regime. For balanced initializations (δ = 0), both µ and ϕ change significantly and S quickly moves away from zero, indicative of the rich regime. For downstream initializations (δ ≪ 0), µ quickly drops to zero, then µ and ϕ slowly climb back to one. Similarly, S remains small before a sudden transition towards one, indicative of a delayed rich regime. See Appendix A.2 for further details.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 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 δ. 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.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,EQUATIONwhere ρ = 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 δ asEQUATIONwhere κ = δ 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, 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,EQUATIONwhereΨ δ (β) = 1 3 δ 2 + 4∥β∥ 2 -2δ δ 2 + 4∥β∥ 2 + δ and ψ δ = δ 2 + 4∥β 0 ∥ 2 -δ.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).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. 
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. We consider the dynamics of a two-layer linear network with h hidden neurons and c outputs, f (x; θ) = A ⊺ W x, 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 [62] . Drawing insights from our analysis of the single-neuron scenario (h = c = 1), we consider the dynamics ofh × h matrix ∆ = η w A 0 A ⊺ 0 -η a W 0 W ⊺ 0 is conserved throughout gradient flowEQUATIONwhere 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,EQUATIONBy 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 δ,EQUATIONwhere 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 isQ(β) = ( 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.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. 
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-colorable7 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 ⊺ σ(W x), 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 ofEQUATIONThe 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 elementsK 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,EQUATION)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 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 ,∂ 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 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. 
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. 
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. 
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). 
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). 
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). 
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. 
 
 Appendix: 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:EQUATION)The gradient flow dynamics are governed by the following coupled ODEs:EQUATIONEQUATION)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.The scaling symmetry between a and w results in a conserved quantity δ ∈ R throughout training, as noted in many prior works [10, 62, 74] , whereEQUATION)This can be easily verified by explicitly writing out the dynamics of δ. Define ρ = (X ⊺ Xaw -X ⊺ y) for succinct notation, such thatδ = 2η w a ȧ -2η a w ⊺ ẇ = 2η w a (-η a w ⊺ ρ) -2η a w ⊺ (-η w aρ) = 0.The conserved quantity confines the parameter dynamics to the surface of a hyperboloid where the magnitude and sign of the conserved quantity determines the geometry, as shown in Fig. 2 . A hyperboloid of the formk i=1 x 2 i - n i=k+1x 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.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 asEQUATIONAs 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 asEQUATIONEQUATION)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 δ,EQUATION)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,EQUATIONWe 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 ϕ.When δ = 0, the dynamics for µ, ϕ are,EQUATIONEQUATIONFirst, 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,EQUATIONwhere c ϕ = tanh -1 (ϕ 0 ). Plugging this solution into the dynamics for µ gives a Bernoulli differential equation,EQUATION)with the solution,EQUATIONwherec µ = 2µ -1 0 cosh 2 (c ϕ ) -(2c ϕ + sinh(2c ϕ )). Note, if ϕ 0 = -1, then φ = 0, and the dynamics of µ will be driven to 0, which is a saddle point.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,EQUATION)The solution is given by,EQUATIONwhereR = 1 2 δ 2 + 4∥β * ∥ 2 .The trajectory of a(t) is given by the Bernoulli equation,EQUATIONwhich can be solved analytically using ν(t). For a 0 > 0, we have thata(t) = 2e tδ/2 ∥β * ∥ √ δ sech 2 (Y (t)) 4e tδ ∥β * ∥ 2 - δ 2 + 4∥β * ∥ 2 ∥β * ∥ 2 δ -a 2 0 + b 2 0 b 2 0 -a 2 0 ∥β * ∥ 2 + a 0 b 0 δ -δe δt δ cosh (2Y (t)) -δ 2 + 4∥β * ∥ 2 sinh (2Y (t)) -1/2where b 0 = w ⊺ 0 β * , andY (t) = 1 2 δ 2 + 4∥β * ∥ 2 t + atanh 2b 0 a 0 +δ √ δ 2 +4∥β * ∥ 2 .From the solutions for ν, a, we can easily obtain dynamics for µ, ϕ.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,EQUATIONEQUATIONConsider the orthant O + = {(a, ω)|a > 0, ω > 0}. The boundary ∂O + is formed by two orthogonal subspaces. On {(a, ω)|a = 0, ω ≥ 0}, ȧ ≥ 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,EQUATIONThese 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,EQUATIONwhich 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 µ, ϕ.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 β * + a0 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:EQUATION)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. ⊺ 0 β * + a0 2 δ + δ 2 + 4∥β * ∥ 2 = 0.For a given δ, this equation describes a hyperplane through the origin. However, a given δ can only be achieved on the surface of some hyperboloid. Thus, the separating surface is the union of the intersections of a hyperplane and a hyperboloid, both parameterized by δ. This intersection is empty if δ > 0. Initializations exactly on the separating surface will travel along the surface to a saddle point where w ⊺ β * = a = 0.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 a = sgn(µ)EQUATION)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) asEQUATIONwhere 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. The true β * is a unit vector pointing in π/4 direction; β(0) is a unit vector pointing towards 3π/2, -π/4, and π/4 directions, respectively, for each of the three rows. δ then defines how a(0) and ∥w(0)∥ are chosen for a particular β(0) where by convention we choose a(0) > 0.(t) 0 1 a(t) = 2 = 1 = 0.2 = 0.02 = 0.2 = 1 = 2 0.5 0.0 0.5 w 1 (t)The network's function is determined by the product β = aw and governed by the ODE,EQUATION)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 = XM X ⊺ . 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 showEQUATIONwhich 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 β,EQUATION)Plugging these expressions into M givesEQUATION)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.The dynamics of the NTK matrix K = XM X ⊺ is determined by Ṁ . From Eq. ( 4), which is derived in this section, we can write Ṁ = 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.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] .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 Φ α (β),EQUATIONwhere ρ = (Xβ -y) is the residual, then the limiting solution of the dynamics is determined by the constrained optimization problem,EQUATIONwhere 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 problemEQUATIONwhereQ(β) = Φ α (β) -∇Φ α (β(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 givesEQUATION)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 equationEQUATIONwhich 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,EQUATIONhas 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.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 asEQUATION)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 asEQUATIONFirst notice the following simplification,u ⊺ A -1 v = √ δ 2 +4∥β∥ 2 -δ √ δ 2 +4∥β∥ 2 +δ. After some algebra, M -1 isEQUATIONTo make notation simpler we will define the following two scalar functions,EQUATIONsuch that we can expressM -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,EQUATIONWe 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 iEQUATIONThus, M -1 is a Hessian map if and only if for all i ̸ = j,EQUATION)while the RHS isEQUATION)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 findEQUATIONwhich 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 ,EQUATIONThus, 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,EQUATIONRearranging terms we find that g δ (x) must solve the ODEEQUATIONUsing our previous expressions (Eq. ( 50) and Eq. ( 55)) we findEQUATIONwhich implies g δ (x) solves the differential equation, ∇g δ (x) =g δ (x) √ δ 2 +4x( √ δ 2 +4x+δ) . The solution is g δ (x) = c √ δ 2 + 4x + δ, where c ∈ R is a constant. Let c = 1.Plugging in our expressions for g δ ∥β∥ 2 , f δ ∥β∥ 2 , h δ ∥β∥ 2 , we get thatg δ ∥β∥ 2 M -1 =   2 δ 2 + 4∥β∥ 2 + δ   I d -     √ δ 2 +4∥β∥ 2 -δ √ δ 2 +4∥β∥ 2 +δ ∥β∥ 2 δ 2 + 4∥β∥ 2     ββ ⊺ (60)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,EQUATIONEquating terms from our expression for g δ ∥β∥ 2 M -1 (equation 60) we get the expression for ∇q(∥β∥) ∇q(∥β∥) = 2∥β∥EQUATIONwhich plugged into the second term gives the expression for ∇ 2 q(∥β∥),EQUATION)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,EQUATIONEQUATIONLetting 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) isEQUATIONand the potential Φ δ (β) isEQUATIONFinally, putting it all together, we can express the inductive bias as in Theorem 3.1.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, β(∞) = arg minEQUATIONwhere q δ (x) =x 2 -δ 2 δ 2 + x 2 + δ 2 4 x 2 + δ 2 4 -δ 2 x and z = -3 2 ∥β 0 ∥ 2 + δ 2 4 -δ 2 β0 ∥β0∥ .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 thatEQUATIONAnother 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.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,EQUATIONWe will also use the expression, ∥β∥ 2 = ∥β * ∥ 2 + α 2 . Pulling these expressions together we get the following equation for α,EQUATION)If we let k = -2z ⊺ v 3 , the solution for α isEQUATION)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.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.We consider the dynamics of a two-layer linear network with h hidden neurons and c outputs, f (x; θ) = A ⊺ W x, 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 ,EQUATIONEQUATION)Extending derivations in [30] , the NTK matrix can be expressed asEQUATIONwhere ⊗ and ⊕ denote the Kronecker product and sum respectively. The Kronecker sum is defined for square matricesC ∈ R c×c and D ∈ R d×d as C ⊕ D = C ⊗ I d + I c ⊗ D.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]:EQUATION)We will also consider a new matrix quantity Q ∈ R h×h with elementsQ kk ′ = w ⊺ k w k ′ ∥w k ∥∥w k ′ ∥ .The resulting dynamics for µ and ϕ can be entirely written in terms µ, ϕ, ∆:EQUATIONEQUATION)whereM = 2 Diag(∆) 2 + 4Diag(µ) 2 -Diag(∆) -1. Using the conserved structure of ∆ we can express Q as a function of µ and M ,EQUATION)This approach yields a coupled nonlinear dynamical system with 2h 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 Ȧ -1 = -A -1 ȦA -1 and our assumption on ∆, we find that the matrix ν evolves according to the matrix Riccati ODE,EQUATION)Additionally, consider the variable C = A ⊺ A, which evolves according to the matrix Bernoulli ODEEQUATION)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,EQUATIONEQUATION)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.We consider the dynamics of β = W ⊺ A ∈ R d×c in function space, which is governed by the ODE,EQUATION)Vectorizing using the identity vec(ABC) = (C ⊺ ⊗ A)vec(B) equation 85 becomesEQUATIONEQUATIONEQUATION)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 . EQUATION)Notice that we can express ∥β k ∥ 2 F asEQUATIONAt each hidden neuron we have the conserved quantity8 η 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 getEQUATION)EQUATION)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 ,EQUATIONEQUATION)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 ,M k = δ 2 k + 4η a η w ∥β k ∥ 2 F + δ k 2 β ⊺ k β k ∥β k ∥ 2 F ⊕ δ 2 k + 4η a η w ∥β k ∥ 2 F -δ k 2 β k β ⊺ k ∥β k ∥ 2 F . (95) B.1.4 Understanding M when there is a single-neuron h = 1When 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 toEQUATIONEQUATION)and we recover Eq. ( 4) presented in Section 3. When h = d = 1, then ββ ⊺ ∥β∥ 2 F = 1 and thus Eq. ( 7) simplifies to,EQUATION)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.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 thatEQUATION)or δ k → -∞ for all k ∈ [h] and there exists a solution such thatEQUATION) 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 Fis 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.As δ k → -∞ for all k ∈ [h], M k → I c ⊗ |δ k | β k β ⊺ k ∥β k ∥ 2 F, and an analogous argument can be made.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.A direct 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 → -∞.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 thatEQUATIONProof. The quantity η w AA ⊺ -η a W W ⊺ = δI h is conserved in gradient flow. Multiplying on the left by W ⊺ and on the right by W we have thatEQUATIONCompleting the square by adding δ 2 4ηa I d to both sides and dividing by η a we get the equality,EQUATIONFor δ ≥ 0, W ⊺ W + δ 2ηa I d ⪰ 0. For δ < 0, then we know from the conserved quantity thatW W ⊺ + δ 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,EQUATIONwhich rearranged gives the final result.Lemma B.4. Let ∆ = δI h . If either δ ≤ 0 or δ > 0 and h ≥ c, we have thatEQUATIONProof. 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 fromEQUATIONPlugging in expressions for W ⊺ W from Lemma B.3 and A ⊺ A from Lemma B.4 we can directly write,EQUATIONFrom this expression for M (β) we can easily consider how it simplifies in limiting settings of δ:EQUATIONAs δ → ±∞, 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.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:EQUATIONIn the last equality we used that fact that for any vector z with a fixed norm, ∥z∥ 2 = 2 ż⊺ 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,EQUATIONwhere λ ∈ 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λ = -M ∇ λ L (111)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 withEQUATIONProof. First note thatEQUATIONwhere we letW ⊺ 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 asEQUATIONwhere 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 thatM ∈ 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(EQUATIONWe 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 flowEQUATIONwhereΦ δ (λ) = min(d,c) i=1q δ (λ i ) and q δ is the hyperbolic entropy potentialEQUATIONProof. 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 β * -we can 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.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. Lemma B.8. For a depth-(L+1) linear network with square width (d = h) and isotropic initialization, then for all t ≥ 0,EQUATIONEQUATIONProof. 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.⊺ = 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 β = -M X ⊺ ρ where M is a positive semi-definite matrix that solely depends on ∥a∥ 2 , W ⊺ 1 W 1 , and δ,EQUATION)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,EQUATIONProof. 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 ββ ⊺ ,EQUATION)Plugged into the previous expression for M results in a positive definite rank-one perturbation to the identity, M = ∥β∥2L L+1 I d + L∥β∥ -2 L+1 ββ ⊺ .(125) Using the Sherman-Morrison formula we find that M -1 isEQUATIONWe 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.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. The input-output map is linear within each non-empty activation region and continuous at the boundary between regions. Linearity implies that every non-empty9 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] .We consider the dynamics of a two-layer piecewise linear network without biases, f (x; θ) = a ⊺ σ(W x), 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 asEQUATIONwhere ρ 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 ofEQUATIONAs 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 ,EQUATION)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 elementsK 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,EQUATIONThese 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 ∥.We used Google Cloud Platform (GCP) nodes to run all experiments. ∼ 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 samplew i i.i.d. ∼ S d-1 and a i i.i.d.∼ Unif({±1}), and then for i ≥ h2 + 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.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.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.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 t1 , K t2 ⟩/ (∥K t1 ∥ F ∥K t2 ∥ 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.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) channelsEQUATIONwhere the Laplacian kernel is defined as However, in B) we see that networks with a small initialization (α < 1) learn much smoother filters, giving quantiative support to results in Fig. 6 . The smoothness is defined as the normalized Laplacian of the filters (see text, eq. 132).EQUATIONWe 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 f f ; 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 RHM with 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.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 [69] and balance α ∈ {0.1, 0.3, 1.0, 3.0, 10} on a regular grid over n = 5 random initializations with a maximal computational budget of m = 30, 000 training steps. B) Same as A), but reporting the number of training steps required until the test performance passes the predefined threshold of 99%. We clearly see the fastest grokking in an unbalanced rich setting.• 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.• 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.Xu and Ziyin [23] presented exact NTK dynamics for a linear model trained with one-dimensional data.The set of saddle points {(a, w)} is the d -1 dimensional subspace satisfying a = 0 and w ⊺ X ⊺ y = 0.We relax this assumption when considering the dynamics of β in function space and their implicit bias.The basin is given by sgn(a0) for δ ≥ 0 or sgn(w ⊺ 0 β * + a 0 2 (δ + δ 2 + 4∥β * ∥ 2 )) for δ < 0. See A.2.5.The Kronecker sum is defined for square matrices C ∈ R c×c and D ∈ R d×d as C ⊕ D = C ⊗ I d + Ic ⊗ D.When h = c = 1 we can recover Eq. (4) presented in the single-neuron setting directly from Eq. (7).To our knowledge, this property has not been recognized before. See Appendix C.1 for a formal statement.As long as c > 1, then the surface of this d + c hyperboloid is always connected, however its topology will depend on the relationship between d and c.While it is trivial to see that for a network f (x; θ) with N hidden neurons there are 2 N distinct activation patterns, not all activation patterns are attainable. See Raghu et al.[65] for a discussion.It is possible for neighboring regions to have the same linear predictor. Some works define linear regions as maximally connected component of input space with the same linear predictor[87]. 


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

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

Reviews: review: paper_summary: paper_summary: The authors derive exact solutions to a minimal model that transitions between lazy and rich learning, elucidating how unbalanced layer-specific initialization variances and learning rates determine the degree of feature learning in a finite-width network. 
They 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.
strength_weakness: strength_weakness: 1. To yield the exact solutions, the authors evoke the assumption of whitened input, which is rather unrealistic. The low-rank case is also addressed but then one needs to additionally assume that the interpolating manifold is one-dimensional to find the solution in terms of $\delta$ exactly. 

missing citations: 
- In line 398, when mentioning grokking and the transition from lazy to rich learning - it is worth noting Ref. [1]. 

[1] https://arxiv.org/abs/2310.03789
questions: questions: 1. In Ref. [1] which you mention, the formula for the NTK with unbalanced LRs is (using your notation) $K = X (\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. 

2. 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? 


[1] https://arxiv.org/abs/2401.07085
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_summary: The paper studies the learning dynamics of a deep network by leveraging the conserved quantities due to symmetries
strength_weakness: strength_weakness: I think the paper could improve by explaining its results in more detail. For example, I find the following points rather confusing

1. Does theorem 4.1 hold after training? Or it holds at any time in training？

2. 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?

3. 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?

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

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

6. 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
questions: questions: See the weakness section
limitations: limitations: NA
paper_summary: paper_summary: The paper studies the dynamics of training in neural networks with the scope of identifying how the variance of weights' initialization together with layer-wise learning rates determines different learning regimes, encompassing lazy, rich, and the transition between them. The paper identifies a conserved quantity $delta$ that is preserved throughout training (i.e. $\dot{\delta} = 0$) and depends on both the weight's magnitude and the learning rates. Thus, the sign and magnitude of $\delta$ at initialization affect learning and the geometry of the landscape that is traversed through gradient flow. The paper finds three different regimes, starting from a solvable model of two-layer linear network. 

1. $\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. 

2. $\delta = 0$ (*Balanced*): corresponds to rich dynamics.

3. $\delta < 0$: initial lazy fitting regime and second rich phase.

The authors then extend these findings to wider and deeper networks and piece-wise linear activations.
strength_weakness: strength_weakness: 1. I would have appreciated more on the NTK dynamics.  There is a clear connection between $\delta$ and (part of) the NTK $K$. I wonder whether this connection can be made more explicit by studying the NTK dynamics $\dot{K}$. This would make the connection between the conserved quantity and the feature learning more explicit. 

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

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

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

Overall, I am in favor of the paper's acceptance because it devised a model where the three different learning regimes are clean and crystallized.
questions: questions: See weaknesses
limitations: limitations: Addressed.
paper_summary: paper_summary: This paper analyses the lazy versus rich learning dynamics of minimal but finite neural networks by deriving exact solutions under arbitrary layerwise initialization and learning rates. The theoretical insights are successively extended to networks of increasing complexity and corroborated by several experiments.
strength_weakness: strength_weakness: Overall the paper leaves little room for criticism. The results being limited to shallow, often linear neural networks trained with small learning rates is justified by the detailed understanding of the training dynamics, and mitigated by experiments. It could have been acknowledged more clearly how the maximal update parameterization by Yang and Hu (2021) already shows that upstream initializations are necessary for non-vanishing feature learning in large models.
questions: questions: - Why do you call the layerwise initialization variances and learning rates *initialization geometry* when learning rate scalings clearly only influence the training dynamics? Similarly, the term geometry is unspecific and does not highlight the interplay between the layers. I encourage the authors to find a more suitable terminology for the choice of hyperparameters.
- Why is the kernel distance in Figure 1 eventually maximized by a downstream initialization $\delta<0$, and not by an upstream initialization $\delta>0$ where feature learning occurs rapidly both directionally and radially?
limitations: limitations: The main limitations are acknowledged in the Conclusion section. It is acceptable that an exact analysis of learning dynamics does not include deep nonlinear networks trained with large learning rates.

Comments: title: Thanks for the reply
comment: Thank you for your reply. We’re glad that our response helped clarify the connection between our work and the scaling limit literature. We are also happy to hear that you favor acceptance of our work.

We certainly agree that the delayed rich regime is intriguing and to our knowledge has not been studied before. While the lazy vs. rich dichotomy is well-studied in infinite-width settings, our analysis (alongside recent work by Xu and Ziyin [reference 17] and analysis of diagonal linear networks [reference 12]) stands out as one of the few analytically tractable models for the transition between lazy and rich learning in a finite-width network. We believe that introducing solvable systems where precise statements can be made, both in terms of dynamics and implicit bias, is a valuable contribution. Additionally, throughout our analysis we highlight the importance of conserved quantities in determining the learning regime, a perspective not considered in these prior works, nor in the work you mentioned (the difference in feature learning between a μP parameterization and an NTK parameterization is due to a change in function scale, not $\delta$, which remains the same in distribution for both parameterizations).

Thank you for the time and effort you put into reading and reviewing our work.
I sincerely thank the authors for the effort spent on this rebuttal. The connection between some of the results of this work and the NTK dynamics is now clear.

**Connections to width-dependent neural network parameterizations.** 

I thank the authors for this clarification. This helps me better understand the connection of their work in the context of the scaling limit literature. 

The reason I am not confident enough to raise my score (which still favors acceptance) is that two of these regimes are already sort of well-understood (lazy vs rich) in the series of existing works that the authors mention. Thus, I believe the crucial contribution lies in the third regime (delayed rich). Having non-vanishing dynamics of the NTK is precisely achieved by correctly setting the initialization variance, learning rate, and output scale of the model. In a sense, $\delta$ is to some extent explicitly already controlled under $\mu$P and related works (e.g. https://arxiv.org/abs/2310.17813).
Thank you for your response, and we agree that the added discussion of NTK dynamics and the connection to infinite-width literature are valuable additions to our work.

To clarify our discussion on width-dependent parametrizations, there are two “knobs” that can determine the degree of feature learning: (1) the overall function scale and (2) the relative scale between layers (this is $\delta$ in our analysis). The difference in feature learning between a mean field parameterization and an NTK parameterization is due to a change in the first knob (the overall function scale). Under both these parameterizations the distribution of $\delta$ is the same even as width $k \to \infty$ (expected $\delta$ is zero, but with a non-vanishing variance). It is possible to have an infinite-width parameterization that uses the second knob to enter into the lazy regime (in this setting the expected $\delta$ would be non-zero and depend on width). We found Reference [14] to be helpful in understanding the connection between our finite-width analysis of feature learning and existing works in the infinite-width setting.

Thank you for clarifying your concern with our use of “initialization geometry”. Your point makes sense and the suggested “optimization geometry” is a good possibility, which we will consider.

Thank you again for your questions, feedback, and constructive criticism!
We are happy to hear that our answers removed your concerns and that you are positive towards our work. Thank you again for your thoughtful and constructive feedback.
Thank you again for your positive and constructive feedback.
Thanks for the detailed explanation. The answer removed my concerns about the work. I will raise to 7. Although I am positive towards the work, I do not find the results significant enough for me to strongly support it.
I thank the authors for their thoughtful response. I think both the added explanation on the dynamics of the NTK matrix and the connection to infinite-width mean-field literature are valuable additions. In particular, the fact that mean-field parameterization is balanced with non-vanishing variance is interesting, and the distinction between the function scale and relative scale mechanism. Can I understand the mean-field parameterization/muP as the unique width-dependent scaling rule that achieves $\delta$ to remain width-independent and hence not degenerate in behaviour with width $k\to \infty$ to for example strictly lazy behaviour in the NTK parameterization?

Concerning the term ‘initialization geometry’, I still believe that per-layer learning rates are not part of the *initialization*, but only of the *optimization*, hence terms in the direction of ‘optimization geometry’ would make more sense to me.

Overall, I believe this paper makes a valuable contribution and I will keep my positive evaluation.
I have read the author's response and will keep my score