Asymptotics of Wide Convolutional Neural Networks

Introduction

Deep neural networks continue to achieve remarkable performance on a diverse range of machine learning tasks, however detailed understanding remains elusive. One of the most promising routes towards understanding is to study very wide neural networks. Wide networks strike an attractive balance between performance and analytic control . Furthermore understanding the performance of networks as the number of parameters is increased is at the heart of the generalization paradox – the observation that over-parameterized deep networks do not over-fit.

In the authors showed that the dynamics of infinitely wide fully connected (FC) neural networks trained under gradient flow simplifies dramatically, and in the limit of infinite width argued that training a deep network is equivalent to training a linear random features model. For models trained with mean squared error (MSE) loss, this infinite width, linear evolution can be written as

where $f(x)$ is the network output on example $x$, $\mathcal{D}_{\textrm{train}}$ is the training dataset, and $\Theta(x,x^{\prime})=\sum_{\mu}\frac{\partial f(x)}{\partial\theta_{\mu}}\frac{\partial f(x^{\prime})}{\partial\theta_{\mu}}$ is the neural tangent kernel (NTK).

In this simplified infinite width evolution was extended to take into account corrections from finite width, giving a controlled prediction for the evolution of wide fully connected networks trained via stochastic gradient descent (SGD).

This formalism has the promise of explaining the training dynamics and predictions of wide neural networks. Empirically, however, there are several results that show convolutional neural networks (CNNs) exhibiting different behavior from fully connected networks (FCs) that have yet to be understood. First, for CNNs, but not FCs, there is evidence that finite width networks trained via SGD outperform their infinite width counterparts . Paradoxically, despite the drop in accuracy for the infinite width predictions, find that the performance of finite width CNNs improves as the width gets larger. These results appear in tension with our understanding of the scaling behavior of large width networks. Second, there is prior empirical evidence suggesting the scaling behavior in CNNs is different from FC networks .

To examine these differences between the observed behavior of CNNs and FCs, we present an extension of the formalism of to convolutional networks. We study correlation functions – ensemble averages over weight configurations of quantities built out of the network map and its derivatives. We conjecture a simple scaling relation for correlation functions for models with convolutional, skip, dense, and global average pooling layers. We prove this relation for deep linear and and single hidden layer networks with smooth activations and check the relation empirically in a broader context including deep non-linear networks, ReLU networks, and networks with max pooling layers. This scaling relation serves as the basis for bounding corrections to linear evolution and allows us to study how model performance depends on network width.

We derive a set of scaling relations for correlation functions, a general class of expectation values of the network map and its derivative. These relations rely on generalizing previous diagramatic methods to convolutional networks.

We apply our scaling relations to the loss and accuracy during training. In particular we argue that the difference between full and linearized test loss scales as $\mathcal{O}(n^{-1})$.

We confirm the predicted loss scaling empirically for deep non-linear networks trained on subsets of CIFAR-10 and find that it is consistent with finite width networks either outperforming or underperforming their infinite width counterparts.

We further apply our relations to derive asymptotically tight bounds for the change of the NTK during training.

Related work

There has been significant progress understanding the behavior of wide neural networks. At initialization, wide networks of any depth behave as Gaussian processes . Finite width corrections to the Gaussian process picture have been discussed in . The infinite width training of NTK parameterized networks was introduced in and studied in . Wide convolutional networks have been studied in . study the infinite width limit of a more general class of network, that includes convolutional networks as a special case. This work is most closely related to which also applied diagramatic techniques to derive asymptotic scaling rules for correlation functions. See also for another use of diagrams in this context. We also discuss finite width corrections to the evolution of convolutional networks, these corrections were studied for fully-connected networks in . Preliminary experiments for this work were run using the neural tangents python library . During the completion of this work appeared which performed extensive empirical studies comparing the performance of finite width CNNs and infinite width kernel methods.

Theory

In this section we present our main theoretical result. Namely, a conjectured class of bounds, Conjecture 1, governing the scaling with respect to width of quantities built out of the network map and its derivatives.

We consider neural networks, $f(x)$, built out of a collection of stacked activations with non-linearity $\sigma$,

For a convolutional layer with kernel size $k_{w}\times k_{h}$ and $n$ channels,

We consider networks terminated by linear transform after a flatten or global average pooling operation,

Following we introduce a class of moments built out of the network map and its derivatives called correlation functions.

A correlation function , $C(x_{1},x_{2},\ldots,x_{m})$, is an ensemble average over products of the network map, $f(x)$, and its derivatives, $\frac{\partial^{k}f(x)}{\partial_{\mu_{1}}\partial_{\mu_{2}}\cdots\partial_{\mu_{k}}}$, subject to the condition that all derivatives are summed in pairs. A general correlation function $C$ takes the form

