The Onset of Variance-Limited Behavior for Networks in the Lazy and Rich Regimes

Introduction

Deep learning systems are achieving state of the art performance on a variety of tasks (Tan & Le, 2019; Hoffmann et al., 2022). Exactly how their generalization is controlled by network architecture, training procedure, and task structure is still not fully understood. One promising direction for deep learning theory in recent years is the infinite-width limit. Under a certain parameterization, infinite-width networks yield a kernel method known as the neural tangent kernel (NTK) (Jacot et al., 2018; Lee et al., 2019). Kernel methods are easier to analyze, allowing for accurate prediction of the generalization performance of wide networks in this regime (Bordelon et al., 2020; Canatar et al., 2021; Bahri et al., 2021; Simon et al., 2021). Infinite-width networks can also operate in the mean-field regime if network outputs are rescaled by a small parameter $\alpha$ that enhances feature learning (Mei et al., 2018; Chizat et al., 2019; Geiger et al., 2020b; Yang & Hu, 2020; Bordelon & Pehlevan, 2022).

While infinite-width networks provide useful limiting cases for deep learning theory, real networks have finite width. Analysis at finite width is more difficult, since predictions are dependent on the initialization of parameters. While several works have attempted to analyze feature evolution and kernel statistics at large but finite width (Dyer & Gur-Ari, 2020; Roberts et al., 2021), the implications of finite width on generalization are not entirely clear. Specifically, it is unknown at what value of the training set size $P$ the effects of finite width become relevant, what impact this critical $P$ has on the learning curve, and how it is affected by feature learning.

To identify the effects of finite width and feature learning on the deviation from infinite width learning curves, we empirically study neural networks trained across a wide range of output scales $\alpha$, widths $N$, and training set sizes $P$ on the simple task of polynomial regression with a ReLU neural network. Concretely, our experiments show the following:

Learning curves for polynomial regression transition exhibit significant finite-width effects very early, around $P\sim\sqrt{N}$. Finite-width NNs at large $\alpha$ are always outperformed by their infinite-width counterparts. We show this gap is driven primarily by variance of the predictor over initializations (Geiger et al., 2020a). Following prior work (Bahri et al., 2021), we refer to this as the variance-limited regime. We compare three distinct ensembling methods to reduce error in this regime.

Feature-learning NNs show improved generalization both before and after the transition to the variance limited regime. Feature learning can be enhanced through re-scaling the output of the network by a small scalar $\alpha$ or by training on a more complex task (a higher-degree polynomial). We show that alignment between the final NTK and the target function on test data improves with feature learning and sample size.

We demonstrate that the learning curve for the NN is well-captured by the learning curve for kernel regression with the final empirical NTK, eNTKf , as has been observed in other works (Vyas et al., 2022; Geiger et al., 2020b; Atanasov et al., 2021; Wei et al., 2022).

Using this correspondence between the NN and the final NTK, we provide a cursory account of how fluctuations in the final NTK over random initializations are suppressed at large width $N$ and large feature learning strength. In a toy model, we reproduce several scaling phenomena, including the $P\sim\sqrt{N}$ transition and the improvements due to feature learning through an alignment effect.

We validate that these effects qualitatively persist in the realistic setting of wide ResNets Zagoruyko & Komodakis (2017) trained on CIFAR in appendix E.

Overall, our results indicate that the onset of finite-width corrections to generalization in neural networks become relevant when the scale of the variance of kernel fluctuations becomes comparable to the bias component of the generalization error in the bias-variance decomposition. The variance contribution to generalization error can be reduced both through ensemble averaging and through feature learning, which we show promotes higher alignment between the final kernel and the task. We construct a model of noisy random features which reproduces the essential aspects of our observations.

Geiger et al. (2020a) analyzed the scaling of network generalization with the number of model parameters. Since the NTK fluctuates with variance $O(N^{-1})$ for a width $N$ network (Dyer & Gur-Ari, 2020; Roberts et al., 2021), they find that finite width networks in the lazy regime generically perform worse than their infinite width counterparts.

The scaling laws of networks over varying $N$ and $P$ were also studied, both empirically and theoretically by Bahri et al. (2021). They consider two types of learning curve scalings. First, they describe resolution-limited scaling, where either training set size or width are effectively infinite and the scaling behavior of generalization error with the other quantity is studied. There, the scaling laws can been obtained by the theory in Bordelon et al. (2020). Second, they analyze variance-limited scaling where width or training set size are fixed to a finite value and the other parameter is taken to infinity. While that work showed for any fixed $P$ that the learning curve converges to the infinite width curve as $O(N^{-1})$, these asymptotics do not predict, for fixed $N$, at which value of $P$ the NN learning curve begins to deviate from the infinite width theory. This is the focus of our work.

The contrast between rich and lazy networks has been empirically studied in several prior works. Depending on the structure of the task, the lazy regime can have either worse (Fort et al., 2020) or better (Ortiz-Jiménez et al., 2021; Geiger et al., 2020b) performance than the feature learning regime. For our setting, where the signal depends on only a small number of relevant input directions, we expect representation learning to be useful, as discussed in (Ghorbani et al., 2020; Paccolat et al., 2021b). Consequently, we posit and verify that the rich network will outperform the lazy one.

