Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel

Diverse aspects of deep learning phenomenology

Recent work has revealed many insights into the shape of loss functions over the high dimensional space of neural network parameters. For example, demonstrates that training even within a random, low-dimensional affine subspace of parameter space can yield a network with low test loss. This suggests that the region of parameter space with low test loss must be a relatively high dimensional object, such that low dimensional random affine hyperplanes can generically intersect it. Moreover, show that different, independently trained networks in weight space with low loss can be connected through nonlinear pathways (found via an optimization process) that never leave this low loss manifold. However, direct linear pathways connecting two such independently trained networks typically always leave the low loss manifold. The loss function restricted to such linear paths then yields a loss barrier at an intermediate point between the two networks. builds and provides evidence for a unifying geometric model of the low-loss manifold consisting of a network of mutually intersecting high dimensional basins (Fig. 1A). Two networks within a basin can be connected by a straight line that never leaves the low-loss manifold, while two networks in different basins can be connected by a piecewise linear path of low loss that is forced to traverse the intersection between two basins. uses these insights to argue that deep ensembles are hard to beat using local subspace sampling methods due to the geometry of this underlying loss landscape. provides further evidence for this large-scale structure by demonstrating that after a very early stage of training of a parent network (but not earlier) two child networks trained starting from the parameters of the parent end up in the same low loss basin at the end of training, and could be connected by a linear path in weight space that does not leave the low loss manifold (Fig. 1B). Furthermore, show that the properties of the final minimum found are strongly influenced by the very early stages of training. Taken together, these results present an intriguing glimpse into the large scale structure of the low loss manifold, and the importance of early training dynamics in determining the final position on the manifold.

Neural tangent kernels, linearized training and the infinite width limit.

The neural tangent kernel (NTK) has garnered much attention as it provides a theoretical foothold to understand deep networks, at least in an infinite width limit with appropriate initialization scale and low learning rate . In such a limit, a network does not move very far in weight space over the course of training, and so one can view learning as a linear process occurring along the tangent space to the manifold of functions $f_{w}$ realizable by the parameters $w$, at the initial function $f^{0}$ (Fig. 1C). This learning process is well described by kernel regression with a certain random kernel associated with the tangent space at initialization. The NTK is also a special case of Taylorized training , which approximates the realizable function space $f_{w}$ to higher order in the vicinity of initialization. Various works compare the training of deep networks to the NTK . In many cases, state of the art networks outperform their random kernel counterparts by significant margins, suggesting that deep learning in practice may indeed explore regions of function space far from initialization, with the tangent space twisting significantly over training time, and hence the kernel being learned from the data (Fig. 1D). However, the nature and extent of this function space motion, the degree of tangent space twisting, and how and when data is infused into a learned tangent kernel, remains poorly understood.

The local geometric structure of neural loss landscapes.

Much effort has gone into characterizing the local geometry of loss landscapes in terms of Hessian curvature and its impact on generalization and learning. Interestingly analyses the Hessian eigenspectrum of loss landscapes at scale, demonstrating that learning leads to the emergence of a small number of large Hessian eigenvalues, and many small ones, bolstering evidence for the existence of many flat directions in low loss regions depicted schematically in Fig. 1A. shows that the gradients of logits with respect to parameters cluster tightly based on the logit over training time, leading directly to the emergence of very sharp Hessian eigenvalues. Moreover, a variety of work has explored relations between the curvature of local minima found by training and their generalization properties , and how learning rate and batch size affect the curvature of the minima found , with larger learning rates generically enabling escape from sharper minima (Fig. 1E). makes a connection between learning rates and the validity of NTK training, showing that for infinitely wide networks, training with a learning rate above a scale determined by the top eigenvalue of the Hessian at initialization results in a learning trajectory that outperforms NTK training, presumably by exploring nonlocal regions of function space far away from initialization.

Towards an integrative view.

