The Neural Covariance SDE: Shaped Infinite Depth-and-Width Networks at Initialization
Authors: Mufan Bill Li
Mihai Nica
Daniel M. Roy
Venue: NeurIPS 2022 Accept (2022)
Keywords: Infinite-depth-and-width
SDE
activation shaping
initialization
NNGP
kernel shaping

Abstract: The logit outputs of a feedforward neural network at initialization are conditionally Gaussian, given a random covariance matrix defined by the penultimate layer. In this work, we study the distribution of this random matrix. Recent work has shown that shaping the activation function as network depth grows large is necessary for this covariance matrix to be non-degenerate. However, the current infinite-width-style understanding of this shaping method is unsatisfactory for large depth: infinite-width analyses ignore the microscopic fluctuations from layer to layer, but these fluctuations accumulate over many layers. 

To overcome this shortcoming, we study the random covariance matrix in the shaped infinite-depth-and-width limit. We identify the precise scaling of the activation function necessary to arrive at a non-trivial limit, and show that the random covariance matrix is governed by a stochastic differential equation (SDE) that we call the Neural Covariance SDE. Using simulations, we show that the SDE closely matches the distribution of the random covariance matrix of finite networks. Additionally, we recover an if-and-only-if condition for exploding and vanishing norms of large shaped networks based on the activation function. 

