Differential Equation Scaling Limits of  Shaped and Unshaped Neural Networks
Authors: Mufan Bill Li
Mihai Nica
Venue: Accepted by TMLR ()

Abstract: Recent analyses of neural networks with shaped activations (i.e. the activation function is scaled as the network size grows) have led to scaling limits described by differential equations. However, these results do not a priori tell us anything about ``ordinary'' unshaped networks, where the activation is unchanged as the network size grows. In this article, we find similar differential equation based asymptotic characterization for two types of unshaped networks. 

Firstly, we show that the following two architectures converge to the same infinite-depth-and-width limit at initialization: 
(i) a fully connected ResNet with a $d^{-1/2}$ factor on the residual branch, where $d$ is the network depth.
(ii) a multilayer perceptron (MLP) with depth $d \ll$ width $n$ and shaped ReLU activation at rate $d^{-1/2}$. 

Secondly, for an unshaped MLP at initialization, we derive the first order asymptotic correction to the layerwise correlation. In particular, if $\rho_\ell$ is the correlation at layer $\ell$, then $q_t = \ell^2 (1 - \rho_\ell)$ with $t = \frac{\ell}{n}$ converges to an SDE with a singularity at $t=0$. 

These results together provide a connection between shaped and unshaped network architectures, and opens up the possibility of studying the effect of normalization methods and how it connects with shaping activation functions.

