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

• Overparameterized networks can remain close to linearization [1] [2] [3] [4]; at infinite width, the NTK becomes deterministic and training resembles NTK kernel regression [6] [7].The introduction describes this as the lazy or kernel regime, associated with minimal parameter movement, static hidden representations, exponential learning curves, and RKHS-aligned implicit biases [8] [9].Source: §1 · p. 1
• Small overall scale produces rich learning, while rapid rich learning and the lowest test loss occur when small scale combines with large relative scale (upstream initialization).The figure’s sweep shows that large overall scale produces lazy learning; at small scale, relative scale can still shift learning between rich and lazy.Source: §1 · p. 2
• Scaling a finite-width model with zero initial output toward large overall scale induces lazy learning, which can worsen test error; small or vanishing initializations are proposed as a basis for rich learning [8].Source: §1 · p. 2
• Rich learning features an evolving NTK, non-convex saddle-traversing dynamics, sigmoidal learning curves, and simplicity biases, but its learned features remain problem-dependent.The paper investigates how layer-specific initialization variances and learning rates affect the lazy-to-rich transition and the inductive biases that support generalization.Source: §1 · p. 2
• Building on Azulay et al. [9], the paper derives exact gradient-flow solutions for a single-neuron, two-layer linear model and finds conserved quantities constrain trajectory geometry and reveal relative-scale directionality.Unbalanced initialization can make earlier layers learn faster (upstream) or later layers learn faster (downstream), producing distinct trajectories.Source: §1 · pp. 2–3
• Extensions find rapid rich learning only from balanced initialization in linear networks, but upstream initialization can accelerate rich learning in nonlinear networks with piecewise-linear activations.Source: §1 · p. 3
• Experiments indicate upstream initializations drive feature learning in deep finite-width networks, promote early-layer CNN interpretability, reduce hierarchical-data sample complexity, and shorten modular-arithmetic grokking.Source: §1 · p. 3

Linear networks

• This work derives exact linear-network solutions using mixed Bernoulli and Riccati equations, exposing separate alignment and fitting phases.Xu and Ziyin [23] also presented exact NTK dynamics for a one-dimensional linear model.Source: §2 · p. 3

Implicit bias

• The paper derives exact inductive-bias expressions for its minimal model and extends Azulay et al. [9]’s time-warped mirror-flow analysis to wide and deep linear networks.Earlier mirror-flow analyses express bias through a Bregman divergence [46], while finding the potential can require a second-order equation that may be unsolvable [47] [48].Source: §2 · p. 4

Two-layer networks

• Unlike prior balanced-initialization analyses [49] [51] [52] [53] [54] [55], this study finds unbalanced starts can eliminate or augment quantization.Earlier work observed quantization in two-layer ReLU networks from small initializations, with alignment preceding task-fitting phases.Source: §2 · p. 4

Infinite-width networks

• Prior infinite-width work develops rich-dynamics scalings and population theories, whereas this study focuses on finite-width networks and discusses connections to width-dependent parameterizations.Related developments include maximal update parametrization (µP) [16] [18].Source: §2 · p. 4

3 A Minimal Model of Lazy and Rich Learning with Exact Solutions

• The invariant δ = η_w a^2 − η_a∥w∥^2 confines trajectories to distinct hyperboloid or cone geometries, and the reduced dynamics are exactly solvable.Except for a measure-zero set converging to saddle points, trajectories reach a global minimum; rescaling symmetry leaves a manifold of parameter minima.Source: §3 · pp. 4–5, Figure 2 · p. 4

Upstream.

• For upstream initialization (δ ≫ 0), ϕ aligns rapidly before µ fits exponentially, while the kernel moves minimally, yielding lazy learning.The kernel is dominated by the η_w a^2 I_d term as w updates.Source: Upstream · p. 5, Figure 3 · p. 5
• For balanced initialization (δ = 0), alignment of w precedes fitting in µ, producing a rank-one NTK update followed by kernel scaling and rich learning.For vanishing initialization, the alignment and fitting phases decouple.Source: Upstream · p. 5, Figure 3 · p. 5