The Neural Covariance SDE: Shaped Infinite
Depth-and-Width Networks at Initialization
Mufan (Bill) Li University of Toronto, Vector Institute
Mihai Nica University of Guelph, Vector Institute
Daniel M. Roy University of Toronto, Vector Institute
Abstract
The logit outputs of feedforward neural network at initialization are condition- ally Gaussian, given random covariance matrix defined by the penultimate layer. In this work, we study the distribution of this random matrix. Recent work has shown that shaping the activation function as network depth grows large is neces- sary for this covariance matrix be non-degenerate. However, the current infinite- width-style understanding of this shaping method is unsatisfactory for large depth: infinite-width analyses ignore the microscopic fluctuations from layer to layer, but these fluctuations accumulate over many layers.
To overcome this shortcoming, shaped aefnite-depph-ande widt limit. We identify the precise scaling of the acti- vation function necessary arrive at non-trivial limit, and show that the random covariance matrix is governed by stochastic differential equation (SDE) that we call the Neural Covariance SDE. Using simulations, we show that the SDE closely matches the distribution of the random covariance matrix of finite networks. Addi- tionally, we recover if-and-only-if condition for exploding and vanishing norms of large shaped networks based on the activation function.
study the random covariance matrix in the
Introduction
Of the many milestones in deep learning theory, the precise characterization of the infinite-width limit of neural networks at initialization as a Gaussian process with non-random covariance ma- trix [MI, 2 was turning point. The so-called Neural Network Gaussian process (NNGP) theory laid the mathematical foundation to study various limiting training dynamics under gradient descent [B-I2]. The Neural Tangent Kernel (NTK) limit formed the foundation for rush of theoretical work, including advances in our understanding of generalization for wide networks [13-15]] Be- sides the NTK limit, the infinite-widtt mean- mean-field limit was developed [IG- [I6-19), where the different parameterization demonstrates benefits for feature learning and hyperparameter tuning [20 22).
Fundamentally, the -iitt paradigm derives results from the assumption that the depth of the network is held fixed while the widths of all layers grow infinity. Unfortunately, this assumption can be problematic for modeling real-world networks, as the microscopic fluctuations from layer to layer are neglected in this limit (see Figure M... particular, infinite-width predictions are shown to be poor approximations of real networks unless the depth is much less than the width [eg [El].
Impressive achievements of deep networks with billions of parameters crystallize the importance of understanding extremely large, deep neural networks (DNNs). An alternative to the infinite-width paradigm the infinite-dpphe -eenite-depth-and-.e. wewptheanddwiieee und-width paradigm. this setting, both the network depth and the width of each layer are simultaneously scaled to infinity, while their relative ratio d/n remains fixed [23, 25-22 29]. Recent work also explores using d/r as an effective perturbation parameter [30
Correspondence: muf li@mail. utwronto.... ca: ni cacamuuuueleph..o ca; daniel. roniilllry@uutoootoo.caa
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Column) Unshaped ReLU DNNs, (2.1)
(Right Column) Shaped ReLU DNNs, Definition
0.55
Median (1-Correlation) 1- 08
Median Correlation p
0.50
0.1
0.45
0.01
0.40
0.35
0.30
100
150
0.25
0.50
0.75
1
Layer F
Time
Markov Chain (Ours)
Infinite Width
-DE (Ours)
Infinite Width
0.20
0.15
Density
Density
2-
0.10
0.05
1
0-
0-
-16
-i2
-8
-0.5
log[l-ooreelttoon) log(1 23
Correlation p
Markov Chain (Ours) i Infinite Width
NN
SDE (Ours)
Infinite Width
NN
(o we
between post-activation vectors in ReLU net- theoretical predictions infinite-width paradigm. Left Col- umn: Pt vs. our Markov chain (2.10) vs. infinite-width update Pt+1 cK)(Pe) (see (2.I0) and ag
Figure 1: Simulations of correlation Ps"
uep lee?l
works, comparing finite NNs
note the log scale and 1 - here). Right Column: o.8 P[mn] vs. our Neural Covariance SDE vs. ODE dp p(pt) dt (see Theorem B Top Row: Median p function of layer. Bottom Row: Full distribution at final layer ( = d. Simulation details: n = = 150, Pa = 0.3, 2 samples for each. In right column: c+ 0, c-
DE step size 1e-2. Densities from kernel density estimation.
33 to study concentration bounds in terms of d/n [S [4]] This limit has the distinct advantage of being incredibly accurate at predicting the output distribution for finite size networks at initialization [27]]] significant improvement over the NNGP theory. Furthermore, it has also been shown that
there is feature learning in this limit [2 in contrast to the linear regime of infinite-widhh limits [S]. Considering the mathematical success of the NNGP techniques, the infinite infinite-depth-an depth-and width limit
hints at the possibility of developing an accurate theory for training and generalization.
An immediate issue of the infinite-depth limit is that this limit predicts that network output becomes degenerate as depth increases: on initialization the network becomes a constant function sending all inputs to the same (random) output BSS.
[3]] While degenerate outputs not necessarily issue in theory, poses a more serious problem in practice: degencrate correlations imply "sharp" input-output Jacobian, and therefore exploding gradients [57, ESJ. Intuitively, the output is not very sensitive to changes in the input, hence the gradient must be very large in the earlier layers.
A promising new attack on this problem is to modify the activation function ("shaping") to reduce to the effect of degeneracy [[8, 39). In this prior work, extensive experiments show that shaping the activation significantly improves training speed without the need for normalization layers. This method has been proven effective for problems as large as standard ResNets on ImageNet data. The authors designed several criteria including reducing estimated output correlation, and numerically

---

Notation
Description
Notation
Description
Table 1: Notation
Tha E I
Input dimension Hidden layer width Base activation
Nhmm E N d E N
Output dimension
E N
Number of hidden layers (depth)
(P()) x" E R"
(P(()) Wo E TR N
Shaped activation
Input for S a S m Network output Neurons activaaiit)io
Weight matrix at layer 0 Weight matrix at final layer
Piur
E R
Won É R R Ws 2
2/ R
Weight matrix layer 1 $ d
for layer < 8 dd Neurons (post-activation for layer < 8
All weights initialized iid (0, 1)
306 E R"
CER
Normalizing constant
e (E for gorg~~V(0.
optimized the shape of activation functions for improved training results. However, their determin- istic estimation of output correlation using the infinite-width limit leads to a poor approximation of real networks, as the additional randomness has both non-zero mean and heavy skew (see Figure right column). Furthermore, numerically searching for the activation shape obscures the picture on how shaping should depend on the network depth and width.
In this paper, we address these problems by providing a precise theory of shaped infinite-depth- and-width networks, extending both the NNGP theories and the activation shaping techniques. In particular, we prescribe an exact scaling of the activation function shape as a function of network width n that leads to a non-trivial nonlinear limit. By keeping track of microscopic O(n 11/2) ran- dom fluctuations in each layer of the network, we show that the cumulative effect is described by stochastic differential equation (SDE) in the limit. contrast to existing infinite-width theory, we able to characterize the random distribution of the output covariance, which matches closely to simulations of real networks. In a similar spirit to how the NNGP theory laid the foundation for studying training and generalization in the infinite-width limit, also see this work as building the mathematical tools for an infinite-depth-and-wee theory of training and generalization.
1.1 Contributions
Similar to the NNGP approach, we use the fact that the output is Gaussian conditional on the penul- timate layer. However, unlike in the infinite-width paradigm, the covariance matrix is no longer deterministic the dnfnite-depth-and--iie covariance matrix. Our main contributions are as follows:
limit. Our focus this paper to study this random
1. We introduce the tool of stochastic Vn-eepansioos and convergence to SDEs for analyzing the distribution of covariances in DNNs.
For unshaped ReLU-like activations, we show that the norm of each layer evolves according to geometric Brownian motion and correlations evolve according to a discrete Markov process. See left column of Figure and Section D
3. For both ReLU-like and large class of smooth activation functions, we derive the Neural Covari- ance SDE characterizing the distribution of the shaped infinte-deeth-andd-iiie limit. See right column of Figure M and Section B
4. We show our prescribed shape scaling is exact, other rates of scaling leads to either degenerate or linear network limits. See Proposition H and Proposition GT00
5. For smooth activations, we derive an if-and-only--f condition for exploding/vanishipg norms based on properties of the activation function. See Proposition B.7 and Section 21
We provide simulations to verify theoretical predictions and help interpret properties real DNNs. See Figures M and A and supplemental simulations in Appendix
Limits for Unshaped ReLU-Like Activations
Using the notation in Table the output of fully connected feedforward network with d hidden layers of width n on input -º is defined by vectors of pre-activations =ē and post-activations y
1
Wox",
28 = =(2)), 24 -
Wepp
2omt = =/
(2.1)
Tim
3

---

Note that factors of ven are equivalent to intializing according to the so-called He initialization [90]. We use Greek indices a, S.... denote multiple different inputs. Note that while our results are all stated for fixed width n in each layer, they can be generalized to layer width in the limit
where all ne
with Σ ni replacing the role of the depth-to width ratio d/n [25].
In this section, we analyze ReLU-like activations by which we mean activations which are linear on the negative and positive numbers given respectively by two slopes 8+ and 8-:
(xx) 54 max(x, 0) 8- min(2; 0) 4+PPReU --- (RRLL((-x) These are precisely the positive homogeneous functions: p((a)) |a (2x) V2, R.
(2.2)
2.1 SDE Limits of Markov Chains
We briefly review the main type of SDE convergence principle used in our main results (see Propo- sition B.G for more precise version). Let Xx, / E IRt, be a continuous time diffusion process obeying an SDE with drift and variance o² as given in 3) Suppose that for each n E N, Y is a discrete time Markov chain é E N whose increments obey
terms of the same functions b, o²:
+ (Y?)) & +0 3/2). (2.3) n
dXx; b(X;) (Xx)
Y - YP" = b({?))
Vn
where & are independent variables with B(G)) = O, Var(E)) 1. With this setup, under technical conditions described precisely in Appendix A,, have convergence of Yz at = [tn] to X or more precisely: with X3" == [tn] we have X" + X as + do in the Skorohod topology. In our applications, n is always the width (i.e., number of neurons in each layer) which may appear implicitly and f always the layer number.
2.2 A Simple SDE: Geometric Brownian Motion Describes |
To motivate our approach of SDE limits, we illustrate the method using the example of the squared norm of the l-th layer, ? where we recall of = ==())) For a single fixed input I. and a ReL.U. like activation yp, the norm of the post-activation neurons | 2 forms a Markov chain in the layer number f. We use the fact that a matrix with iid Gaussian entries applied to any unit vector gives a Gaussian vector of iid (0, 1) entries. Hence, in each layer, can define the Gaussian vector g" as follows, and use (2.I) with the positive homogeneity of yp to write the Markov chain update rule:
| 1- = 118 Σ cp())) where g" = We
In)
(2.4)
At this point, the widtti approach applies the law of large numbers (LLN) to conclude lim |r2+1| 1v8 1 [g)) 1 a.s. by definition of c. However, the LLN cannot n-+00
be applied when depth d is diverging with n, the cumulative effect of the fluctuations over d lay- ers does not vanish! Instead, we keep track of the ((1/VT) fluctuations in each layer by introducing the zero mean finite variance random variable n7 Vu Ex= (ou(ge)) - 1).. This allows us to rewrite this Markov chain update rule as
1 2 = 1 1
- 1700
1
(2.5)
Vn
which allows us to see that the Markov chain Yr = , 11 is now in the form of (2.33) with YI 1, 12 b(Y) (Y)²(Y) = Var Var(R R22 = Var(cp(g)>)Y² Consequently, we have that the squared norm Markov chain converges to a geometric Brownian motion d.x, = B, or more precisely
|[tm] = X, d N( N(-
lim ==0
(2.6)
where the convergence in the Skorohod topology (see Appendix A).. When y is the ReLU function
(++ = 1, 8_ 0). we have 2 and o² = 5, which recovers known results in [25, 27-29]
We remark again this simple Markov chain example illustrates the main technique we use later sections to establish SDE convergence for shaped networks in Section B

---

2.3 Non-SDE Markov Chains: the Gram Matrix (F8, (p?) and Correlation Pt
We can generalize Section 2
to a collection of inputs [a")m= la=1 by looking at the entire Gram matrix K(p ( 80 where we again recall Pi (p(=)) We note that the convergence of Markov chains to SDEs in (2.5)) can be generalized to Y E R- by considering Cov(Ex) = Ix. b RN + RN, and a : IR + N The Gram matrix is of particular interest because the neurons in any layer are conditionally Gaussian when conditioned the previous layer, with covariance matrix proportional to the Gram matrix:
F N (o, "Kpt. 90 8 8In) Fa d
[2 1f+1la=1
n (o, C [Kp8 (. &a) @ Inn.
(2.7)
n
where Fz denotes the sigma-algebra generated by the f-th layer [=f|==1, and @ denotes the Kro- necker product (here indicating conditionally independent entries in each vector). With this prop- erty in mind, we will introduce [E = [F[] denote the conditional expectation, and Varz(.) Cove(-) similarly to denote the conditional variance and covariance. If we define g" as in (2.d), we see that the g are all marginally N NV In). Similar to (24)) we can write the update
rule for the a; B-entry of the Gram matrix:
( (9tt 108 11
1 n
" Σ
oploft)(oe))
(2.8)
Just as we did in (2.5), we can define Rºd = 1 (cleep(o))
(1lpplp((pp(() and write
v
(fi (( = 148 11 [E (cplof)(lo"))
Vn J/
1
(2.9)
where R28 are mean zero with covariance [o R R7$1 Covellp(o")))))o (( cp(o))
(By the Central Limit Theorem, 12 will be approximately Gaussian for large n.)
However, unlike the simple single-data- point case from Section 2.2 we do not have convergence to a continuous time SDE. This is because the differences (42) ( - (59 v/) ~ 0 as n + 00.
Instead, (2.3) is discrete recursion update with additive noise of the form Y f(Y?) + J
for some function f. and consequently Y Y does not vanish as > 3c.
For a clarifying example, we can consider the one-dimensional Markov chain of hidden layer correla- tions. More precisely, we can define as (vp / /l 11 which we observe can be extracted from the entries of the Gram matrix. In fact, we can write down an approximate recursion update a9 for Pe (see Appendix and Proposition BN for details):
Pet R cKi ( + 1 (
+ Ee oRaLL(0)))
(x x
(2.10)
Vn
nin
where = =x(p)==Elp(0)) El(())((pp + w p=)]==p) for iid N 1) random variables, and & are iid first calculated in [II]]. In fact, we can observe that as n - 0o, P[m]] converges to the fixed point of 1.) at p = for all t > 0. We note this limiting behaviour cannot be described by an SDE, as the solution must jump from the initial condition to the fixed point at
[O, 1). For the ReLU case, 2 and cK2 (p) (vI -
prccos(-p)) pprccppparccos(ppp)
Despite not having an SDE limit, we observe that the approximate Markov chain (2.10) already provides a much better approximation to finite size networks compared to the infinite-width theory (see left column of Figure M.. This is because the infitite-width approach discards the terms in (2.10) that vanish as n + oc and consider only the update PE+1 = cK1(p a²)) Analysis of this
deterministic equation leads to the prediction that Pt a and new bound in Appendix E)).
101100 O for e >>>>> (see (4.8) in
Furthermore, we observe that in this case, the microscopic OIn-1) and O(n -1/2) terms (2.10)
accumulate to macroscopic differences! For the examples in Figure MI we see their net effect is that
5

---

0.3 - faster than the infinite-widde prediction. Heuristically, the reason for this discrepancy is Pe due )R-LU(D)
p + 1. This means that the randomness can push o closer 1, but becomes "trapped" when 2/ is close to I because RELU is so small here. In the next section, we will see that are just one step away from achieving limiting SDEs.
3 Neural Covariance SDEs: Shaped Infinite-Deptt-and--e Width Limit
In this section, we follow the ideas of [38 39] to reshape the activation function Reshaping means to replace the base activation function yp in (2.I) with ys that depends on width n. We will also replace the normalizing constant
will choose o to depend on n such that in the limit as n + oo, we have that yp is approximately an identity function, Sa + Id. Recalling from (2.7) that the output is conditionally Gaussian with covariance determined by the Gram matrix [(p8 P 8 [2),),,a= therefore recover complete characterization by describing the random covariance matrix.
(E (g)
for g NN0,, 1). Specifically,
3.1 Neural Covariance SDE for Shaped ReLU-Like Activations
Definition 3.1. We shape the ReLU-like activation (2(2x) = 84 max(x, 0) + s_. min(2; 0), by setting the slopes to depend n according $+ == 1 + et for some given constants C4, R. v
We will also set c = (Bv,(g)))) for g ~ 10.
We will show that with shaping of Definition B one gets non-trivial SDEs that describe the covari- ance (Theorem ( and correlations (Theorem ( of the network. The precise scaling is shown to
be the critical scaling for non-trivial limit Proposition 3.4. All proofs for results this section appear in Appendix Q
Remark. Note that in the statement of our theorems, abuse notation and the same letter to denote the pre-limit Markov chain and the limiting SDE. For example, in Theorem 6.7 we use Vi for the covariance layer& and V, to denote the limiting SDE at time t.
Theorem 3.2 (Covariance SDE, ReLUU), Let Va == f. (et :P)) and define V== == [V2 15a5g== to be the upper triangular entries thought of as vector in 1mm(+1)/2 Then, with $ = 1 + =
v
as in Definition B.l in the limit as n + 00; + T, the interpolated process V VImn converges in distribution the Skorohod topology D Dz Bm(m+ to the solution of the SDE
1
dV22 b(V)) dt E(V)) dB..
Vo
x
(3.1)
niin
Isasssm
where v(p)=
(c) - )
(V1- p2
arccossp p) Pt ai
yog
2x
b(V.) = ( (020) V
and E(Vi) = V. + Va05ve
15a58$m
(3.2)
Furthermore, the output distribution can be described conditional Vr evaluated at final time T
==odda=1 Vr d N (0, [V. ( )=,-=1)-
(3.3)
Here remark that v11 0, and therefore the drift component of diagonal entries (Voa) zero, as they are geometric Brownian motion. However, we emphasize that the m m-point joint out- put distribution is not characterized by the marginal for each of the pairs, as the output 2 is not Gaussian. particular, we observe the diffusion matrix entry corresponding to V.a V.7 involves other processes V.ay, V,95 V. V. This implies that the Neural Covariance SDE limit cannot be described by kernel, unlike stacking random features or NNGP.
That being said, is still instructive to study the marginal for pair of data points. More specifically, it turns out in the generalized ReL.U case, we can derive the marginal SDE for the correlation process.

---

Theorem 3.3 (Correlation SDE, ReLU), Let PE = (p7 |
where o2 po(zzf. In the limit as
==
n - do and 81 =1+ tn the interpolated prucess PLin] converges in distribution to the solution of the following SDE in the Skorohod topology of D
aB
dp:
as
[v(p= +
dt + olpe
a0 dBB,
Po
ag
f2 xx
x |r
(3.4)
where
V1 12 - v(p) = (c+ - c_
((() = -2 - p²) o(0) = 11)==1pp1 p²_ (3.5)
2x
To help interpret the SDE, we observe that 4 a are entirely independent of the activation function. In other words, these terms will be present in this limit even for linear networks. At the same time, describes the influence of the shaped activation function in this limit. [69] has derived a related ordinary differential equation (ODE) of dp: p(px) dt in the sequential limit of n + oo then d d 00, where the activation is shaped depending on depth. Here we also note that v(p) closely related to the J1 function derived in [_I]]. See Appendix C for the m-point joint version of the correlation SDE, and Appendix for an empirical measure of convergence in the Kolmogorov- Smirnov distance.
One immediate consequence of the correlation SDE is that nition I. is the only case where the limit neither degenerate nor a linear network.
show the n -122 scaling in Defi-
Proposition 3.4 (Critical Exponent, ReLU), Let PE aa = 107 1ø/l where ot = (o(2))) Consider the limit + and 8+ 1 + et ny for some p O. Then depending the value of p. the interpolated process Plenj converges in distribution w.r.t. the Skorohod topology of D to
(i) the degenerate limit: Pt = for all t > 0, if0 S < 2. and c # C...
(ii) the critical limit: the SDE from Theorem
i
(iii) the linear network limit: i > 1 the following SDE, with μ. a as defined in (3.))
00 dB,, Po o8 (x xx dp. p( ) (pi
(3.6)
[ [200
Here we remark that the unshaped network case (p 0) is contained by the above in case (i). At the same time, we observe that case (iii) is equivalent to the correlation SDE in Theorem B. except with v = 0. In particular, we observe this limit is also reached when e_ = c... which implies (x) s is linear, which the reason we call this the linear network limit. Furthermore. without much additional work, the same argument also implies the joint covariance SDE also loses the drift component, i.e,, dV, = 1/2 dB,
3.2 Neural Covariance SDE for Shaped Smooth Activations
In this section, we consider smooth activation functions and derive a similar covariance SDE, All the proofs for results in this section can found in Appendix D
Assumption 3.5. &p E C4 (R), (0) = 0, (0) = L, and | (x) & C(1+(xP) some C, >
We note that for any non-constant function a C' and zo R such that o' o'(x)) 0. we can always define *xx) == a(z+ (zo) such that it satisfies (0 0, 4'(0) 1. The choice of xo
d'(2g)
will be discussed further in Section HL The fourth derivative growth condition is used to control the Taylor remainder term in expectation, but any control over the remainder will suffice.
Following the ideas of [38], we consider the following shaping smooth activation function y. Definition 3.6. For some constant 0. we set (p.()) = syp (3) with avTi, and (E po)))) for g (0, 1).
7

---

Observe that in the limit + oc, we will achieve that Pe - Id as desired. We also observe that the shaping factor outside the activation cancels out with the next layer's factor, therefore it is equivalent shape the entire network. More precisely, if we view Zout as an input-output map f: TR IR. Foee of unshaped network, then shaping the smooth activation functions equivalent to the modification sf (#).
In this regime, we can similarly characterize the joint output distribution, however the limiting SDEs are not always well behaved. In particular, they can have finite time explosions as described by the Feller test for explosions [42, Theorem 5.5.29]. Here the SDE in Proposition B is exactly the Vi marginal of the Neural Covariance SDE, with the parameter b determined by the activation function o and controls whether or not finite time explosions happen (see (I.
Proposition 3.7 (Finite Time Explosion). Let X, E 1 be a solution to the following SDE
dX, -Xx(x - 1) dl + V2X, dB:,,
Xo= zo 0, I.
(3.7)
Let 7" sup inf(l X 2 M or Xx S M-1] } be the explosion time, and say X has finite time explosion if < xo. For this equation, P([*=xc] x if and only if I O.
=
Technically speaking, the main culprit behind finite time explosions is the non-Lipschitzness of the drift coefficient. This issue requires us to weaken the sense of convergence in this section; the ordinary convergence in the Skorohod topology is in general not true when the diffusion has finite time explosions. A weakened type of convergence is the best we can hope for. To this goal, we introduce the following definition.
Definition 3.8. We say sequence of processes X" converge locally to X in the Skorohod topology for any O. we define the following stopping times
T" = {t 0: [X]' 2 ==
0: 0:
(3.8)
and we have that X E converge to XLAT in the Skorohod topology.
This weakened sense of convergence essentially constrains the processes X", _X in a bounded set by adding an absorbing boundary condition. Not only do these stopping times rule out explosions, the drift coefficient is now also Lipschitz on compact set. With this notion of convergence, we can now state precise Neural Covariance SDE result for general smooth activation functions.
v. = 5. (V2 ( where off = (2f)) ) and define Ve == [V lıgas=m to be the upper triangular entries thought of as a vector in mm(m+1)/22 Then, with = ayn as in Definition B.d, in the limit as - xo, - T, the interpolated process Vitn| converges locally in distribution the solution of the following SDE in the Skorohod topology of DR,
Theorem 3.9 (Covariance SDE, Smooth). Let ip satisfy Assumption M
R Rm(m+1)/2
1
dV. = b(V)) dt + 12V)) dB. Vo =
[x .23) niin
(3.9)
15as3sm
where E(Vi) is the same as Theorem and 608 (08(V)) Furthermore, if VT is fintte then the output distribution can be described conditional on Vr as
(0)²
(v. + V -3)) +
(0((0)
yai (Vaa + v.8 22). (3.10)
Aa²
2a
[2 IVg - N (0. [V ( 10,
(3.11)
and otherwise the distribution of [2 1m is undefined.
We also have a similar critical scaling result for general smooth activations. Proposition 3.10 (Critical Exponent, Smooth). Let ip satisfy Assumption vaB [1p :0)) where op (2() with anP for some p O, and define Vi [V, I5a5s=m be the upper triangular entries thought of vector. Then in the limit
co, d T, the interpolated process V/en converges locally in distribution w.rlt the Skorvhod topology of Dz R1/ttt//22 to V, which depending on the value of p
We want to thank Boris Hanin for observing this equivalent parameterization.
8

---

(a) Unstable Centre 2o 0
(b) Stable Center dgg log
Magnitude Sample Paths 1/Ma
1e17
Magnitude Sample Paths V
1e11
1e5
1e-1
0.1-
0
25
50 Layer
75
100
25
50
75
100
Layer
Figure 2: Simulation of 10 shaped softplus networks in Example BW with d 100,
1. Vi 1- = centred at two different values. "Stable" here means the Neural Covariance
SDE guaranteed not to have finite time explosions; unstable networks can explode on initializa- tion!
(i) the degenerate limit: if0 < <
Vaa = 0 or oc,
is 122"
> 0 and Vo #0,
Vas const.
if
(3.12)
for all 0 and 1 S a B m,
(i) the critical limit: the solution of the SDE from Theorem B.9 is
(iii) the linear network limit: the stopped solution to the SDE dV; E(V) dB, with coefficient
E defined Theorem if > 2"
Here we observe that in case when )0) + (0) O, we also have a constant (in time) cor- relation aB similar to the ReLU case in Proposition 3.4, however in this case p. a5 is not necessarily equal the same time, the linear network limit case (iii) also has the same covariance SDE
as Proposition 3.4
4 Consequences, Discussion, and Future Directions
So far, we have derived the Neural Covariance SDE. Analysis of this SDE reveals important be- haviour of the network on initialization. Here we lay out one concrete example and provide some discussion and future directions.
Exploding and Vanishing Norms. Here we consider the behaviour of shaping smooth activation functions, it is done in the experiments of [G8]. While the authors here avoided exploding and vanishing norms by numerically optimizing shaping parameters, we can actually describe the precise behaviour a priori with the Neural Covariance SDE. Recall the shaping parameter from Defini- tion B.R. Let VÚ be the solution to the SDE in (B.9). We can write down the marginal SDE for V
dVac (3 (0)) (0)
Vax
1) dt V2V. dB,, (V
(4.1)
which implies by Proposition 3.7 that VE has finite time explosion (with non-zero probability) This criterion can be used to help choose how activation
and only if 3 ==" (0) (0(()
functions should be centered for shaping; below are two examples.
Example 4.1 (Sigmoid and tanh at 20 = =)) We start with the sigmoid activation oxx) Ite
then we can define (xx) = 40 - to satisfy Assumption HR which leads to (" 0, (0) and therefore leads stable network. It nurns (p(x) tanh(x) already satisfies Assump-
tion m which leads to (0(()) (0)(0) 22, and therefore is also stable.

---

More generally, if a behaves like a cumulative distribution function for a symmetric unimodal den- sity, we will have that (0) 0 and (0)(() < O as desired.
Example 4.2 (Soft Plus General To E IR). Let us consider so E R and ((x) log(1 (1 + ea+20), which implies ((x) = (1 + e === log 1+c satisfies Assumption B. This gives us 4" 1te'0
1 Ite (d((0) -- (I+e*ayz, and therefore 30"(0)² + (0) = (Ite=ojz ( -e"o). In other words, 5 the shaped network is stable if and only if ao 2 log ā (see Figure R). We note that the authors of
[38] numerically found a shift of zo = 0.41, which is in the stable regime of xu log 3 0.097.
Relationship to Edge of Chaos. The finite time explosion example above resembles the Edge of Chaos (BOC) analysis of gradient stability [83 G5, Eel, ESI, where the weight and bias variance at initialization determines a stability criterion. However, we note that the EOC regime is sufficiently different that the results are not directly comparable. More precisely, the EOC analysis in the sequential limit of infinite- width and then infinite-depth, which also leaves the activation function unchanged. Under very weak assumptions, the variance (diagonal of Vi) will not explode in this regime; instead, the gradient can explode due to the covariance (off diagonals). On the other hand, our finite explosion result is in the joint limit of depth and width, where the variance (diagonal of Vi) can explode instead.
Posterior Inference. ate a prior over functions f II + Ro Consequently, an interesting future direction would be to study the posterior distribution. i.c. the output m conditioned on æ==1 and a training dataset (x 2gm y However, to our best knowledge, it is not straightforward to explicitly compute or sample from the conditional distributions for this SDE structure. It would be desirable to extend
Similar to the NNGP setting, we can use the Neural Covariance SDE to gener-
existing approaches in the perturbative regime [GO, BII to our setting.
Extension to Other Architectures. The key step to deriving the covariance SDE is the conditional Gaussian distribution in (2.2),, which directly leads to Markov chain. It follows immediately that ResNets [86] admit a similar conditional structure. With bit more work for convolutional networks, we can obtain = 1 IFee (0, A(V)) @ In) where A is an affine transformation and Vz is the previous layer's 's Gram matrix [_2]]. We note that recurrent networks will not lead to a Markov chain or SDE limit, as the weight matrix is reused from layer to layer.
Simulating SDEs. Both the Markov chains and SDEs predict neural networks at initialization very well (see Figure I1,, but the SDE is significantly faster to simulate. In particular, we can view the Markov chain as an approximate Euler discretization of the SDE, but with very small step size n=1 In contrast, to simulate the SDE we should only need a step size that is small on the scale of depth-to-width ratio T = n/,, which is independent of width n. Therefore, practitioners using the shaping techniques of [E8, [9] can now simulate the covariance SDEs at low computational cost to significantly improve estimates of the output correlation (see Figure II and additional simulations in Appendix E).
Analytical Tractability of SDEs. Besides numerical tractability, the SDEs are also far more tractable to analyze. For example, in the one input case, we arrive at geometric Brownian mo- tion (2.6), which is known to have a log-normal distribution at fixed times. Similarly, our finite time explosions hinge on the fact we identified an SDE limit. In the same way that NNGP theory played a major role in the mffnite-widthh regime, the Neural Covariance SDEs and the techniques developed here also serve as mathematical foundation for studying training and generalization.
Acknowledgement
We would like to thank Sinho Chewi, James Foster, Boris Hanin, Cameron Jakub, Jeffrey Negrea, Nuri Mert Vural, Guodong Zhang, Matthew S. Zhang, and Yuchong Zhang for helpful discussions and draft feedback. We would like to thank Sam Buchanan and Soufiane Hayou for pointing out a gap in the proof of Proposition B. ML is supported by Ontario Graduate Scholarship and the Vector Institute. MN is supported by an NSERC Discovery Grant. DMR is supported in part by Canada CIFAR AI Chair funding through the Vector Institute, NSERC Discovery Grant, Ontario Early Researcher Award, stipend provided by the Charles Simonyi Endowment, and a New Frontiers in Research Exploration Grant.
10

---

References
[1] R. M. Neal. Bayesian learning for neural networks. Vol. 118. Springer Science & Business
Media, 1995.
[2] J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. "Deep Neural Networks as Gaussian Processes." Int. Conf. Learning Representations (ICLR). 2018.
[3] A. Jacot, F. Gabriel, and C. Hongler. "Neural tangent kernel: Convergence and generalization
in neural networks." Advances Information Processing Systems (NeurIPS). 2018, arXiv:
1806.07572.
[4] S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. "Gradient descent finds global minima of deep neural networks." Int. Conf. Machine Learning (ICML). PMLR. 2019, pp. 1675-1685.
[5] Z. Allen-Zhu, Y. Li, and Z. Song. "A convergence theory for deep learning via over- parameterization." " Int. Conf. Machine Learning (ICML). PMLR. 2019, pp. 242-252.
[6] D. Zou, Y. Cao, D. Zhou, and Q. Gu. "Gradient descent optimizes over-parameterized deep ReLU networks." In: Machine Learning 109.3 (2020), pp. 467-492.
[7] L. Chizat, E. Oyallon, and F. Bach. "On Lazy Training in Differentiable Programming." In: Advances in Neural Information Processing Systems 32 (2019), pp. 2937-2947.
[8] J. Lee, L. Xiao, S. S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington.
Wide neural networks of any depth evolve as linear models under gradient descent. 2019.
arxiv: 1902.06720..
[9] G. Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behav- ior, gradient independence, and neural tangent kernel derivation. 2019. arXiv: 1902. 04/60 /60.
[10] G. Yang. Tensor programs il: Neural tangent kernel for any architecture. 2020. arXiv: 2006
14548.
[11] S. Arora, S. S. Du, W. Hu, Z. Li, R. Salakhutdinov, and R. Wang. "On exact computation with infinitely wide neural net." Proceedings of the 33rd International Conference on Neural Information Processing Systems. 2019, pp. 8141-8150.
[12] Z. Chen, Y. Cao. D. Zou, and Q. Gu. "How Much Over-parametteizatton Is Sufficient to Learn Deep Re[LU] Networks?" International Conference on Learning Representations. 2021. URL:
uZa6.
[13] Z. Ji and M. Telgarsky. "Polylogarithmic width suffices for gradient descent to achieve ar- bitrarily small test error with shallow relu networks." In: arXiv preprint arXiv: 1929.1222 (2019).
[14] J. Ba, M. Erdogdu, T. Suzuki, D. Wu, and Zhang. "Generalization of two-layer neural networks: An asymptotic viewpoint." International conference on learning representations. 2019.
[15] P. L. Bartlett, A. Montanari, and Rakhlin. "Deep learning: statistical viewpoint." In: Acta numerica 30 (2021), pp. 87-201.
[16] G. M. Rotskoff and E. Vanden-Eijnden, Trainability and Accuracy of Neural Networks: An Interacting Particle System Approach. 2018. arXiv: 1805500915
[17] L. Chizat and F. Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. 2018. arXiv:
[18] J. Sirignano and K. Spiliopoulos. Mean Field Analysis of Neural Networks: A Law of Large Numbers. 2018 arXiv: 1805_ 0T00X
[19] S. Mei, A. Montanari, and P.-M. Nguyen. "A mean field view of the landscape of two- layer neural networks." In: Proceedings of the National Academy of Sciences 115.33 (2018), E7665-E7671. ISSN: 0027-8424. DOI: 10. 1073//0aa 1806579115, eprint: httpp::/wwwr mmas.org/ wwwooneennnn//////eooo
full. w.l.lppp.. fdf.
[20] G. Yang and E. J. Hu. "Feature Learning in Infinite-Wird Neural Networks." Int. Conf. Ma- chine Learning (ICML). 2021. arXiv: 2011. 14522
[21] G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, Pachocki, W. Chen,
and J. Gao. "Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparam- eter Transfer." In: arXiv preprint arXiv:2203 03366 (2022).
[22] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. "High-dimensions Asymp- totics of Feature Learning: How One Gradient Step Improves the Representation." " In: arXiv preprint 2rXiv:2205.014 (2022).
11

---

[23] B. Hanin and M. Nica. "Finite Depth and Width Corrections to the Neural Tangent Kernel."
Int. Conf. Learning Representations (ICLR). 2019.
[24] M. Seleznova and G. Kutyniok. Analyzing Finite Neural Networks: Can We Trust Neural
Tangent Kernel Theory? 2020. arXiv: 2012.04471.
[25] B. Hanin and M. Nica. "Products of many large random matrices and gradients in deep neural networks." In: Communications in Mathematical Physics (2019), pp. 1-36.
[26] Z. Hu and H. Huang. "On the Random Conjugate Kernel and Neural Tangent Kernel." Inter- national Conference on Machine Learning. PMLR. 2021, pp. 4359-4368.
[27] M. Li, M. Nica, and D. Roy. "The future is log- gg-Gaassian: ResNets and their infinte-depthh and-width limit at initialization." In: Advances in Neural Information Processing Systems 34 (2021).
[28] J. Zavatone- Veth and C. Pehlevan. "Exact marginal prior distributions of finite Bayesian neu- ral networks." " In: Advances in Neural Information Processing Systems 34 (2021).
[29] L. Noci, G. Bachmann, K. Roth, S. Nowozin, and T. Hofmann. "Precise characterization of the prior predictive distribution of deep ReLU networks." In: Advances in Neural Information Processing Systems 34 (2021).
[30] S. Yaida. "Non-Gaussian processes and neural networks at finite widths." " Mathematical and Scientific Machine Learning. PMLR. 2020, pp. 165-192.
[31] D. A. Roberts, Yaida, and B. Hanin. The principles of deep learning theory. Cambridge University Press, 2022.
[32] J. Zavatone- Veth, A. Canatar, B. Ruben, and C. Pehlevan. "Asymptotics of representation learning in finite Bayesian neural networks." In: Advances in Neural Information Processing Systems 34 (2021).
[33] B. Hanin. "Correlation Functions in Random Fully Connected Neural Networks Finite Width." In: arXiv preprint arXiv:2204.0100 (2022).
[34] S. Buchanan, D. Gilboa, and J. Wright. "Deep Networks and the Multiple Manifold Problem." International Conference on Learning Representations. 2021. URL: https:/ pet/forum/id-Due
d Q=.
[35] G. Yang and S. S. Schoenholz. "Mean field residual networks: on the edge of chaos." Ad- vances in Neural Information Processing Systems. 2017, pp. 2865-2873,
[36] S. Hayou, F. Clerico, B. He, G. Deligiannidis, A. Doucet, and J. Rousseau. "Stable ResNet." Int. Conf. Artificial Intelligence and Statistics (AISTATS). PMLR. 2021, pp. 1324-1332.
[37] B. Hanin and D. Rolnick. "How to Start Training: The Effect of Initialization and Architec- ture." In: Advances in Neural Information Processing Systems 31 (2018).
[38] J. Martens, A. Ballard, G. Desjardins, G. Swirszez, V. Dalibard, J. Sohl-Dickstein, and S. S. Schoenholz. "Rapid training of deep neural networks without skip connections or normaliza- tion layers using Deep Kernel Shaping." " In: arXiv preprint 10.0122 (2021).
[39] G. Zhang, A. Botev, and J. Martens. "Deep Learning without Shortcuts: Shaping the Kernel with Tailored Rectifiers." In: arXiv preprint arXiv:2203.08122 08100 (2022).
[40] K. He, X. Zhang, S. Ren, and J. Sun. "Delving deep into rectifiers: Surpassing human- level performance on imagenet classification." Proc. IEEE Int. Conf. Computer Vision. 2015,
pp. 1026-1032
[41] Y. Cho and L. K. Saul. "Kernel methods for deep learning." Advances in Neural Information Processing Systems (NeurIPS). 2009, pp. 342-350.
[42] I. Karatzas and S. Shreve. Brownian motion and stochastic calculus. Vol. 113. Springer Sci-
ence & Business Media, 2012.
[43] S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. "Deep information propaga- tion." In: arXiv preprint 161..0111..0123 (2016).
[44] S. Hayou, A. Doucet, and J. Rousseau. "On the impact of the activation function on deep neu- ral networks training." International conference on machine learning. PMLR. 2019, pp. 2672- 2680.
[45] M. Murray, V. Abrol, and J. Tanner. "Activation function design for deep networks: linearity and effective initialisation." In: Applied and Computational Harmonic Analysis 59 (2022). pp. 117-154.
12

---

[46] K. He. X. Zhang, S. Ren, and J. Sun. "Deep residual learning for image recognition." Pro- ceedings of the IEEE conference on computer vision and pattern recognition. 2016. pp. 770- 778.
[47] R. Novak, L. Xiao, J. Lee, Y. Bahri, G. Yang, J. Hron, D. A. Abolafia, J. Pennington, and J. Sohl-Dickstein. "Bayesian deep convolutional networks with many channels are gaussian processes." In: arXiv preprint arXiv: 1810.05148 (2018).
[48] O. Kallenberg. Foundations of Modern Probability. Probability theory and stochastic mod- elling. Springer, 2021. ISBN: 9783030618722
[49] S. N. Ethier and T. G. Kurtz. Markov processes: characterization and convergence. John
Wiley Sons, 2009.
[50] D. W. Stroock and S. S. Varadhan. Multidimensional diffusion processes. Vol. 233. Springer
Science & Business Media, 1997.
[51] A. Meurer, C. P. Smith, M. Paprocki, O. ertík, S. B. Kirpichev. M. Rocklin, A. Kumar, Ivanov, J. K. Moore, S. Singh, T. Rathnayake, S. Vig, B. E. Granger, R. P. Muller, F. Bonazzi, H. Gupta, S. Vats, F. Johansson, F. Pedregosa, M. Curry, A. R. Terrel,, Rouka, A. Saboo, I. Fernando, S. Kulal, R. Cimrman, and A. Scopatz. "SymPy: symbolic computing in Python." In: PeerJ Computer Science (Jan. 2017), e103. ISSN: 2376-5992, DOI: 10. 7717/peerr
103. URL: https: /ttp:://dooooo//////// Ntrg/10.1e 1.1..7177/peer-eee
103,
[52] T. Campbell and T. Broderick. "Automated scalable Bayesian inference via Hilbert coresets."
In: The Journal of Machine Learning Research 20.1 1 (2019), pp. 551-588.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] The contributions list directly references the theorems and sections corresponding them.
(b) Did you describe the limitations of your work? [Yes] In introduction, discussed that our theory is only at initialization, and does not describe training dynamics and generalization.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] The only assumption required is in Assumption SMS
(b) Did you include complete proofs of all theoretical results? [Yes] Please see the rele- vant sections in Appendices Al C and D
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material as URL)? [Yes] All of our simulations (at initialization only) are contained the file Correlation.ipyns.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A] Our experiments only required simulations at initialization, and therefore no training.
(c) Did you report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [N/A] We plotted the histogram and full density of the samples instead.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [N/A] Our simulations were small enough that it did not require GPUs.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
13

---

(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you' re using/curating? [N/A]
(e) Did you diseuss whether the data you using/curating contains personally identifi- able information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "jjf2GogXDy4", "reviewer": "Reviewer_stwd", "summary": "The paper studies the random covariance matrix of infinite depth and width limit of initialized deep neural networks. Using recent work on deep kernel shaping, the authors show that shaping activation function leads to non-trivial stochastic differential equations (SDE) of neural covariance matrix. Authors show that simulation of finite networks match the distribution predicted by the SDE. ", "questions": "One limitation I see from the current approach is how to use the current framework to do **inference** In the conventional infinite-width limit, due to deterministic covariance matrix inference leads to conventional Gaussian process inference. With non-deterministic covariances, does the author have ways to perform posterior inference to make predictions? Interesting related work is [Yaida, Non-Gaussian processes and neural networks at finite widths, MSML 2020] (also covered in 6.4 of ref [30] Roberts, Yaida and Hanin 2022) which looks at finite width correction to NNGP and obtain non Gaussian process as correction to the infinite-width NNGP but still retains ways to perform inference as a correction to NNGP inference. \n\nAlso for the first attempt of a new theoretical approach, simplicity is a virtue and I appreciate it. However, it would be beneficial for the readers to see applicability / limitation of neural covariance SDE to general architecture. For example, DKS[37] or TAT[38] both main practical applications are to Convolutional Networks, setup in Section 2 deals just with fully connected networks. Could authors expand more on generalization to more structured architecture? \n\nWhat is the practical benefit of having SDEs vs just Markov Chains (as in section 2.3)? I understand there's a nice \"trainability\" benefit coming from activation shaping and SDE maybe coming along with the choice. But my read from Figure 1 in terms of predictability of the theory of neural covariances, both are as good matching the simulations.\n\nIn the setup, neural networks weights are initialized with 0-mean variance 1 Gaussian without any bias terms. Often exploding/vanishing of input norms are phrased in terms of initialization variances of weights and biases(e.g. Schoenholz et al., Deep Information Propagation, ICLR 2017 or [34]) and denoted Chaotic/Critical/Ordered phase. Wonder if the condition in section 4, eq (24) is directly related to the phases if accounted for weight/bias variances. Or should I consider if and only if condition in (24) as totally different conditions due to simultaneous width-depth limit?\n\nSimple Nit: L200 should be either PDE or derivative on $\\rho_t$ should be full derivative? Maybe the point is to say it's not a stochastic DEs? From the first read, it was slightly confusing. \n", "limitations": "Mostly a theoretical analysis of deep neural networks, so there's no foreseen negative societal impact. \n\nFew limitations as raised in the Questions section: (not necessarily needs resolving but would be good if mentioned as future directions etc)\n- Methods to perform inferences\n- Generalization beyond fully connected networks\n", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper presents an interesting theoretical framework for understanding simultaneously the large depth + width limit of random neural networks at initialization. Much of conventional large-width limit has focused on fixed depth and large width limit but there is evidence as presented by the authors as well that if one simultaneously takes large width and large depth limit, the behavior of neural networks can be systematically different. \n\nInteresting contribution of this paper is that even in this seemingly complicated limit, there's tractable theory determining the random covariance matrix in terms of stochastic differential equations. This is allowed by utilizing activation shaping which was used to show that one could train normalization-layer and skip-connection free very deep feed-forward networks.\n\nPaper does a very thorough job of covering related works, with the only exception I see being [Yaida, Non-Gaussian processes and neural networks at finite widths, MSML 2020] mentioned in the Questions section. \n\nExposition is clear. The exposition split of starting from unshaped ReLU networks in Section 2 with SDE example in section 2.2 then leading to shaped activation case of Section 3 was quite helpful.  However, to understand the full details of the work, it does require readers to have some prior exposure to modern probability and stochastic processes. Appendix A attempts to provide some background but this serves more of a refresher than an introduction. This may limit the target audience. Having said that, probably simple exposure to stochastic processes and stochastic differential equations may be sufficient to understand the main paper's message. \n\nAuthors did a great job of providing, ipython notebook of the empirical simulations. Empirical results do show solid support for correctness and generality of derived results. One thing to be cautious about is to consider removing unnecessary outputs in the notebook, since there is some identifying information left by mistake. This doesn't seem to be a severe violation of anonymity but needs to be careful for future submissions. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "iWoP880D9Ex", "reviewer": "Reviewer_SK2F", "summary": "This paper study the random covariance matrix in the shaped infinite-depth-and-width limit to model the layer to layer fluctuations in random feature neural networks.\n\nThe authors claim that the modeling of covariance and depth has a significant improvement over the NNGP regime.", "questions": "Is there any further interpretation when incorporating this technique to neural tangent kernel?\n", "limitations": "The analysis is still in the linearization regime, which is known underestimate the power of neural networks. The evidence is not strong enough to justify the necessity of “reshaping” in NN literature at this moment.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n\nThe idea of depth and covariance modeling is intuitive and non-trivial. This paper analyzes infinite-depth-and-width limit by reshaping the activation function, resulting non-trivial SDE structure.\n\nWeakness\n\nAlthough the “reshaping” idea is theoretically elegant, the practical interpretation is still based on the linearization regime, which is known to underestimate the neural networks.\n\nBesides, the gain from “reshaping” is not significant enough to replace NNGP/NTK (at this moment).\nThe theoretical modeling is based on the evolution of random features with the depth. However, deep random features are not better than two-layer one in real practice. More interpretation might be needed to justify the significance.\n\nThe authors raise some examples such as the Exploding and Vanishing Norms. However, it is an artificial setting which is not even a problem in real practice. A more practical example indicated by this framework is appreciated. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "1n0tXnjkUBj", "reviewer": "Reviewer_xxTF", "summary": "Infinite-depth-and-width limits for the covariances of activation layers are derived for neural networks with ReLU-like and smooth activation functions at initialization (under the iid Gaussian weights assumption). To obtain a nontrivial limit, leaky ReLU-like activations are assumed with the slopes differing by a factor decreasing as $O(n^{-1/2})$. This factor is shown to be critical for this limit. For smooth activation functions, a criterion is obtained on the activation function to ensure that the limiting covariance SDE is nonexplosive. Sigmoidal, tanh, and some classes of softplus activation functions are determined to satisfy this criterion. Some experiments are conducted which show agreement between the distribution of correlations for initialized finite NNs and those of the SDE predictions.", "questions": "How does the kernel derived from these neural covariance SDEs behave in a GP relative to e.g. the arc-cosine kernel?\n\nDoes the typical ReLU setup $s_+ = 1$, $s_- = 0$ not reveal an interesting limiting SDE? This seems like a surprising omission, given that almost any other activation function seems to exhibit nontrivial limiting behaviour. ", "limitations": "While the authors do not discuss limitations of their own work, I find this to be reasonable, as there are few apparent limitations here. Perhaps the most significant is the focus on initialized networks (requiring all weights to be iid normal), which has surprisingly little discussion in the main document. ", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "NNGPs and NTKs are among the most useful theoretical tools for analyzing neural networks, and provide a powerful link between the versatile Gaussian process paradigm and deep learning. However, it has been shown time and again that these infinite-width approximations do not mimic the performance of modern neural networks. Here, the authors have instead considered the infinite-depth-and-width regime, and while analogues of NNGPs and NTKs are not yet available, the proposed building blocks are provided here for neural networks at initialization. Comfortingly, the behaviour of these approximations seem to mimic those of finite NNs --- as in double descent theory, the large n,d regime appears to be the correct one. The theoretical analysis is extensive, and the presentation is (mostly) excellent, although I did not have the opportunity to inspect the proofs in detail. \n\nWith this in mind, and given that this work is novel to my knowledge, I believe this is a valuable contribution and am excited by potential follow-up work here. However, on first read, I was having a difficult time drawing conclusions from these particular results. In particular, I believe the numerical results and their details deserve their own section in the main document. To make room for this, I would suggest shrinking Section 2.3 (I found this section to be more confusing than enlightening, although I understand why it is here), and moving the critical exponent results to the Supplementary Material, as I believe they can be easily summarized in a sentence or two. It would also have been nice to see what these covariances could do as part of a GP, but I would not expect any significant improvements over known kernels. \n\nSome minor points:\n- The first paragraph of the abstract is confusing as written; I would strongly suggest rewriting it, starting with 'Recent work...' in the first sentence\n- line 72: 'according a' -> 'according to a'\n- line 90: 'chains' should probably be capitalized to be consistent\n- line 103: 'of, the' -> 'of the'\n- line 106: 'Gaussian random matrix' - Define. There should also be explicit mention of the Gaussian weights assumption here.\n- line 108: Consider reversing the order of these two equations, as the significance of g is unclear when reading from left-to-right.\n- line 119-121: last sentence is unnecessary\n- line 161: replace colon with full stop\n- line 162-163: last sentence is repetitive; consider removing\n- Can different symbols be used to describe the Markov chain and the limiting SDE (e.g. \\mathcal{V})? The results are difficult to parse as is.\n- line 187: 'they are' -> 'they follow'\n- line 232: I'm not sure the exclamation mark is warranted here...\n- line 274: Proposition 7 should probably be moved here to avoid confusion", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "00JhUqh791d", "reviewer": "Reviewer_TjDp", "summary": "This paper is theoretical and studies the behavior of neural networks in the shaped infinite depth/width limit. As stated in the appendix, the logits at the end of a feedforward network are conditionally Gaussian at initialization with a random covariance matrix. This work studies that matrix at infinite depth/width.\n\nCurrent theory does not work at large depths because small fluctuations between layers are disregarded. When properly accounted for, this paper demonstrates that the fluctuations add up so that the covariance matrix is governed by a stochastic differential equation (SDE). This work therefore provides a stronger theoretical understanding about networks at initialization. \n\nThe distribution of the covariance matrix closely follows the distribution predicted using the SDE theory developed here, this is shown by simulating the initialization of many neural networks.", "questions": "- Is there a way this work can be used outside of our understanding? For example to find a valuable initialization scheme?\n- Is it possible to extend this to understand training dynamics?", "limitations": "The societal impact is not addressed. However, this is a theoretical paper and a broader impact statement is not required (though it is always appreciated). \n\nThe limitations are discussed in the introduction. Briefly, the paper only applies to initialization and not training dynamics.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe paper provides a creative theoretical contribution. Furthering our understanding of large neural networks, from a novel angle using SDEs. The work is also tested briefly, leading to positive results like those in Figure 1.\n\nWeaknesses:\n\nThe only weakness I can find is in the presentation, which is mainly down to personal preference. In my opinion the paper feels very math heavy, and possibly not as accessible as it could be. The paper may benefit from writing more intuitive statements in the main paper and what their consequences are, essentially walking through the work. Then the more rigorous proofs and statements can be placed in the appendix. For example, equation 2 for the ReLU-like activations is currently:\n\n$\\phi(x) = s_{+}\\max(x, 0) + s_{-}\\min(x, 0) = s_{+}\\text{ReLU}(x) - s_{-}\\text{ReLU}(-x)$\n\nThese are both technically correct, but an easier way of writing may be:\n\n$\\phi(x) = s_{+}x ~~ \\text{if} ~~ x\\geq 0,\n~~ \\text{or} ~~ \ns_{-}x ~~ \\text{if} ~~ x< 0$\n\nAdditionally more figures/tables showing visually how well the Covariance SDE works always help.\n\nMinor Typos (do not affect the quality of the paper):\n\n- Line 60: \"matches matches\" should be \"matches\"\n- Line 246: \"Not only do these stopping time rules out\" should be \"Not only do these stopping times rule out\" or \"Not only does this stopping time rule out\"", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept