# Feature learning as alignment: a structural property of gradient descent in non-linear neural networks

arXiv: https://arxiv.org/abs/2402.05271 · HTML: https://mlsanity.com/paper/2402.05271/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 60 references that are in the map.
- First posted on arXiv 7 Feb 2024.

## Builds on (most cited first, 30 of 36)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [Linearized two-layers neural networks in high dimension](https://mlsanity.com/paper/1904.12191/) (2019)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)
- [A Short Note on Concentration Inequalities for Random Vectors with SubGaussian Norm](https://mlsanity.com/paper/1902.03736/) (2019)
- [On the Power and Limitations of Random Features for Understanding Neural Networks](https://mlsanity.com/paper/1904.00687/) (2019)
- [The Neural Tangent Kernel in High Dimensions: Triple Descent and a Multi-Scale Theory of Generalization](https://mlsanity.com/paper/2008.06786/) (2020)
- [What Can ResNet Learn Efficiently, Going Beyond Kernels?](https://mlsanity.com/paper/1905.10337/) (2019)
- [Limitations of Lazy Training of Two-layers Neural Networks](https://mlsanity.com/paper/1906.08899/) (2019)
- [Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking](https://mlsanity.com/paper/2311.18817/) (2023)
- [Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed](https://mlsanity.com/paper/2102.11742/) (2021)
- [Neural Networks Efficiently Learn Low-Dimensional Representations with SGD](https://mlsanity.com/paper/2209.14863/) (2022)
- [A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks](https://mlsanity.com/paper/2310.07891/) (2023)
- [Spectrum of inner-product kernel matrices in the polynomial regime and multiple descent phenomenon in kernel ridge regression](https://mlsanity.com/paper/2204.10425/) (2022)
- [Second-order regression models exhibit progressive sharpening to the edge of stability](https://mlsanity.com/paper/2210.04860/) (2022)
- [Knowledge Distillation in Wide Neural Networks: Risk Bound, Data Efficiency and Imperfect Teacher](https://mlsanity.com/paper/2010.10090/) (2020)
- [Temperature check: theory and practice for training models with softmax-cross-entropy losses](https://mlsanity.com/paper/2010.07344/) (2020)
- [Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning](https://mlsanity.com/paper/2306.04815/) (2023)
- [Layer-Specific Adaptive Learning Rates for Deep Networks](https://mlsanity.com/paper/1510.04609/) (2015)
- [Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks](https://mlsanity.com/paper/2305.06986/) (2023)

## Built on by (most cited first, 10 of 10)

- [Emergence in non-neural models: grokking modular arithmetic via average gradient outer product](https://mlsanity.com/paper/2407.20199/) (2024)
- [On the Complexity-Faithfulness Trade-off of Gradient-Based Explanations](https://mlsanity.com/paper/2508.10490/) (2025)
- [Weibull Weight-Scale Parameter Evolution under AdamW Training Dynamics](https://mlsanity.com/paper/2606.19367/) (2026)
- [From Coefficients to Directions: Rethinking Model Merging with Directional Alignment](https://mlsanity.com/paper/2512.00391/) (2025)
- [Breaking Data Symmetry is Needed For Generalization in Feature Learning Kernels](https://mlsanity.com/paper/2604.00316/) (2026)
- [Why Geometric Continuity Emerges in Deep Neural Networks: Residual Connections and Rotational Symmetry Breaking](https://mlsanity.com/paper/2605.04971/) (2026)
- [The Model Shape Behind Scaling Laws: A Gradient Superposition Perspective](https://mlsanity.com/paper/2605.27989/) (2026)
- [Dual-Manifold Geometry Guided Representation Learning: Adaptive Coupling between Kernel and Data Spaces](https://mlsanity.com/paper/2608.12737/) (2026)
- [A Function-Space Approach to the Statistical Mechanics of Learning Dynamics](https://mlsanity.com/paper/2609.09589/) (2026)
- [Can Representation Learning Decouple from Loss Minimization? Polar Updates Have an Answer](https://mlsanity.com/paper/2609.36240/) (2026)

## Abstract (arXiv)

Understanding the mechanisms through which neural networks extract statistics from input-label pairs through feature learning is one of the most important unsolved problems in supervised learning. Prior works demonstrated that the gram matrices of the weights (the neural feature matrices, NFM) and the average gradient outer products (AGOP) become correlated during training, in a statement known as the neural feature ansatz (NFA). Through the NFA, the authors introduce mapping with the AGOP as a general mechanism for neural feature learning. However, these works do not provide a theoretical explanation for this correlation or its origins. In this work, we further clarify the nature of this correlation, and explain its emergence. We show that this correlation is equivalent to alignment between the left singular structure of the weight matrices and the newly defined pre-activation tangent features at each layer. We further establish that the alignment is driven by the interaction of weight changes induced by SGD with the pre-activation features, and analyze the resulting dynamics analytically at early times in terms of simple statistics of the inputs and labels. We prove the derivative alignment occurs almost surely in specific high dimensional settings. Finally, we introduce a simple optimization rule motivated by our analysis of the centered correlation which dramatically increases the NFA correlations at any given layer and improves the quality of features learned.