Research map / Architectures and efficiency / RNNs and normalization
Three Mechanisms of Feature Learning in a Linear Network
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Solves gradient-flow learning dynamics of a one-hidden-layer linear network on 1D data at any finite width, identifies alignment, disalignment and rescaling as feature-learning mechanisms, and gives kernel/feature-learning phase diagrams validated on nonlinear nets.
Problem
Existing theory of neural network learning dynamics (NTK, mean-field) mostly covers infinite width and perturbative finite-width corrections; there is no analytically solvable finite-width model that exhibits both kernel and feature-learning dynamics.
Approach
Derive an exact closed-form gradient-flow solution for a one-hidden-layer linear network with 1D data at any width and initialization, with separate layer learning rates; use it to show that layer alignment/disalignment and output rescaling are feature-learning mechanisms absent in the kernel regime, and to derive scaling-limit phase diagrams.
Key results
- Exact solution for the one-hidden-layer linear network shows the alignment zeta evolves monotonically from 0 to (alpha_+^2-1)/(alpha_+^2+1) in the feature learning phase, while staying near zero in the kernel phase.
- In a two-layer ReLU net on CIFAR-10, turning Kaiming into the feature learning phase improves test accuracy by about 5%, while turning Xavier into the kernel phase decreases it by about 10%.
- Measured weight-evolution exponents delta=0.47, 1.08, 1.93 match predicted 0.5, 1, 2 for NTK, Xavier(-) and Kaiming(-) scalings.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use feature-learning rather than kernel-phase scaling for Kaiming initialization | Two-layer ReLU FCN · CIFAR-10 | test accuracy: approximately 5% | Helped |
| Use kernel-phase rather than feature-learning scaling for Xavier initialization | Two-layer ReLU FCN · CIFAR-10 | test accuracy: approximately 10% | Hurt |
| Use feature-learning rather than kernel-phase scaling | Two-layer ReLU FCN, d=2000 · CIFAR-10 | test accuracy: about 5% higher | Helped |
| Use feature-learning rather than kernel-phase scaling | Two-layer ReLU FCN, d=2000 · CIFAR-10 | training accuracy: 1.0 | Helped |
| Increase initialization scale σ (ablation) | ResNet-18 · CIFAR-10 | test accuracy | Hurt |
| Increase initialization scale σ (ablation) | ResNet-18 · CIFAR-10 | weight norm | Hurt |
| Use feature-learning rather than kernel-phase scaling | Two-layer ReLU FCN · CIFAR-10 | relative weight evolution: δ=0; δ=0.47,1.08,1.93 | Mixed |
| Vary dataset size and input dimension (ablation) | ReLU network alignment experiment · Alignment experiment | layer alignment | No clear effect |
| Increase learning rate to 15 times the reference setting (ablation) | Linear network alignment experiment · Alignment experiment | layer alignment: 15 times | No clear effect |
| Use data not confined to a one-dimensional subspace (ablation) | Linear network alignment experiment · Alignment experiment | layer alignment | No clear effect |
| Use a three-layer instead of one-hidden-layer linear network (ablation) | Linear network alignment experiment · Alignment experiment | layer alignment | No clear effect |
| Use nonlinear activations instead of linear activation (ablation) | Network alignment experiments · Alignment experiment | layer alignment | No clear effect |
Surprising findings
- Under gradient descent, alignment between layers is strongly initialization-dependent, whereas SGD converges to aligned solutions regardless of initialization.
- Larger initialization can hurt generalization through disalignment, which is distinct from the usual NTK explanation.
Limitations
- Analytical results hold only for inputs lying in a one-dimensional subspace.
- Analysis is limited to gradient flow (deterministic), not discrete SGD with finite step sizes.
- Closed-form solution is only tractable for beta=1 (linear network); general power activations are reduced to an ODE.
Experimental setup
Datasets: MNIST CIFAR-10 synthetic 1D regression data
Benchmarks: CIFAR-10 MNIST
Metrics: alignment angle zeta test accuracy train accuracy relative weight distance from initialization weight norm
Baselines: NTK parameterization mean-field parameterization Xavier initialization Kaiming initialization lazy training
Builds on: neural tangent kernel mean-field parameterization tensor programs
Closest papers by meaning
- From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning (ICML 2025 · similarity 0.90)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.89)
- Feature Learning Dynamics in Infinite-Depth Neural Networks (NeurIPS 2026 · similarity 0.89)
- Learning dynamics in linear recurrent neural networks (ICML 2025 · similarity 0.89)
- Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks (NeurIPS 2026 · similarity 0.89)
- From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.