Above, we have reviewed previously distinct strands of inquiry into deep learning phenomenology that have made little to no contact with each other. Indeed, we currently have no understanding of how local and global loss geometry interacts with the degree of kernel learning in state of the art architectures and training regimes used in practice. For example, at what point in training is the fate of the final chosen basin in Fig. 1 A,B irrevocably determined? Does the kernel change significantly from initialization as in Fig. 1D? If so, when during training does the tangent kernel start to acquire knowledge of the data? Also, when does kernel learning finally stabilize? What relations do either of these times have to the time at which basin fate is determined? How does local geometry in terms of curvature change as all these events occur? Here we address these questions to obtain an integrative view of the learning process across a range of networks and datasets. While we only present results for ResNet20 trained on CIFAR10 and CIFAR100 in the main paper, in Appendix C we find similar results for a WideResNet, variations of Resnets and a Simple CNN trained on CIFAR10 and CIFAR100, indicating our results hold generally across architectures, datasets and training protocols. Many experimental details are covered in our Appendix.

Definition of measurement metrics for geometry and training

Consider also a second-order Taylor expansion to approximate the change to the logits for input $x_{i}$:

Note, that for an infinitesimal $\eta$, the dynamics in Eq. 1 are those of gradient flow, and terms higher than order 1 in Eq. 2 vanish. In this case, steepest descent in the parameter space corresponds to steepest descent in the function space using a neural tangent kernel (NTK),

Hierarchical exploration of the loss landscape through parents and children.

In order to explore the loss landscape and the stability of training dynamics in a more multiscale hierarchical manner than is possible using completely independent training runs, we employ a method of parent-child spawning (shown schematically in Fig. 1B). In this process, a parent network is trained from initialization to a spawning time $t_{s}$, yielding a parent weight trajectory $\{w_{t}\}_{t=0}^{t_{s}}$. At the spawn time $t_{s}$, several copies of the parent network are made, and these so-called children are then trained with independent minibatch stochasticity, yielding different child weight trajectories $\{w^{t_{s},a}_{t}\}_{t=t_{s}}^{T}$, where $a$ indexes the children, and $T$ is the final training time. We will be interested in various measures of the distance between children after training, as a function of their spawn time $t_{s}$, as well as measures of the distance between the same network (either parent or child) at two different training times. We turn to these various distance measures next.

Kernel distance.

For finite width networks, the kernel $\kappa_{t}(S)=\kappa_{w_{t}}(S)$ changes with training time $t$. We compare two Kernel gram matrices in a scale-invariant manner by computing a kernel distance:

Kernel velocity.

We further track the speed at which the kernel changes. As discussed above, in non-linear neural networks, we do not expect Eq. 3 to vanish. In order to capture the evolution of the quantity in Eq. 3, we compute the kernel velocity $v(t)\equiv S(w_{t},w_{t+dt})/dt$, i.e. the rate of change of kernel distance. We use a time separation of $0.4$ epochs to capture appreciable change.

Error barrier between children.

ReLU activation pattern distance.

Function space distance.

To compute the distance between the two functions $f_{w}$ and $f_{w^{\prime}}$, parameterized by weights $w$ and $w^{\prime}$, we would ideally like to calculate the degree of disagreement between their outputs averaged over the whole input space $x$. However, since this is computationally intractable, we approximate this distance by the normalized fraction of test examples on which their predicted labels disagree. Let $S^{\textrm{test}}$ denote the test set. Then, $\lVert f_{w}(x)-f_{w^{\prime}}(x)\rVert_{S^{\textrm{test}}}=\frac{1}{Z|S^{\textrm{test}}_{x}|}\sum_{x\in S^{\textrm{test}}_{x}}\left(f_{w}(x)\neq f_{w^{\prime}}(x)\right)$ l where $S^{\textrm{test}}_{x}$ are test inputs and $Z$ is a normalizing constant chosen to aid comparison. In particular, we define $Z$ to be the expected number of examples on which two classifiers would disagree assuming each made random independent predictions with the same error rates, $p$ and $p^{\prime}$, as their error rates on the test set. This quantity is used also by , and is given by $Z=p(1-p^{\prime})+p^{\prime}(1-p)+pp^{\prime}\frac{K-2}{K-1}$, where $K$ is the number of classes. A unit distance indicates two networks make uncorrelated errors.

