The Weight Gram Matrix Captures Sequential Feature Linearization in Deep Networks
Not on the map. Why?
Abstract (arXiv)
Understanding how deep neural networks learn representations remains a central challenge in machine learning theory. In this work, we propose a feature-centric framework for analyzing neural network training by relating weight updates to feature evolution. We introduce a simple identity, the Feature Learning Equation, which identifies the weight Gram matrix as the key object capturing feature dynamics. This enables us to interpret gradient descent as implicitly inducing a hypothetical evolution of features, whose covariance structure - termed the Virtual Covariance - characterizes how representations evolve during training. Building on this perspective, we introduce Target Linearity, a measure quantifying the linear alignment between features and targets. By analyzing the training and layer-wise dynamics, we show that deep networks learn to sequentially transform representations toward target-linear structure. This linearization perspective provides a unified interpretation of several empirical phenomena, including Neural Collapse and linear interpolation in generative models.
Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Builds on 35 of its 52 references that are in the map. First posted on arXiv 7 May 2026.
Builds on
35 of 52 references.
References · 35
- 2017Attention Is All You Needcited by 59,932
- 2015Deep Residual Learning for Image Recognitioncited by 51,809
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
- 2018GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understandingcited by 4,854
- 2013Efficient Estimation of Word Representations in Vector Spacecited by 4,346
- 2016Gaussian Error Linear Units (GELUs)cited by 3,431
- 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
- 2019Similarity of Neural Network Representations Revisitedcited by 1,695
- 2016Understanding deep learning requires rethinking generalizationcited by 1,657
- 2016Understanding intermediate layers using linear classifier probescited by 1,279
- 2023Representation Engineering: A Top-Down Approach to AI Transparencycited by 1,239
- 2019BERT Rediscovers the Classical NLP Pipelinecited by 1,046
23 more
- 2020Prevalence of Neural Collapse during the terminal phase of deep learning trainingcited by 713
- 2022Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasetscited by 599
- 2018A Mean Field View of the Landscape of Two-Layers Neural Networkscited by 463
- 2025Muon is Scalable for LLM Trainingcited by 381
- 2014On the Number of Linear Regions of Deep Neural Networkscited by 378
- 2019Gradient Descent Maximizes the Margin of Homogeneous Neural Networkscited by 308
- 2018Implicit Bias of Gradient Descent on Linear Convolutional Networkscited by 251
- 2016Benefits of depth in neural networkscited by 248
- 2022Towards Understanding Grokking: An Effective Theory of Representation Learningcited by 229
- 2020Directional convergence and alignment in deep learningcited by 172
- 2022Neural Networks can Learn Representations with Gradient Descentcited by 161
- 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 146
- 2022The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networkscited by 141
- 2020Neural collapse with unconstrained featurescited by 140
- 2017The Riemannian Geometry of Deep Generative Modelscited by 81
- 2025Toward universal steering and monitoring of AI modelscited by 29
- 2024Average gradient outer product as a mechanism for deep neural collapsecited by 21
- 2023Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networkscited by 18
- 2024Linear Recursive Feature Machines provably recover low-rank matricescited by 16
- 2025Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometrycited by 9
- 2023ReLU Neural Networks with Linear Layers are Biased Towards Single- and Multi-Index Modelscited by 7
- 2024Early Directional Convergence in Deep Homogeneous Neural Networks for Small Initializationscited by 6
- 2025Neural Collapse under Gradient Flow on Shallow ReLU Networks for Orthogonally Separable Datacited by 2
Built on by
Citing papers per year: 2026: 3.
3 papers.
- 2026Learning Dynamics Reveal a Hierarchy of Weight-Induced Layerwise Gram Metricscited by 0
- 2026Dual-Manifold Geometry Guided Representation Learning: Adaptive Coupling between Kernel and Data Spacescited by 0
- 2026Refusal geometry reflects refusal training: diverse refusal prefixes can raise stable rank and weaken refusal vector ablation attackscited by 0
- 2026Refusal geometry reflects refusal training: diverse refusal prefixes can raise stable rank and weaken refusal vector ablation attacks26 Aug 2026
- 2026Dual-Manifold Geometry Guided Representation Learning: Adaptive Coupling between Kernel and Data Spaces13 Aug 2026
- 2026Learning Dynamics Reveal a Hierarchy of Weight-Induced Layerwise Gram Metrics8 Jun 2026