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

TL;DR

The paper addresses how initialization scale and layer-wise learning rates influence feature learning beyond the lazy NTK regime. It derives exact gradient-flow solutions for a minimal two-layer linear model and extends the analysis to wide/deep linear and shallow nonlinear networks, introducing a conserved quantity delta and a time-warped mirror-flow framework. The work provides precise trajectory characterizations, regime classifications (upstream, balanced, downstream), and explicit inductive-bias potentials, supported by experiments showing practical benefits like improved interpretability, reduced sample complexity, and faster grokking. These results offer principled initialization strategies to promote rapid feature learning in finite-width networks and deepen understanding of depth- and layer-scale interactions in training dynamics.

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.

Figures (12)

• Figure 1: Unbalanced initializations lead to rapid rich learning and generalization. We follow the experimental setup used in Fig. 1 of chizat2019lazy -- a wide two-layer student ReLU network $f(x;\theta) = \sum_{i=1}^h a_i \max(0,w_i^\intercal x)$ trained on a dataset generated from a narrow two-layer teacher ReLU network. The student parameters are initialized as $w_i \sim \text{Unif}(\mathbb{S}^{d-1}(\frac{\tau}{\alpha}))$ and $a_i = \pm\alpha\tau$, such that $\tau > 0$ controls the overall scale of the function, while $\alpha > 0$ controls the relative scale of the first and second layers through the conserved quantity $\delta = \tau^2 (\alpha^2 - \alpha^{-2})$. (a) Shows the training trajectories of $|a_i|w_i$ (color denotes $\mathrm{sgn}(a_i)$) when $d = 2$ for four different settings of $\tau, \delta$. 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 $\delta > 0$ shows striking alignment to the teacher (dotted lines), while a downstream initialization $\delta < 0$ shows no alignment. (b) Shows the test loss and kernel distance from initialization computed through training over a sweep of $\tau$ and $\delta$ when $d=100$. Lazy learning happens when $\tau$ is large, rich learning happens when $\tau$ is small, and rapid rich learning happens when both$\tau$ is small and $\delta$ is large -- an upstream initialization. This initialization also leads to the smallest test loss. See \ref{['fig:two-layer-supporting']} in \ref{['app:experimental-details-two-layer']} for supporting figures.
• Figure 2: Balance determines geometry of trajectory. The quantity $\delta = \eta_wa^2 - \eta_a\|w\|^2$ is conserved through gradient flow, which constrains the trajectory to: (a) a one-sheeted hyperboloid for downstream initializations, (b) a double cone for balanced initializations, and (c) a two-sheeted hyperboloid for upstream initializations. Gradient flow dynamics for three different initializations $a_0, w_0$ with the same product $\beta_0 = a_0w_0$ are shown. The minima manifold is shown in red and the manifold of equivalent $\beta_0$ initializations in gray. The surface is colored according to training loss, with blue representing higher loss and red representing lower loss.
• Figure 3: Exact solutions for the single hidden neuron model. Our theoretical predictions (black dashed lines) agree with gradient flow simulations (solid lines, color-coded based on $\delta$ values), shown here for three key metrics: $\mu$ (left), $\phi$ (middle), and $S(0,t)$ (right). Each metric starts at the same value for all $\delta$, but varying $\delta$ has a pronounced effect on the metric's dynamics. For upstream initializations ($\delta \gg 0$), $\mu$ changes only slightly, $\phi$ exponentially aligns, and $S$ remains near zero, indicative of the lazy regime. For balanced initializations ($\delta = 0$), both $\mu$ and $\phi$ change significantly and $S$ quickly moves away from zero, indicative of the rich regime. For downstream initializations ($\delta \ll 0$), $\mu$ quickly drops to zero, then $\mu$ and $\phi$ slowly climb back to one. Similarly, $S$ remains small before a sudden transition towards one, indicative of a delayed rich regime. See \ref{['app:single-neuron-exact-solutions']} for further details.
• Figure 4: Balance modulates $\beta$ dynamics and implicit bias. Here we show the dynamics of $\beta = a w$ with different values of $\delta$, but the same initial $\beta_0$. When $X^\intercal X$ is whitened (left), we can solve for the dynamics exactly using our expressions for $\mu, \phi$ (black dashed lines). Upstream initializations follow the trajectory of gradient flow on $\beta$, downstream initializations first move in the direction of $\beta_0$ before sweeping around towards $\beta_*$, and balanced initializations take an intermediate trajectory between these two. When $X^\intercal X$ is low-rank (right), then we can only predict the trajectories in the limit of $\delta = \pm \infty$. If the interpolating manifold is one-dimensional, then we can solve for the solution in terms of $\delta$ exactly (black dots). See \ref{['app:single-neuron-inductive-bias']} for details.
• Figure 5: Rapid feature learning is caused by large activation changes with minimal parameter movement. (a) We show the surface of a two-layer ReLU network trained on an XOR-like task, starting with a near-zero input-output map, $f(x;\theta_0)\approx0$. The surface consists of convex conic regions, each with a distinct activation pattern, colored by the parity of active neurons. A lazy initialization (bottom) maintains a fixed activation partition throughout training, reweighting the hidden neurons to fit the data. In contrast, a rich balanced or upstream initialization (top) features an initial alignment phase where the partition map changes rapidly while the input-output map remains close to zero, followed by a data-fitting phase. (b) We show the evolution of Hamming distance in activation patterns and parameter distance, relative to $t=0$, as a function of overall and relative scales (same experiments as in \ref{['fig:two-layer-relu']}(b)). Rapid feature learning occurs from a small-$\tau$ 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-$\tau$ downstream initializations require large parameter movement to fit the data in the delayed rich regime.

Theorems & Definitions (26)

• Theorem 3.1: Extending Theorem 2 in azulay2021implicit
• Theorem 4.1
• Theorem 4.2
• Lemma A.1
• proof
• Theorem A.2: Theorem 2 from azulay2021implicit
• Theorem B.1
• proof
• Lemma B.3
• proof