The Onset of Variance-Limited Behavior for Networks in the Lazy and Rich Regimes
This work addresses the limitations of infinite-width approximations for neural networks in practical settings, which is important for researchers and practitioners in machine learning, though it is incremental in nature.
The paper investigates the transition from infinite-width to variance-limited behavior in wide neural networks as training set size increases, finding that finite-size effects become relevant at a critical sample size scaling as the square root of network width, and that feature learning or ensemble averaging can mitigate these effects.
For small training set sizes $P$, the generalization error of wide neural networks is well-approximated by the error of an infinite width neural network (NN), either in the kernel or mean-field/feature-learning regime. However, after a critical sample size $P^*$, we empirically find the finite-width network generalization becomes worse than that of the infinite width network. In this work, we empirically study the transition from infinite-width behavior to this variance limited regime as a function of sample size $P$ and network width $N$. We find that finite-size effects can become relevant for very small dataset sizes on the order of $P^* \sim \sqrt{N}$ for polynomial regression with ReLU networks. We discuss the source of these effects using an argument based on the variance of the NN's final neural tangent kernel (NTK). This transition can be pushed to larger $P$ by enhancing feature learning or by ensemble averaging the networks. We find that the learning curve for regression with the final NTK is an accurate approximation of the NN learning curve. Using this, we provide a toy model which also exhibits $P^* \sim \sqrt{N}$ scaling and has $P$-dependent benefits from feature learning.