Here, $0\leq k_{1}\leq\cdots\leq k_{m-1}\leq k_{m}$ are integers,We adopt the convention that $k_{a}=k_{a-1}$ represents a factor of $f$ with no derivatives acting. $m$ and $k_{m}$ are even, $\pi\in S_{k_{m}}$ is a permutation, and $\Delta_{\mu_{1}\dots\mu_{k_{m}}}^{(\pi)}=\delta_{\mu_{\pi(1)}\mu_{\pi(2)}}\cdots\delta_{\mu_{\pi(k_{m}-1)}\mu_{\pi(k_{m})}}\,.$ We use $\delta$ to denote the Kronecker delta.

We refer to paired summed indices as contracted derivatives and refer to factors of the network map with such contracted derivatives as being contracted in $C$.

For every such correlation function we define the associated cluster graph.

The cluster graph, $G_{C}(V,E)$, associated to a correlation function $C(x_{1},x_{2},\ldots,x_{m})$, is a graph consisting of $m$ vertices, one corresponding to each factor of the network map in Equation (8). The graph has a single edge between vertices corresponding to factors of the network map sharing a pair of contracted derivatives.

In the authors argue for a set of bounds on correlation functions for fully connected networks based on the cluster graph. For a correlation function $C$ with a cluster graph containing $n_{e}$ even components (connected components with an even number of vertices), $n_{o}$ odd components, and $m$ vertices, they argue that the correlation function satisfies $C=\mathcal{O}(n^{n_{e}+n_{o}/2-m/2})$. The authors prove these bounds for deep linear and one-hidden non-linear fully connected networks and check empirically in a variety of contexts. Subsequently established this bound for deep non-linear networks with polynomial activations. Here we extend this to networks with convolution, skip, and global average pooling layers. We present a proof in the deep-linear and 1-hidden layer non-linear case.

Let $C(x_{1},\dots,x_{m})$ be a correlation function with cluster graph, $G_{C}$. Suppose that $G_{C}$ has $n_{e}$ connected components with an even size, and $n_{o}$ components of odd size, then $C(x_{1},\dots,x_{m})=\mathcal{O}(n^{s_{C}})$, where

We have tested this conjecture empirically in a variety of contexts, and these results appear in Section 4. We are also able to prove Conjecture 1 for deep linear and one-hidden-layer non-linear networks.

Conjecture 1 holds for deep linear and one-hidden-layer networks with smooth activations.

The argument for the deep linear case is summarized below. A detailed proof for deep linear networks and one-hidden-layer non-linear networks appears in the Supplement.

In the deep linear case, the proof follows from two lemmas.

Let $f(x)$ be a deep linear network of depth $d$ made up of convolutional, skip, dense, and GAP layers. Then the network function can be written as a finite sum over $\mathcal{N}_{f}$ functions $f_{I}(x)$, where each function $f_{I}$ has the topology of a fully connected network with depth $d_{I}\leq d$.

Furthermore, let $\{\theta_{I}\}$ be the set of weights of each $f_{I}(x)$ and $\{\theta\}$ the weights of $f$. Then, $\{\theta_{I}\}\in\{\theta\}$.

The decomposition of the network map motivates a generalization of the definition of correlation functions to expectations involving the maps $f_{I}(x)$. We dub such expectation values mixed correlation functions.

A mixed correlation function, $C_{I_{1},I_{2},\ldots,I_{m}}(x_{1},x_{2},\ldots,x_{m})$, is an ensemble average over products of the functions, $f_{I}(x)$, and its derivatives, $\frac{\partial^{k}f_{I}(x)}{\partial_{\mu_{1}}\partial_{\mu_{2}}\cdots\partial_{\mu_{k}}}$, subject to the condition that all derivatives are summed in pairs.

The decomposition of the network map in Lemma 1 allows us to write correlation functions of CNNs in terms of finite sums over correlation functions of fully connected networks and thus reuse much of the technology introduced in . This is formalized in the following lemma.

Let $C_{I_{1},I_{2},\ldots,I_{m}}(x_{1},x_{2},\ldots,x_{m})$ be a mixed correlation function. Let $G_{C}$ be a cluster graph associated to $C_{I_{1},I_{2},\ldots,I_{m}}(x_{1},x_{2},\ldots,x_{m})$ via Definition 2 ignoring the labels $\{I_{1},I_{2},\ldots,I_{m}\}$. Let $n_{e}$, $n_{o}$ be the number of even, odd clusters in $G_{C}$, then $C_{I_{1},I_{2},\ldots,I_{m}}(x_{1},x_{2},\ldots,x_{m})=\mathcal{O}(n^{s_{C}})$ with