An integrative view of learning dynamics

Figs. 4 and 4 plot the full range of metrics defined in Section 2 for two SOTA networks. Panel A presents standard training curves. Panel B confirms the results of , that the error barrier on a linear path between two children decreases rapidly with spawning time, falling close to within two to five epochs. Panel C and D indicate that the NTK changes rapidly early in training, and more slowly later in training, as quantified by ReLU activation distance (C) and kernel distance (D) measured on a parent run at different pairs of times in training. Finally, Panel E shows that function, kernel and ReLU distances between children at the end of training also drop as a function of spawn time.

We note that the SOTA training protocols in Figs. 4 and 4 involve learning rate drops later in training, which alone could account for a slowing of the NTK evolution. Therefore we ran a constant learning rate experiment in Fig. 4. We see that all tracked metrics still exhibit the same key patterns: the error barrier drops rapidly within a few epochs (B), the NTK evolves very rapidly early on, but continues to evolve at a constant slow velocity later in training (C,D), and final distances between children drop at an early spawn time and remain constant thereafter (E).

Overall, these results provide an integrative picture of the learning process, which reveals an early, extremely short chaotic period in which the final basin chosen by a child is highly sensitive to SGD noise and the NTK evolves very rapidly, followed by a later more stable phase in which the basin selection is determined, the NTK continues to evolve, albeit more slowly, and the final distance between children remains smaller. In the next few sections we explore these results in more detail.

Equivalent results for other networks are shown in Appendix C in Figs. 14, 17, 17, 17, 14 and 14, together with additional network properties tracked over epochs.

The local and global geometry of the loss landscape surrounding children

We first explore how both the global and local landscape geometry surrounding two child pairs and their spawning parent depend on the spawn time $t_{s}$ in Fig. 5. These three networks define a $2D$ affine plane in weight space and a curved $2D$ manifold in function space.

The first two columns of Fig. 5 clearly indicate that two children spawned at an early time $t_{s}$ in the chaotic training regime arrive at two different loss basins that are well separated in function space (top row), while two children spawned at a later time $t_{s}$ in the stable training regime arrive at the same loss basin, though this loss basin can still exhibit non-negligible function diversity (albeit smaller than the diversity between basins). Furthermore, the right two columns of Fig. 5 indicate that the test error as a function of position along the tangent plane to the 2D curved manifold in function space (either at the spawn point or a child point) is insufficient to describe the error along the full curved $2D$ manifold in function space when the children are in different basins (top row), but can approximately describe the loss landscape when the children are in the same basin (bottom row). Thus Fig. 5 constitutes a new direct data-driven visualization of loss landscape geometry that provides strong evidence for several aspects of the conceptual picture laid out in Fig. 1: the existence of multiple basins (Fig. 1A), the chaotic sensitivity of basin fate to SGD choices early in training (Fig. 1B), and the twisting of tangent planes in function space that occur as one travels from one basin to another (Fig. 1ACD). See also Fig. 11 for a t-SNE visualization of the bifurcating evolution of all parents and children in function space that further corroborates this picture of loss landscape geometry and its impact on training dynamics.

In Fig. 6 we explore more quantitatively the relationship between the final function space distance between children, spawn time of children, and the error barrier. This figure demonstrates that the error barrier drops to zero rapidly within 2-3 epochs, and then after that, the later two children are spawned, the closer they remain to each other. Since these experiments were done by training parents and children at a constant learning rate over 200 epochs such that child distances stabilized, the reduction in achievable function space distance between children as a function of spawn time cannot be explained either by learning rate drops or by insufficient training time for children (see Fig. 10).

NTK velocity slows down and stabilizes after basin fate is determined

We next explore the relation between error barrier and kernel velocity in Fig. 7 by zooming in on the early epochs, compared to the full training shown in Fig. 4-4 panels B and D. This higher resolution view clearly reveals that the early chaotic training regime is characterized by a tightly correlated reduction in both error barrier and kernel velocity, with the latter stabilizing to a low non-zero velocity after the error barrier disappears. Thus the NTK evolves relatively rapidly until basin fate is sealed.

