Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Critical feature learning in deep neural networks
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.
Derives a finite-width Bayesian theory of deep non-linear networks in which the network prior is a superposition of Gaussian processes, and uses a large-deviation approach to obtain forward-backward self-consistency equations for MAP kernels, linking feature learning to criticality and output scale.
Problem
Understanding feature learning in deep neural networks at finite width (proportional limit P=alpha N), beyond the lazy/NNGP regime, and what drives kernel adaptation to data.
Approach
Writes the network prior exactly as a superposition of Gaussian processes with fluctuating kernels of variance O(1/N); a large-deviation principle then gives forward-backward equations for MAP kernels across all layers, perturbatively solved to linear order in conjugate kernels, showing feature learning is strongest slightly away from criticality and increases with output feature scale.
Key results
- Feature-corrected kernels from the theory show larger centered kernel alignment with the target YY^T than the NNGP on XOR and MNIST (0 vs 3), matching Langevin-trained networks.
- Kernel adaptation is strongest slightly away from edge-of-chaos criticality, peaking at weight variance g_f well below g_crit.
- Downscaling the output weight variance by a non-extensive feature scale gamma_0 increases the conjugate kernel in the input layer roughly linearly in gamma_0, enhancing feature learning.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use feature-corrected kernels instead of the NNGP kernel | Bayesian feature-learning theory · XOR | CKA with target kernel | Helped |
| Use feature-corrected kernels instead of the NNGP kernel | Bayesian feature-learning theory · MNIST digits 0 vs 3 | kernel alignment with target kernel | Helped |
| Reduce hidden-layer weight variance | Feature-corrected kernels · XOR | CKA with target kernel | Helped |
| Increase feature scale by reducing output-layer weight variance | Deep neural networks · Training-data kernel adaptation | feature learning | Helped |
| Increase feature scale by reducing output-layer weight variance | Deep neural networks · Training-data kernel adaptation across hidden-layer weight variances | feature learning | Helped |
| Move hidden-layer weight variance toward g_f from above | Fully trained networks · Kernel adaptation | kernel adaptation | Helped |
| Reduce hidden-layer weight variance below g_f | Deep neural networks · Kernel adaptation in the vanishing-gradient regime | kernel adaptation | Hurt |
| Increase network width so that P/N approaches zero with P fixed | Feature-corrected kernels · XOR | difference from NNGP limit | Hurt |
Surprising findings
- Maximal kernel adaptation occurs at prior weight variances significantly below the critical point predicted by criticality alone.
- Feature learning is not monotone in criticality: adaptation drops suddenly below g_f as gradients vanish.
Experimental setup
Datasets: MNIST XOR
Benchmarks: MNIST XOR
Metrics: centered kernel alignment
Baselines: NNGP Langevin-sampled finite-width networks
Builds on: NNGP neural tangent kernel Langevin dynamics Gärtner-Ellis theorem
Closest papers by meaning
- From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning (ICML 2025 · similarity 0.90)
- Adaptive kernel predictors from feature-learning infinite limits of 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.88)
- Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.88)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.87)
- Deep Kernel Posterior Learning under Infinite Variance Prior Weights (ICLR 2025 · similarity 0.87)
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.