Together Lemma 1 and Lemma 2 imply Theorem 1 for deep linear networks. Lemma 1 follows from the definitions of our layers above. Here we show this for convolution layers, and leave the details of skip and GAP layers to the Supplement.

We prove Lemma 2 in the Supplement. We rely on the graphical Feynman diagram techniques introduced in .

Conjecture 1 has important implications for the training dynamics of convolutional networks at both infinite and finite width.

Consider a network trained via gradient flow with mean squared error (MSE) loss.

In general this equation can describe quite complicated training dynamics, as a result of the time dependence of $\Theta(t)$. Empirically we find some non-trivial late time behavior for certain CNN models which has to be treated with care. We elaborate on this in the Supplement. At infinite width, however, $\Theta$ is constant and the dynamics reduce to that of training a linear model.

As we will explain, Conjecture 1 bounds the change in the kernel as

In Figure 5 and Table 3 we see evidence that this bound is saturated in a variety of CNNs.

To understand this analytically, we begin with an illustrative example. Consider the time derivative of the NTK,

The constancy of the NTK is a striking feature of infinite width networks. However in practice we mostly consider finite width networks and the connection between infinite and finite width evolution is not immediately clear. In the authors take steps towards understanding finite width networks. In particular they show that the scaling relations, Conjecture 1, imply a systematic expansion for the evolution of the network map.

Here, $f^{(0)}(x;t)$ is the linearized infinite width evolution, Equation (1), and the higher order terms can be iteratively solved for in terms of the network map at initialization. The derivation of this result relies only on Conjecture 1 and so applies here as well. Some important consequences of this expansion are scaling relations for the loss and accuracy during training, which we now describe.

2 Performance scaling

One natural application of the large width evolution in Equation (19) is understanding the dynamics of the loss and accuracy during training of wide networks. This question is at the heart of the generalization paradox, the observation that over-parameterized networks suffer no degradation in performance as they become larger . It also describes the asymptotic behavior of the so called double descent curve .

The expansion of the network map in powers of $1/n$ leads to a corresponding expansion in the test loss

Thus we expect the difference between the full model and linearized test loss to scale as

Here $L_{\textrm{test}}^{\textrm{lin}}:=L_{\textrm{test}}(f^{(0)})$. This scaling has been observed empirically in fully connected networks and is the same scaling predicted in for linear models. Here we see good agreement with this asymptotic behavior of the loss in deep convolutional networks (see Figure 8).

This relation between the loss of a finite width network and the infinite width loss is of particular interest for CNNs. As mentioned above, convolutional networks often exhibit a gap in performance between infinite width networks and their finite width counterparts . It is natural to ask whether we can understand this performance gap within the framework of the perturbative expansion around large width.

The relation, Equation (21), implies that the full and linearized loss approach each other at infinite width, however the relative ordering is not dictated. We will see in the models studied below, that either ordering is possible depending on the time during training. In particular, in the setup studied here, finite width models outperform their linear counterparts if training is stopped at the non-linear early stopping time. We expand further on the dynamics and performance of the particular models studied in Section 4.3.

Numerical Experiments

To support the theoretical predictions in Section 3, we present numerical results for the scaling of the NTK at initialization as well as at convergence (100% training accuracy) for one- and three-hidden-layer convolutional neural networks with layers of the types defined in Section 3. All models are trained on 2-class MNIST (0’s and 1’s) with 10 examples per class, with the exception of Section 4.3 which is trained on 2-class CIFAR (airplane and automobile) with 100 examples per class and tested on the full 2-class test dataset. Training was done with full-batch gradient descent and all models achieved 100% training accuracy. The learning rate used (unless otherwise specified) was $0.25\cdot\frac{1}{\max{\lambda_{i}}}$, where $\lambda_{i}$ are the eigenvalues of the NTK.

Next we consider the expectation of the variance of the NTK. This variance can be written as the difference of two correlation functions,

From Conjecture 1, each of these correlation functions is $\mathcal{O}(n^{0})$, thus we can bound $\text{Var}_{\theta}\left[\Theta\right]=\mathcal{O}(n^{0})$. In the Supplement we show that for deep linear networks and for one-hidden-layer networks, we can actually do better giving $\text{Var}_{\theta}\left[\Theta\right]=\mathcal{O}(n^{-1})$. In the spirit of Conjecture 1, we predict this scaling more generally. Note that the suppression of the variance of the NTK with width is crucial for $\Theta$ to have a well defined infinite width limit. In particular, this implies typical realizations of the kernel will be close to the mean as the width increases. Sub-panel 4b of Figure 4 corroborates this predicted scaling for three-hidden-layer CNNs with more examples in sub-panel 4d.