The data-dependent NTK rapidly learns features useful for performance

NTK and nonlinear training remain different even at low learning rates

In the NTK limit, which involves both infinite widths and infinitesimal learning rates, linearized training and full nonlinear training dynamics provably coincide. However, the persistent performance gap up to 30 to 90 epochs between linearized and full nonlinear training (green curves versus purple baselines in Fig. 8) indicates the NTK limit does not accurately describe training dynamics used in practice, at finite widths and large learning rates. We remove one of the two reasons for the discrepancy by comparing the same linearized training dynamics to extremely low learning rate nonlinear training dynamics (blue curves in Fig. 8). In this finite width low learning rate regime, we find, remarkably, that a significant performance gap persists between linearized and nonlinear training (the red nonlinear training advantage region in Fig. 9, left), but only during the first few epochs of training, corresponding precisely to the chaotic regime before basin fate is sealed. Indeed the disappearance of this low learning rate nonlinear advantage is tightly correlated with the disappearance of the error barrier (Fig. 9, right). This indicates that while the data-dependent NTK limit can describe well the low (but not high) learning rate dynamics after the first few epochs, this same NTK limit cannot accurately describe the full nonlinear learning dynamics during the highly chaotic early phase prior to basin fate determination, even when the full nonlinear training uses very low learning rates, and when the NTK is learned from the data. We present additional experiments with Taylor expansions of order 2 on ResNet in Fig. 19 and linear order for WideResNet in Fig. 19. In Section C.3, we also perform the same set of linearized training experiments on networks trained with no batch normalization to ensure that we observe the same effect.

Summary of contributions and discussion

In summary, we have performed large scale simultaneous measurements of diverse metrics (Figs. 4, 4 and 4) finding a strikingly universal chaotic to stable training transition across datasets and architectures that completes within two to three epochs. During the early chaotic transient: (1) the final basin fate of a network is determined (Fig. 7 left); (2) the NTK rapidly changes at high speed (Fig. 7 middle and Fig. 4D); (3) the NTK rapidly learns useful features in training data, outperforming the standard NTK at initialization by a factor of 3 within 3 to 4 epochs (Fig. 8 green curves); (4) even low learning rate training retains a nonlinear performance advantage over linearized NTK training with a learned kernel (Fig. 9 red regions); and (5) the error barrier, kernel velocity, and low learning rate nonlinear advantage all fall together in a tightly correlated manner (Fig. 7, right) and (Fig. 9, right). After this rapid chaotic transient, training enters a more stable regime in which: (6) SGD stochasticity allows more limited child exploration in terms of function space distance, leading to smaller function diversity within basins compared to between basins (Figs. 5 and 6); (7) the kernel velocity stabilizes to a fixed nonzero speed (Fig. 7 middle and Fig. 4D); (8) the data dependent kernel performance continues to improve, matching that of full network training by 30 to 90 epochs, of training, representing 15% to 45% of the full 200 epochs (Fig. 8 green curves).

The empirical picture uncovered by our work is much richer than what any theory of deep learning can currently capture. In particular, the NTK theory attempts to describe the entire nonlinear deep learning process using a fixed random kernel at initialization. While this description is provably accurate at infinite width and low learning rate, our results show it is a poor description of what occurs in practice at finite widths and large learning rates (Figs. 7 and 8). More interestingly, the NTK theory is even a poor description of nonlinear training at finite width and extremely low learning rates, especially during the early chaotic training phase (Fig. 9).

This rich phenomenological picture of the rapid sequential nature of the learning process could potentially yield practical dividends in terms of a theory for the rational design of learning rate schedules. For example, the timing of optimized learning rate drops coincide with the time when the data-dependent tangent kernel can achieve high accuracy. Indeed our observations are consistent with findings in . But more generally, we hope that our empirical measurements of such a rich phenomenology may serve as an inspiration for developing an equally rich unifying theory of deep learning that can simultaneously capture these diverse phenomena.

Broader Impact

The goal of our work is to gain a better understanding of deep neural networks. This could potentially make machine learning applications more reliable and transparent in the long run.

