Feature learning from one gradient step

Ba et al.'s results. The step is rank one up to $\log^2n/\sqrt n$ (Proposition 2). For $\eta=\Theta(1)$, with $\theta_1=\mu_1\eta\sqrt{\bar\mu^2/\psi_1+\mu_1^{*2}}$, $\theta_2=\mu_1\mu_1^*\eta$ and $\bar\mu^2=\|f^*\|_{L^2}^2+\sigma_\varepsilon^2$ the energy of the labels, label noise included, a spike appears iff $\theta_1\gt\psi_2^{1/4}$ (Theorem 3), with $$s_1(W_1)\to\frac{\sqrt{(1+\theta_1^2)(\psi_2+\theta_1^2)}}{\theta_1},\qquad|\langle u_1,\beta^*\rangle|^2\to\frac{\theta_2^2}{\theta_1^2}\Big(1-\frac{\psi_2+\theta_1^2}{\theta_1^2(\theta_1^2+1)}\Big).$$ Label noise strengthens the spike and dilutes its overlap, and the effect is large: the printed Theorem 3 drops $\sigma_\varepsilon^2$ from $\bar\mu^2$ although its own proof carries it, and at $d=1000$ a simulation gives $s_1=3.70$, overlap 0.167, against 3.92 and 0.171 from the formula with the noise term and 2.59 and 0.44 without (my check). After any fixed number of such steps, for odd $\sigma$, the risk equals that of the Gaussian-equivalent features $\mu_1W_t^\top x+\mu_2z$ (Theorem 5), so $R_1\ge\mu_2^{*2}$; the gain $R_0-R_1\to\delta\ge0$ is explicit through fixed-point equations, positive iff $\mu_1\mu_1^*\eta\ne0$, and vanishes as $\psi_2\to\infty$ (Theorem 7). For a large step $\eta=\bar\eta\sqrt N$ and bounded $\sigma$, with $\tau^*=\inf_\kappa\mathbb E_{\xi_1}\big(\sigma^*(\xi_1)-\mathbb E_{\xi_2}\sigma(\kappa\xi_1+\xi_2)\big)^2$, Theorem 11 gives, for $n/d$ above a constant, $$R_1(\lambda)\le10\tau^*+C\Big(\sqrt{\tau^*d/n}+d/n\Big),$$ so one large step beats every kernel when $10\tau^*\lt\mu_2^{*2}$; for $\sigma=\sigma^*=\operatorname{erf}$, $\tau^*=0$ and $R_1=O(d/n)$.

The first gradient is rank one, because a Hadamard product with a rank-one matrix loses a factor $\sqrt n$. Split $\sigma'=\mu_1+\sigma'_\perp$, so $G_0=A+B+C$ with $A=\frac{\mu_1}{n\sqrt N}X^\top ya^\top$, $B$ the $\sigma'_\perp$ part, and $C$ the network's own output. The estimate used is $$\|(ya^\top)\odot M\|=\|\operatorname{diag}(y)\,M\operatorname{diag}(a)\|\le\|y\|_\infty\|a\|_\infty\|M\|:$$ where $A$ gets the $\ell^2$ norms $\|y\|\approx\sqrt n$, $\|a\|\approx1$, $B$ gets only the sup norms, so $\|B\|=O(\log^2n/n)$ against $\|A\|\asymp1/\sqrt n$, and $C$ is small under mean-field scaling. The same induction keeps each neuron within $O(\log^2N/\sqrt N)$ of its start for any fixed number of steps at $\eta=\Theta(1)$.

The spike is a BBP transition, and its overlap with $\beta^*$ is inherited from $u$. $W_1\approx W_0+ua^\top$ with $u=\frac{\eta\mu_1}{n}X^\top y$, which is $\eta\mu_1\mu_1^*\beta^*$ plus an isotropic fluctuation of squared norm of order $d/n$. Sylvester's determinant identity reduces the outlier equation to a $2\times2$ determinant of resolvent quadratic forms, $$\det\Big(I_2+\begin{pmatrix}\|a\|^2&1\\1&0\end{pmatrix}\begin{pmatrix}u^\top\\a^\top W_0^\top\end{pmatrix}Q_0(x)\begin{pmatrix}u&W_0a\end{pmatrix}\Big)=0,\qquad Q_0=(W_0W_0^\top-x)^{-1},$$ whose entries converge to Marchenko-Pastur transforms, leaving $1+\theta_1^2(1+xm(x))=0$, with a root beyond the edge iff $\theta_1^4\gt\psi_2$. The overlap is the classical spike overlap with $u/\|u\|$ times $\cos^2(u,\beta^*)=\theta_2^2/\theta_1^2$, so even as $\eta\to\infty$ the spike points along $\beta^*$ only as well as $X^\top y/n$ does.