2 Change in the NTK

Next, we consider how the deviation in the NTK from initializaion depends on the number of color channels. Equation (17) predicts that this difference scales as $\mathcal{O}(n^{-1})$. As discussed above, this constancy of the kernel is what underlies the large width linear dynamics.

Figure 5 shows that the scaling of the expected deviation from initialization of the NTK for a three-hidden-layer convolutional networks for different layer types and ReLU activation function. For each architecture, the fit is done at the time step when all the models for the 10 initializations have hit 100% training accuracy. Further results are recorded in Table 3.

Figure 6 shows the evolution of $\Theta(t)$ during training with gradient descent for a single hidden layer convolutional network with ReLU activation function and global average pooling over a range of widths, $n$. Equation (17) is an asymptotic statement. In practice, to see this scaling it is often necessary to go to $n\gg 32$. As an example of the sensitivity to width, we look at the $1/n^{\alpha}$ fit over large and small ranges of $n$ in Figure 7 during training. We find greater deviations from $\mathcal{O}(n^{-1})$ scaling when including smaller widths. More generally, experimentally testing convolutional networks with enough channels to convincingly confirm or rule out Conjecture 1 can be a challenge and we detail this in the Supplement.

3 Loss scaling

In this section we empirically study the performance scaling discussed in Section 3.2 for a four-hidden-layer CNN trained on 2-class CIFAR with 100 examples per class. These models exhibit the now familiar property that their performance does not get worse as the number of channels is increased.

We compare the training dynamics of a non-linear model and its linearized counterpart. We find that if both models are stopped at the optimal early stopping time for the non-linear model, then the non-linear model outperforms the linear model. This gap in performance fits well with the predicted $\mathcal{O}(n^{-1})$ scaling. However this does not represent a true performance gap between linear and non-linear model, but rather the fact that the linear and non-linear models achieve their maximum accuracy at different times. Indeed for late times we see the non-linear model under-performs the linear model, again with a gap scaling as $\mathcal{O}(n^{-1})$. As both linear and non-linear models over-fit, there is also a gap between the optimally stopped non-linear model and the infinite time predictions of the linear model. These results are summarized in Figure 8.

Discussion

We have presented a simple relation, Conjecture 1, for the asymptotic scaling of correlation functions with width and tested the predictions in a variety of CNN architectures. We used these scaling relations to study the training dynamics and performance of wide convolutional networks. At infinite width, CNNs evolve as linear models, with training controlled by the constant NTK. Our conjecture gives an asymptotically tight bound on the approach to constancy.

Away from infinite width the NTK is no longer constant, but the dynamics can still be systematically approximated in a large width expansion. We use this to predict an asymptotic $\mathcal{O}(n^{-1})$ scaling for the difference between the linearized loss and the full non-linear loss Equation (21). In Section 4.3 we presented evidence corroborating this prediction. Though the linearized and non-linear loss approach each other at infinite width, we found that their order can depend on the training time for which they are being compared. This same sensitivity is reflected in the accuracy. In particular, depending on the stopping criteria, the full finite width model can be either better or worse then the linear model.

One motivation for extending the analysis of to convolutional networks is to bridge the gap between our analytic understanding at infinite width and the empirically best performing finite width networks used in practice. Our analysis of the performance gap between linearized and non-linear models is a step in this direction. We hope the tools developed here can be extended to increasingly realistic scenarios.

Acknowledgements

The authors wish to thank Yasaman Bahri, Guy Gur-Ari, Jaehoon Lee, Aitor Lewkowycz, Sam Schoenholz, and Jascha Sohl-dickstein for useful discussions during the completion of this work.

References

Appendix A Expansion of deep-linear network maps

In this section, we complete the proof of Lemma 1. In the body we discussed convolution layers. Here we extend the analysis to skip and global average pooling layers.

Global average pooling

Here each $f_{I}(x)$ has the topology of a deep network without a the global average pooling layer.

Lemma 1 follows from repeated applications of (3), (23), and (24) to each such layer, and relabeling of indices so that $I$ runs over all terms in the resulting sums. ∎

Appendix B Feynman diagrams for deep-linear networks

We are interested in computing correlation functions (Definition 1). This involves computing expectation values over the multi-variate Gaussian initial weight values. A central tool in computing these expectation values is Isserlis’ Theorem (sometimes referred to as Wick’s theorem). Which establishes that higher moments under a Gaussian distribution can be computed as sums of products of second moments. For example,

In general, keeping track of all the moments be cumbersome and Feynman diagrams are a useful book keeping tool. These diagrams were used in to compute the ensemble averages over initial weights for fully-connected networks. We begin by reviewing the key definitions and results and then extend the technology as is needed to prove Theorem 1.