Published in Transactions on Machine Learning Research (04/2024)
Differential Equation Scaling Limits of Shaped and Unshaped Neural Networks
Mufan (Bill) Li Princeton University
edu
Mihai Nica
nicam Sicram@huoguepph..o
University of Guelph and Vector Institude
Reviewed on OpenReview: https: // openreview. net/ forum? id= +RDuUXYsSJ
Abstract
Recent analyses of neural networks with shaped activations (i.e. the activation function is scaled as the network size grows) have led to scaling limits described by differential equations. However, these results do not a priori tell us anything about "ordinary" unshaped networks, where the activation is unchanged as the network size grows. this article, we find similar differential equation based asymptotic characterization for two types of unshaped networks.
Firstly, show that the following two architectures converge to the same infinito-depth- and-width limit at initialization: (i) fully connected ResNet with d-1/2 factor the residual branch, where d is the network depth. (ii) multilayer perceptron (MLP)
with depth d & width n and shaped RcLU activation at rate d-1/2_
Secondly, for unshaped MLP at initialization, correction to the layerwise correlation. In particular, if pe is the correlation at layer f, then = (1 px) with t = converges to an SDE with singularity at t= 0.
derive the first order asymptotic
These results together provide a connection between shaped and unshaped network architectures, and opens up the possibility of studying the effect of normalization methods and how connects with shaping activation functions.
1 Introduction
Martens ct all (2021);; Zhang et all (2022) proposed transforming the activation function to be more lincar as
the neural network becomes larger size, which significantly improved the speed of training deep networks without batch normalization. Based on the infintteddeph--ad--iide
limit analysis of Li all (2022), the
principle of these transformations can be roughly summarized as follows: choose an activation function Gs R + R as a perturbation of the identity map depending the network width (or depth = n p 0)
M(xx + O(n -2)))))- = x (x(x) = - n²
Vd h(()) + O(d-))
(1.1)
where for simplicity we will ignore the higher order terms for now. Li all (022)) also showed the limiting multilayer perceptron (MLP) can be described by a Neural Covariance stochastic differential equation (SDE). Furthermore, it appears the choice of 2 is necessary to reach a non-degenerate nor trivial limit, at least when the depth-to-width ratio d converges to positive constant. (sce (L all 2022, Proposition 3.4)).
Recently, Fayou & Yang (2023): Cirone el al (0123 also studied the infinite-depth-and-wiie specific ResNet architecture (He et all 201F). Most interestingly, the limit is described by an ordinary differential equation (ODE) very similar to the neural covariance SDE. Furthermore, Hayou & Yang (2023) showed the width and depth limits commute, i... no dependence on the depth-to- tiith ratio It then natural to consider more careful comparison and understanding of the two differential equations.
limit of a

---

Published in Transactions on Machine Learning Research (04/2024)
0.20
0.15
Density
0.10
0.05.
0-
-4
Transformed Correlation ft
SDE (Ours)
NN
Figure 1: Empirical distribution of the transformed correlation r log(L pe)) for unshaped ReLU
MLP,, SDE sample density computed via. kernel density estimatiion Simulated with n d = =150, Po 0.3, To log(1 Pa) = loog(.7)), SDE step size 10-2, and 213 samples.
At the same time, via Markov chain. lakub & Nica (212k) further studied the large width asymptotics of the Markov chain
(2022) demonstrated the unshaped network also has an accurate approximation
updates, where the transition kernel still depends on the width. Since the Markov chain quickly converges to
fixed point, it does not immediately appear to have a scaling limit. However, this motivates us to consider modified scaling around the fixed point, so that we can recover a first order asymptotic correction term
In this note, we provide two technical results that address both of the above questions. Both of the results are achieved by considering a modification of the scaling, which leads to the following results.
Firstly, we demonstrate that shaping the activation has close connection to ResNots, and the covariance ODE is in fact just the deterministic drift component of the covariance SDE, Furthermore, in the limit the scaled ratio of प to converges to positive constant and E (0, 1), the shaped MLP covariance also converges to the same ODE.
Secondly, we analyze the correlation of an unshaped MLP: providing derivation of the first order asymptotic correction. The correction term arises from rescaling the correlation PE in layer by ae = f - pr), and we show it closely approximated by an SDE.
The rest of this article is organized as follows. Firstly, we will provide a brief literature review the rest of this section. Next, we will review the most relevant known results on this covariance SDEs and ODEs Section Ø. Then in Section D. will make the connection between shaping and ResNets precise. At the same time, will also provide derivation of the unshaped regime in Section HI where show that by modifying the scaling yet again, we can recover another SDE related to the correlation of ReLU MTP.
1.1 Related Work
On conceptual level, the main difficulty of analyzing neural networks is due to the lack of mathematical tractability. In seminal work, Neal (1485) showed that two layer neural networks at initialization converges to Gaussian process. Beyond the result itself, the conceptual brcakthrough opened up the field to analyzing
large size asymptotics of neural net works. In particular, this led to a large body of work on large or infinite
width neural networks (Lee et all 20S Jacot et al, 20L Dn et all 2014 Mei et all 2015; Sirignano &
Spiliopoulas 2018; Yang. 2019; |Barflett et all 2021). However, majority of these results relied on the network converging to kernel limit, which are known to perform worse than neural networks (Ghorbani ef al 2020)
The gap in performance is believed to be primarily due to lack of feature learning (Yang & Hi. 2021: Abbe

---

Published in Transactions on Machine Learning Research (04/2024)
Notation
Description
Notation
Description
Table I: Notation
Pin E N
Input dimension Hidden layer width
Reut E E d d p.()) Wo R
Output dimension
VEN
Number of hidden layers (depth)
p(()
Base activation Iuput for S a Network output Neurons (pre-activation)
Shaped activation
E ER
Weight matrix layer 0 Weight matrix at final layer
Pout R
Wont E R
Ri R
We Rux
Weight matrix layer 1 <
for layer S & S d Neurons (post-activation) for layer 6
All weights initialized iid ~ (((, 1) Normalizing constant
PF I
CER
(B(p(g)²) for g N c=
voB E R
Covariance ;
00 E [-1,11]
Correlation / VVpav
While this motivated the study of several alternative scaling limits, in this work we are mostly interosted in the infnite-depth--nd-wittt limit.
all 212 Ba et all
First investigated by Hamin & Nical (Ν0ии)) was shown that not only does this regime not converge to a Gaussian process at initialization, it also learns features [amin & Nical (zuнa). This limit has since been analyzed with transform based methods (NNcii et all, 2021) and central limit theorem approaches (Li all, (021). As we will describe in more detail soon, the result of most interest is the covariance SDE limit of al (2022) The MLP results
also further extended to the transformer setting ([Noci all, 202)
The d-1/2 scaling for ResNets was first considered by Faaouu et all (2021), with the depth limit carefully studied afterwards (Hayou, 2022, Hayou & Yang, 2023: Havon 2723) Fischer al (2122) also arrived at the
same scaling through different theoretical approach. This scaling has found applications for hyperparameter
tuning (Bordelon et all 2023; Yang et al, 2023] when used in conjunction with the uP scaling (Yang & Hu. 2021).
Batch and layer normalization methods introduced remedy for unstable training (loffe & Szegedy, 20E Ba ct all, 2016), albeit theoretical analyses of these highly discrete changes per layer has been challenging. A recent promising approach studies the isometry gap, and shows that batch normalization methods achieves similar effect as shaping activation functions (Meterez et all 2л23). Theoretical connections between these approaches using a differential equation based description remains an open problem.
2 Background on Shaped Networks and ResNets
Let {=") j== be set of input data points in R and let 2 En denote the f-th hidden layer with respect to the input z". We consider the standard width-n depth-d MLP architecture with Ho-initialization (Ho atl 201) defined by the following recursion
Zou - - WoW ((=3),
C [Weps(zf), o4 2441 f41=1 -
Win x
(2.1)
/niin
where P8 IR - R is the activation function to be specified, c- E for g~ N(,, 1), of E ER" =uu e Bout, and the matrices Wont E Rous We Rox., Wi E m are initialized with iid N(O, 1) entries.
The main structure used to study neural networks at initialization, such as the noural network Gaussian processes (NNGP) (INaal, иииа, Lee et all 201)), on the conditional Gaussian property. More precisely, if we condition on the previous layers F, = a( ((=f)))(m))sss)) we have that
[=f m |F N (0, (1))) @ Im):
(2.2)
where we use the notation H vertically stack vectors, let oo (2f) be the post activation hidden layer, and let be the Kronecker product.
3

---

Published in Transactions on Machine Learning Research (04/2024)
This naturally leads us to define the covariance matrix V2 12) The NNGP results essentially reduces down to applying the Law of Large Numbers inductively to show the covariance Ve converges to its expect value. More precisely, if [=] ~ N(0, Ve-1 @ In), then in the limit as - x
Ve
c Σ P. (x(()))))))))))) 4 (p. where g" ~N
o.
Ve 4-1
Vi
(2.3)
op
E-1
However, since the next layer covariance Ve a deterministic function of the previous Vt-1, this forms a fixed point type iteration V/2 = f(V--)). Indeed, if we observe the correlations aB
which is bounded 1
in 1, 1], it does in fact converge to fixed point at Pso =1 1 for ReLU activations (see e.g. Proposition 3.4 (i) it et al (2022))
This degeneracy of correlations also causes unstable gradients, which led to a proposal by Martens et aL (2021); Zhang et. al. (2022) to modify the shape of activation functions yex depending on the size of the network, leading to improved training speeds without using normalization methods. However, as both of these works computed the activation shape based on set of criteria, it is unclear what is the appropriate modification in the scaling limit depth and width both approaches infinity.
2.1
Shaped Limit of Neural Networks
To this end, all (2022) is the first to describe the limit as d, n - oc with d + 0 and the activation n function iPs is shaped at precise rate to be closer to the identity as increase n. In particular, the covariance matrix Ve forms Markov chain,
it satisfies the definition of Ve+1lF; V++1lo(V,) (see e.g.
The main result describes the scaling limit of the Markov chain Vital via a stochastic differential equation (SDE), which we can intuitively interpret as the Euler discretization converging to the differential equation
(2
V4+1 = Vie b(V))
E(V2) + +Onn 3/3) n-00, dV. = b(v)) dt + (VV2) dB,,
(2.4)
Vn
where & are iid. zero mean identity variance random vectors, and we interpret as the step size of the discretization.
The activation functions modified as follows. For a ReLU-like activation, choose
(x) = + max(x, 0) + min 0), 84 = 1 + np E K.pp2 00,
(2.5)
or for a smooth activation & C* C+(R) (R) such that 4(0) 0' (0) 1 and g( (x) bounded by a polynomial,
choose
(x) up
an a + 0 0.
(2.6)
We will first recall one of the main results of Li all
A127 which is stated informally below.
Informal). Let p 1. Then in the limit d, n + oc, d T>> O, and sss defined as above, we have that the upper triangular entries of V (fattened to vector) converges to the following SDE weakly
Thcorem 2.1 (Theorem 3.2 and 3.9 of Li ct all
(202))
dV. = b(V)(V)) dt 12 dB. Vo = nin
1100 x 1 Ma,
(2.7)
where E(V) 885 + padybr and if is ReLU-like activation we have
Va8
v(p)= (ct cc
b(V)las= v Vaa 8
Pl--p-parccoss (2.8)
V35
2x
The statement "Informal" in the sense that have stated what the final limit is, but not the precise sense of the convergence. See Appendix for rigarous treatment of the convergence result.
4

---

Published in Transactions on Machine Learning Research (04/2024)
else if 4 is smooth activation have
8ao (V.)
"²)())
(vaov + v (2v -3)) + 4a
(0)(0) v. (Vea + V, 2). 2). 2a²
(2.9)
While the formulae may seem overwhelming, there is actually a fairly straight forward interpretation of both the drift and the diffusion coefficient 11/2, In particular, for the unshaped ReLU network, that is (x(x)) = max(x, 0), the deterministic component of the update for the covariance matrix compared to the shaped network are follows
Unshaped: EV41 - Ve x b(V2).
Shaped: E VV+1 Ve x b(V))
(2.10)
Effectively, the drift component got slowed down by a multiplicative factor of 1, which can interpreted an Euler discretization step. In order to achieve stable limit as we take depth d to infinity, we would also require each layer to contribute infinitesmmally as a rate proportional to 2. therefore this is a desired rescaling.
The diffusion coefficient is much more interesting, Since we can interpret this diffusion be on the manifold of symmetric positive definite matrices, we would expect the diffusion coefficient of Brownian motions this manifold to correspond to Riemannian metric. Indeed, this shown in Li Li all (20221) Σ-1 corresponds to the afline invariant metric. More precisely, for all V SPD(m), the inner product corresponding to E-1
(A, ()
E Ar TV-VV-BV-1) for all A, Sym(m). ISiSism
(2.11)
Furthermore, this gives the linear network SDE dV, (VV))2/2 dB, an interpretation the dual Brownian motion in information geometry, which uses pair of dually flat affine connections instead of the standard Levi-Civita connection.
2.2 ODE Limit of Residual Networks
At the same time, Hayou & Yang (2023): Cirone et all (2IZk) found an ordinary differential equation (ODE) limit describing the covariance matrix for infniite-depth-and-rie ResNet architecture with va factor on their residual branch, more preciscly their recursion is defined as follows (in our notation and convention)
ResNets. The authors considered
2 = 2 +
We (2))) where ppx) = max(x, 0).
(2.12)
Vdn
Intuitivoly, this is similar to shaping activations, also weakens the effect of cach layer as we take d infinity. This will be discussed in more details in the following section.
One of their main results can be stated informally as follows.
Theorem 2.2 (Theorem of Hayou & Yang (2023), Informal). Let d,n - (in any onder), and the covoriance process yag converges to the following ODE [td]
1 1 a3
de
VoB
=
a3
VoB
(2.13)
Pt
where f(p) = (paresin + /1- 10-ppp -2) + p.
Here, we observe this ODE
is exactly the drift component of the covariance SDE (27)) i.e.
1 1 a3
yol u(p: ad
V.av if (c+ 2
d
1.
(2.14)
2
Pt
5

---

Published in Transactions on Machine Learning Research (04/2024)
To see this, we just need to use the identity arcsin p 2 arccos | to get that ff(p) v(p,, and that yau 2nd = Voaa,, is exactly the definition of Pt We also note the correlation ODE 4 aB = u was Pi
first derived in (Zhang et all, 2022, Proposition 3), where they considered the sequential width then depth limit with fixed initial and terminal condition.
In the next section, we will describe another way to recover this ODE from an alternative scaling limit of the shaped MLP.
3 An Alternative Shaped Limit for p E (0,
We start by providing some intuitions on this result. The shaped MLP can be seen as layerwise perturbation of the linear network
21+1 = - --eeede)) / + Va
Jav
e W. h(ze)
(3.1)
where e-1 = B (g) for g (0, 1) corresponds to the He-initialization He et all MIBB and We = has iid 10, entries.
On an intuitive level (which we will make precise in Remark 8 if we take the infinitewidtth limit first, then this removes the effect the random weights. In other words, if we replace the weights 1 We with the identity matrix In in each hidden layer, we get the same limit at initialization. Therefore, we can heuristically write
v
se+1 28 + hl(z),
(3.2)
vd
where we also used the fact e in the limit.
Observe this is resembling ResNet, where the first ze the skip connection. In fact, we can again heuristically add back in the weights on the residual branch to get
21-1 RJ + Weh(ze),
(3.3)
which exactly recovers the ResNet formulation of Hayou et all (2021);; Hayou & Yang (2023) where the authors studiod the case when h(x) = max(x, 0) the ReLU activation.
Remark 3.1. On a heuristic level, this implies that whenever the width limit is taken first (or equivalently d d= n for E (0, 1/2)), the shaped network with shaping parameter d-1/2 has the same limiting distribution at initialization as a ResNet with a d-2 weighting on the residual branch.
However, we note that despite having identical ODE for the covariance at initialization, this does not imply the training dynamics will be the same it will likely be different. Furthermore, since Hayou & Yang (2023) showed the width and depth limits commute for ResNets, this provides the additional insight that noncommuntatiiitty of limits in shaped MLPs arises from the product of random matrices.
3.1 Precise Results
The core object that forms a Markov chain the post-activation covariance matrix A # 1f)m, To see this, we will use the property of Gaussian matrix multiplication, where we let W E Roxn with iid entries Wi ~ (0, 1), and {um) E Rn be a collection of constant vectors, which gives us
Wumm d N(, [(w", u/)m & In),
(3.4)
where use the notation [ m to stack the vectors vertically. This forms Markov chain because we can condition on Fz = a1
[=]]m la=1) to get
1 F lo NOO
(3.5)
and we see that Ve+iFFe (V+)) which is exactly the definition of Markov chain.
6

---

Published in Transactions on Machine Learning Research (04/2024)
We will start by deriving the precise Markov chain update up term of size O(n (n)) which will be slight modification of the Euler discretization we saw in (
Lemma 3.2 (Covariance Markov Chain for the Shaped MLP). Let of be the MLP in defined in (2. with shaped R activations defined in (2.3) For p E (0,2) and d n?p, the Markov chain satisfies
b(V2) d
+ 2V + 0/d
V4+1 = V2
(3.6)
Vn
where (Ee) are iid zero mean and idenlity covoriance random veclors, and in the limit as n, d ++ oo we have that M, + D, and the coefficients are defined as
b(V)a = v(pa8) o(p0))/vaav333 I(V)a8.7 varves + vogy8,
(3.7)
where pa = VVaayaz and v(p) = (c+-c-)² Proof.
Voo
(V1- + parccos
2n
To start, we will observe that conditioned on Fe. we have that
1e Σ pa(= Pe(e (=(+111,)) = 1 Σ (.(o)f) P.(o))
(3.8)
n
1=1
where used the fact that 2441 are jointly Gaussian, and we have that
0, [ Pe
pi
0 i=1 ~ N
x
In
(3.9)
for pad Vap VVaavas" At this point, we can define Ki (ppa))= E (9f) (a(o,)
and the random variable
R
Vn Σ
cKi (pt (o,
(3.10)
which allows us to decompose this into a deterministic and a random component as
Val Fe = 148 1 (<Ki(08") + Vn R2
(3.11)
We can Taylor expand (p) terms of n-P get the following result from Lemma (p) p
v(p) n O(n "))
v(p) cc)) (V1 (VI-p² p² + parccos p)
(3.12)
2T
Similarly from Lemma also have the approximation
E R2 R7 Pi 95
pi pt + O(n
(3.13)
Putting everything logether, we can write
1 Vn
v(pr.)
cKi
Pi
Ra
OOnn
(3.14)
ndp
vn
which implies we can write
b(Vi)as
Ras,
Vag = vas + 2+1
O(n O(n-pp).
(3.15)
ndp
where b(V)= = (())=((p+))
v Now taking the upper triangular entries of V as a vector, we have that
b(V2)
(Ve)
V4+1 Vt
+ +O 3p),
(3.16)
2p
Vn
where (V = 1Vaa/33(007) 80 + ) + On = 88 + vad + ((n and & is zero mean identity covariance random vector. We recover the exact desire result from writing d = n 29

---

Published in Transactions on Machine Learning Research (04/2024)
Remark 3.3. We note that the drift term arising from the activation depends only on the depth d and the random term only depends on the width n. If we decouple the dependence on d and n, and take the infihitewwidth limit first, we arrive at
(V2) O(n 8/2), 0n2))
Ve+ Ve +
(3.17)
d
which is equivalent to removing the randomness of the weights.
We note this Markov chain behaves like a sum of two Euler updates with step sizes ndp and 1 where the 1 corresponds to the random term with coefficient J However since E (0, 2), the first term with step size ndp will dominate, which is the term that corresponds to shaping the activation function. Therefore, we expect the random term to vanish in the limit, and hence leaving with ODE only. We make this result. precise in the following Proposition.
n
Proposition 3.4 (Covariance ODE for the Shaped ReLU MLP). Let p (0, 2)) Then in the limit d, n 0G, ग. 1, and yps is the shaped ReLU defined in (2a1), we have that the upper triangular entries of Vlnnn (flattened to vector) converges to the following ODE weakly with respect to the Skorohod topology of DR_,R.m 13/2
dV. = b(V)) dt,
Vo = 11--16 1(0", xx "x Na,
(3.18)
Rén
where is defined in Theorem and Lemma FF
Proof. Starting with the Markov chain in Lemma S part of the drift instead. More precisely, we let
we will treat the random term of order ((n) as
B(V) b(V) np
+ (V)
(3.19)
2D
Vn
which in expectation just equals to 6(V)n 2p Since there is no random term at the order of n n we apply Proposition A.7 to the special as if the diffusion coefficient is equal to zero. At the same time, since the higher order term in the Markov chain is at the desired order of O(n ((n,, which will vanish in the limit, we get the desired result.
At this point it's worth pointing out the regime of p E (0, 1) was studied Li all (2022), but the scaling limit was taken to be - T instead of n This led to a "degenerate" regime where Pt = for all & > 0. The above ODE result implies that the degenerate limit can be characterized a more refined way the scaling chosen carefully.
d
In the next and final section, we show that actually even when the network is unshaped (i.c. p p= 0), there exists scaling such that we can characterize the limiting Markov chain up to the first order asymptotic correction.
4 An SDE for the Unshaped ReLU MLP
In this section, we let px(z) = =x())= = max(x, 0), and we are interested in studying the correlation
(49 v() 1o8l kop
a. Pt.
(4.1)
V
From this point onwards, we will only consider the marginal over two inputs, so we will drop the superscript a,3. Similar to the previous section, we will also start by providing an intuitive sketch.
Many existing work has derived the rough asymptotic order of the unshaped correlation to be Pt 1-0(E-2) where & is the layer (see for example Appendix E. of Li al all (2022) and Dakub & Nical (202)) Firstly, this
8

---

Published in Transactions on Machine Learning Research (04/2024)
implies that Taylor expansion of all functions of p in the Markov chain update around p p= will be very accurate. At the same time, it is natural to magnify the object inside the big O by reverting the scaling, or more precisely consider the object
qe = (1))
(4.2)
which will hopefully remain at size ((1).
For simplicity, we can consider the infinite-width update of the unshaped correlation (which corresponds to the zeroth-order Taylor expansion in :))
PE+1 Pe ci (1- pc) + O((1 - (p))pp))/3) /2)
(4.3)
where for the sake of illustration we will take and drop the big O term for now. Substituting in ge. we can recover the update
3/2
2q g1+1 qe t
q
(4.4)
While this doesn't 't quite look like an Euler update just yet, we can substitute in t for the time scale, which will lead us to have
ge+1 = qe
tn (2qe
1
2/2)
(4.5)
hence (heuristically) giving us the singular ODE
dg: 2qt: 3/2 t
(4.6)
To recover the SDE, will simply include the additional terms of the Markov chain instead taking the infinite-width limit first.
4.1 Full Derivation
In the rest of this section, we will provide derivation of an SDE arising from appropriate scaling of pe-
Theorem 4.1 (Rescaled Correlation). Let qe = f pt). Then for all to > 0, the process {q h1z0 converges to a solution of the following SDE weakly in the Skorohod topology (see (la all 2122, Appendir A))
1
2 1/2
4t
-1) dt 2-2222q dBr
dq. 2qi
(4.7)
The above statement holds only when ty > O, and there interesting technicality that must be resolved to interpret what happens as L 0+. In particular, the Markov chain is not time homogeneous, and the limiting SDE has singularity at = 0. The contribution of the singularity needs to be controlled in order to establish convergence for all t 2 0. Furthermore, due to the singularity issuc, it is also unclear what the initial condition of 4 should be.
In our simulations for Figure M, we addressed the time singularity by shifting the time evaluation of 1 to the next step of T where A, > 0 is the time step size. More preciscly, we first consider the log version of rt log qt
dr. -2 (1-
v2 ax
exp((i) di - 2V2dd.
(4.8)
Then we choose the following discretization
1
/2 exp(#) 3r
TH4 rt 2 1
A, 21+2V2EVV
(4.9)
t A
where & N(. 1). For initial conditions, we also noticed that since the initial correlation must be contained in the interval [-1,1], 1], the end result. was not very sensitive to the choice of ro-
9

---

Published in Transactions on Machine Learning Research (04/2024)
Proof of Theorem Firstly, will introduce the definitions
(p) = (o
(4.10)
where 9, w N(, 1) and we define Pg + qw with -i- We will also use the short hand notation to write Kp == p.p. Here we will recall several formulae calculated in Cho & Sanl (2009) and (Li et all, 2022, Lemma B.4)
= arccos(-p)
Ku(p) = 1 [0>0) 1l (po+qw>0)
2π
Ki(p) = E ((g)p(pg + qw) = 9 parccos(-1)
2π
(4.11)
K2(p) = E ( + qu)² 3pg + arccos arccos -p) + 2p²))
2m
K3,1(p) (pg))² p(g qu) q(2 p²) + arccos(-PP))
2n
Furthermore, will define
M2 = E [cp(g)) 11] 5.
(4.12)
Using the steps of (Li aL 2122
Proposition B.8), we can establish an approximate Markov chain
Pr+1 = cK)( + (p((
ox(ce Vn + O(n O(23/)))
(4.13)
where & are iid with zero mean and unit variance, and
pr(e)) (pprpp) [
+M2 + 3) 4cK3,1
(4.14)
2² (p(()) = 2 [K + M2 1) 4cK + +2K2]
and we write K. = (pt).
Here we use the big (n E)) to denote a random variable X such that for p 2 1
E
S CCp <0,
(4.15)
f(n Eyp
for some constants C > 0 independent of and l.
In view of the SDE convergence theorem Proposition A.M if we eventually reach an SDE, will only need to keep track of the expected drift instead of the random drift. We can then Taylor expand the coefficients in terms of Pe about P = (from the negative direction) using SymPy (Meurer et all am7)) which translates to the following update rule
2/2
22
PE+1 = Pe +
(1
1pp)3// +
(1 5/)
30m
4
73v2
+ -2(1 - ppr) +
(1 pe)
p0)
3/2
+ +3(1 - pp)²
(1 -pp)s/2
(4.16)
15t
&
56
+ 2//((1 pt)
po)
O(111 - px))* + +n 3/2)
Vn
15th
We note up this point, a similar approach was taken in lakub Nical (212% here by consider the following scaling
However, we will diverge
l pt)
(4.17)
10

---

Published in Transactions on Machine Learning Research (04/2024)
The choice of l scale is motivated by the infinite-width limit, where (1 -p)) shown to be of order (
in eL all 20222 Appendix E).. This implies the following update
f (L + 1)2 qp
-2 222/2 (1
/2 (1 - p2) p=)
= qe
3x
30m
e
4v2 (1 pt) 3/2 + +311 -po)² 73v2
-2(1 -pe) +
(1 -p))0/2
(4.18)
11
15t
Ee
22/2(1 pa)
56
(1 (1- po) +0( de + +02
+
vn
15
Next, we will drop all higher order terms in e, then using the fact that (2+1)² = 1 20(E-), we can write
2/2 qr
2qe +
& 2v + 0 + 1112
g++1 = qr(1 (1 + 1)
(4.19)
3x E
Vn
where we dropped 8 8 3/2 in the big O since it gets dominated by 8-2 Choosing the time scaling t = & then gives us
2 qt 1
2 1/2 41
2 qe 3 +
4++1 Qe +
3x
-1
+ +21
(4.20)
n
Finally, will use the Markov chain convergence result to SDE result Proposition A 77 which leads to the desired SDE for 2 to > 0
/2 1/2
3n 9t
dg. = 29
-1)
di + 2v dB,-
(4.21)
5 Discussion
In this section, we provide some discussion on the potential impact of this work, from both practical and a thcoretical point of view.
Stable Training via Depth Scaling. Martens et all (2021) make the key observation that as depth increases, the increased instability of training dynamics is due to the key role of nonlinear activation functions. Since then, it is better understood that to achieve stable training in large depth, it is necessary to weaken the nonlinearities of each layer (Noci et all, 2023; [Bordelon et all, 2023; Yang et all, 2023). Our results observe key connection between two strategies, either weakening the activation function, or weakening the entire layer via skip connections directly. From practical point of view, since both the shaping and ResNet approaches can lead to the same covariance ODE at initialization, we understand that the key to preventing unstable gradients is weakening nonlinearities. To choose between the two regimes, left to study the role of wcakening the weight matrix or otherwise, and how this affects training dynamics. In particular, we note that the shaped limit admits feature learning without modifying the scaling (Hanin & Nical 2m)) which is fundamentally different than the pP regime (Yang & HU, 2)2
therefore
Analysis of Normalization Methods. Since the correlation Markov chain has large jumps, and quickly to degenerate fixed point, it is intuitive to assume this chain does not admit a continuous time limit, and therefore difficult to analyze from an analytical approach. Indeed, this is the approach taken by Meterez all (2023), which yielded one of the first analysis of normalization methods in deep network. However, our result provides counter-intuitie understanding: it remains possible to analyze seemingly large discrete jumps in Markov chain if rescaled appropriately. In our case, since the Markov chain converges quickly to
11

---

Published in Transactions on Machine Learning Research (04/2024)
the fixed point at p -= 1, zooming in around the fixed point function of time (or layer f) allows view the dynamics at the correct scale. This overcomes previously known technical hurdle, and opens up the possibility of analyzing normalization methods, which is still lacking theoretical progress. More specifically, we would like to understand how normalization methods may help or hurt performance in practice, and how to best choose and lune these methods given their many varianis, all of which are now made easier by our scaling approach.
Foundation for Studying Training Dynamics. Until recently, almost all of the work infinite- depth neural networks remain at initialization. This is not due to lack of attempts, but rather due to lack of mathematical techniques available. In particular, we also emphasize the seminal work of neural langent kernels for training dynamics (Jacol el all, 2 is entirely built on the same techniques studying initialization (Neal, 1898 Lee et all, 2015). For this reason, it is important to slowly yet firmly build up a foundation of theoretical results, that understands as much structure as possible at initialization. To this goal, the key distinction from the anfinitewwithh regime is that cach layer must be viewed as an infinitesimal discretization of a continuous "layer time." This approach is what helped yield some of the first. characterization of training dynamics for infinite-depth ResNets (Bordelon et all, LIZK: Yang et all (IzK). In fact, any theory of training dynamics must also account for this infinitesimal treatment of layers, otherwise the limit cannot possibly be stable. Therefore, we view this line of work as building towards a theory of training dynamics, and eventually generalization as well.
References
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4887. PMLR, 2022.
Jimmn Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Demny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. arXiv preprint (:Xiv:225.0144:4 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and GeofTrey E Hinton. arXiv: arXiv:1607.06; 607.06450. 2016.
Layer normalization. arxiv preprinl
Peter T. Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: statistical viewpoint. Acta numerica, 30:87-2011 2021.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit, 2023.
Youngmin Cho and Lawrence K Saul. Kernel methods for deep learning. In Advances in Neural Information Processing Systems (NeurIPS), pp. 342-350, 2009.
Nicola. Muca Cirone, Maud Lemercier, and Cristopher Salvi. Neural signature kernels as infinite-width-depth- limits of controlled resnets. arXiv preprint arXiv:2303.1167. 2023.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Int. Conf. Machine Learning (ICML), pp. 1675 1685. PMLR, 2019.
Stewart N Ethier and Thomas G Kurtz. Markov processes: characterization and convergence. John Wiley & Sons, 2009.
Kirsten Fischer, David Dahmeen, and Moritz Helias. Optimal signal propagation in resnets through residual scaling. arXiv preprint arXio:2305.0711. 2023.
Bchrooz Ghorbani, Song Mci, Thoodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Newral Information Processing Systems, 33:14820 18:14820-14800, 2020.
Boris Hanin and Mihai Nica. Finite depth and width corrections the neural tangent kernel. In Int. Conf. Learning Representations (ICLR), 2019a.
12

---

Published in Transactions on Machine Learning Research (04/2024)
Boris Hanin and Mihai Nica. Products of many large random matrices and gradients in deep neural networks. Communications in Malhematical Physics, pp. 1 1-36, 2019b.
Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Transactions on Machine Learning Research, 2022.
Soufiane Hayou. Commutative width and depth scaling in deep neural networks, 2023.
Soufiane Hayou and Greg Yang. Width and depth limits commute in residual networks. orXiv preprint arXiv:2302.00m 50453, 2023.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Roussean. Stable ResNet. In Int. Conf. Artificial Intelligence and Statistics (AISTATS), pp. 1324 1332. PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. In Proc. IEEE Int. Conf. Computer Vision, pp. 1026-1034 2015.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual Icarning for image recognition. In
Proceedings of the IEEE conference on compuler vision and pallern recognilion, pp. 770 778, 2016.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pp. 448-456. pmlr, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Adoances in Information Processing Systems (NeurIPS), 2018.
Cameron Jakub and Mihai Nica. Depth degeneracy in neural networks: Vanishing angles in fully connected relu networks on initialization. arXiv preprint 0at:2302.09712. 2023.
O. Kallenberg. Foundations of Modern Probability. Probability theory and stochastic modelling. Springer, 2021. ISBN 9783030618728.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl- Dickstein. Deep neural networks as gaussian processes. In Int. Conf. Learning Representations (ICLR), 2018.
Mufan Li, Mihai Nica, and Dan Roy. The future is log-gaussian Resnets and their infnite-deepth-and-wiihhh limit at initialization. Advances in Neural Information Processing Systems, 34, 2021.
Mufan Li, Jaume de Dios Pont, Mihai Nica, and Daniel M. Roy. Geometric dyson brownian motion and the free log- normnolaal for minor of products of random matrices. In Preparation, 2024.
Mufan Bill Li, Mihai Nica, and Daniel M Roy. The neural covariance sde: Shaped infinite depth-and-widd networks at initialization. orXiv preprint Xix:222006 Kk0.2226.0,768, 2022.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl- Dickstein, and Samuel Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using deop kernel shaping. arXiv preprint arXiv:2110.01765, 2021.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7671, 2018.
Alexandru Meterez, Amir Joudaki, Francesco Orabona, Alexander Immer, Gunnar Rätsch, and Hadi Daneshmand. Towards training without depth limits: Batch normalization without gradient explosion, 2023.
13

---

Published in Transactions on Machine Learning Research (04/2024)
Aaron Meurer, Christopher P. Smith, Mateusz Paprocki, Ondřej Certik, Sergey B. Kirpichev, Matthew Rocklin, AMiT Kumar, Sergiu Tvanov, Jason K. Moore, Sartaj Singh, Thilina Rathnayake, Sean Vig, Brian E. Granger, Richard P. Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Podregosa, Matthew J. Curry, Andy R. Terrel, Stěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. Sympy: symbolic computing python. PeerJ Computer Science, 3cc103, January 2017. ISSN 2376-5992. doi: 10.7717//eerj-cs.103 103. URL https://dol.
28g/10.7710 hTg/10.7/A17/perr-ccs.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Modia, 1995.
Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofınnn. Precise characterization of the prior predictive distribution of deep relu networks. Advances in Neurol Information Processing Systems, 34, 2021.
Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel M Roy. The shaped transformer: Attention models in the infinite depth-and-width limit. arXiv preprint arXiv:236..17755 2023.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers, 2018.
Daniel W Stroock and SR. Srinivasa Varadhan. Multidimensional diffusion processes, volume 233. Springer Science Business Media, 1997.
Greg Yang. Tensor programs i: Wide feedforward or recurrent neural networks of any architecture are gaussian processes, 2019.
Greg Yang and Edward J. Hu. Feature learning in infinite-width neural networks. In Tnl. Conf. Machine Learning (ICML), 2021.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-depth neural networks, 2023.
Guodong Zhang, Aleksandar Botev, and James Martens. Deep learning without shortcuts: Shaping the kernel with tailored rectifiers. arXiv preprint arXiv:2203.08 08120, 2022.
14

---

Published in Transactions on Machine Learning Research (04/2024)
A Background on Markov Chain Convergence to SDEs
In this section, we will review the main technical results used for Markov chain convergence to SDEs. In particular, we will consider piccewise constant interpolations of Markoy chains in contimuous time defined by V(n) which converges to a continuous process V, Due to measurability concerns of the Markov chain with jumps, in order to define weak convergence on the space of processes, we need to define a precise space along with its equipped topology first introduced by Skorohod. We will mostly follow the references Kallenberg (202T); Ethier & Kuriz (2009); Stroock &r. Varadhan (1997) in the rest of this section, and the content will be essentially a rephrasing of Appendix A in Li 21122)
To start, we will let S be a complete separable metric space, where the stochastic processes takes value in (in our case, it's the upper triangular entries Rm((+1)/22). Lct DR.,S be the space of cådlàg functions, i.c. functions that are right continuous with left limits, from R+ to S. For In E R4, we use In " x to denote locally uniform convergence, or uniform convergence on compact subsets of R., We also consider class of bijections À on Rt+ such that is strictly increasing, and A O. We define Skorohod convergence, written as In + x on DR..s; if there exists a sequence of such bijections An such that
An w1 Id. In An
(A.1)
Intuitively, this allows us to side step any discont in the notion of convergence, as Xn allowed perform "time change" so we always land on the "correct side" of the discontinuity.
Importantly, we note that DR_,S with the above notion of convergence forms well behaved probability space.
Theorem A.1 (Theorem A5.3, Kallenberg (111 For any separable complete metric space S, there exists a topology T on D such that
(i) T induces the Skorohod convergence In 4 x,
(i) DR_. Polish (separable completely metrizable topological space) under T,
(18) T generales the Borel o-field generaled by the evaluation maps Te, L 2 0, where TT(x)) = Tt.
To be fully self-contained, we will also introduce the definitions of a Feller semi-group, generator, and the core. Let S be locally compact separable metric space, and let Co Co(S) be the space of continuous functions that vanishes at infinity, equipped with the sup norm, hence making Co Banach space. We say T : Co + Cp is a positive contraction operator if for all 0 < f S 1, we have that 0 < Tf S 1. A semi-groop (Ti) of positive contraction operators on Co is called a Feller semi-group if it also satisfies
Ti C Co,
t20 0,
(A.2)
Тe fx) + x, as t - 0, f x
Let D C Co and A D Co- We say the pair D) is generator of the semi-group (T,) if D is the maximal set such that
T f Af.
lim t02
(A.3)
L
We say an operator with domain D on Banach space B is closed, if its graph G = (U, / DJ is closed subset of B x B. If the closure of G is the graph of an operator A, we say that A is the closure of A.
We say a linear subspace D C D is a core of A, if the closure of AlD is A. Intuitively, all important properties of A can be recovered via a limit point of ASn, for fn E D. Furthermore, if (A, DD is a geneator of a Feller semi-group, every dense invariant subspace D C D is a core of (Kallenberg 2121, Proposition 17.9). In particular, the core we will work with is the space C of smooth functions that vanishes infinity.
The following is a sufficient condition for semi-group to be Feller.
15

---

Published in Transactions on Machine Learning Research (04/2024)
Theorem A.2 (Section 8, Theorem 2.5, Ethier & Kurtz (2m0)) Let a (RR))) with Og be bounded for all i, k e [d), and let b: Rd - Rd be Lipschilz. Then the genembtor defined by the elliptic operator
AS - Σ a 000, Σ RO..
(A.4)
generates Feller semi-group on Co-
Next, will state a set of equivalent conditions for convergence of Feller processes.
Theorem A.3 (Theorem 17.25, Kallenberg (202T)). Let , x². be Feller processes in S with semiigroups (Tn,) and generalors (A,D), (An, Dn), respectively, and for core D for A. Then these conditions are equivalent:
(i) for any E TED, there crists some fn E Da with fn - and AnSa + AS,
(ti) Tn, - Tr strongly for each t > 0,
(tit) Tn.f + Tef for every f E Co, uniformly for bounded t >> O,
(iv) Xg $ Xu in S X 4 X in the Skohorod topology of DR..S.
We note that it is common to choose the core D = C8 and that checking the first condition is sufficient for convergence in the Skorohod topology. The next result will help us apply the above criterion to continuous time interpolated Markov chains.
Theorem A.4 (Theorem 17.28, Kallenberg (2022)). Let YI, y2, y3.... .... be discrele lime Markoe chains
S with transition operators U1, U2, Uz...., and let a Feller process with semi-group (Tt) and generator A. Pir D for A, and let 0 hn Then conditions (i) (iv) of Theorem remain equivalent for the operators and processes
An \ 1), T, = U((hm X [ithon
(A.5)
Here we note that in our applications, we will always choose h, n 2p, which is essentially dependent on the width of the neural network.
At this point, still need to check that the generators A, converges to A with respect to the core D C For this goal, we will use lemma from Stroock & Varadhar (1997). Here we will define IT (1, dy) to be the Markov transition kernel of y and let
a (x) hn Jy-ulS1
(yi- - r.) xz)) (x dy),
1
(x)= lnn /1-0|/1
(ys- - xz) II (x,
(A.6)
1
(x
II (x, RR B(x,e)).
hn
Lemma A.5 (Lemmn 11.2.1, Stroock & Varadhan (1977) ). The following two conditions are equivalent:
(i) For any R R> 0, c > 0 we have that
lim sup [[an(x) axx))lop + (bx(()) +( 0,
(A.7)
n- |
(i) For each f C (RR)), we have that
1 Arf + Af.
(A.8)
h
uniformly compact sets of IRd, where A is defined
16

---

Published in Transactions on Machine Learning Research (04/2024)
We will also need to weaken the definition slightly for non-Lipschitz coefficients. Definition A.6. We say sequence of processes X converge locally to X in the Skorohod topology if for any 0, we define the following slopping times
= [t20: T= [t: |xel 2 r),
(A.9)
and we have that iir converge to XEAT in the Skorohod topology-
Finally, to summarize everything into a useful form for this paper, we will state the following proposition.
Proposition A.7 (Convergence of Markov Chains to SDE, Proposition A.6, ef all (2022)) Let y be a discrele time Markov chain on R.N defined by the following updale for p, >
b bn((F,07) nP
on(X?") E2 O -2---)
Y = YE
(A.10)
nP
where E E R are üid random variahles with zero mean, identtty covariance, and moments uniformly bounded in n. Furthermore, w/ are also tid random variables such that Ei,,, (Y7", =f) yl = bn (v) and bn w/) has uniformly bounded moments in n. Finally, on is deterministic function, and the remainder terms in O(n-2p- have uniformly bounded moments in
Suppose ba. on are uniformly Lipschitz functions in n and converges to b, uniformly on compact sets, then in the limit as n + oc, the process XE = [6n2o] converges in distribution to the solution of the following SDE in the Skorohod topology of DR,
dX, = +x
Xo= lim Yo x
(A.11)
Suppose otherwise bas on are only locally Lipschitz (but still uniform in n), then X converges locally to X in the same lopology (see Desomition A...). More precisely, for any fired O, we consider the slopping limes
inf (1 2 0: 21) T= inf (1 2 0 2
(A.12)
then the stopped process X converges in distribution to the stopped solution Xear of the above SDE in the same topology.
B Technical Lemmas for Shaped Activations
Here we will recall and slightly modify a collection of definitions and technical results from Appendix B and C of ef all (2022), which are related to the shaped activation that we will use in the main theorems. To start, we will let (p(x) max(æ, 0), and recall the ReI.U-like activation function as
(x) 8 max(x, 0) +
+__minnm min 0) (45((z) s_x(xx)
(B.1)
Let g NN0, 1), then we will restate Li et al
2122 Lemma B.3, B.6)
1
Ey((g) 3
By(g)
E((g)
2'
2
(B.2)
87 +
E (g)
84 S_
2g,((g)
= (44+
V2T
We will also recall the definitions
Jp.(pp) = E())))(()'' Kpr(p) = E..(
(B.3)
where g, w are iid A 1) and we define 9 pg + qw with g = VI- p². We will also use the short hand notation to write Jn = Jppp Kp = Kpp-
17

---

Published in Transactions on Machine Learning Research (04/2024)
Here recall from Cho &r Saul (2009) and Li ct. all (2022, Lemma B.7) that
J((p) = /1 - p² + (T arccos p)p
Ki(p) = (6 + (z) 2848-J((-p))
(B.4)
2
Next we will slightly modify Taylor expansion result from (Li
2022. Lemma C.1)
Lemma B.1 (Taylor Expand Shaping Correlation). Recall (a) 84 max(æ, 0) min(x, 0). Let 8
Ct for p> > O, then we have the following Taylor ezpansion
cKi(p) = = v(p) n2p + O(n ()), v(p)
cc)3 (VI p² + parceos
(B.5)
2m
Proof. We start by expanding the formula for (p) to get
cKi(p) = 5 + s- 2x
2
(s7 + +2) (V1 p2 + (w arccosp)p) 2848- 8- (V1-ph pp)c (B.6)
At this point, we can plug in 82 =
f and Taylor expanding gives us nes
cki(p) = P parccos (p)
+ (x - arccos (p))
[po arccos (p) + 2pc c_ arccos (p) - -pc2. arccos (p)
+ (n==)]
2m
(B.7)
+ / p² 22
V p² 0² V 1 p²
2x
+ +0 (n
Simplifying the expressions gives us the desired result of
cK((p) p ((p) np + O(n (P).
(B.8)
Similar to Li of all (2022, Lemma C.2), can also approximate the fourth moment of shaped activation functions via the fourth moment of Gaussians.
Lemma B.2 (Fourth Moment Approximation). Considers the jointly Gaussian random variables
1 N (0, (00
(B.9)
where paa for all a. Let po(x) = 8t max 0) + s_ min(x, 0) with coefficients s= = 1 + have
then we
no.
II p.(gg) = E II 9" + O(n P) = 12 + p + p2 + +(n P). a=1
(B.10)
a=1
Proof. Observe that we can write (x)) as a perturbation of identity
(x(x) x (c4p(z) -app())-c-((=))) O(n
(B.11)
This means we can approximate the product of shaped activations without the activation
4
4 II gd + O(n
Il
(.(g))
(B.12)
a=l
a=l
18

---

Published in Transactions on Machine Learning Research (04/2024)
Finally, use the Isscrlis Theorem to write
E II Eg B +E " - 222 1p0002 014 1,2
(B.13)
a=l
which is the desired result.
Furthermore, we can characterize the covariance of the following random variables
RaB
Σ (Po(of)) (k1())))
(B.14)
vn
where the random vector [gft "f") are iid copies of the same Gaussian vector. This follows from a modification of LLetta all (2n22, Lemma C.3).
Lemma B.3 (Covariance of Roo). Let Rod be defined above. Then if 8+ 1 + for the shaped activation, we have that
E Rad R = pm 9م + pa or + O(n
(B.15)
Proof. Firstly, we note since each entry of the sum in Rad are iid, we will only need to compute the moments for single entry. This means
a R =E6² (p.(of))) Ki (0()0)) (pal9)))x
Ki (o²)))
(B.16)
this point, we observe that c 1 On ki(p) = p 1+00 O(n On),)),, and we can write
E E 1255 = E (P.c()))) pa8) (o (of) /² + P).
(B.17)
Finally, the fourth moment approximation Lemnaa
gives us the desired result
8 1716 = -00 908 + 1o 0 p0 1 pm - 60 03 0 n6 + O(
(B.18)
- 00 + pab p0r P)n
19


Reviews:
[{"review_id": "qvOBfT7dr9", "reviewer": "Reviewer_YWHM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \nTwo key results - \n1. Shaping Activation and ResNet connection: The authors note that the shaping of the activation function in neural networks is closely related to ResNet architectures. In particular, they note that the covariant ordinary differential equation (ODE) is merely the deterministic drift part of a broader covariant stochastic differential equation (SDE). Moreover, for the shaped multilayer perceptron (MLP), they find that if the scaled ratio $\\frac{d}{n^{2p}}$ converges to a positive constant, where $p$ belongs to the interval $(0, \\frac{1}{2})$, the covariance of the shaped MLP also converges to this same ODE. This result emphasises the connection between shaped activations and ResNet structures in neural networks.\n\n2. Correlation Analysis in Unshaped MLP: The second result deals with the correlation dynamics within an unformed MLP. They derive the first-order asymptotic correction for the layer-wise correlation. This correction is introduced by rescaling the correlation $\\rho_{\\ell}$ in the layer $\\ell$ using the term $q_{\\ell} = \\ell^2(1 - \\rho_{\\ell})$. Their analysis shows that this rescaled correlation is approximated by a stochastic differential equation (SDE), indicating a relationship between layer depth, width and correlation in unshaped MLPs.\n\nThese findings contribute to our understanding of the behavior of neural networks. They not only provide a link between shaped and unshaped network architectures, but also offer a new perspective on the dynamics of correlations in neural networks and the principles governing them.\n\nWeakness:\n- The complexity of the article arises from the extensive recourse to external references and the brevity of the proof, which presupposes knowledge of numerous previous results. This approach makes it difficult for readers to understand the arguments put forward. To improve comprehensibility, it would be beneficial to construct the theoretical objects with due precision, formulate the assumptions more clearly. Such a detailed presentation would help readers understand the complicated concepts and logical flow of the arguments, making the material more accessible and informative.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This article explores the behavior of neural networks with unshaped activations, where the activation function remains constant regardless of network size. Two main findings are presented:\n\n1. Convergence of Two Architectures: It demonstrates that two different network architectures converge to the same limit when they are infinitely deep and wide. These architectures are:\n   - A fully connected ResNet with a scaling factor inversely proportional to the square root of the network depth ($d^{-1/2}$).\n   - A multilayer perceptron (MLP) with depth significantly less than width, employing a shaped ReLU activation scaled at a rate of $d^{-1/2}$.\n\n2. Layerwise Correlation in Unshaped MLPs: For an unshaped MLP, the research derives a first-order asymptotic correction to the layerwise correlation. Specifically, if $\\rho_{\\ell}$ is the correlation at a certain layer $\\ell$, then $q_t = \\ell^2(1 - \\rho_{\\ell})$ with $t = \\frac{\\ell}{n}$ converges to a stochastic differential equation (SDE) with a singularity at $t = 0$.\n\nThese insights bridge the gap between shaped and unshaped neural network architectures. They also pave the way for investigating how normalization methods might be related to shaping activation functions.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Restate clearly Theorem 2.1 - with all definitions and notations, conclusions. \n- Make in an appendix a list of definitions [when reading results, would be great to know where the symbols are defined]\n- Restate clearly Theorem 2.2\n- Would like to have a detailed derivation of Theorem 4.1 - The treatment of complicated mathematical concepts, especially the convergence of processes, suffers from the imprecision of definitions and constructions. A detailed proof from first principle is essential. Such clarity would not only improve the understanding of the subject, but also maintain the precision expected in mathematical presentation.\n- The article would benefit from a clearer discussion of the significance and practical implications of its findings. In particular, it should clarify how these results influence our understanding of learning mechanisms and guide the selection of activation functions in neural networks. At present, the relevance and scope of these results for practise remain rather vague. A more explicit link between the theoretical results and their practical applications would greatly enhance the reader's understanding of the implications and relevance of the study in this area.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "MZBKoeG9T0", "reviewer": "Reviewer_BR7V", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Unfortunately, I am not an expert on this topic, so please regard my review accordingly.\n\nBased on the authors' claims, the main strength of this submission is the relation between shaped and unshaped networks.\n\nThe main weakness is the exposition and clarity of writing. In particular, the paper is written for people working on this problem, and it is not accessible to a larger audience. Heavy use of jargon and prior results made it difficult to follow the general plot. In addition, the manuscript contains (or at least, what seems to be) several unfinished sentences. There is no conclusion section. The authors do not motivate their research problem, nor its significance and the impact of their results.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper relates between scaling limits of shaped and unshaped neural networks. Specifically, the authors show that certain unshaped ResNets converge to shaped MLP. In addition, the authors derive correlation results for unshaped MLPs.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I believe the authors should strive to make their submission more accessible to general ML audience.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "b1XJxTSlE0", "reviewer": "Reviewer_mCoA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well written.\n- The organization is clear. \n- I appreciated the authors' efforts to provide intuition with the mathematics.  \n- The contributions are clearly specified in relation to previous work. \n\n\nWeaknesses:\n- As a non-expert in this particular field, I find it difficult to judge how useful and interesting these contributions are to this line of work. It is clear to me that they *could* be interesting to some researchers in this field, but I am not familiar enough with the area to be confident. \n- It would be helpful to have a conclusion that summarizes why these contributions will be interesting to other researchers studying scaling limits and machine learning more generally.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper contributes two ideas to the asymptotic analysis of infinitely wide and deep neural networks. In this line of work, it is common to analyze what happens to the \"covariance\" in each layer, that is the expectation of how the activations in layer L for input x1 covary with those for input x2. This type of analysis has been useful in understanding neural network hyperparameters including weight initializations, learning rates, activation functions, batch norm, and layer norm. \n\nIn the present paper, the authors first draw a connection between two previously-published analyses, showing that the ODE describing the depth-evolution of the covariance in a infinitely wide and deep ResNet (Hayou and Yang 2023) is the drift term in the SDE describing the depth-evolution of the covariance in a fully-connected multilayer perceptron with shaped ReLU activation (Li et al 2022). Secondly, the paper provides an analysis of the covariance for an unshaped ReLU MLP.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have no major requested changes. \nSome typos: \n- Section 1, in the description of the first contribution, \"Furthermore in the limit requires the scale ratio...\"\n- Section 1.1 ends with \"...description remains an open.\"", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper analyzes the infinite-width-and-depth limit of certain neural network architectures at initialization, finding connections between residual networks and shaped activations, as well as the leading-order correction to the asymptotic activation covariance in standard ReLU MLPs.

The reviewers generally appreciated the results and the intuition provided to help readers understand the findings. The derived SDEs provide a precise new basis for our understanding of the dynamics of signal propagation at initialization, and will surely be of interest to the community. As the authors note, conclusions about initialization do not necessarily carry over during training, but I agree that building a deep understanding at initialization is a prerequisite for the much more ambitious goal of characterizing training dynamics. All together, the scope and analysis seem entirely appropriate for TMLR.

The reviewers noted that the original submission was challenging to read for someone without expertise in the area. In the revision, the authors provided significantly more background, discussion, and details that make the paper considerably more approachable. As such, I believe the paper is now appropriate to publish, and recommend accepting the paper as is.
Decision: Accept (poster)