Our toy model is inspired by the literature on random feature models. Analysis of generalization for two layer networks at initialization in the limit of high dimensional data have been carried out using techniques from random matrix theory (Mei & Montanari, 2022; Hu & Lu, 2020; Adlam & Pennington, 2020a; Dhifallah & Lu, 2020; Adlam & Pennington, 2020b) and statistical mechanics (Gerace et al., 2020; d’Ascoli et al., 2020; 2020). Several of these works have identified that when $N$ is comparable to $P$, the network generalization error has a contribution from variance over initial parameters. Further, they provide a theoretical explanation of the benefit of ensembling predictions of many networks trained with different initial parameters. Recently, Ba et al. (2022) studied regression with the hidden features of a two layer network after taking one step of gradient descent, finding significant improvements to the learning curve due to feature learning. Zavatone-Veth et al. (2022) analyzed linear regression for Bayesian deep linear networks with width $N$ comparable to sample size $P$ and demonstrated the advantage of training multiple layers compared to only training the only last layer, finding that feature learning advantage has leading correction of scale $(P/N)^{2}$ at small $P/N$.

Problem Setup and Notation

Empirical Results

In this section, we will study learning curves for ReLU NNs trained on polynomial regression tasks of varying degrees. We take our task to be learning $y=Q_{k}(\bm{\beta}\cdot\bm{x})$ where $\bm{\beta}$ is random vector of norm $1/D$ and $Q_{k}$ is the $k$th gegenbauer polynomial. We will establish the following key observations, which we will set out to theoretically explain in Section 4.

Both eNTK0 and sufficiently lazy networks perform strictly worse than NTK∞ , but the ensembled predictors approach the NTK∞ test error.

NNs in the feature learning regime of small $\alpha$ can outperform NTK∞ for an intermediate range of $P$. Over this range, the effect of ensembling is less notable.

Even richly trained finite width NNs eventually perform worse than NTK∞ at sufficiently large $P$. However, these small $\alpha$ feature-learning networks become variance-limited at larger $P$ than lazy networks. Once in the variance-limited regime, all networks benefit from ensembling over initializations.

For all networks, the transition to the variance-limited regime begins at a $P^{*}$ that scales sub-linearly with $N$. For polynomial regression, we find $P^{*}\sim\sqrt{N}$.

These findings support our hypothesis that finite width introduces variance in eNTK0 over initializations, which ultimately leads to variance in the learned predictor and higher generalization error. Although we primarily focus on polynomial interpolation tasks in this paper, in Appendix F we provide results for wide ResNets trained on CIFAR and observe that rich networks also outperform lazy ones, and that lazy ones benefit more significantly from ensembling.

In this section, we first investigate how finite width NN learning curves differ from infinite width NTK regression. In Figure 1 we show the generalization error $E_{g}(f^{*}_{\theta_{0},\mathcal{D}})$ for a depth 3 network with width $N=1000$ trained on a quadratic $k=2$ and quartic $k=4$ polynomial regression task. Additional plots for other degree polynomials are provided in Appendix F. We sweep over $P$ to show the effect of more data on generalization, which is the main relationship we are interested in studying. For each training set size we sweep over a grid of 20 random draws of the train set and 20 random network initializations. This for 400 trained networks in total at each choice of $P,k,N,\alpha$. We see that a discrepancy arises at large enough $P$ where the neural networks begin to perform worse than NTK∞ ​.

We probe the source of the discrepancy between finite width NNs and NTK∞ by ensemble averaging network predictions $\bar{f}_{\mathcal{D}}(\bm{x}):=\langle f_{\theta_{0},\mathcal{D}}^{*}(\bm{x})\rangle_{\theta_{0}}$ over $E=20$ initializations $\theta_{0}$. In Figures 1b and 1d, we calculate the error of $\bar{f}_{\mathcal{D}}(\bm{x})$, each trained on the same dataset. We then plot $E_{g}(\bar{f}_{\mathcal{D}})$. This ensembled error approximates the bias in a bias-variance decomposition (Appendix B). Thus, any gap between 1 (a) and 1 (b) is driven by variance of $f_{\theta,\mathcal{D}}$ over $\theta$.

In addition to initialization variance, variance over dataset $\mathcal{D}$ contributes to the total generalization error. Following (Adlam & Pennington, 2020b), we discuss a symmetric decomposition of the variance in Appendix B, showing the contribution from dataset variance and the effects of bagging. We find that most of the variance in our experiments is due to initialization.

We show several other plots of the results of these studies in the appendix. We show the effect of bagging (Figure 7), phase plots of different degree target functions (Figures 10, 9), phase plots over $N,\alpha$ (Figure 11) and a comparison of network predictions against the initial and final kernel regressors (Figures 18, 19).

2 Final NTK Variance leads to Generalization Plateau

In this section, we show how the variance over initialization can be interpreted as kernel variance in both the rich and lazy regimes. We also show how this implies a plateau for the generalization error.

To begin, we demonstrate empirically that all networks have the same generalization error as kernel regression solutions with their final eNTKs. At large $\alpha$, the initial and the final kernel are already close, so this follows from earlier results of Chizat et al. (2019). In the rich regime, the properties of the eNTKf have been studied in several prior works. Several have empirically demonstrated that the eNTKf is a good match to the final network predictor for a trained network (Long, 2021; Vyas et al., 2022; Wei et al., 2022) while others have given conditions under which such an effect would hold true (Atanasov et al., 2021; Bordelon & Pehlevan, 2022). We comment on this in appendix C.4. We show in Figure 3 how the final network generalization error matches the generalization error of eNTKf . As a consequence, we can use eNTKf to study the observed generalization behavior.