We refer to $T$ as a derivative tensor. As above, if two derivative tensors in a correlation function, $C$, have $k$ paired summed indices, we say that the tensors are contracted $k$-times in $C$.

We now describe the Feynman diagrams associated to a correlation function.

Let $C(x_{1},\dots,x_{m})$ be a correlation function for a network with $d$ hidden layers. The family $\Gamma(C)$ is the set of all graphs that have the following properties.

There are $m$ vertices $v_{1},\dots,v_{m}$, each of degree $d+1$.

Each edge has a type $t\in\{U,W^{(1)},\dots,W^{(d-1)},V\}$. Every vertex has one edge of each type.

The graphs in $\Gamma(C)$ are called the Feynman diagrams of $C$.

Some example Feynman diagrams are shown in Figure 9.

For one hidden layer networks, the Feynman diagrams allow one to easily compute the scaling of a correlation function. Deep networks require additional technology, the double line graph.

Let $\gamma\in\Gamma(C)$ be a Feynman diagram for a correlation function $C$ involving $k$ derivative tensors for a network of depth $d$. Its double-line graph, $\textrm{DL}(\gamma)$ is a graph with $kd$ vertices of degree 2, defined by the following blow-up procedure.

Each vertex $v_{i}$ in $\gamma$ is mapped to $d$ vertices $v^{(1)}_{i},\dots,v^{(d)}_{i}$ in $\textrm{DL}(\gamma)$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $U$ is mapped to a single edge $(v^{(1)}_{i},v^{(1)}_{j})$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $W^{(l)}$ is mapped to two edges $(v^{(l)}_{i},v^{(l)}_{j})$, $(v^{(l+1)}_{i},v^{(l+1)}_{j})$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $V$ is mapped to a single edge $(v^{(d)}_{i},v^{(d)}_{j})$.

The number of faces in $\gamma$ is given by the number of loops in the double-line graph $\textrm{DL}(\gamma)$.

Some example double line graphs are shown in Figure 10. With these definitions correlation functions of deep linear networks satisfy a theorem originally due to .

Let $C(x_{1},\dots,x_{m})$ be a correlation function of a deep linear network with $d$ hidden layers, and let $\gamma\in\Gamma(C)$ be a Feynman diagram. The diagram represents a subset of terms that contribute to $C$, and its asymptotic behavior is determined by the Feynman rules: the subset is $\mathcal{O}(n^{s_{\gamma}})$ where $s_{\gamma}=l_{\gamma}-\frac{dm}{2}$, and $l_{\gamma}$ is the number of loops in the double-line diagram $\textrm{DL}(\gamma)$. Furthermore, the correlation function is $C=\mathcal{O}(n^{s})$, where $s=\max_{\gamma\in\Gamma(C)}s_{\gamma}$.

Below we generalize this construction to accommodate networks with convolution, skip, and GAP layers.

B.2 Extension

We must extend the above technology to our mixed correlation functions (Definition 3). Mixed correlation functions are expectations of maps of the form

Let $C_{I_{1},\ldots,I_{m}}(x_{1},\ldots,x_{m})$ be a mixed correlation function for a collection of maps, $\{f_{I_{1}},\ldots,f_{I_{m}}\}$ with depths $\{d_{1},\ldots,d_{m}\}$. The family $\Gamma(C_{I_{1},\ldots,I_{m}})$ is the set of all graphs that have the following properties.

There are $m$ vertices $v_{1},\dots,v_{m}$, where each $v_{i}$ has degree $k_{i}=d_{i}+1$.

Each edge has a type $t$. Every vertex, $v_{i}$ has one edge of each type $t\in\{V_{\alpha^{I_{i}}_{d_{I_{i}}}},W_{\alpha^{I}_{d_{I}-1}}^{(\beta^{I}_{d_{i}-1})},\ldots,W_{\alpha^{I}_{1}}^{(\beta^{I}_{1})},W_{\alpha^{I}_{0}}^{(0)}\}$.

Let $\gamma\in\Gamma(C_{I_{1},\ldots,I_{m}})$ be a Feynman diagram for a mixed correlation function $C_{I_{1},\ldots,I_{m}}$ involving $m$ derivative tensors. Its double-line graph, $\textrm{DL}(\gamma)$ is a graph with vertices of degree 2, defined by the following blow-up procedure.

Each vertex $v_{i}$ in $\gamma$ of degree $k_{i}$ is mapped to $k_{i}-1$ vertices $v^{(1)}_{i},\dots,v^{({k_{i}}-1)}_{i}$ in $\textrm{DL}(\gamma)$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $W_{\alpha}^{(\beta)}$ is mapped to two edges, $(v^{(e_{i})}_{i},v^{(e_{j})}_{j})$ and $(v^{(e_{i}+1)}_{i},v^{(e_{j}+1)}_{j})$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $W_{\alpha}^{(0)}$ is mapped to a single edge $(v^{(1)}_{i},v^{(1)}_{j})$.

