Asymptotics of feature learning in two-layer networks after one gradient-step

TL;DR

This paper provides an exact asymptotic analysis of how two-layer neural networks learn features after a single gradient step, revealing the importance of data adaptation for learning non-linear functions.

Contribution

It introduces a spiked Random Features model to describe trained networks and derives precise asymptotics for generalization error in high dimensions, connecting theory with actual network learning curves.

Findings

The model accurately predicts the learning curves of two-layer networks.

Adapting to data is essential for learning non-linear functions.

One gradient step significantly improves feature learning.

Abstract

In this manuscript, we investigate the problem of how two-layer neural networks learn features from data, and improve over the kernel regime, after being trained with a single gradient descent step. Leveraging the insight from (Ba et al., 2022), we model the trained network by a spiked Random Features (sRF) model. Further building on recent progress on Gaussian universality (Dandi et al., 2023), we provide an exact asymptotic description of the generalization error of the sRF in the high-dimensional limit where the number of samples, the width, and the input dimension grow at a proportional rate. The resulting characterization for sRFs also captures closely the learning curves of the original network model. This enables us to understand how adapting to the data is crucial for the network to efficiently learn non-linear functions in the direction of the gradient -- where at initialization it can only express linear functions in this regime.