Funding Sources

DMR was supported, in part, by an NSERC Discovery Grant, Ontario Early Researcher Award, and a stipend provided by the Charles Simonyi Endowment. SG thanks the Simons Foundation, James S. McDonnell Foundation, NTT Research, and an NSF Career award for support. This research was in part carried out while GKD and DMR participated in the Special Year on Optimization, Statistics, and Theoretical Machine Learning at the Institute of Advanced Studies.

Acknowledgements

The authors would like to thank Jonathan Frankle, Ekansh Sharma, and Mufan Li for feedback on drafts, and Shems Saleh for helping to produce Fig. 1.

References

A Function distance between children runs

The function distance between the children runs is shown in Fig. 10 (right column) for ResNet20 on CIFAR-10 (top row) and CIFAR-100 (bottom row). The signal is relatively noisy from iteration to iteration, we therefore overlay the raw data with a smoothed out version with a window of $\pm 5$ epochs in Fig. 10.

We also produced a t-SNE visualization of parent and children evolution in the function space. To do that, we took predictions of the parent and children runs at different stages of their training on the test set. We then flattened the vector of predicted probabilities for all images and all of their classes into a single long vector, one for each stage of training of a network. We then used the t-SNE embedding to embed all parent and children runs into a 2D space. For the individual panels in Fig. 11, we highlighted the relevant embedded points, however, due to the nature of t-SNE all predictions had in fact been embedded together.

B Definitions of additional metrics

The logit gradient centroid for each class $k$ is defined as $\mu_{w}^{k}=\frac{1}{m}\sum_{i=1}^{m}\nabla_{w}(f_{w}(x_{i}))_{k}$. Previous work has shown that the span of the $K$ logit gradient centroids approximately tracks an important local quantity: the span of the top $K$ directions of maximal Hessian curvature (see also see Section B.2). Thus the span of the logit gradient centroids describes the orientation of the walls of the basins shown schematically in Fig. 1AB, and two trained networks with highly dissimilar logit gradient centroids likely lie in differently oriented basins. In order to evaluate how logit gradient centroids compare at $w$ and $w^{\prime}$, we compute average cosine similarity

which we refer to as logit gradient centroid alignment.

B.2 Logit gradient centroids and the top Hessian eigenvectors

Here we note qualitative relations between the space of logit gradient centroids and the principal Hessian subspace. First, without loss of generality, assume $K=1$ and consider a single data point $(x,y)$. Consider mean squared error (MSE) loss, with empirical risk term $R_{S}(w)=(f_{w}(x)-y)^{2}$. Then

where $H_{w}=\nabla^{2}f_{w}(x)$ (as defined in Section 2).

Since $(J_{w}(x))_{k}=\nabla_{w}(f_{w}(x))_{k}$, and logit gradient centroids are defined as

we decompose $(J_{w}(x))_{k}=\mu_{w}^{k}+e_{w}^{n,k}$, for $w=w_{0}$. This decomposition of $J_{w}$ has been previously studied in . Then $\nabla^{2}R_{S}(w)=\frac{1}{K^{2}m}\sum_{k,k^{\prime}}\mu_{w}^{k}(\mu_{w}^{k^{\prime}})^{T}+\mathcal{O}(e)$.

According to our empirical results as well as other literature , the logit gradient centroids are mutually almost orthogonal. Therefore, $\nabla^{2}R_{S}(w)\approx\sum_{k}\|\mu_{w}^{k}\|^{2}$. For mutually orthogonal gradient centroids $\mu_{w}^{k}$, this amounts to a singular value decomposition with $K$ non-zero singular values $\|\mu_{w}^{k}\|^{2}$ associated with singular vectors $\mu_{w}^{k}$.

While the relative length of $\mu_{w}^{k}$ can be changing with training time, the empirically observed stability of their directions $\mu_{w}^{k}$ makes the Hessian eigenvector associated with the highest eigenvalue constrained to lie primarily within the vector space defined by the span of $\{\mu_{w}^{k}\}_{k\in\{1,...,K\}}$. This subspace of dimension $K$ ($K=10$ for CIFAR-10, $K=100$ for CIFAR-100) has a significantly lower dimensions than the typical weight space of the network.

