Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks

Technical Summary: Sharp Feature-Learning Transitions and Bayes-Optimal Neural Scaling Laws in Extensive-Width Networks

Problem Setup
The paper investigates the information-theoretic limits of learning a one-hidden-layer teacher neural network with hierarchical features from noisy queries, specifically in the context of knowledge transfer to a smaller student model. The study operates in the extensive-width regime, where the teacher width $k$ scales linearly with the input dimension $d$ ($k = \Theta(d)$). This setting is positioned between the well-studied multi-index regime ($k = \Theta(1)$) and the Neural Tangent Kernel (NTK) regime ($k \gg d$), capturing the "large-but-finite" width of practical networks.

The teacher network generates labels $y_\mu = v^\top \sigma(W_0 x_\mu)$ from inputs $x_\mu \sim \mathcal{N}(0, I_d)$, with additive Gaussian noise. The readout vector $v$ follows a hierarchical structure (e.g., power-law decay $v_i \propto i^{-\beta}$), implying that features have geometrically decaying importance. The authors analyze two types of learners:

1. Bayesian Student: Knows the data generation mechanism (activation $\sigma$, noise variance $\Delta$, readout $v$) but not the specific teacher weights $W_0$. It predicts via the posterior mean, achieving the Bayes-optimal error ($\varepsilon_{BO}$).
2. SGD Student: A one-hidden-layer network of width $k_s \leq k$ trained with ADAM.

The primary questions addressed are: (Q1) Given a fixed dataset, which teacher features are statistically recoverable? and (Q2) How does the recoverable subset of features constrain the scaling law of $\varepsilon_{BO}$ with respect to data size $n$?

Methodology
The authors employ a heuristic leave-one-out (cavity) decoupling argument, a technique standard in statistical physics, to derive asymptotically sharp characterizations of the learning process.

• Decoupling: The method isolates the contribution of a single feature $w_i$ by treating the contributions of all other features as an effective noise term.
• Gaussian Assumption: The authors assume that the "cavity fields" (the effective noise from other features) behave as independent Gaussian variables. This reduces the multi-feature inference problem to a sequence of independent scalar generalized linear models (GLMs).
• Fixed-Point System: By analyzing the overlap between the true feature direction and the posterior estimate, the authors derive a closed system of fixed-point equations. These equations relate the overlap $q_i$ of feature $i$ to the signal-to-noise ratio (SNR) and the global generalization error $\varepsilon$.
• Validation: The theoretical predictions are validated numerically using Hamiltonian Monte Carlo (HMC) simulations across various parameter regimes, showing excellent agreement even at moderate dimensions ($d \sim 10^2$).

Key Results

1. Sharp Staircase of Feature-Learning Transitions:
   The analysis reveals that feature learnability is governed by a sharp phase transition. A feature $i$ is recoverable (overlap $q_i > 0$) if and only if its effective SNR, $snr_i = \frac{(n/d)v_i^2}{\Delta + \varepsilon_{BO}}$, exceeds a critical activation-dependent threshold $\lambda_\sigma$.

   • As data size $n$ increases, features cross this threshold sequentially, starting from the strongest (largest $v_i$).
   • This results in a "staircase" of discontinuous jumps in feature overlaps and corresponding drops in generalization error, rather than a smooth transition.

2. Effective Width ($k_c$) and Error Decomposition:
   The number of learnable features at a given data budget is defined as the effective width $k_c = |\{i : q_i > 0\}|$. The Bayes-optimal error decomposes into two terms:
   $$\varepsilon_{BO} = \Theta\left(\frac{k_c d}{n}\right) + \Theta\left(\sum_{j > k_c} v_j^2\right)$$
   The first term represents the estimation error on learnable features (scaling as $d/n$ per feature), while the second term is the irreducible bias from unlearnable features.

3. Unified Scaling Laws:
   For power-law readouts with exponent $\beta > 1/2$, the paper derives two distinct scaling regimes that unify under the relation $\varepsilon_{BO} = \Theta(k_c d / n)$:

   • Feature-Learning Regime ($n \ll k^{2\beta}d$): Here, $k_c \sim (n/d)^{1/(2\beta)}$. The error scales as $\varepsilon_{BO} \sim n^{1/(2\beta) - 1}$. In this regime, the model is limited by the number of features it can learn.
   • Refinement Regime ($n \gg k^{2\beta}d$): All features are learnable ($k_c = k$). The error scales as $\varepsilon_{BO} \sim n^{-1}$, recovering the classical parametric rate.
     The crossover point $n^\star = \Theta(k^{2\beta}d)$ marks the transition where the weakest feature becomes learnable.

4. Implications for SGD-Trained Students:
   Empirical results show that an ADAM-trained student with width $k_s$ near the theoretical effective width $k_c$ achieves these optimal scaling laws (up to a small algorithmic gap).

   • If $k_s < k_c$, the student underfits, limited by its capacity.
   • If $k_s \gg k_c$, the student overfits non-learnable directions, degrading performance.
   • The optimal student width scales with data size according to the feature hierarchy, providing a theoretical basis for hyperparameter transfer.

Significance and Claims
The paper claims to provide the first analytical characterization of Bayes-optimal learning in the extensive-width regime ($k = \Theta(d)$) for networks with hierarchical features and vanishing second Hermite coefficients ($\mu_2 = 0$).

• Unification: It unifies two previously distinct empirical scaling regimes (feature-learning and refinement) into a single framework driven by the effective width $k_c$.
• Mechanism: It offers a mechanistic explanation for the "staircase" learning dynamics observed in practice, attributing them to sequential phase transitions in feature recoverability.
• Methodological Contribution: The work extends the cavity method beyond dense readouts to arbitrary readout structures (including sparse and power-law), recovering known results for dense readouts while generalizing to sparse settings where error decay rates differ significantly.
• Practical Insight: The derived effective width $k_c$ provides a first-principles prescription for sizing student models based on data budgets, suggesting that model capacity should scale with the number of learnable features rather than the total number of teacher features.

Limitations
The authors acknowledge that their derivations rely on unproven assumptions, specifically the Gaussianity of cavity fields and the concentration of overlaps. Additionally, the analysis is restricted to activations with a vanishing second Hermite coefficient ($\mu_2 = 0$), which includes odd activations but excludes others like ReLU (though the paper notes that ReLU can be analyzed via the complementary regime $\mu_2 \neq 0$ in other works). The leave-one-out reduction is heuristic, though numerically validated.