The Weight Gram Matrix Captures Sequential Feature Linearization in Deep Networks

TL;DR

This paper introduces a feature-centric framework using the weight Gram matrix to analyze how deep neural networks learn representations, revealing sequential feature linearization during training.

Contribution

It proposes the Feature Learning Equation and Virtual Covariance to interpret gradient descent as evolving features, offering a new perspective on deep learning dynamics.

Findings

Deep networks transform features towards target-linear structure during training.

The framework explains phenomena like Neural Collapse and linear interpolation in generative models.

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.