Downstream.

• For downstream initialization (δ ≪ 0), fast fitting in µ is followed by slow alignment in ϕ, so the initially lazy trajectory becomes rich after a δ-dependent delay.In function space, β first moves toward β* along β0's direction; if β0 is misaligned, later alignment changes the NTK.Source: Downstream · pp. 5–6, Figure 4 · p. 6
• For low-rank data, δ controls the trade-off between minimum-norm interpolation and preserving β0's direction; with vanishing initialization, upstream and balanced biases are functionally equivalent.In the downstream setting, preservation of the initialization direction dominates the implicit bias.Source: Downstream · pp. 6–7

Wide linear networks

• Theorem 4.1 expresses the wide-network preconditioner as a sum of neuron-specific terms determined by rank-1 contributions β_k and conserved quantities δ_k.The preconditioner characterizes the NTK and depends on quadratics of the layer parameters.Source: §4 · p. 7
• When h ≥ max(d, c), uniform-sign conserved quantities yield lazy learning, while δ_k → 0 yields rich learning regardless of dimensions.For min(d, c) ≤ h < max(d, c), networks narrowing from input to output are lazy when all δ_k ≫ 0, and expanding networks are lazy when all δ_k ≪ 0; opposite signs produce a delayed rich regime under the stated initialization condition.Source: §4 · p. 7
• The singular values of β follow a mirror flow with a hyperbolic entropy potential, interpolating between ℓ1 in the rich limit and ℓ2 in the lazy limits.The potential is δ-dependent.Source: §4 · p. 8

Deep linear networks

• For deep linear networks, the rich-regime inductive bias balances attaining the minimum-norm solution with preserving the initialization direction.This generalizes the inductive bias derived for rich two-layer linear networks to depth-(L + 1) networks.Source: §4 · p. 8

Piecewise Linear Networks

• Piecewise-linear networks partition input space into convex regions with linear maps, and the activation regions form a 2-colorable partition, a property the paper says was previously unrecognized.The networks use activations σ(z) = max(z, γz); with L hidden layers and width h, the input-output map can have O(hdL) regions.Source: §5 · p. 8

Two-layer network

• When neuron directions remain fixed, the activation partition, activation patterns, preconditioning matrices, and NTK stay fixed as slopes change; insufficient width can delay rich alignment, with |δk| setting the delay.The output weights change to fit the data while the kernel remains constant in this regime.Source: §5 · p. 8
• At vanishing scale, balanced training separates into partition alignment, where directions align to data-dependent fixed points while outputs stay small, and data fitting, where radial dynamics dominate.During alignment, each neuron's driving vector is approximately independent of the other hidden neurons; later, residuals depend on all neurons and directions may continue to change.Source: §5 · p. 9
• With upstream initialization, neuron directions change activation patterns and the NTK while the preconditioning matrix stays constant, enabling rich learning with little parameter movement.When δk ≫ 0, the directional dynamics are approximately ∂t β̂k = −δk μk^-1 ξk; vanishing initializations produce an extremely fast alignment phase.Source: §5 · p. 9, Figure 5 · p. 9

Connections to infinite-width

• Finite-width learning regimes align with an infinite-width phase portrait in Luo et al. [17], but the role of the distribution of per-neuron conserved quantities in finite-width feature learning remains open.Across width-dependent parametrizations, random initialization induces per-neuron conserved quantities with zero expectation and non-vanishing variance.Source: §5 · pp. 9–10

Unbalanced initializations in practice

• Upstream initialization links activation-pattern changes to early NTK evolution, enhances or degrades CNN filter alignment, changes hierarchical-task sample complexity, and reduces grokking time.The grokking experiment reduces the embedding initialization variance in a single-layer transformer, which has fewer than 6% of the parameters.Source: §5 · p. 10, Figure 6 · p. 10