Next, we relate the variance of the final predictor $f_{\theta_{0},\mathcal{D}}^{*}$ to the corresponding infinite width network $f^{\infty}_{\mathcal{D}}$. The finite size fluctuations of the kernel at initialization have been studied in (Dyer & Gur-Ari, 2020; Hanin & Nica, 2019; Roberts et al., 2021). The variance of the kernel elements has been shown to scale as $1/N$. We perform the following bias-variance decomposition: Take $f_{\theta_{0},\mathcal{D}}$ to be the eNTK0 predictor, or a sufficiently lazy network trained to interpolation on a dataset $\mathcal{D}$. Then,

We demonstrate this equality using a relationship between the infinite-width network and an infinite ensemble of finite-width networks derived in Appendix B. There we also show that the $O(1/N)$ term is strictly positive for sufficiently large $N$. Thus, for lazy networks of sufficiently large $N$, finite width effects lead to strictly worse generalization error. The decomposition in Equation 2 continues to hold for rich networks at small $\alpha$ if $f^{\infty}$ is interpreted as the infinite-width mean field limit. In this case one can show that ensembles of rich networks are approximating an infinite width limit in the mean-field regime. See Appendix B for details.

3 Feature Learning delays variance limited transition

We now consider how feature learning alters the onset of the variance limited regime, and how this onset scales with $\alpha,N$. We define the onset of the variance limited regime to take place at the value $P^{*}=P_{1/2}$ where over half of the generalization error is due to variance over initializations. Equivalently we have $E_{g}(\bar{f}^{*})/E_{g}(f^{*})=1/2$. By using an interpolation method together with bisection, we solve for $P_{1/2}$ and plot it in Figure 4.

We can understand the delay of the variance limited transition, as well as the lower value of the final plateau using a mechanistic picture similar to the effect observed in Atanasov et al. (2021). In that setting, under small initialization, the kernel follows a deterministic trajectory, picking up a low rank component in the direction of the train set targets $\bm{y}\bm{y}^{\top}$, and then changing only in scale as the network weights grow to interpolate the dataset. In their case, for initial output scale $\sigma^{L}$, eNTKf is deterministic up to a variance of $O(\sigma)$. In our case, the kernel variance at initialization scales as $\sigma^{2L}/N$. As $\sigma\to 0$ the kernel’s trajectory becomes deterministic up to a variance term scaling with $\sigma$ as $O(\sigma)$, which implies that the final predictor also has a variance scaling as $O(\sigma)$.

Signal plus noise correlated feature model

In Section 3.2 we have shown that in both the rich and lazy regimes, the generalization error of the NN is well approximated by the generalization of a kernel regression solution with eNTKf ​. This finding motivates an analysis of the generalization of kernel machines which depend on network initialization $\theta_{0}$. Unlike many analyses of random feature models which specialize to two layer networks and focus on high dimensional Gaussian random data (Mei & Montanari, 2022; Adlam & Pennington, 2020a; Gerace et al., 2020; Ba et al., 2022), we propose to analyze regression with the eNTKf for more general feature structures. This work builds on the kernel generalization theory for kernels developed with statistical mechanics (Bordelon et al., 2020; Canatar et al., 2021; Simon et al., 2021; Loureiro et al., 2021). We will attempt to derive approximate learning curves in terms of the eNTKf ​’s signal and noise components, which provide some phenomenological explanations of the onset of the variance limited regime and the benefits of feature learning. Starting with the final NTK $K_{\theta_{0}}({\bm{x}},{\bm{x}}^{\prime})$ which depends on the random initial parameters $\theta_{0}$, we project its square root $K^{1/2}_{\theta_{0}}({\bm{x}},{\bm{x}}^{\prime})$ (as defined in equation 32) on a fixed basis $\{b_{k}(x)\}_{k=1}^{\infty}$ orthonormal with respect to $p({\bm{x}})$. This defines a feature map

To gain insight into the role of feature noise, we characterize the test error associated with a Gaussian covariate model in a high dimensional limit $P,M,N_{\mathcal{H}}\to\infty$ with $\alpha=P/M,\eta=N_{\mathcal{H}}/M$.

This model was also studied by Loureiro et al. (2021) and subsumes the classic two layer random feature models of prior works (Hu & Lu, 2020; Adlam & Pennington, 2020a; Mei & Montanari, 2022). The expected generalization error for any distribution of ${\bm{A}}(\bm{\theta}_{0})$ has the form

where $\alpha=P/M$ and $\gamma=\frac{\alpha}{(\lambda+q)^{2}}\text{Tr}{\bm{G}}^{2}[{\bm{A}}\bm{\Sigma}_{M}{\bm{A}}^{\top}+\bm{\Sigma}_{\epsilon}]^{2}$. Details of the calculation can be found in Appendix D. We also provide experiments showing the predictive accuracy of the theory in Figure 6. In general, we do not know the induced distribution of ${\bm{A}}(\theta_{0})$ over disorder $\theta_{0}$. In Appendix D.5, we compute explicit learning curves for a simple toy model where ${\bm{A}}(\bm{\theta}_{0})^{\prime}s$ entries as i.i.d. Gaussian over the random initialization $\bm{\theta}_{0}$. A similar random feature model was recently analyzed with diagrammatic techniques by Maloney et al. (2022). In the high dimensional limit $M,P,N_{\mathcal{H}}\to\infty$ with $P/M=\alpha,N_{\mathcal{H}}/M=\eta$, our replica calculation demonstrates that test error is self-averaging (the same for every random instance of ${\bm{A}}$) which we describe in Appendix D.5 and Figure 16.

