High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

Abstract

We study the first gradient descent step on the first-layer parameters W in a two-layer neural network:, where W ∈ R d×N , a ∈ R N are randomly initialized, and the training objective is the empirical MSE loss: 1 n n i=1 (f (xi) − yi) 2 . In the proportional asymptotic limit where n, d, N → ∞ at the same rate, and an idealized student-teacher setting, we show that the first gradient update contains a rank-1 "spike", which results in an alignment between the first-layer weights and the linear component of the teacher model f * . To characterize the impact of this alignment, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on W with learning rate η, when f * is a single-index model. We consider two scalings of the first step learning rate η. For small η, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large η, we prove that for certain f * , the same ridge estimator on trained features can go beyond this "linear regime" and outperform a wide range of random features and rotationally invariant kernels. Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.

The Neural Covariance SDE: Shaped Infinite Depth-and-Width Networks at Initialization

“…∂ t U αα t = bU αα t (U αα t − 1) , (D. 22) where we observe if b > 0 this ODE is "mean avoiding" as it will drift towards 0 or ∞. And since the V t time scale is on the order of 1 n , for all t > 0 we have that…”

“…The Neural Tangent Kernel (NTK) limit formed the foundation for a rush of theoretical work, including advances in our understanding of generalization for wide networks [13][14][15]. Besides the NTK limit, the infinite-width mean-field limit was developed [16][17][18][19], where the different parameterization demonstrates benefits for feature learning and hyperparameter tuning [20][21][22].…”

On the symmetries in the dynamics of wide two-layer neural networks

“…Learning representations. An existing line of work (Yehudai and Shamir, 2019;Allen-Zhu et al, 2019;Abbe et al, 2021;Damian et al, 2022;Ba et al, 2022) studies in depth the representations learned by neural networks trained with (S)GD at finite-width from a different perspective focusing on the advantages of feature-learning in terms of performance comparatively to using random features. In contrast, our aim is to describe the representations themselves in relationship with the symmetries of the problem.…”