6 Conclusion

• Unbalanced initializations that promote faster learning in earlier layers can accelerate rich learning, according to the paper’s theory and experiments.The theory is difficult to extend to deeper nonlinear networks, and omits discretization and stochastic SGD effects; future work should study how layerwise learning-speed profiles affect feature learning, inductive bias, and generalization.Source: §6 · p. 10

A Single-Neuron Linear Network; A.1 Conserved quantity

• Scaling symmetry between a and w yields a conserved quantity δ that confines parameter dynamics to a hyperboloid whose geometry depends on its magnitude and sign.The appendix also relates this symmetry to the single-neuron linear network’s parameterization and training trajectory.Source: §A.1 · p. 17, A Single-Neuron Linear Network · p. 17

A Single-Neuron Linear Network; A.2 Exact solutions; A.2.1 Deriving the dynamics for µ and ϕ; A.2.2 Balanced δ = 0; A.2.3 Upstream δ > 0

• For whitened inputs (X⊺X = Id) with δ = 0 and a0 > 0, signed magnitude µ and alignment ϕ describe training: µ keeps its sign, while ϕ decouples and evolves from ϕ0 to 1.Source: §A.2 · p. 17, §A.2.1 · p. 17, §A.2.2 · p. 18

A.2.4 Downstream δ < 0; A.2.5 Basins of attraction

• For δ < 0, a may change sign, and non-saddle trajectories intersect the boundary aω = 0 at most once; initializations exactly on the separating surface travel to a saddle point where a = wᵀβ* = 0.Source: §A.2.4 · p. 19, §A.2.5 · p. 20

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

• The model parameters can be recovered from µ and ϕ because w moves only in the span of its initialization w0 and β∗.The coefficient of w along β∗ is determined by ∥w∥ϕ.Source: §A.2.6 · p. 21

A.3 Function space dynamics of β

• In function space, β follows preconditioned gradient flow with a time-dependent matrix M that also determines the NTK, and its dynamics can be expressed entirely in β when a does not cross zero.This characterization applies to δ ≥ 0 except at the origin saddle, and to δ < 0 when the sign of a remains unchanged.Source: §A.3 · p. 22

A.3.1 Kernel dynamics

• The NTK is nearly unchanged for δ ≫ 0, changes mainly with β’s direction for δ ≪ 0, and changes with both magnitude and direction when δ = 0.These regimes correspond respectively to lazy, delayed-rich, and rich learning.Source: §A.3.1 · p. 22

A.4 Deriving the inductive bias

• In overparameterized settings with infinitely many interpolating solutions, δ can determine which solution the dynamics reach, motivating a time-warped mirror-flow analysis.The discussion extends prior work by Azulay et al. [9].Source: §A.4 · p. 22

A.4.1 Overview of time-warped mirror flow analysis

• Mirror-flow dynamics for a linear predictor can be related to a constrained optimization characterized by the potential’s Bregman divergence.The overview identifies the resulting strictly convex function Q(β) as the implicit bias.Source: §A.4.1 · p. 23
• The KKT argument derives the constraint multiplier by integrating the mirror-flow dynamics, assuming the residual integral exists and the required second-order ODE has a solution.The overview states that residual decay must be fast enough for the integral to be well defined.Source: §A.4.1 · p. 23
• The required second-order ODE solution may fail to exist even for extremely simple Jacobian maps [47].Source: §A.4.1 · p. 23
• Azulay et al. [9] extend the interpretation using a smooth positive scalar function g(β) that warps time along the predictor trajectory.The warped-time integral must be finite, and τ(∞) = ∞ is required to avoid stalling and preserve the interpolating solution.Source: §A.4.1 · p. 23

A.4.2 Applying time-warped mirror flow analysis

• For initializations with ∥β(t)∥ > 0 for all t ≥ 0, M is positive definite and has a unique inverse computable with the Sherman–Morrison formula.The analysis applies to the β dynamics satisfying ∇θβ∇θβ⊺ = M.Source: §A.4.2 · p. 23

Proving

• The Hessian integrability condition discussed by Gunasekar et al. [47] provides a test for whether M⁻¹ is the Hessian of a potential.Source: Proving · p. 24
• The componentwise test reduces to 2∇fδ(x) + hδ(x) = 0, which fails, so M⁻¹ is not a Hessian map.Source: Proving · p. 24
• Although M⁻¹ is not a Hessian map, a scalar function gδ(∥β∥²) can rescale it into one.The analysis derives a differential equation for gδ and specifies a solution with constant c = 1.Source: Proving · pp. 24–25
• A radial potential can be chosen so its Hessian satisfies the required conditions; it is convex for x ≠ 0 and minimized at x = 0.The construction chooses the constant cδ so that Φδ(β) > 0 for β ≠ 0 and Φδ(0) = 0.Source: Proving · p. 25
• Combining the construction yields the inductive bias stated in Theorem 3.1.Source: Proving · p. 25

A.4.3 Connection to Theorem 2 in Azulay et al. [9]

• The inductive biases in Theorem 3.1 and Azulay et al. [9] are functionally equivalent despite being derived in different forms, while this work allows a broader set of initializations.Both analyses ensure that M remains invertible along the trajectory under their stated assumptions; extending the analysis to trajectories where M becomes low-rank is only a possibility discussed by the authors.Source: §A.4.3 · p. 26

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

• When the interpolating manifold is one-dimensional, all interpolating solutions can be parameterized by a scalar, yielding an exact analytic solution for the inductive-bias optimization.The authors also report that a solution with the opposite sign appears to work for previously unconsidered initializations, including trajectories passing through the origin where a changes sign.Source: §A.4.4 · p. 26

B Wide and Deep Linear Networks

• The analysis techniques for a single-neuron linear network are extended to linear networks with multiple neurons, outputs, and layers.This section introduces the broader linear-network setting before focusing on wide networks.Source: B Wide and Deep Linear Networks · p. 27

B.1 Wide linear networks

• For a two-layer linear network with h ≥ min(d, c), the initialization-dependent matrix Δ = A0A0⊺ − W0W0⊺ is conserved throughout gradient flow.The condition on h lets the parameterization represent all linear maps from Rd to Rc.Source: §B.1 · p. 27

B.1.1 Parameter space dynamics

• For wide-network parameter dynamics, the analysis uses coordinate transformations; with c = 1, it produces a coupled nonlinear system with 2h variables.Additional initialization assumptions, such as permutation invariance among hidden neurons, can simplify the system.Source: §B.1.1 · p. 27
• For c = h, assuming Δ = δIh and that A stays invertible, transformed variables ν and C evolve according to matrix Riccati and Bernoulli equations.For whitened input, β* = X⊺Y is the unique minimum of the function-space dynamics.Source: §B.1.1 · pp. 27–28
• The transformed parameter-space system is difficult to solve exactly unless ν and β*⊺β* share an eigenspace, so the authors turn to the more tractable function-space dynamics.That eigenspace condition would decouple ν into scalar Riccati equations; the function-space approach requires fewer assumptions.Source: §B.1.1 · p. 28

B.1.2 Function space dynamics

• In function space, β = W⊺A follows gradient flow preconditioned by M, a matrix built from quadratic forms of A and W.The displayed dynamics express M using a Kronecker sum involving A⊺A and W⊺W.Source: §B.1.2 · p. 28

B.1.3 Proving Theorem 4.1

• Each hidden neuron contributes a matrix β_k = w_k a_k^T to the network map β, which is the sum of these contributions.The proof considers a single hidden neuron with parameters w_k and a_k.Source: §B.1.3 · p. 28
• For each hidden neuron, η_w∥a_k∥^2 − η_a∥w_k∥^2 = δ_k is conserved.The passage states that this quantity is used to invert the expression for ∥β_k∥_F^2.Source: §B.1.3 · p. 28
• Using the inverted norm expression, the proof solves for the outer products and derives the representation for M in Theorem 4.1.The provided passages do not include the full expression for M.Source: §B.1.3 · pp. 28–29

• For h=min(d,c)=1, Theorem 4.1’s M simplifies; h=c=1 recovers Eq. (4), while h=d=1 reduces Eq. (7) using ββ⊺F=1.Source: §B.1.4 · p. 29
• As δ→∞, M is proportional to I_d for h=c=1 and β⊺β for h=d=1; as δ→−∞, these become ββ⊺ and I_c, so M is β-independent in both limits.In both settings, M is a weighted sum of the identity matrix and a rank-one projection matrix.Source: §B.1.4 · p. 29
• When δ=0, M depends on β, and δ’s effect on lazy, rich, or delayed-rich dynamics depends crucially on the relative sizes of d, h, and c.Source: §B.1.4 · p. 29

B.1.5 Interpreting M in different limits and architectures

• When all δk tend to +∞ or −∞ and a least-squares solution satisfies the corresponding initial span condition, Theorem B.1 gives lazy dynamics with static Mk and NTK.The parameters follow projected gradient descent in their row spaces for the positive limit, while the negative limit has an analogous column-space condition.Source: §B.1.5 · pp. 29–30
• With independent initializations, narrowing networks become lazy as all δk → ∞ when h ≥ c, expanding networks as all δk → −∞ when h ≥ d, and h ≥ max(d, c) suffices for either limit.These dimensional guarantees follow because the initialized rows span the required output or input space.Source: §B.1.5 · p. 30
• If no least-squares solution lies in the span of the initial βk rows, the +∞ limit begins lazily but later enters a delayed rich phase; the magnitude of δk sets the delay.The transition occurs when fitting requires the row directions to change; an analogous column-space statement holds for the −∞ limit.Source: §B.1.5 · p. 30

B.1.6 Simplifying M through assumptions on ∆

• Under isotropic initialization, ∆ = δIh simplifies M, but for general δ this assumption requires h ≤ min(d, c).For square networks with equal input, hidden, and output dimensions and the stated Gaussian initializations, isotropy is naturally satisfied as h → ∞.Source: §B.1.6 · p. 30
• As δ → ±∞, M → δIdc and learning is lazy with a constant NTK, whereas δ = 0 yields rich dynamics in which the NTK changes in magnitude and direction.In the infinite-|δ| limit, β follows the gradient-flow trajectory of linear regression.Source: §B.1.6 · p. 31

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

• For ∆ = δIh, h ≥ min(d, c), and δ ≠ 0, the singular values of β follow a mirror flow with a δ-dependent hyperbolic-entropy potential.This potential smoothly interpolates between ℓ1 and ℓ2 penalties in the rich (δ → 0) and lazy (δ → ±∞) limits, respectively.Source: §B.1.7 · p. 32
• The mirror-flow potential does not generally establish an inductive bias at interpolation because singular vectors evolve; whitening and task-aligned initialization keep them fixed, but make the final solution independent of δ.Under those additional assumptions, the dynamics decouple and applying an ℓ1 or ℓ2 singular-value penalty makes no difference.Source: §B.1.7 · p. 33

B.2 Deep linear networks

• The deep-linear analysis assumes a depth-(L + 1) network with square width, isotropic initialization, and shared learning rate η = 1.The setup takes input and hidden dimensions equal.Source: §B.2 · p. 33
• Conserved relations among layers preserve the initialization structure, while first-layer weights significantly influence predictor dynamics.The analysis expresses the squared norm and outer product of the predictor using parameter quantities and conserved layer relations.Source: §B.2 · p. 33
• As δ→0, interpolating deep-linear solutions follow a depth-dependent inductive bias balancing minimum norm with preservation of the initialization direction [9].The theorem assumes square width, isotropic initialization, and a solution that interpolates the data.Source: §B.2 · p. 34

C Piecewise Linear Networks