Each edge $(v_{i},v_{j})$ in $\gamma$ of type $V_{\alpha}$ is mapped to a single edge $(v^{(k_{i}-1)}_{i},v^{(k_{j}-1)}_{j})$.

Here, $e_{i}$ take values in $\{2,\ldots,k_{i}-2\}$. As in Definition 5, the number of faces in $\gamma$ is given by the number of loops in the double-line graph $\textrm{DL}(\gamma)$.

Example Feynman and double-line diagrams for mixed correlation functions are shown in Figure 11.

With the double line diagram defined, we can generalize Theorem 2 for the mixed correlation functions.

Let $C_{I_{1},\ldots,I_{m}}(x_{1},\dots,x_{m})$ be a mixed correlation function. Let $\gamma\in\Gamma(C_{I_{1},\ldots,I_{m}})$ be a Feynman diagram. The diagram represents a subset of terms that contribute to $C_{I_{1},\ldots,I_{m}}$, and its asymptotic behavior is determined by the Feynman rules: the subset is $\mathcal{O}(n^{s_{\gamma}})$ where $s_{\gamma}=l_{\gamma}-\sum_{i=1}^{m}\frac{k_{i}-1}{2}$, with $k_{i}$ the degree of the $i$-th vertex, and $l_{\gamma}$ the number of loops in the double-line diagram $\textrm{DL}(\gamma)$. Furthermore, the mixed correlation function satisfies $C_{I_{1},\ldots,I_{m}}=\mathcal{O}(n^{s})$, where $s=\max_{\gamma\in\Gamma(C_{I_{1},\ldots,I_{m}})}s_{\gamma}$.

We now use the Feynman rules (Theorem 3) to bound the scaling of a correlation function by the maximal number of connected components appearing in any single line Feynman diagram.

Let $C_{I_{1},\ldots,I_{m}}(x_{1},\dots,x_{m})$ be a mixed correlation function. Let $c_{\gamma}$ be the number of connected components of a graph $\gamma\in\Gamma(C_{I_{1},\ldots,I_{m}})$. Then $C_{I_{1},\ldots,I_{m}}=\mathcal{O}(n^{s})$, where

It is enough to show that each connected component $\gamma^{\prime}$ in $\gamma$ is bounded as $\mathcal{O}(n^{s_{\gamma^{\prime}}})$ where $s_{\gamma^{\prime}}\leq 1-\frac{v_{\gamma^{\prime}}}{2}$. The graph $\textrm{DL}(\gamma^{\prime})$ is a triangulation of a Riemann surface with $f$ faces $e$ edges, and $v_{\gamma^{\prime}}$ vertices of degrees $k_{1},\ldots,k_{v_{\gamma^{\prime}}}$. The Feynman rules give

Using the relation $e=\sum_{i=1}^{v_{\gamma^{\prime}}}\frac{k_{i}}{2}$ and the definition of the Euler character, $\chi=v_{\gamma^{\prime}}-e+f$ we have

The diagram $\textrm{DL}(\gamma^{\prime})$ is a triangulation of a Riemann surface with at least one boundary, thus $\chi\leq 1$ and $s_{\gamma^{\prime}}\leq 1-\frac{v_{\gamma^{\prime}}}{2}$. ∎

We are now ready to prove Theorem 1 for deep-linear networks.

Let $C(x_{1},\ldots,x_{m})$ be a correlation function for a deep-linear network built out of dense, convolution, skip, and GAP layers. By Lemma 1 we can write

Thus, by Lemma 3 $C(x_{1},\ldots,x_{m})=\mathcal{O}(n^{s})$ where

We now show that $c_{\gamma}\leq n_{e}+\frac{n_{o}}{2}\ \forall\,\gamma$. Firstly note that the cluster graph $G_{C}$ (Definition 2) is a sub-graph of all $\gamma\in\Gamma(C_{I_{1},\ldots,I_{m}})$, thus $c_{\gamma}\leq n_{e}+n_{o}$ as each cluster in $G_{C}$ can form at most one connected component in $\gamma$. Furthermore, note that each connected component in $\gamma$ contains an even number of vertices, thus even clusters in $G_{C}$ can form there own connected components in $\gamma$, but odd clusters must be paired in the connected components of $\gamma$. Thus, $c_{\gamma}\leq n_{e}+\frac{n_{o}}{2}$. ∎

Appendix C One-hidden-layer non-linear Networks

