Investigating the Early Signatures of Feature Learning in Deep Linear Networks

Shaunak Bhandarkar, Boris Hanin

Deep neural networks learn data-dependent features, a capacity widely credited for their empirical success. However, explicitly characterizing what features emerge -- and when they improve generalization -- is challenging, since training is nonlinear and non-convex and the data is often anisotropic. Here, we give a closed-form account of feature learning for a deep linear network in a mean-field parameterization, trained by one gradient step from initialization on high-dimensional data carrying a low-dimensional (spiked) signal. In the infinite-width mean-field limit, one step updates each layer's feature kernel by a target-aligned rank-one spike, amplified by both the learning rate and the layer depth. We next assess how these features impact prediction through kernel ridge regression (KRR), finding that KRR on the learned feature kernel yields scalar inflation of the KRR estimator at initialization -- an "anti-shrinkage" mechanism that reinflates task-relevant directions suppressed by ridge, while also amplifying noise. To study when this inflation improves generalization, in the overparameterized proportional limit we derive closed-form expressions for the bias, variance, and…