2 Explaining Feature Learning Benefits and Error Plateaus

Using our model, we can also approximate the role of feature learning as enhancement in the signal correlation along task-relevant eigenfunctions. In Figure 6 (d) we plot the learning curves for networks trained with different levels of feature learning, controlled by $\alpha$. We see that feature learning leads to improvements in the learning curve both before and after onset of variance limits. In Figure 6 (e)-(f), we plot the theoretical generalization for kernels with enhanced signal eigenvalue for the task eigenfunction $y({\bm{x}})=\phi_{k}({\bm{x}})$. This enhancement, based on the intuition of kernel alignment, leads to lower bias and lower asymptotic variance. However, this model does not capture the fact that feature learning advantages are small at small $P$ and that the slopes of the learning curves are different at different $\alpha$. Following the observation of Paccolat et al. (2021a) that kernel alignment can occur with scale $\sqrt{P}$, we plot the learning curves for signal enhancements that scale as $\sqrt{P}$. Though this toy model reproduces the onset of the variance limited regime $P_{1/2}$ and the reduction in variance due to feature learning, our current result is not the complete story. A more refined future theory could use the structure of neural architecture to constrain the structure of the ${\bm{A}}$ distribution.

Conclusion

We performed an extensive empirical study for deep ReLU NNs learning a fairly simple polynomial regression problems. For sufficiently large dataset size $P$, all neural networks under-perform the infinite width limit, and we demonstrated that this worse performance is driven by initialization variance. We show that the onset of the variance limited regime can occur early in the learning curve with $P_{1/2}\sim\sqrt{N}$, but this can be delayed by enhancing feature learning. Finally, we studied a simple random-feature model to attempt to explain these effects and qualitatively reproduce the observed behavior, as well as quantitatively reproducing the relevant scaling relationship for $P_{1/2}$. This work takes a step towards understanding scaling laws in regimes where finite-size networks undergo feature learning. This has implications for how the choice of initialization scale, neural architecture, and number networks in an ensemble can be tuned to achieve optimal performance under a fixed compute and data budget.

References

Appendix A Details on Experiments

We used JAX (Bradbury et al., 2018) for all neural network training. We built multi-layer perceptrons (MLPs) of depth 2 and 3. Most of the results are reported for depth 3 perceptrons, where there is a separation between the width of the network $N$ and the number of parameters $N^{2}$. Sweeping over more depths and architectures is possible, but because of the extensive dimensionality of the hyperparameter search space, we have not yet experimented with deeper networks.

We considered MLPs with no bias terms. Since the Gegenbauer polynomials are mean zero, we do not need biases to fit the training set and generalize well. We have also verified that adding trainable biases does not change the final results in any substantial way.

As mentioned in the main text, we consider the final output function to be the initial network output minus the output at initialization:

Here, only $\theta$ is differentiated through, while $\theta_{0}$ is held fixed. The rationale for this choice is that without this subtraction, in the lazy limit the trained neural network output can be written as

We trained this network with full batch gradient descent with a learning rate $\eta$ so that

Each network was trained to an interpolation threshold of $10^{-6}$. If a network could not reach this threshold in under 30k steps, we checked if the training error was less than $10$ times the generalization error. If this was not satisfied, then that run of the network was discarded.

For each fixed $P,k$, we generated 20 independent datasets. For each fixed $N,\alpha$ we generated 20 independent neural network initializations. This $20\times 20$ table yields a total of 400 neural networks trained on every combination of initialization and dataset choice.

The infinite width network predictions were calculated using the Neural Tangents package (Novak et al., 2020). The finite width eNTK0 ​s were also calculated using the empirical methods in Neural Tangents. They were trained to interpolation using the gradient_descent_mse method. This is substantially faster than training the linearized model using standard full-batch gradient descent, which we have found to take a very long time for most networks. We use the same strategy for the eNTKf ​s.

For the experiments in the main text, we have taken the input dimension to be $D=10$ and sweep over $k=1,2,3,4$. We swept over 15 values $P$ in logspace from size $30$ to size 10k, and over 6 values of $N$ in logspace from size $30$ to size $2150$. We then swept over alpha values $0.1,0.5,1.0,10.0,20.0$. Depending on $\alpha,N$, we tuned the learning rate $\eta$ of the network small enough to stay close to the gradient flow limit, but allow for the interpolation threshold to be feasibly reached.

For each of the 1800 settings of $P,N,\alpha,k$ and each of the 400 networks, 400 eNTK0 ​s, 400 eNTKf ​s, and 20 NTK∞ ​s, the generalization error was saved, as well as a vector of $\hat{y}$ predictions on a test set of 2000 points. In addition, for the neural networks we saved both initial and final parameters. All are saved as lists of numpy arrays in a directory of about 1TB. We plan to make the results of our experiments publicly accessible, alongside the code to generate them.