In this section we prove Theorem 1 for the case of networks with a single hidden layer. In this case, there are no skip connections, so we consider networks with a single convolutional layer terminated by either a gap or flatten layer.

It is convenient to adopt a notation that highlights the scaling with respect to the network width (number of convolutional channels). To this end, we write the network function as,

To write $f$ in this form, we have juggled the indexing over input pixels, channels, and kernels as follows.

$\mathbf{X}$ is a $\mathcal{W}\mathcal{H}\times k_{w}k_{h}c_{\textrm{in}}$ matrix. $\mathbf{X}_{\{r,s\},\{a,b,j\}}:=x_{r+a,s+b;j}$.

For each $i$, $\mathbf{U}_{i}$ is a $k_{w}k_{h}c_{\textrm{in}}$ vector. $\left(\mathbf{U}_{i}\right)_{\{a,b,j\}}:=W^{(0)}_{a,b;ij}$.

For each $i$, $\mathbf{V}_{i}$ is a $\mathcal{W}\mathcal{H}$ vector. For flatten models, $\left(\mathbf{V}_{i}\right)_{\{r,s\}}=V_{r,s;i}$, while for GAP models $\left(\mathbf{V}_{i}\right)_{\{r,s\}}=V_{i}\delta_{rs}$.

Here $\delta_{rs}$ is the Kronecker delta and $c_{\textrm{in}}$ is the number of input channels. We have also dropped the normalization over input channel number, kernel size, and image size as they do not effect the scaling with respect to width.

If we adopt the notation $R=\{r,s\}$, $A=\{a,b,j\}$, and let $\mathbf{U}_{A,i}$ and $\mathbf{V}_{i,R}$ be Gaussian distributed with unit variance.

With all of this notation out of the way, we can prove Theorem 1. An outline of the argument is the following. A general correlation function can be written as a sum of terms,

This form of the correlation function follows from computing the expectation over the weights $\mathbf{V}$ and evaluating all derivatives in the correlation function. These operations generate a sum of $K$ expectation values over $\mathbf{U}$ and $\alpha$ indexes the terms in this sum. Before completing the proof, as an explicit example, let us evaluate Equation (38) for the NTK,

In the general case, we argue that the maximum number of sums, $r_{\textrm{max}}=\max_{\alpha}r_{\alpha}$ is bounded as $r_{\textrm{max}}\leq n_{e}+\frac{n_{o}}{2}$, where $n_{e}$ and $n_{o}$ are the number of even and odd clusters in the cluster graph $G_{C}$. We further argue that the $\mathcal{S}^{(\alpha)}_{i_{1}\ldots i_{r_{\alpha}}}$ are bounded by an $n$-independent constant. These two statements establish Theorem 1.

To prove these results we again take a graphical approach. We introduce a new class of graphs to keep track of the index sums in $C$.

Let $C(x_{1},\dots,x_{m})$ be a correlation function for a one-hidden-layer non-linear network. The family $\Gamma^{\prime}(C)$ is the set of all graphs that have the following properties.

There are $m$ vertices $v_{1},\dots,v_{m}$, each of degree at least one.

Each edge has a type $t\in\{U,V\}$. Every vertex has one edge of type $V$.

There is an edge $(v_{i},v_{j})$ for $f(x_{i})$, $f(x_{j})$ contracted in $C$.

Each graph $\gamma_{\alpha}\in\Gamma^{\prime}(C)$ corresponds to a term in the $\alpha$ sum (38). The number of index sums, $r_{\alpha}$ is the number of connected components in $\gamma_{\alpha}$. To see this, note that we have one index sum for each factor of the network map $f(x)$ (Equation (36)) in the correlation function $C$. A contracted derivative between pairs of network maps in $C$, results in a delta function eliminating one index sum. Similarly, the $V$ edges in $\gamma_{\alpha}$ correspond to using Isserlis’ theorem to evaluate the Gaussian expectation over the $\mathbf{V}_{i}$ in $C$. A single covariance as in (37) also eliminates an index sum due to the Kronecker delta factor. The result is that each connected component corresponds to a single sum.

We now argue for the bound on the maximal number of connected components.

Let $C(x_{1},\dots,x_{m})$ be a correlation function for a one-hidden-layer non-linear network, let $\gamma_{\alpha}\in\Gamma^{\prime}(C)$, and let $c_{\alpha}$ denote the number of connected components in $\gamma_{\alpha}$. Further, let $G_{C}$ be the cluster graph of $C$ and $n_{e}$ ($n_{o}$) denote the number of even (odd) clusters. Then $c_{\alpha}$ satisfies.

