Disentangling feature and lazy training in deep neural networks
Mario Geiger, Stefano Spigler, Arthur Jacot, Matthieu Wyart
cs.LG, stat.ML
2019-06-19
Varying output scale α maps two infinite-width limits: frozen NTK when α√h ≫ 1, evolving features when α√h ≪ 1. Which wins depends on architecture.
Infinitely wide nets have two clean limits. Last-layer weights scaled as h^{-1/2} freeze the Neural Tangent Kernel (NTK) and make training linear in the parameters. Scaling as h^{-1} yields mean-field dynamics: the distribution of per-neuron parameters follows a PDE and features move. Both limits are theoretically worked out. Which one describes nets people actually train, and which one generalizes better, was still a fight.
Chizat and Bach already noted that multiplying the output by a large factor pushes training into a lazy regime where weights barely move. This paper turns that knob into a sweep: train F(w,x)=α[f(w,x)-f(w0,x)] over the (α,h) plane, and track test error, output fluctuations, and kernel drift.
The main setup is a 3-hidden-layer fully connected net with Softplus activations and continuous-time gradient flow, binary classification on 10k Fashion-MNIST images. The loss is a soft hinge, so training stops in finite time once the data are fitted. α is varied over 11 orders of magnitude. Each point is an ensemble of 10–20 random initializations, so single-model and averaged performance can be separated.
Lazy training is operationally the region where the kernel at the end of training is essentially the kernel at initialization, and test error sits on the frozen-NTK curve. Feature training is the region where the kernel changes by a relative O(1). The control parameter is α√h. MNIST, CIFAR-10, EMNIST, and a CNN trained with ADAM are used as checks.
The two regimes are split by α=Θ(h^{-1/2}). After an affine rescaling of test error, curves for different widths collapse versus α√h. For large α√h, error converges to frozen NTK and stops improving if α is increased further.
Fluctuations of the output around the ensemble mean decay as Var F ∼ 1/h in both regimes, i.e. δF ∼ h^{-1/2}. A large part of the gain from width is just killing initialization noise: the ensemble-averaged test error is essentially flat in h, and that plateau is the single-net h→∞ limit. At fixed α, growing h pushes α√h across 1 and slides the run from feature training into lazy training. That is why ensemble error was previously seen to worsen past some finite width.
The dynamics differ. In feature training there is a time t1∼α√h: for t≪t1 the dynamics stay approximately linear; at t∼t1 the output has grown by a factor of about √h and the kernel starts to move. Eventually ‖ΔΘ‖/‖Θ0‖ ∼ (α√h)^{-a} with a<2 for ReLU and Softplus, and a→2 as depth grows. The measured a≈1.3 for three hidden layers; a scaling argument gives a=2/(1+1/L)=1.66. In the lazy regime the kernel change scales as 1/(α√h), as NTK fluctuation theory predicts. ReLU has an extra intermediate window where the kernel change goes as α^{-1/2}.
Which regime wins depends on architecture and data. Fully connected nets with gradient flow prefer lazy training on Fashion-MNIST, MNIST, EMNIST, and CIFAR-10. A 4-layer CNN with ADAM prefers feature training on CIFAR-10; the same CNN on Fashion-MNIST peaks at intermediate α. Projecting MNIST onto its first 10 principal components slightly favors feature training.
Any claim about infinite width or about NTK explaining deep learning needs to say which side of α√h it sits on. That axis changes the conclusion. Practically, widening a fully connected net and ensembling usually lands in the lazy regime; a conv net that wants feature learning should push the output scale down near α∼h^{-1/2}. The h^{-1/2} fluctuations hold on both sides, so averaging independently trained copies at moderate width buys much of what extra width would have bought.
This is a phase diagram, not a new algorithm. The diagram is the contribution: kernel motion becomes a measurable function of α√h.
The backbone is fully connected, Softplus, gradient flow, binary labels, 10k points, far from ImageNet with cross-entropy and ADAM. CNN numbers are checks, not a systematic scan. The predicted a=2/(1+1/L) still misses the measured 1.3; equal amplification across layers is a simplification. Below α√h∼10^{-4} the variance blows up and is left unstudied. Later work found that as dataset size p grows, a CNN's learning-curve exponent can favor feature training; the 10^4-sample ranking here does not automatically extend to full datasets.