We apply the same methodology of centering the network and allowing $\alpha$ to control the degree of laziness by redefining

We consider the task of binary classification for CIFAR-10. In order to allow $P$ to become large we divide the data into two classes: animate and inanimate objects. We choose to subsample eight classes and superclass them into two: (cat, deer, dog, horse) vs (airplane, automobile, ship, truck). Each superclass consists of 20,000 training examples and 4,000 test examples retrieved from the CIFAR-10 dataset.

On subsets of this dataset, we train wide residual networks (ResNets) Zagoruyko & Komodakis (2017) of width $64$ and block size $1$ with the NTK parameterization Jacot et al. (2018) on this task using mini-batch gradient descent with batch size of 256 and MSE loss. Step sizes are governed by the Adam optimizer Kingma & Ba (2014) with initial learning rate $\eta_{0}=10^{-3}.$ Every network is trained for 24,000 steps, such that under nearly all settings of $\alpha$ and dataset size the network has attained infinitesimal train loss.

We sweep $\alpha$ from $10^{-3}$ to $10^{0}$ and $P$ from $2^{9}$ to $2^{15}$. For each value of $P$, we randomly sample five training datasets of size $P$ and compute ensembles of size 20. For each network in an ensemble the initialization and the order of the training data is randomly chosen independently of those for the other networks.

Appendix B Fine-grained bias-variance decomposition

Let $\mathcal{D}$ be a dataset of $(\bm{x}^{\mu},y^{\mu})_{\mu=1}^{P}\sim p(\bm{x},y)$ viewed as a random variable. Let $\theta_{0}$ represent the initial parameters of a neural network, viewed as a random variable. In the case of no label noise, as in section 2.2.1 of Adlam & Pennington (2020b), we derive the symmetric decomposition of the generalization error in terms of the variance due to initialization and the variance due to the dataset. We have

$V_{\mathcal{D}}$ and $V_{\theta_{0}}$ give the components of the variance explained by variance in $\mathcal{D},\theta_{0}$ respectively. $V_{\mathcal{D},\theta_{0}}$ is the remaining part of the variance not explained by either of these two sources. As in the main text, $\bar{f}^{*}_{\mathcal{D}}(\bm{x})$ is the ensemble average of the trained predictors over initializations. $E_{\mathcal{D}}[f^{*}_{\theta_{0},\mathcal{D}}(\bm{x})|\theta_{0}]$ is commonly referred to as the bagged predictor. In the next subsection we study these terms empirically.

B.2 Empirical Study of Dataset Variance

Using the network simulations, one can show that the bagged predictor does not have substantially lower generalization error in the regimes that we are interested in. This implies that most of the variance driving higher generalization error is due to variance over initializations. In figure 7, we make phase plots of the fraction of $E_{g}$ that arises from variance due to initialization, variance over datasets, and total variance for width 1000. This can be obtained by computing the ensembled predictor, the bagged predictor, and the ensembled-bagged predictor respectively.

B.3 Relating Ensembled Network Generalization to Infinite Width Generalization

Making use of the fact that at leading order, the eNTKf (either in the rich or lazy regime) of a trained network has $\theta_{0}$-dependent fluctuations with variance $1/N$, one can write the kernel Gram matrices as

Here, $\bm{\delta K}_{\theta_{0}},\bm{\delta k}_{\theta_{0}}$ are the leading order fluctuations around the infinite width network. Because of how we have written them, their variance is $O(1)$ with respect to $N$. Using perturbation theory (Dyer & Gur-Ari, 2020), one can demonstrate that these leading order terms have mean zero around their infinite-width limit.

The predictor for the eNTK0 (or for a sufficiently large $\alpha$ neural network) for a training set with target labels $\bm{y}$ is given by:

Upon taking the ensemble, because of the mean zero property of the deviations, we get that

We can now bound the generalization error of the ensemble of networks in terms of the infinite-width generalization:

By equation 18, the second term yields a positive contribution going as $O(N^{-2})$. The last term can be bounded by Cauchy-Schwarz:

After we enter the variance limited regime by taking $P>P_{1/2}$ we get $E_{g}^{\infty}\leq O(1/N)$ so this last term is bounded by $N^{-3/2}$. Consequently, the difference in generalization error between the infinite width NTK and an ensemble of lazy network or eNTK0 predictors is subleading in $1/N$ compared to the generalization gap, which goes as $N^{-1}$.

The same argument can be extended to any predictor that differs from some infinite width limit. In particular Bordelon & Pehlevan (2022) show that the fluctuations of the eNTKf in any mean field network are asymptotically mean zero with variance $N^{-1}$. The above argument then applies to the predictor obtained by ensembling networks that have learned features. This implies that in the variance limited regime, ensemble averages of feature learning networks have the same generalization as the infinite-width mean field solutions up to a term that decays faster than $N^{-3/2}$.

Appendix C Feature Learning

After appropriately rescaling learning rate to $\eta=\sigma^{-2L}$ we get

Under the assumption that $\sigma^{L}\ll 1$ and $y^{\mu}=O(1)$ so that the error term is $O(1)$ we get that the output changes in time as $df/dt=O(1)$.

On the other hand, using the chain rule one can show that the features change as a product of the gradient update and the features in the prior layer, yielding the scaling