First note that $G_{C}$ is a sub-graph of $\gamma_{\alpha}$. As each vertex in $\gamma_{\alpha}$ must have one $V$ edge, connected components in $\gamma_{\alpha}$ have an even number of vertices, thus the even clusters in $G_{C}$ can form their own connected components in $\gamma_{\alpha}$ while an odd clusters in $G_{C}$ must pair with at least one other odd cluster to form a connected component in $\gamma_{\alpha}$. ∎

An immediate consequence of this lemma is the bound $r_{\textrm{max}}\leq n_{e}+\frac{n_{o}}{2}$. We now prove Theorem 1 for one-hidden-layer non-linear networks.

We can bound the correlation function $C(x_{1},\ldots,x_{m})$ as,

Here we have introduced $s_{\textrm{max}}$ as a bound on the expectation values, $s_{\textrm{max}}=\max_{\alpha,i_{1},\ldots,i_{r_{\alpha}}}|\mathcal{S}^{(\alpha)}_{i_{1}\ldots i_{r_{\alpha}}}|$. The expectations $S^{(\alpha)}_{i_{1}\ldots i_{r_{\alpha}}}$ can only take $\mathcal{O}(1)$ different values, as the $\mathbf{U}_{i}$ are i.i.d. $s_{\textrm{max}}$ is the maximum over these $\mathcal{O}(1)$ options. Thus $C=\mathcal{O}(n^{n_{e}+\frac{n_{o}}{2}-\frac{m}{2}})$. ∎

Appendix D Late Time Behavior of GAP Networks

In this section we will briefly discuss some non-trivial features we see during training of convolutional neural networks with global average pooling. When training these models until convergence, we see a discontinuity in the NTK and the network function at late times that has to be avoided to see the theoretically predicted scaling behavior with width.

One such example is shown in Figure 12a, where the evolution of the NTK and the training loss and accuracy exhibits a discontinuity after about 250 steps. The evolution of the NTK for different widths is shown in Figure 12b, which shows the discontinuity and a high frequency oscillation that is width dependent. We also note that the effect gets smaller with larger width. The number of steps until this feature appears does depend on the hyperparameters in a systematic way, so they can be chosen such that this behavior can be avoided. It occurs earlier for deeper networks, and later for wider networks; later for smaller learning rates; and increasing the precision from float-32 to float-64 slightly delays it.

Appendix E Variance of NTK

Here we establish the claim that the variance of the NTK is $\mathcal{O}(n^{-1})$ for deep-linear networks with convolutional, skip, GAP layers as well as for one-hidden-layer non-linear CNNs.

In this case we can argue using Lemma 1 and the Feynman diagrams. The variance of the kernel takes the form

We are now in a position to use the Feynman diagrams to compute the scaling of $\textrm{Var}_{\theta}\left[\Theta_{I}(x,x^{\prime})\right]$. If we denote by $C_{1;I}$ the expectation of $\Theta_{I}$ and $C_{2;I}$ the expectation of the square.

$C_{1}$ can be computed from the double line diagrams associated to each graph in $\Gamma(C_{1})$ and $C_{2}$ can be computed from the graphs in $\Gamma(C_{2})$. For the variance, we actually need to compute $C_{1}^{2}$. The value of $C_{1}$ is given by summing the result of applying the Feynman rules to each diagram in $\Gamma(C_{1})$. Thus $C_{1}^{2}$ is given by applying the Feynman rules to every graph in $\Gamma(C_{1})\times\Gamma(C_{1})$.

From the rules for drawing Feynman diagrams, Definition 6, every diagram $\gamma\in\Gamma(C_{1})\times\Gamma(C_{1})$ is also a valid diagram in $\Gamma(C_{2})$. The scaling of the variance can thus be read off by considering the double-line diagrams $\textrm{DL}(\gamma)$ for $\gamma\in\Gamma(C_{2})\setminus\Gamma(C_{1})\times\Gamma(C_{1})$. These diagrams have four vertices, but a single connected component, thus the Feynman rules (Theorem 3) give $\textrm{Var}_{\theta}[(x,x^{\prime})]=\mathcal{O}(n^{-1})$

One-hidden-layer non-linear convolutional networks

Here, we proceed by brute force calculation. Adopting the notation of Equation (36), the NTK can be written as

To simplify notation, we write the above as $\Theta(x,x^{\prime})=\frac{1}{n}\sum_{i=1}^{n}\phi[\mathbf{U}_{i},\mathbf{V}_{i}]$. The variance of the NTK can be written as

In arriving at the last line we have used the fact that $\mathbf{U}_{i},\mathbf{V}_{i}$ are i.i.d. and that $\textrm{Var}_{\theta}\left[\phi[\mathbf{U},\mathbf{V}]\right]=\mathcal{O}(1)$.