The Weight Gram Matrix Captures Sequential Feature Linearization in Deep Networks

Abstract: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.

NeurIPS 2026, 32 pages including appendix. Added Proposition 2 (pre-conditioning) and extra whitening experiments (Table 1-3) for V2

Machine Learning (cs.LG); Artificial Intelligence (cs.AI)

arXiv:2605.06258 [cs.LG]

(or arXiv:2605.06258v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2605.06258 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.