This gives us that the change in the features scales as $(\alpha\sqrt{N})^{-1}$ while the change in the output scales as $O(1)$. Thus, for $\alpha\sqrt{N}$ sufficiently small, the features can move dramatically.

C.2 Output Rescaling without Rescaling Weights

C.3 Kernel Alignment

In this section we comment on our choice of kernel alignment metric

For kernels that are diagonally dominant, such as those encountered in the experiments, this metric is related to another alignment metric

Here $|\bm{K}|_{F}$ is the Frobenius norm of the Gram matrix of the kernel. This metric was extensively used in Baratin et al. (2021). The advantage of the first metric over the second is that one can quickly estimate the denominator of $A(\bm{K})$ via Monte Carlo estimation of $\langle\bm{u}^{\top}\bm{K}\bm{u}\rangle_{\bm{u}\sim\mathcal{N}(0,\bm{1})}$.

We use $A(\bm{K}_{f})$ as a measure of feature learning, as we have found that this more finely captures elements of feature learning than other related metrics. We list several metrics we tried that did not work.

One option for a representation-learning metric involves measuring the magnitude of the change between the initial and final kernels, $\bm{K}_{i},\bm{K}_{f}$:

However, this is more sensitive to the raw parameter change than any task-relevant data. If one instead were to normalize the kernels to be unit norm at the beginning and the end, the modified metric

This metric however remains remarkably flat over the whole range of $\alpha,P$, as does the centered kernel alignment (CKA) of Cortes et al. (2012)

Here $C$ is the centering matrix that subtracts off the mean components of the kernel for a $P\times P$ kernel. This alignment metric has been shown to be useful in comparing neural representations (Kornblith et al., 2019). For our task, however, because the signal is low-dimensional, only a small set of eigenspaces of the kernel align to this task. As a result, the CKA, which counts all eigenspaces equally, appears to be too coarse to capture the low-dimensional feature learning that is happening.

On the other hand, we find that $A(\bm{K}_{f})$ (with $\bm{K}_{f}$ given by the eNTKf evaluated on a test set) can very finely detect alignment along the task relevant directions. This produces a clear signal of feature learning at small $\alpha$ and large $P$ as shown in Figure 2c.

$A(\bm{K}_{f})$ can be related to the centered kernel alignment between the eNTKf and the (mean zero) task kernel $\bm{y}\bm{y}^{\top}$, where $\bm{y}$ is a vector of draws from the population distribution $p(\bm{x},y)$.

C.4 Relationship Between Trained Network and Final Kernel

In general, the learned function contains contributions from the instantaneous NTKs at every point in the training. Concretely, following Atanasov et al. (2021) we have the following formula for the final network predictor $f(x)$

where $[\bm{k}(x,t)]_{\mu}=K(x,x_{\mu},t)$ and $[\bm{K}(s)]_{\mu\nu}=K(x_{\mu},x_{\nu},s)$ and $[\bm{y}]_{\mu}=y_{\mu}$. In general there are contributions from earlier kernels $\bm{k}(x,t)$ for $t<\infty$ and so the function $f$ cannot always be written as a linear combination of the final NTK $K_{f}$ on training data: $f=\sum_{\mu}\alpha_{\mu}K_{f}(x,x_{\mu})$. However, as Vyas et al. (2022); Atanasov et al. (2021) have shown, the final predictions of the network are often well modeled by regression with the final NTK. We verify this for our task in section 3.2.

Appendix D Generic Random Feature Model

For a random kernel, $K({\bm{x}},{\bm{x}}^{\prime};\theta)$, we first compute its Mercer decomposition

From the eigenvalues $\lambda_{k}$ and eigenfunctions $\phi_{k}$, we can construct the square root

Lastly, using $K^{1/2}$, we can get a feature map by projecting against a static basis $\{b_{k}\}$ giving

These features reproduce the kernel so that $K({\bm{x}},{\bm{x}}^{\prime};\theta)=\sum_{k}\psi_{k}({\bm{x}})\psi_{k}({\bm{x}}^{\prime})$. This can be observed from the following observation

where the last line follows from the orthogonality of $U_{km}$ and recovers $K(\bm{x},\bm{x}^{\prime};\theta)$.

D.2 Decomposition of Finite Width Features

We now attempt to characterize the variance in the features over the sample distribution. We will first consider the case of a fixed realization of $\bm{\theta}_{0}$ before providing a typical case analysis over random $\bm{\theta}_{0}$. For a fixed initialization $\bm{\theta}_{0}$ we define the following covariance matrices

where $\bm{\psi}_{M}$ are the truncated (but deterministic) features induced by the deterministic infinite width kernel. We will mainly be interested in the case where $M\to\infty$ and where the target function can be expressed as the linear combination $y({\bm{x}})={\bm{w}}^{*}\cdot\bm{\psi}_{M}({\bm{x}})$ of these features. For example, in the case of our experiments on the sphere, $\bm{\psi}_{M}$ could be the spherical harmonic functions. Further, in the $M\to\infty$ limit, we will be able to express the target features $\bm{\psi}$ as linear combinations of the features $\bm{\psi}_{M}$

The matrix $\bm{A}(\bm{\theta}_{0})$ are the coefficients of the decomposition which can vary over initializations. Crucially $\bm{A}(\bm{\theta}_{0})$ projects to the subspace of dimension $N_{\mathcal{H}}$ where the finite width features have variance over ${\bm{x}}$. The population risk for this $\bm{\theta}_{0}$ has an irreducible component