Small steps keep the neurons nearly orthogonal, which is the condition Gaussian equivalence needs. Hu and Lu's Lindeberg proof for random features uses the weights only through $\|W\|=O(1)$ and pairwise overlaps (with each other and with $\beta^*$) of order $\mathrm{polylog}/\sqrt d$, plus an $\ell^\infty$ bound on the ridge solutions, which the paper reproves for $W_0+\Delta$ by leaving one neuron out. Step 1's columnwise bound gives the first: each neuron moves no more than the random overlaps it already had. In the equivalent model the prediction is a multiple of $\xi_1=\langle\beta^*,x\rangle$ plus independent Gaussian noise, so the risk is at least $\mathbb E(\sigma^*(\xi_1)-c\,\xi_1)^2\ge\mu_2^{*2}$.

The small-step gain is a rank-one perturbation of the random-features risk, computed without rotation invariance. The variance term is unchanged, since $W_1^\top W_1-W_0^\top W_0$ has Frobenius norm $O(1)$; only the bias moves, and averaging $\beta^*$ over the sphere is no longer allowed because $W_1$ is aligned with it. So split the feature covariance into the random part plus a rank-two term and separate it by Sherman-Morrison-Woodbury. The remaining quadratic forms, such as $u^\top\tilde X^\top\Phi_0R_0\Phi_0^\top\tilde Xu$, concentrate on $\theta_1^2$, $\theta_2$ or $\theta_2^2$ times normalized traces, because $u$ and $a$ are independent of the second sample, and those traces are evaluated with linear pencils, so $\delta$ is a rational function of their limits.

A step of size $\sqrt N$ shifts each neuron along $\beta^*$ by its own output weight, and an average of neurons is a Gaussian smoothing of $\sigma$. With $\eta=\bar\eta\sqrt N$ the signal part of column $i$ of the update is $$\eta\sqrt N[A_1]_i=\bar\eta\mu_1\mu_1^*\big(\tfrac1nX^\top X\big)\beta^*\cdot\sqrt Na_i\ \approx\ c_i\beta^*,\qquad c_i=\bar\eta\mu_1\mu_1^*\sqrt Na_i\sim N\big(0,(\bar\eta\mu_1\mu_1^*)^2\big),$$ order one, so near-orthogonality and Gaussian equivalence are lost. Keep the neurons with $\sqrt Na_i\approx\alpha$; they share $\kappa=\bar\eta\mu_1\mu_1^*\alpha$ while their $w_i^0$ are independent, and a law of large numbers turns their average into $\bar f(x)=\mathbb E_{\xi_2}\sigma(\kappa\xi_1+\xi_2)$. The rest of the update adds an error of order $d/n$: every selected neuron inherits the error of $X^\top y/n$ as an estimate of $\mu_1^*\beta^*$. Since the Gaussian $a_i$ take every value, $\kappa$ is free, which is why the bound involves $\tau^*$ and not $\bar\eta$; for erf, $\mathbb E\operatorname{erf}(\kappa\xi_1+\xi_2)=\operatorname{erf}(\kappa\xi_1/\sqrt3)$, so $\kappa=\sqrt3$ reproduces the teacher.

Ridge on fresh data does as well as the oracle, up to a constant. Ridge on the trained features is kernel ridge regression in the finite-dimensional RKHS of $\frac1{\sqrt N}\sigma(W_1^\top x)$, and once the empirical feature covariance concentrates (bounded $\sigma$, $n/d$ above a constant) ordinary KRR bounds give $R_1\le B_1+B_2+2\sqrt{B_1B_2}+o_P(1)$, with both terms controlled by Step 5's $\|f^*-f_r\|^2$ and a norm $\|f_r\|^2_{\mathcal H}=\Theta(N^r)$ that the ridge range makes negligible. That gives Theorem 11 with the unoptimized constant 10.

The first step replaces each neuron by a function of itself. Because $f_{\theta_0}\equiv0$, the gradient of neuron $j$ involves only the preprocessed labels, $\nabla_{w_j}L=-2a_jg_n(w_j)$ with $g_n(w)=\frac1n\sum_i\hat y_ix_i\sigma'(w\cdot x_i)$, and the weight decay $\lambda_1=\eta_1^{-1}$ exactly cancels $W^{(0)}$: $$w_j^{(1)}=2\eta_1a_j\,g_n(w_j).$$ No trace of the random initialization is left, and by homogeneity the question is what $w\mapsto g_n(w)$ does to a uniform direction.

