Research map / Optimization / Implicit bias and generalization

The Weight Gram Matrix Captures Sequential Feature Linearization in Deep Networks

NeurIPS 2026 · Optimization › Implicit bias and generalization · approach: Neural tangent kernel analysis · 0 citations (OpenAlex)

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.

Proposes a feature-centric view in which the weight Gram matrix captures a virtual feature-space update (Virtual Covariance), introduces Target Linearity (R^2 of ridge regression onto targets) and shows deep networks sequentially linearize representations, linking this to Neural Collapse and VAE latent interpolation.

Problem

How feature representations evolve during deep network training is poorly explained by general theories; existing analyses are parameter-centric or limited to structured case studies, and phenomena like Neural Collapse and linear interpolation lack a unified account.

Approach

The Feature Learning Equation links weight updates to feature gradients via the weight Gram matrix, so the Gram shift under GD equals the Virtual Covariance Shift of a hypothetical feature-space gradient update. This yields Target Linearity, a surrogate Y^T H^T W^T W H Y whose growth explains sequential layer-wise linearization.

Key results

• Virtual Covariance correlates with the actual weight Gram update at rho=0.98 versus rho=0.71 for AGOP on a two-layer CIFAR-10 network.
• Whitened gradient updates (preserving the Gram shift) match or exceed standard GD on MNIST, SVHN and CIFAR-10 and track GD with a small gap on CIFAR-100 for 3-layer MLPs.
• Layer-wise Target Linearity increases with depth and training for a 4-layer MLP on CIFAR-10/100, with the surrogate growing roughly exponentially while TL grows roughly logarithmically.

Surprising findings

• Pre-trained BERT's Target Linearity profile differs by task, with syntactic tasks peaking in early-to-middle layers, while fine-tuning makes all tasks near-monotone in depth.
• Gradient updates that preserve the weight Gram matrix recover most of standard GD's performance.

Limitations

• Dynamics rely on first-order Taylor and m-positive homogeneity approximations and assumptions.
• Layer-wise results depend on the Neural Feature Ansatz, which lacks full theoretical justification and does not hold in all settings.
• Gradient-whitening results are limited to small fully connected and CNN models on standard benchmarks.

Experimental setup

Datasets: MNIST SVHN CIFAR-10 CIFAR-100 SST-2 MNLI CoNLL-2003

Benchmarks: MNIST SVHN CIFAR-10 CIFAR-100 SST-2 MNLI CoNLL-2003

Metrics: test accuracy Target Linearity (R^2) Pearson correlation

Baselines: Neural Tangent Kernel regime Average Gradient Outer Product

Builds on: Neural Feature Ansatz Neural Tangent Kernel Neural Collapse Muon optimizer

Closest papers by meaning

• From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks (ICLR 2025 · similarity 0.89)
• Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.88)
• Feature Learning Dynamics in Infinite-Depth Neural Networks (NeurIPS 2026 · similarity 0.87)
• How are linear representations learned? Exact solutions to the dynamics of abstraction (NeurIPS 2026 · similarity 0.86)
• Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.86)
• On the Emergence of Cross-Task Linearity in Pretraining-Finetuning Paradigm (ICML 2024 · similarity 0.86)