where the bound is tight for the optimal weights ${\bm{w}}=\left({\bm{A}}(\bm{\theta}_{0})\bm{\Sigma}_{M}{\bm{A}}(\bm{\theta}_{0})^{\top}\right)^{-1}{\bm{A}}(\bm{\theta}_{0})\bm{\Sigma}_{M}{\bm{w}}^{*}$. The irreducible error is determined by a projection matrix which preserves the subspace where the features $\bm{\psi}({\bm{x}},\bm{\theta}_{0})$ have variance: ${\bm{I}}-{\bm{A}}(\bm{\theta})^{\top}\left({\bm{A}}(\bm{\theta}_{0})\bm{\Sigma}_{M}{\bm{A}}(\bm{\theta}_{0})^{\top}\right)^{-1}{\bm{A}}(\bm{\theta}_{0})\bm{\Sigma}_{M}$. In general, this will preserve some fraction of the variance in the target function, but some variance in the target function will not be expressible by linear combinations of the features $\bm{\psi}({\bm{x}},\bm{\theta})$. We expect that random finite width $N$ neural networks will have unexplained variance in the target function on the order $\sim 1/N$.

D.3 Gaussian Covariate Model

Following prior works on learning curves for kernel regression (Bordelon et al., 2020; Canatar et al., 2021; Loureiro et al., 2021), we will approximate the learning problem with a Gaussian covariates model with matching second moments.

The features $\bm{\psi}_{M}({\bm{x}})$ will be treated as Gaussian over random draws of datapoints. We will assume centered features. We decompose the features in the orthonormal basis ${\bm{b}}({\bm{x}})$, which we approximate as a Gaussian vector ${\bm{b}}\sim\mathcal{N}(0,{\bm{I}})$.

We refer readers to Hu & Lu (2020) for a discussion of this equivalence between random feature regression and this Gaussian covariate model.

D.4 Replica Calculation of the Learning Curve

To analyze the typical case performance of kernel regression, we define the following partition function which is dominated

For proper normalization, we assume that $\left<\bm{\psi}_{M}\bm{\psi}_{M}^{\top}\right>=\frac{1}{M}\bm{\Sigma}_{M}$ and $\left<\bm{\epsilon}\bm{\epsilon}^{\top}\right>=\frac{1}{M}\bm{\Sigma}_{\epsilon}$. We note that in the $\beta\to\infty$ limit, the partition function is dominated by the unique minimizer of the regularized least squares objective (Canatar et al., 2021; Loureiro et al., 2021). Further, for a fixed realization of $\bm{\theta}_{0}$ the average generalization error over datasets $\mathcal{D}$ can be computed by differentiation of the source term $J$

Thus the $\beta\to\infty$ limit of the above quantity will give the expected generalization error of the risk minimizer. We see the need to average the quantity $\ln Z$ over realizations of datasets $\mathcal{D}$. For this, we resort to the replica trick $\left<\ln Z\right>=\lim_{n\to 0}\frac{1}{n}\ln\left<Z^{n}\right>$. We will compute the integer moments $\left<Z^{n}\right>$ for integer $n$ and then analytically continue the resulting expressions to $n\to 0$ under a symmetry ansatz. The replicated partition function thus has the form

We now need to perform the necessary average over the random realizations of data points $\mathcal{D}=\{{\bm{b}}_{\mu},\bm{\epsilon}_{\mu}\}$. We note that the scalar quantities $h^{a}_{\mu}={\bm{w}}^{a}\cdot\bm{\psi}_{\mu}-{\bm{w}}^{*}\cdot{\bm{\psi}}_{M,\mu}$ are Gaussian with mean zero and covariance

We further see that the generalization error in replica $a$ is $E_{g}({\bm{w}}^{a})=Q_{aa}$. Performing the Gaussian integral over $\{h^{a}_{\mu}\}$ gives

To take the $n\to 0$ limit, we make the replica symmetry ansatz

which is well motivated since this is a convex optimization problem. Letting $\alpha=P/N$, we find that under the RS ansatz the replicated partition function has the form

In a limit where $\alpha=P/M$ is $O(1)$, then this $S$ is intensive $S=O_{M}(1)$. We can thus appeal to saddle point integration (method of steepest descent) to compute the set of order parameters which have dominant contribution to the free energy.

The order parameters $q^{*},q_{0}^{*},\hat{q}^{*},\hat{q}_{0}^{*}$ are defined via the saddle point equations $\frac{\partial S}{\partial q}=\frac{\partial S}{\partial q_{0}}=\frac{\partial S}{\partial\hat{q}}=\frac{\partial S}{\partial\hat{q}_{0}}=0$. For our purposes, it suffices to analyze two of these equations

We can now take the zero temperature ($\beta\to\infty$) limit to solve for the generalization error

We see that we need to compute the $J$ derivatives on $\hat{q}$. We let $\kappa=\lambda+q$ and note

We see that this recovers the minimal possible error in the $P\to\infty$ limit. The derived learning curves depend on the instance of random initial condition $\bm{\theta}_{0}$. To get the average case performance, we take an additional average of this expression over $\bm{\theta}_{0}$

This average is complicated since $\gamma,\hat{q},{\bm{G}}$ all depend on $\bm{\theta}_{0}$. In the next section we go beyond this analysis to try average case analysis for random Gaussian ${\bm{A}}$.