The population feature map is $Hw/\sqrt{2\pi}$ once the low-order parts are subtracted. Stein's lemma and the Hermite expansion of $\sigma'$ give, for unit $w$, $$g(w)=\mathbb E[\hat f^*(x)\,x\,\sigma'(w\cdot x)]=\frac{C_1-\beta}{2}+\frac{w(C_0-\alpha)}{\sqrt{2\pi}}+\frac{Hw}{\sqrt{2\pi}}+\sum_{k\ge2}\Big(c_{k+1}\frac{C_{k+1}(w^{\otimes k})}{k!}+c_{k+2}\,w\,\frac{C_k(w^{\otimes k})}{k!}\Big),$$ an asymptotic series in $d^{-1/2}$ whose later terms carry more information about $f^*$ but are smaller by further powers of $d^{-1/2}$. Without preprocessing the order-one $C_1/2$, the same for every neuron, would move all neurons one way; subtracting $\alpha+\beta\cdot x$ leaves $Hw/\sqrt{2\pi}$, of order $d^{-1/2}$, as the leading term. Since $H$ maps into $S^*$, every neuron lands in the principal subspace, and Assumption 2 makes the images span it. Ba et al.'s mechanism is the first term of this same expansion.

Estimating the Hessian term takes $d^2$ samples. The signal has size $\|Hw\|\gtrsim1/(\kappa\sqrt d)$, and the empirical map must match the population one uniformly over neurons, $$\sup_{w\in S^{d-1}}\|g_n(w)-g(w)\|\lesssim\sqrt{\frac{d\,\iota^{p+1}}{n}},$$ so $n\gtrsim d^2\kappa^2$ up to logs. The uniform bound combines truncation by hypercontractivity, a hyperplane-arrangement count (the sign patterns $\sigma'(w\cdot x_i)$ take at most $n^d$ values) and subgaussian tails; a leave-one-out split shows the new preactivations are at most $\sqrt{\iota^3/d}$ on the training points, so $\eta_1\asymp\sqrt d$ puts them in $[-1,1]$ and exactly undoes the $d^{-1/2}$ size of $Hw$.

Random biased ReLUs reproduce monomials on $[-1,1]$ with bounded weights. For $a\sim\mathrm{Unif}\{\pm1\}$ and $b\sim\mathrm{Unif}[-1,1]$, integration by parts gives $\mathbb E[2(1-a)f''(b)\sigma(ax+b)]=f(x)$ plus an affine function, so for every $k$ there is $v_k(a,b)$ with $$\mathbb E[v_k(a,b)\,\sigma(ax+b)]=x^k\quad(|x|\le1),\qquad\sup|v_k|\lesssim1,$$ transferred to Gaussian biases by importance weighting. That is why the biases are redrawn: a ReLU with random bias acts, in expectation, like any power of its preactivation.

Random directions $Hw$ span every polynomial on $S^*$, at a cost set by $\kappa$. Write each monomial $\langle T,x^{\otimes k}\rangle$ as an average $\mathbb E_w[z_T(w)(g_n(w)\cdot x)^k]$ with the minimum-norm coefficient from the moment matrix of $g_n(w)^{\otimes k}$. That works once the moment matrix is nondegenerate on symmetric tensors over $S^*$, $$\mathbb E_w\langle T,(\Pi^*g_n(w))^{\otimes k}\rangle^2\ \gtrsim\ (rd\kappa^2)^{-k}\|T\|_F^2,$$ which holds for the leading term because $Hw$ is a nondegenerate random vector in $S^*$ of smallest scale $\lambda_{\min}(H)/\sqrt d$, with $\lambda_{\min}$ the smallest nonzero eigenvalue in absolute value ($H$ need not be positive semidefinite); the remainder is absorbed using $n\gtrsim d^2r\kappa^2$ and $d\gtrsim\kappa r^{3/2}$. With Step 10, the dimension drops out, $\eta_1^{-2k}$ cancelling the $d^k$, and summing over the Taylor tensors of $f^*$ represents it with second moment $\tilde O(r^p\kappa^{2p})$.

Finite width, ridge and generalization. Sampling $m$ neurons realizes the representation with training error $\tilde O(r^p\kappa^{2p}/m)$ and $\|a^*\|^2\lesssim r^p\kappa^{2p}/m$ (Bernstein); ridge on $a$ is norm-constrained least squares, which gradient descent approximates; and a Rademacher bound gives $\|a^*\|\sqrt{md/n}=\sqrt{dr^p\kappa^{2p}/n}$. The factor $d$ is there because $W^{(1)}$ was built from the same sample, so the bound must hold uniformly over first layers; with fresh data it disappears. The $n^{-1/4}$ is the square root of the label-noise fluctuation of the training loss.

Transfer follows from the polynomial representation, and the lower bound from targets with $H=0$. Steps 10 to 12 used $f^*$ only as a degree-$p$ polynomial on $S^*$, so a new head fits any $\tilde g(\Pi^*x)$ with $N,m\gtrsim r^p\kappa^{2p}$. For the lower bound take $\mathrm{He}_p(v\cdot x)/\sqrt{p!}$ over $e^{c\epsilon^2d}$ unit vectors with pairwise overlap at most $\epsilon$: $\mathbb E f_vf_w=(v\cdot w)^p\le\epsilon^p$, and Szörényi's argument eliminates at most $2/(\tau^2-\epsilon^p)$ candidates per query, which forces the tolerance bound. These functions have $H=0$ for $p\ge3$, so the gradient's expansion starts at order $d^{-(p-1)/2}$: information exponent $p$.