When evaluating the cosine similarity of the logit gradient centroids as in Eq. 5, we are approximately estimating the overlap between the low dimensional subspaces to which the sharpest directions of the Hessian are constrained between the two networks.

B.3 Escape threshold

Consider iterative optimizers that at each iteration perform an update on the weights of the form

where $\bar{S}$ is a minibatch, i.e., a random subset $\bar{S}\subset S$.

Let $L_{t}=\nabla^{2}R_{S}(w_{t})$ be the Hessian of the empirical loss. Then, by a second order Taylor expansion of $R_{S}(w_{t+1})$ around $w_{t}$, we have

The loss after one iteration decreases if the difference $R_{S}(w_{t+1})-R_{S}(w_{t})$ is negative. Under the second order approximation, the condition for non-increasing loss is equivalent to

We refer to left hand side term in Eq. 10 as the escape threshold. If the escape threshold is below zero, then the trajectory will be descending in the quadratic basin, under the assumption that the local quadratic approximation is accurate. If the escape threshold is positive, the loss will increase or the trajectory will escape the quadratic basin.

Gradient descent (GD) update is $\Delta_{t}(S)=\nabla R_{S}(W_{t})$. Combining this with the bound $\langle\Delta_{t},L_{t}\Delta_{t}\rangle\leq\lambda_{t}\|\Delta_{t}\|^{2}$, where $\lambda_{t}$ is the spectral norm of $L_{t}$, Eq. 10 simplifies to $2-\eta\lambda_{t}$. We refer to this term as the escape threshold for GD.

C Additional results

Here we present additional experiments similar to the ones in Figs. 4, 4 and 4 that include measurements of logit gradient centroid clustering, Hessian spectral norms and escape-time threshold analysis (for the networks with no batch norm).

We compute all of above metrics for SOTA networks in Figs. 4, 4, 14, 17 and 17 (see Appendix D for details of training and hyperparameters). The top rows of Figs. 14, 17 and 17 describe the dynamics of parent training. From left to right, the test and training error drop, as well as the top Hessian eigenvalue drops over time, and all three distances (ReLU pattern, mean logit gradient, and kernel distance) computed on a parent run between pairs of training epochs reveal a rapid change around a very early point at about 5 epochs, followed by slow freezing of all three quantities. In particular, there is significant kernel learning. Moreover, there is a period of time early on where the top Hessian eigenvalue $\lambda>2/\eta$ where $\eta$ is the learning rate. This large learning rate condition would be necessary for gradient descent to escape a quadratic minimum (Fig. 1E, top).

The bottom row computes various distances between pairs of children at the end of training as a function of the common time $t_{s}$ at which the pairs of children were spawned. These plots indicate that at an extremely early spawn epoch (x-axis) of around $t_{s}=1$, the basin fate of the two children is sealed by their parent. Beyond $t_{s}=1$, the two children end up in the same basin, as evidenced by the lack of a loss barrier along a linear path between them, and are much closer to each other, as measured by distances in (from left to right), weight space, ReLU pattern, logit gradient, function space, and kernel space. In contrast, before this early spawn epoch of $1$, the final basin choice of the children displays a chaotic sensitivity to SGD steps, as evidenced by a large loss barrier and larger distances. Intriguingly, significant, albeit slow kernel learning continues after basin fate selection.

C.2 Further discussion

The integrative analysis of diverse measurements during neural network training (Figs. 4, 4, 14, 17, 17 and 17) and the results of linearized training (Fig. 8) reveal a uniform and striking story. First, very early in the training process—between 1 and 10%—of training time, the final basin fate of the neural network is determined. After this point, large scale motion of the neural network is no longer influenced by SGD noise and the networks trained with different SGD noise converge to low loss points in the same (linearly connected) basin. This is indicated by the fact that children spawned after this point are linearly connected through low error networks: the error barrier between spawned children goes below 0 at around 10 epochs for all networks. Additionally, the escape threshold and spectral norm analysis reveals that, after this point, the learning rate is small enough to keep the network within the local quadratic approximation of the loss surface, further supporting the hypothesis that the network does not leave the basin. Once the network is in the basin, we find that various metrics of network distance start decreasing.