D.5 Quenched Average over Gaussian A

In this section we will define a distribution of features which allows an exact asymptotic prediction over random realizations of disorder $\bm{\theta}_{0}$ and datasets $\mathcal{D}$. This is a nontrivial extension of the result of Loureiro et al. (2021) since the number of necessary saddle point equations to be solved doubles from two to four. However, this more complicated theory allows us to exactly compute the expectation in equation 55 under an ansatz for the random matrix ${\bm{A}}$. We construct our features with

We will now perform an approximate average over both datasets $\mathcal{D}$ and realizations of $\bm{A}$

As before, we first average over ${\bm{b}}_{\mu},\bm{\epsilon}_{\mu}|A$ and define order parameters $Q_{ab}$ as before.

where we defined the fields ${\bm{g}}^{a}=\frac{1}{\sqrt{N}}{\bm{A}}{\bm{w}}^{a}$ which are mean zero Gaussian with covariance $\left<{\bm{g}}^{a}{\bm{g}}^{b\top}\right>=V_{ab}{\bm{I}}$ where $V_{ab}=\frac{\sigma^{2}}{N}{\bm{w}}^{a}\cdot{\bm{w}}^{b}$. Performing the Gaussian integral over ${\bm{G}}=\text{Vec}\{{\bm{g}}^{a}\}$, we find

Next, we need to integrate over ${\bm{W}}=\text{Vec}\{{\bm{w}}^{a}\}$ which gives

Now the replicated partition function has the form

Now we make a replica symmetry ansatz on the order parameters ${\bm{Q}},\hat{{\bm{Q}}},{\bm{V}},\hat{{\bm{V}}}$

We introduce the shorthand for normalized trace of a matrix ${\bm{G}}$ as $\text{tr}\ {\bm{G}}=\frac{1}{M}\text{Tr}{\bm{G}}$. Under the replica symmetry ansatz, we find the following free energy

Letting $F=2M^{-1}\left<\ln Z\right>$, the saddle point equations read

Now the generalization error can be determined from

We see that it is necessary to compute $\partial_{J}\hat{q}$ and $\partial_{J}v$ in order to obtain the final result. For simplicity, we set $\sigma^{2}=1$. The equations for the source derivatives are

Once the value of the order parameters $(q,\hat{q},v,\hat{v})$ have been determined, these source derivatives can be obtained by solving a $4\times 4$ linear system. Examples of these solutions are provided in Figure 16.

We can compute the asymptotic ($\alpha\to\infty$) generalization error due to the random projection ${\bm{A}}$ in the limit of $\bm{\Sigma}_{\epsilon}=0$. First, note that if $\hat{v}\to O_{\alpha}(1)$, then the asymptotic error would be zero. Therefore, we will assume that $\hat{v}\sim a\alpha^{c}$ for some $a,c>0$. The saddle point equations give the following asymptotic conditions

For $0<\eta<1$, this equation can only be satisfied as $\alpha\to\infty$ if $c=1$ so that $\hat{v}$ has the same scaling with $\alpha$ as $\hat{q}$. If $c<1$ then we could get the equation $\eta=1$. If $c>1$, then the equation would give $\eta=0$. The constant $a$ solves the equation

Using this fact, our order parameters satisfy the following large $\alpha$ scalings

The source derivative equations simplify to $\partial_{J}\hat{q}\sim 1\ ,\ \partial_{J}q\sim 0$ and

We note that $\partial_{J}\hat{v}$ only depends on the product $a\lambda$ which is an implicit function of $\eta$ and $\bm{\Sigma}_{M}$. The generalization error is $E_{g}=\frac{1}{M}\partial_{J}\hat{q}(1+\hat{v}){\bm{w}}^{*}[(1+\hat{v}){\bm{I}}+\hat{q}\bm{\Sigma}]^{-1}\bm{\Sigma}_{M}{\bm{w}}^{*}$

We see that in the generic case, the asymptotic error has a nontrivial dependence on the task ${\bm{w}}^{*}$ and the correlation structure $\bm{\Sigma}_{M}$. To gain more intuition, we will now consider the special case of isotropic features $\bm{\Sigma}_{M}={\bm{I}}$. In this case, we have $\eta=\frac{1}{1+\lambda a}$ so that $\lambda a=\frac{1-\eta}{\eta}$. This results in the following generalization error

We see that as $\eta=\frac{N_{\mathcal{H}}}{M}\to 1$, the asymptotic error converges to zero since all information in the original features is preserved.

D.5.2 Simplified Isotropic Feature Noise

We can simplify the above expressions somewhat in the case where $\sigma^{2}=1$ and $\bm{\Sigma}_{\epsilon}=\sigma^{2}_{\epsilon}{\bm{I}}$. In this case, the order parameters become

Letting ${\bm{G}}=[(1+\hat{v}+\sigma^{2}_{\epsilon}\hat{q}){\bm{I}}+\hat{q}\bm{\Sigma}_{M}]^{-1}$, the source derivatives have the form

For each $\alpha$, we can solve for $\partial_{J}\hat{q}$ and $\partial_{J}\hat{v}$ to get the final generalization error with the formula

An example of these solutions can be found in Figure 16, where we show good agreement between theory and experiment.

Appendix E ResNet on CIFAR Experiments

Appendix F Additional Experiments