• For piecewise-linear networks, the section focuses on input-output-map dynamics rather than the inductive bias of interpolating solutions.It extends tools from the linear-network analysis while examining their limitations.Source: C Piecewise Linear Networks · p. 35
• Directly extending mirror-flow analysis to non-trivial piecewise-linear networks is difficult or provably impossible [9] [80].The section proceeds from properties of piecewise-linear maps to two-layer dynamics, then discusses challenges and future directions for deeper networks.Source: C Piecewise Linear Networks · p. 35

C.1 Surface of a piecewise linear network

• The network is linear within each non-empty activation region and continuous across boundaries, whose union forms a convex partition of input space.Each region has a predictor vector; boundaries occur where a neuron’s pre-activation is zero.Source: §C.1 · p. 35
• Without redundant neurons, neighboring activation regions can be two-colored by the parity of their active-neuron counts.The authors use this property for visualization and connect it to paper folding [82].Source: §C.1 · p. 35

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

• For a two-layer, bias-free piecewise-linear network, each hidden neuron has conserved quantity δ_k = η_w a_k^2 − η_a∥w_k∥^2.The activation σ(z) = max(z, γz) includes ReLU and Leaky ReLU; a neuron’s contribution depends on whether an input lies in its active halfspace.Source: §C.2 · p. 36
• Unlike in the linear setting, each neuron’s dynamics have a distinct driver dependent on its activations, and the NTK depends on both M_k and C.The matrix C records neuron activations on training points.Source: §C.2 · p. 36
• A signed spherical coordinate change separates each neuron’s contribution β_k into direction β̂_k and slope μ_k, which determine its active halfspace orientation and linear-region slope.The coordinates are expressed in parameters as β̂_k = w_k/∥w_k∥ and μ_k = a_k∥w_k∥.Source: §C.2 · p. 36

D Experimental Details

• Experiments ran on GCP, with Figure 1 using a 360-core AMD Genoa node, Figure 5 on single A100 GPUs, and estimated compute totaling about 200 A100 hours plus 100 node hours.Figure 5 runtimes ranged from four to 20 hours per experiment; Figures 2–4 ran on a personal computer.Source: D Experimental Details · p. 37

D.1 Figure 1: Teacher-Student with Two-layer ReLU Networks

• Between 10^5 and 10^6 training steps, training loss continues to fall while test loss and other measured distances remain largely unchanged.Training loss saturates at 1e-10; the supporting experiment trains for ten times longer than Fig. 1(b).Source: Figure 10 · p. 38

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

• The single-neuron linear-network figures simulate gradient flow for full-rank and low-rank data cases.The full-rank setup uses X^T X = I2; the passage also specifies parameter choices for the low-rank case.Source: §D.2 · p. 38

Kernel Distance

• Across 500 epochs of MNIST training with learning rate 1e-4 and batch size 512, parameter distance is L2 over all parameters, activation distance is normalized Hamming distance over binarized hidden activations, and kernel distance measures NTK similarity over time.Hamming and kernel distances use a 10% MNIST subsample; curves are averaged over 8 runs.Source: Kernel Distance · p. 39

Gabor Filters

• Networks with small initialization (α < 1) learn smoother convolutional filters while achieving comparable training and test accuracy.The experiment uses a small ResNet on CIFAR10 with 15 × 15 convolution kernels and no weight decay.Source: Figure 11 · p. 39, Gabor Filters · p. 39

Random Hierarchy Model

• The random hierarchy model experiment trains a width-64 MLP on hierarchical classification data while varying first-layer initialization scale and training-set size.The setup uses three hidden layers, L = 3, eight classes, and averages test accuracy over six random seed configurations.Source: Random Hierarchy Model · pp. 39–40

Grokking

• The fastest grokking occurs in an unbalanced rich setting on the modular arithmetic task.Grokking is measured as reaching 99% accuracy; the experiments vary embedding balance α and logit scaling τ.Source: Figure 12 · p. 40, Grokking · p. 40