Our linearized training analysis (Figs. 19 and 19) reveals further interesting features about this early period of training. Very early in the training process, before an epoch is completed, the data-dependent NTK—the first order approximation of the neural network—rapidly starts learning useful information about the data as evidenced by the fall off of the green lines in Figs. 8, 19 and 19. This occurs uniformly across networks and datasets. When a kernel machine is trained with these data dependent features, it performs significantly better than the NTK at random initialization. In fact, less than halfway through training, the data-dependent neural tangent kernel machine performs nearly as well as as the full network.

These observations have two important implications. First, while training the NTK at random initialization may not be very representative of training finite sized networks, across a range of networks, the data-dependent NTK obtained from a small amount of training is. Second, the features built up by the NTK relatively early in the training process are sufficient for achieving low errors competitive with the full non-linear networks.

While deep learning research often focuses on improving accuracy towards the end of training, our results show that the early phase of training is important for determining the final fate of the network. A better understanding of this phase may provide us with tools to diagnose and improve networks early on in the training process, thus decreasing the cost of training neural networks.

C.3 Nonlinear advantage and function space distance without batch norm

In Fig. 6 we show the function space distance between children runs and its correlation with the error barrier for a network with batch normalization. The nonlinear advantage with batch normalization is also shown in Fig. 9. While state-of-the-art models use batch normalization and we therefore focused on it in the main text, we wanted to confirm that the same phenomena are observed for networks without batch normalization. Fig. 20 shows the function space results, and Fig. 21 shows the nonlinear advantage for ResNet20v1 without batch normalization trained on CIFAR-10 with Adam at learning rate $10^{-3}$.

D Experimental details

SimpleCNN: SimpleCNN is a 6 layer fully convolutional neural network. Each convolution has a 3x3 kernel, stride 1 and bias. The layers have 32, 32, 64, 64, 128, and 128 channels from the first to the last layer. The weights are initialized using Kaiming initialization and the biases are initialized to 0. There is a 2D maxpooling with a 2x2 kernel and stride 2 after layers 2, 4, and 6. Layer 6 is followed by a 2d global average pool which results in a 1x128 unit feature vector from which the classes are linearly predicted.

ResNet20: We use the ResNet20 used with CIFAR-10/100 data in the original paper .

ResNet20 without batchnorm (RN no BN): Same as ResNet20 but with batchnorm turned off.

WideResNet-16-4 (WRN-16-4): We use the WideResNet-16-4 described in (16 layers, widen factor 4, no dropout).

D.2 Training Details

See Table 1 for training details and hyperparameters.

D.3 Extended training time and learning rate ablations

In Figs. 8 and 9, we train networks with a lower learning rate to compare to linearized training and measure the nonlinear advantage. This is done as follows: we first train a parent network at a constant learning rate of 0.1. At certain epochs (the x-axis of Fig. 8 and labels of Fig. 9) we spawn a child network and train it with a learning rate of 0.001 and independent SGD noise until convergence (the learning rate is chosen as the smallest learning rate at which the network converges in a reasonable amount of time, 1000 epochs in our case.) We then use the learning rate dropped accuracy is the final accuracy that the child converges to. For stability, we repeat this process and average the two independent runs. The linearized training is performed using the package and implemented in JAX . It trains for additional 200 epochs at learning rate 0.001 that we choose based on a small-scale grid search for linearized training specifically. The nonlinear advantage is computed as the error difference between the final test performance of the low learning rate child network, and the final test performance of the linearized neural network.

D.4 Linearized training details

The linearized training was performed using the Taylor expansions tools described in and implemented in . We trained for 200 epochs using SGD with Momentum, implementing the training loop in JAX . The learning used was $0.001$. We did a small-scale grid search for the best performing learning rate between $0.0001$ and $0.1$, choosing $0.001$ in the end.