Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit
Authors: Blake Bordelon
Lorenzo Noci
Mufan Bill Li
Boris Hanin
Cengiz Pehlevan
Venue: ICLR 2024 poster (2024)
Keywords: Deep Learning Theory
Hyperparameter Transfer
Training Dynamics
Residual Networks
Primary area: optimization
TL;DR: A study of hyperparameter transfer and training dynamics in large depth residual neural networks.

Abstract: The cost of hyperparameter tuning in deep learning has been rising with model sizes, prompting practitioners to find new tuning methods using a proxy of smaller networks. One such proposal uses $\mu$P parameterized networks, where the optimal hyperparameters for small width networks *transfer* to networks with arbitrarily large width. However, in this scheme, hyperparameters do not transfer across depths. As a remedy, we study residual networks with a residual branch scale of $1/\sqrt{\text{depth}}$ in combination with the $\mu$P parameterization. We provide experiments demonstrating that residual architectures including convolutional ResNets and vision transformers trained with this parameterization exhibit transfer of optimal hyperparameters across width and depth on CIFAR-10 and ImageNet.  Furthermore, our empirical findings are supported and motivated by theory. Using recent developments in the dynamical mean field theory (DMFT) description of neural network learning dynamics, we show that this parameterization of ResNets admits a well-defined feature learning joint infinite-width and infinite-depth limit and show convergence of finite-size network dynamics towards this limit.

Published as a conference paper at ICLR 2024
DEPTHWISE HYPERPARAMETER TRANSFER IN RESID- UAL NETWORKS: DYNAMICS AND SCALING LIMIT
Blake Bordelon" 1, Lorenzo Noci Mufan (Bill) Li Boris Hanin & Cengiz Pehlevan
Harvard University
: ETH Zürich
8 Princeton University
ABSTRACT
The cost of hyperparameter tuning in deep learning has been rising with model sizes, prompting practitioners to find new tuning methods using proxy of smaller networks. One such proposal uses μP parameterized networks, where the optimal hyperparameters for small width networks transfer to networks with arbitrarily large width. However, in this scheme, hyperparameters do not transfer across depths. As remedy, study residual networks with residual branch scale of 1/ /epth in combination with the μP parameterization. We provide experiments demonstrating that residual architectures including convolutional ResNets and Vi- sion Transformers trained with this parameterization exhibit transfer of optimal hyperparameters across width and depth on CIFAR-IO and ImageNet. Further- more, our empirical findings are supported and motivated by theory. Using recent developments in the dynamical mean field theory (DMFT) description of neural network learning dynamics, we show that this parameterization of ResNets ad- mits well-defined feature learning joint infinite-wiie and infinite-depe depth limit and show convergence of finite-size network dynamics towards this limit.
INTRODUCTION
Depth
1.0
1.0
6
0.8
0.8
12
Train Loss
Train Loss
0.6
0.6
Width
128
1024
0.4
0.4
10-3
10-2
10-1
10-3
10-2
17
Learning Rate
Learning Rate
(a) ResNet, μP
(b) This work: J ReeNNet + JPP
Figure 1: The optimal learning rate 17" transfers across both depth and width in our proposed param- eterization but not in μP or standard parameterization (Fig. B.2). Loss is plotted after 20 epochs on CIFAR-10. 10. All the missing datapoints indicate that the corresponding run diverged.
Increasing the number of parameters in a neural network has led to consistent and often dramatic improvements in model quality (Kaplan et al. (2020) Hoffmann et al.) 2022 Zhai et al. [022]] Klug & Heckel 2023 OpenAl 2023). To realize these gains, however, it is typically necessary to conduct by trial-and-eme grid search for optimal choices of hyperparameters, such as learning rates. Doing so
"Equal first authors. I Equal senior authors.

---

Published as a conference paper at ICLR 2024
-318
0.45
268
Train Loss
218
0.35
168
1188
0.25
68
0.20
10
10
10
10-1
Learning Rate
Learning Rate
(a) ResNet+ JPP BatchNorm
ResNet JPP, BatchNorm
Figure 2: Effect of batch normalization standard Resnet18 8-tppe architecture (He et al. 2016 The above examples are taken after 20 epochs on CIFAR-10. Normalization layers can slightly improve consistency and trainability across depths in standard ResNets, but consistency is much more reliable with the 1 VL scaling (precisely, at finite depth). Runs that exceed target loss of 0.5 are removed from the plot for visual clarity.
Bg 3/ 3/vL. here to increase feature learning
directly in SOTA models, which may have hundreds of billions (Brown et al. 2020) or even trillions of parameters (Fedus al. 2022), often incurs prohibitive computational costs.
To combat this, an influential recent line work by Yang & Hu (2021) proposed the so-called μP parameterization. which seeks develop principles by which optimal hyperparameters from small networks can be reused - or transferred prescription focuses transfer from narrower wider models, but does not always transfer across depth (see Figure a)). This leads us to pose the following problem:
to larger networks (Yang et al., 2021). The jP
Question: Can we transfer hyperparameters simultaneously across depth and width?
In this work, we provide an affirmative answer for a particular flexible class of residual architectures (including ResNets and transformers) with residual branches scaled like 1 /eyttphh (see Equation and Table [ for the exact parameterization). More specifically, our empirical contributions can be summarized follows:
We show that a simple variant of the μP parameterization allows for transfer of learning rates simultaneously across depth and width in residual networks with 1/v /dytepth scaling on the residual branches (see Equation and Figure lbb).. Morcover, we observe transfer not only for training with a fixed learning rate, but also with complex learning rate schedules (Figure a-dd)..
We show that the lack of learning rate transfer across depth in μP ResNets can be partly overcome by adding normalization layers (Figure 2) However, we find on CIFAR-10, Tiny ImageNet, ImageNet, that the transfer can be significantly improved with our 1/ depth residual branches (Figures and B.6
While most of our experiments concern residual convolutional networks, such as wide ResNets, we also present evidence that our prescription leads to learning rate transfer in Vision Transformers (ViTs), both with and without normalization (Figure 3
We demonstrate transfer for learning rates, as well as other hyperparameters such momentum coefficients and regularization strengths in convolutional residual networks (Figure(4).
The reason hyperparameters can transfer is due to the existence of scaling limit where the net- work predictions and internal features of the model move rate independent of the size of the network. The key underlying idea is that two neural networks will exhibit similar training dynamics (and hence similar optimal hyperparameters) if they are both finite size approximations to, or dis- cretizations of, such consistent scaling limit. In fact, the success of the mean field jP scaling (Yang& & Hu| [202]] Bordelon & Pehlevan (022b) is precisely due to the width-independent scale of teature updates. In the same spirit, propose our parameterization that the rate of feature up-
2

---

Published as a conference paper at ICLR 2024
SP (PyTorch)
Mean Field
uP+1/VL- L-Rssidul (Ours)
N N-
{ =
N [N/1/2 8>
Branch Scale Be
1
(LN)-112
0 < f <
D-11/2 t=
D DD-1/2
8 = 0
Output Scale Y
1
No.11/2
No22
I.R Schedule 17(1)
(1)(1)
mo(4))
[N-1, 8>0 D-11 t=
Weight Variance of
1
Table 1: In the framework of Equation different choices of coefficients leads standard param- eterization (SP, PyTorch default (Paszke et al. 2019)), JPP Mean Field parameterization, and our /uPts/VTrrsiidull branch scaling which transfers over width and depth.
dates and prediction updates are also independent the depth of the network. Our main theoretical contributions can be summarized as follows:
Extending the infinite-widd width limit results Bordelon & Pehlevan (002b), we show the infinite- width-and-depth limit of residual network with 1/ /epth residual branch scaling and uP initial- ization can be characterized by a dynamical mean field theory (DMFT) (Section 4 4.2) where the layer index takes on continuous "layer time." In general, the network's internal features evolve during training in this limit. We show that a lazy limit the DMFT dynamics, where features are frozen, also be analyzed
special case. See Section 4.3.1 and Proposition 2
We provide exact solutions of the DMFT dynamics in the rich (i.e. non-kernel) regime for simple deep linear network. See Section 4.3.2
The rest of the article organized as follows. In Section 2 we define our parameterization and scaling rules. In Section demonstrate our empirical results on hyperparameter transfer. Then in Section 4 we proceed to characterize the limiting process at infinite width and depth and the approximation error finite N, L. Next we provide a brief overview of related works Section before concluding with a discussion of limitations and future directions in Section
FRAMEWORK FOR DEPTHWISE HYPERPARAMETER TRANSFER
We consider the following type of residual network (and suitable architecturalloptinizzooo gener- alizations, see Appendix K of width N and depth L with inputs N V-dimensional preactivations ht E IRN and outputs f(æ)= E IR:
IR which mapped to
BL "0-=" - (x)) h h (x)) h hxx)) + 3 W ((h(())), h h²(x)) 30 Wo, (1) 7
[(x)
where 7 is a scaling factor, o(-) the activation function, and the weights are initialized W
(0, of) with corresponding learning rate v/((). See Table 1
Compared to what is commonly done practical residual architectures, such as ResNets and Trans- formers, our parameterization uses the VL factor proposed in Hayou et al. 2021) Fischer et al. 2023). and the coefficient 7 from jPP, but do not normalization layers (Table I. The parame- ters 0 = (Wo ...., w are updated with a gradient based learning rule. For concreteness consider SGD which we use in many experiments (we also allow for gradient flow and other optimizers in Appendix K
@t O(t) -nt)) EZ-D, Vo (ff(0,))]
(2)
where 2, is the minibatch at iteration i, v)() is the learning rate schedule, and the loss function L. depends on parameters through the predictions of the model.
The choice of how (7, (t)) should be scaled as width N and depth L go to infinity determines the stability and feature learning properties of the network. If 7. O(1 (kernel regime), then the
We use (1) denote constant independent of width N and depth I..
3

---

Published as a conference paper at ICLR 2024
Accuracy
Dept
Depth
40
Depth
10
20
Learning Rate
Loarning
Learning
(a) VL
VTT + PP
(b) V
ViT + p.P. LayerNorm
(c) ViT J.P, LayerNorm
90
Depth
Train Accuracy
0.0
20
60
Trals
Train
10
50
80
30
10-3
12
Learning Rate
Epochs
Learning
(d)
VIT + JaP, Cosine Annealing
(c)
1 VI
ViT ImageNet Dynamics
(f) ViT p.P, ImageNet
VI
Figure 3: ViTs trained with Adam also exhibit learning rate transfer, with or without LayerNorm and 1/V scalling. The above examples are taken after 20 epochs on Tiny ImageNet for (a)-(d) and ImageNet after 10 epochs for (e)-(f). In (a)-(c), the learning rate is linearly increased to the target in the first 1000 steps. In (d), after 2000 warm-up steps, the learning rate is decreased to zero with cosine schedule in 100 epochs. Notice that in addition the transfer across all the settings, the 1/VL models show stronger non-saturating benefit of deeper models. (e)-(f) ViTs on first few epochs of ImageNet. Also, we note (a), (d), (e) and (f) refer to models without normalization layers.
network evolution is stable under gradient descent, but the network does not learn features infinite width (Jacot et al.) 2018) Lee et al., 2019 Yang & Hu 2021). In this work, we study the richer feature learning (mean field) parameterization = To v N, rn(() no (1) N, where 70, (() are independent of N. In this regime, the scaling factor of 1/7 in the last layer enables feature learning at infinite width, and the constant "Yo controls the rate of feature learning. At fixed depth L, this mean field parameterization capable generating consistent feature learning dynamics across network widths N Yang & Hu 2021 Bordelon & Pehlevan 20226).
One contribution of this work identifying how yoo "o should be scaled with depth in order to obtain a stable feature-learning large L limit. We argue that after the factors have been introduced the correct scaling is simply
1 = 7o = 7oV,, 110, To 0(1.
(3)
We show both theoretically and experimentally that this parameterization exhibits consistent learn- ing dynamics and hyperparameter optima across large network depths L. To justify this choice, we will characterize the large width and depth N, >> co limit of this class of models.
3 EXPERIMENTAL VERIFICATION OF HYPERPARAMETER TRANSFER
We report here a variety of experiments on hyperparameter transfer in residual networks. We begin with basic modification of equation by replacing fully connected layers with convolutional layers. Apart from the proposed 1/VL scaling, notice that this convolutional residual model is similar to widely adopted implementations of ResNets (He et al. 2016), with the main difference being the absence of normalization layers, and a single convolution operator per residual block (details in Appendix
B.2 and B.I We also experiment with a practical ResNet (in figures] 2 B.4 3.4.B... B.6). which has two convolutional layers per block, and normalization layers (see Appendix A. Finally, also perform selected experiments on Vision Transformers (Dosovitskiy et al. 2020), to demonstrate the broad applicability of our theoretical insights (Figure 3). Below list the the main takeaways from our experiments.
A)). We this residual convolutional architecture in figures 4
μP dues not Transfer at Large Depths. While μP exhibits transfer over network widths, does not immediately give transfer network depths. In Figure [a), we show the train loss after 20
4

---

Published as a conference paper at ICLR 2024
2.0
Mom:
1.2
16
0.05
Train Loss
0.4
0.1
1.5
0.8
0.5
1.0
Depth
1.0
0.6
Depth
12
30
10 Epochs
Epoch
(a) Momentum
(b) Feature Learning Rate Yo
Figure 4: Other hyperparameters also transfer. This example shows training dynamics during the first 20 epochs on CIFAR-10 (architecture details in Appendix A The dynamics at two different depths are provided for (a) momentum (b) feature learning rate Y0-
epochs on CIFAR-10. Notice how the network transfers over widths but not depths for p.. For completeness, in Fig. B.2 we show that also SP does not transfer.
1/VL-scaled Residual Networks Transfer. We repeat the same experiment, but now scale the residual branch by I/vL. The networks transfer over both width and depth (Figure 1 (b)).
The Role Normalization Layers. Normalization layers, such as LayerNorm and BatchNorm, are commonly used instead of the 1/ //VT-scalnn in residual architectures. In Fig. we repeat the learning rate transfer experiment with BatchNorm placed after the convolutional layer. In Fig. B.4 we repeat the same experiment with LayerNorm. In both cases, while is known that the presence of normalization layers certainly helps the trainability of the model at larger depth (Daneshmand al.,
2021 Joudaki et al. 2023), we observe that the 1/v L/vLlld version exhibits more consistent transfer across both width and depth.
Empirical Verification Vision Transformers with Adam. Despite our theory being primarily on (S)GD-based training, we empirically test whether the transfer results can be extended to Vision Transformers (ViTs) trained with yıP- parameterizzed Adam optimizer (Yang et al. (2021). In particu- lar, we consider Pre-LN Transformers (Xiong et al.) 2020). and adapt to our setting by scaling all the residual branches by 1/VL. (as in Noci et al. 2022)), and make the architecture compatible with the /uP framework of Yang al. (222) (see Appendix A.3 for details). We test the variants both with and without LayerNorm. Results for 20 epochs training with the Tiny ImageNet dataset are shown in Fig. (a-c), where we also see good transfer across all the settings. The observation of a good transfer with depth for Pre-LN Transformers was empirically observed for translation tasks Yang et al.] (202T). Here, we additionally notice that while the performance benefits of Pre-LN Transformers saturates large depth, the models with 1
exhibits more consistent im- provement with depth. Finally, in Fig. e-f) we successfully test scaling the learning rate transfer experiment of // /L-scaled ViTs (without LayerNorm) to ImageNet (see Appendix B for ImageNet experiments with ResNets).
Other Hyperparameters and Learning Rate Schedules also Transfer. Given the standard prac- tiee of using learning rate schedulers optimizing Transformers (Dosovitskiy et al. 2020 Touvron al.), show that linear warm-up (Fig a-c)),, optionally followed by cosine decay (Fig d))
also transfer consistently. Finally, Figure we plot the learning dynamics for momentum, weight decay, and the feature learning rate "Yo, which interpolates between the kernel and feature learning regimes (Bordelon & Pehlevan, 2022b). Despite the usual correct transfer, we highlight the consis- tency of the dynamics, in the sense that learning curves with the same hyperparameter but different depths are remarkably similar, especially at the beginning of training.
CONVERGENCE TO THE LARGE WIDTH AND DEPTH LIMIT
We argued that width/depth-inwann invariant feature and prediction updates crucial for hyperparameter transfer. To formalize this, let (t) be the network output after steps of optimization, where explicitly write the dependence width N and depth L. For hyperparameter transfer want
5

---

Published as a conference paper at ICLR 2024
10-4
|fx f
2.0
64
128
1.6
256
512
200
10a
10
Steps
N 512 SGD Dynamies
(b) Convergence L,
Figure 5: Approximation of the joint N co, limit requires both sufficiently large N and L. CNNs are compared after 250 steps on CIFAR-10 with batchsize 32 with To 1.0. Since cannot compute the infinite width and depth predictor, use proxy f= the ensemble averaged (over random inits) predictor standard deviation computed over 15 random initializations. (a) SGD dynamics large width are strikingly consistent across depths. (b) The convergence of fN.L to a large width and depth proxy f000,00 is bottlenecked by width N for small N while for large N, it decays like O(L-1).
networks with N = 512 and
128. Error bars are
(t) to be close to fN , (t). One strategy to achieve this is to adopt a parameterization which makes finite networks close possible to a limit fo0 (t). Following this line of thought, desire that a parameterization:
Have well-defined large depth and width limit, i.e. limn,L-4so fN,L(t).
Have internal features kernels evolve asymptotically independent of width or depth L.
3. Have minimal finite size approximation error.
These desiderata motivated the - -uP ResNet in Equation and Table The existence of the limit for the forward pass at initialization = O was shown in Hayou & Yang (202)) Cirone et al. (2023). Here, we study the learning dynamics in this model with μPP paramcterization, establish that the limit exists and show that feature updates are width and depth independent. Further, we characterize the limit throughout training in Appendices limit in Figure 5) and Appendix Figure B. See Appendix H and for theoretical analysis of finite size approximation error and Figure B. for additional tests of the sources of finite width/depth approximation error.
We attempt to visualize convergence to the
4.1 PRIMER ON THE DMFT IN THE N -> co LIMIT
To derive our theoretical results, we use a technique known as dynamical mean field theory (DMFT) Bordelon & Pehlevan 2022b). DMFT describes the dynamics of neural networks at (or near) infinite width (preserving feature learning). this setting one tracks a set kernel order parameters
1 きh (' a/:1, @ (a,(),))) (1'((a)))) ('r, = gar, (
(x', 8),
(4)
where gf (x. t) = Noo (f(m.)
(a.(). are the back-propagation signals. DMFT reveals the following facts
about the infinite width limit:
1. The dynamical kernels (Ø', Gz} and network outputs f concentrate over random initializations of network weights throughout training.
The preactivations and gradient fields of each neuron {hā,g() become i.i.d random variables drawn from single-site (i.e. neuron marginal) density which depends on the kernels.
3. The kernels be computed averages over this single-site density.
In addition, an ideal parameterization would also have the property that scaling up the width and depth of the model would monotonically improve performance, which empirically seems hold in our parameterization.

---

Published as a conference paper at ICLR 2024
The last two facts provide a self-consistency interpretation: given the kernels, one compute or sample the distribution of preactivations {h², Given the density of {h², g }. one can compute the kernels. Prior work on this limit has demonstrated that these equations be approximately solved at finite I. with a Monte Carlo procedure (but with computational complexity cubic in the training iterations) (Bordelon & Pehlevan 2022b).
4.2 THE LARGE-WIDTH--NE WIDTHI-AND-DEPTI LIMIT DMFT EQUATIONS
From the finite-L DMFT equations (App- D we calculate the large depth 1. + co limit, generating
continuum process over the layer time I [0, 1], a concept introduced in Li et al. (2022
Proposition 1 (Informal) Consider network training dynamics of a ResNet as in Equation ! with our pP J.. scaling and the appropriate choice of learning rate and "/o in the infinite width and depth limit. The preactivation h(T; x, t) drawn from the marginal density of neurons in layer at layer-time t, data point x and training time obeys a stochastic integral equation
h(T; ax; 1) = h(0;=;;)) + S( du(t'; x); + 170010 S. dr' S( ds da' Tac'C((T'ss x'; 1,8888)) s)g(T' : x';
(5)
where du( is zero mean Brownian motion. The covariance of du is the feature kernel and the deterministic operator Cn can be computed from the deterministic limiting kernels 4. G. The gradients g(T) satisfy an analogous integral equation (See App. E.7 for formulas for Cn and g). The stochastic variables h,g g and the kernels 4 G satisfy a closed system of equations.
A full proposition can be found in Appendix E.7 equations challenging solve in the general case, the h and g distributions are Gaussaan Its main utility in the present work is demonstrating that there is a well-defined feature-learning infinite width and depth limit and that, in the limit, the predictions. and kernels evolve by O(1).
We derive this limit in the Appendix E These
4.3 SPECIAL EXACTLY SOLVEABLE CASES OF THE N, L - DO DMFT EQUATIONS
Though the resulting DMFT equations for the kernels are intractable in the general case, we can provide analytical solutions in special cases to verify our approach. We first write down the training dynamics in the lazy limit by characterizing the neural tangent kernel (NTK) at initialization. Next, we discuss deep linear ResNets where the feature learning dynamics gives ODEs that can be closed while preserving Gaussianity of preactivations.
4.3.1 THE LAZY LEARNING LIMIT
A special case of the above dynamics potentially of interest is the limit of n0 - 0 where the feature kernels are frozen. In this limit the NTK K is static and can be computed at initialization and used to calculate the full trajectory of the network dynamics (see Appendix G.I.
Proposition 2 (Informal) Consider a ResNet with the - uP parameterization in the N, L + co limit. In the % - limit, the feature kernel (TT), gradient kernel G(T) and NTK K do not evolve during training and satisfy the following differential equations over layer time T
0, H H, x') = Φ(T; x'), ((T:0, = (6(h))((h))) (,.
0,G(T; n, x.)) = -G(T; x²) (((1))(())
H(T K(x, = ('' dr G(T, ((),,, H(T; H(T;2',x))
H(T
Hn (T)
(6)
The dynamics of the neural network outputs f(x,t) t) can be computed from K (x, x').
The forward ODE for & at initialization was studied by Hayou & Yang (2023); ; Cirone et al. (2023).
Figure [ shows this limiting NTK for a ReLU network. The finite depth kernels converge to the limiting large T. kernel (blue line in (a)) at a rate (T.-22 in square error.
A Atternativel, one can also take the large-I. limit to get limiting distribution over order parameters (ker- nels, output logits, etc) fixed N, which yields the same limiting equations when N oc (see App. F).
7

---

Published as a conference paper at ICLR 2024
0.20
1-4
1-
10-1
K(0)
512
Lue
0.05
0.00
-3
(a) ReLU NTKs at Initialization
(b) Initial NTK Convergence in L
Figure 6: Convergence of the depth L NTK to the limiting kernel at initialization. (a) The infinite width NTK for a ReLU residual multilayer perpectron (MLP) for two inputs 2, x'' separated by angle 0 E [- π over varying depth L. The kernel converges as + oo. (b) The convergence rate of the initialization averaged NTK (K) with network widths N and network depths L at initialization. For sufficiently large N, the initial NTK will converge at rate (L -) in square emror
2.75
dasth -
10
desth 10
depth
107
2,25
Hgl
2.00
10
Septh 10
depth =
depth
10-2
Sc5th 20
Inf Depth Theery
depth 50
1.59
1.25
desth
Inf Depth Theory
0.8
1.0
Kernels at initialization
(b) Kernel Updates after 1 Step
Figure 7: We plot the feature kernels H()) = H l- in width N = 1000 residual linear networks before and after taking gradient step in parameterization. The kernel profiles at initialization Ho(r) and the kernel changes (T) Ho(T) converge for large L.. We compare to the theoretical result expected from the + oo limit of the DMFT equations (dashed black).
4.3.2 DEEP LINEAR RESNETS
Next we explore the dynamics of feature learning in a solveable infinite depth ResNet with linear activations. In this case, the preactivations h and gradient fields g remain Gaussian over the full time-course of learning, rendering the DMFT equations tractable without any need for Monte Carlo methods. As an example, in Figure 7 we show the kernel updates in deep linear ResNet after single training step on single training example (x, y). The theoretical solution to the kernel ODE initialization HT( (0) and after one step H1 (7) is (see Appendix G.2.2 for derivation)
T2 + (1 ()+ - -e" + e -n] -
Hi(T) = e H (0)
(7)
3
For networks with depths on the order of ~ 100 the infinite depth solution gives an accurate descrip- tion of the profile of the kernel updates. More general equations for infinite depth linear networks are provided in Appendix G.2
RELATED WORKS
The 1 1/V L-scaleed residual model (He et al. 2016 adopted this work belongs to a rapidly expand- ing research domain. This approach has several advantageous characteristies in signal propagation, in terms of dynamical isometry (Tarnowski et al. 2019), as well as stable backward and forward ker- nels (Balduzzi et al. 2017 Hayou al. 2021 Fischer al. 2023). As result, this has paved the
8

---

Published as a conference paper at ICLR 2024
way for the development of reliable initialization methods Taki 2017 Zhang et al. 2019). including
Transformers without normalization layers (Huang et al. 2020 Noci et al. 2022)
The study of neural network scaling limits started with MLPs in the infinite-width kernel regime (Jacot et al. 2018 Neal, 2012 Matthews et al. 2018 Lee et al. 2017), where the preactivations also Gaussian with deterministic covariance. However, the joint depth and width limit has been shown to be non-Gaussian (Hanin & Nica, 2020; Noci et al. (2021) and have a stochastic covariance (Li et al. 2021). In particular, the feature covariance of particular classes of networks have limit described by stochastic differential equations (SDEs) (Li et al., 2022 Noci et al. (023). where the layer time the depth-to- towth ratio in this case. On the other hand, Hayou & Yang (2023) study the forward signal propagation in 1/ 1/vLe L-scaled ResNets at initialization, where the preactivations converge to a Gaussian distribution with a deterministic covariance described by an ODE.
Transfer across width and depth was achieved on non-residual networks with the Automatic Gradient Descent algorithm proposed by |Brnsttin et al. (2020) 2023). Successful hyperparameter transfer
across network widths μrP was shown in Yang et al. (021). Width-consistent early time predic- tion and representation dynamics were also observed in an empirical study of Vyas et al. (2023).
The infinite width limits of feature learning networks have been characterized with a mean-field
PDE in two layer networks (Chizat & Bach 2018; Mei et al. 2019; Chizat & Bach 2020). Tensor programs Yang 2020 Yang & Hu 2021), and DMFT techniques (Bordelon & Pehlevan [022b a
2023], which have been proven useful understand the gradient dynamics in other settings well
iMignacco et al. 2020 Gerbelot et al.] 2022). This present work extends the DMFT technique to
ResNets to analyze their feature learning behavior at large depth L.
Understanding feature learning in non-residual feedforward networks close to the asymptotic regime has also been studied with finite width and depth corrections from the corresponding infinite limit, both in the lazy (Chizat et al. 2019) NTK (Hanin) [018 HHanin& & Nica, 2019) and NNGP limits
Antognnin 2019; Yaida 2020 Li & Sompolinsky 2021 Zavatone- Veth et al. [222]] Segadlo al.
2022 Hanin 2022), and more recently in the rich Teature learning regime (Bordelon & Pehlevan 2023 At any finite width and depth the analysis is even more challenging and has been studied in special cases using combinatorial approaches and hypergeometric functions (Noci et al. 2021
Zavatone- Veth & Pehlevan 2021 Hanin & Zlokapa 2023).
DISCUSSION
In this work, we have proposed simple parameterization of residual networks. In this parameteri- zation, we have seen empirically that hyperparameter transfer remarkably consistent across both width and depth for variety of architectures, hyperparameters, and datasets. The experimental ev- idence is supported by theory, which shows that dynamics of all hidden layers are both non- con-trivial and independent of width and depth while not vanishing in the limit. We believe that these results have the potential to reduce computational costs of hyperparameter tuning, allowing the practitioner to train a large depth and width model only once with near optimal hyperparameters.
Limitations. While we demonstrate the existence of an infinite width and depth feature learning limit of neural network dynamics using statistical field theory methods and empirical tests, our results are derived at the level of rigor of physics, rather than formal proof. We also do not consider joint scaling limits in which the dataset size and number of optimization steps can scale with width N or depth L. Instead, we treat time and sample size as ((1) quantities. Due to computational limits, most of our experiments are limited 10-20 epochs on each dataset.
Future Directions. This paper explores one possible scaling limit of infinite width and depth net- works. Other parameterizations and learning rate scaling rules may also be worth exploring in future work (Jelassi et al., 2023) Li et al.) 2022 Noci et al.) 2023). It would be especially interesting to see if an stochastic limit for the kernels can be derived for the training dynamics in ResNets without the 1/VL scaling factor. This parameterization makes it especially convenient to study depth-width trade-off's while scaling neural networks since hyperparameter dependencies are consistent across widths/depths (see Figure B.I. It would also be interesting to study our limiting dynamics in fur- ther detail to characterize what kinds of features are learned. Finally, it remains an open problem theoretically why hyperparameters transfer across widths and depths even when the actual predic- tions/losses differ across width and depth.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
BB is supported by a Google PhD research fellowship. BB and CP are also supported by NSF Award DMS-2134155 This work has been made possible in part by gift from the Chan Zucker- berg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. BH and MBL are supported by NSF grant DMS-2133806 BH is further supported by NSF CAREER grant DMS- 2143754, NSF grant DMS-1855684 and an ONR MURI on Foundations of Deep Learning. We would like thank Jacob Zavatone- Veth, Mary Letey, Alex Atanasov for comments on early versions of this manuscript.
REFERENCES
Joseph M Antognini. Finite size corrections for neural network gaussian processes. arXiv preprint arXiv: /908. 10030, 2019.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Unceriainty in Artificial Intelligence. pp. 1352-1361. PMLR, 2021.
David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets the answer, then what is the question? In Interna- tional Conference on Machine Learning, pp. 342-350. PMLR, 2017.
Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning, Advances in Neural Information Processing Systems, 33: 21370-21381, 2020.
Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic gradient descent: Deep learning without hyperparameters. arXiv preprint arXiv:2304.05187, 2023.
Blake Bordelon and Cengiz Pehlevan. The influence of learning rule on representation dynamics in wide neural networks. arXiv preprint w..iv:22/0.02/11 2022a.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022b.
Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. arXiv preprint arXiv:2304.03400, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pp. 1305-1338. PMLR, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. Advances in neural informa- tion processing systems, 22, 2009.
Nicola Muca Cirone, Maud Lemercier, and Cristopher Salvi. Neural signature kernels as infinite- width-depth-limits of controlled resnets. arXiv preprint arXiv:2303.17671 2023.
Hadi Daneshmand, Amir Joudaki, and Francis Bach. Batch normalization orthogonalizes repre- sentations in deep random networks. Advances in Neural Information Processing Systems, 34: 4896-4906, 2021.
10

---

Published as a conference paper at ICLR 2024
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. /929, 2020.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270, 2022.
Kirsten Fischer, David Dahmen, and Moritz Helias. Optimal signal propagation in resnets through residual scaling. arXiv preprint arXivv2305.0777 07715, 2023.
Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zdeborova. Rigorous dynamical mean field theory for stochastic gradient descent methods. arXiv preprint arXiv:2210.06591 2022.
Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? Ad- vances neural information processing systems, 31, 2018.
Boris Hanin. Correlation functions in random fully connected neural networks at finite width. arXiv preprint arXiv:2204. 20000008 2022.
Boris Hanin and Mihai Nica. Finite depth and width corrections the neural tangent kernel, arXiv preprint arXiv:1909 00999, 2019.
Boris Hanin and Mihai Nica. Products of many large random matrices and gradients in deep neural networks. Communications in Matheratical Physics, 376(1):287-322 2020.
Boris Hanin and Alexander Zlokapa. Bayesian interpolation with deep linear networks. Proceedings of the National Academy of Sciences, 120(23):e23013455122 2023.
Soufiane Hayou. On the infinite-depth limit of finite-width neural networks. Transactions on
Machine Learning Research, 2023. ISSN 2835-8856. URL httpp:://peeneevve..
Forum?1d=RbLsyz10
Soufiane Hayou and Greg Yang. Width and depth limits commute residual networks. arXiv preprint arXiv:2302.00 00005333000003, 2023.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousscau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pp. 1324-1332, PMLR, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint 2rXiv:2203.15566 2022.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer opti- mization through better initialization. In International Conference on Machine Learning, pp- 4475 4483. PMLR, 2020.
Kiyosi Itō. 109. stochastic integral. Proceedings of the Imperial Academy, 20(8):519 524,, 1944.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Samy Jelassi, Boris Hanin, Ziwei Ji, Sashank J Reddi, Srinadh Bhojanapalli, and Sanjiv Kumar. Depth dependence of up learning rates in relu mlps. arXiv preprint arXiv:2305.07810 2023.
Amir Joudaki, Hadi Daneshmand, and Francis Bach. On the impact of activation and normalization in obtaining isometric embeddings at initialization. arXiv preprint arXiv:2305. 18399, 2023.

---

Published as a conference paper at ICLR 2024
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
Tobit Klug and Reinhard Heckel. Scaling laws for deep learning based image reconstruction. In The Eleventh International Conference on Learning Reprexentations, 2023. URL https://
openreview net/ ferss/forunte wtw.:m/ddooo-o-eeeeeeeeeepeeeeep
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.001655, 2017.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Mufan Li, Mihai Nica, and Dan Roy. The future is log-gaussian: Resnets and their infinite-depth- and-width limit at initialization. Advances in Neural Information Processing Systems, 34:7852- 7864, 2021.
Mufan Li, Mihai Nica, and Dan Roy. The neural covariance sde: Shaped infinite wedth-and-width networks at initialization. Advances Neural Information Processing Systems, 15:05 2022.
Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The back- propagating kernel renormalization. Physical Review X, 11(3):031059, 2021.
Paul Cecil Martin, ED Siggia, and HA Rose. Statistical dynamics of classical systems. Physical Review A, 8(1):423, 1973.
Alexander G de Matthews, Mark Rowland, Jiri Hron, Richard Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv: 11271,
2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory two-layers neural networks: dimension-free bounds and kerne limit. In Conference on Learning Theory, pp. 2388- 2464. PMLR, 2019.
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540-9550, 2020.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofmann. Precise characterization of the prior predictive distribution of deep relu networks. Advances in Neural Information Processing Systems, 34:20851-20862, 2021.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical persecttives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198 2,2221198-2211,,222 2022.
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 and-width limit. arXiv preprint arXiv:2306. 17759, 2023.
OpenAl. Gpt-4 technical report, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Walter Rudin. Principles of mathematical analysis. 1953.
12

---

Published as a conference paper at ICLR 2024
Kai Segadlo, Bastian Epping, Alexander Meegen, David Dahmen, Michael Krämer, and Moritz Helias. Unified field theoretical approach to deep and recurrent neuronal networks. Journal of Statistical Mechanics: Theory and Experiment. 2022(10):00301, 2022.
Masato Taki. Deep residual networks and weight initialization. arXiv preprint arXiv:1709.02956 2017.
Wojciech Tarnowski, Piotr Warchot, Stanistaw Jastrzebski, Jacek Tabor, and Maciej Nowak. Dy- namical isometry is achieved in residual networks in a universal way for any activation function. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 2221-2230. PMLR, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: open and efficient foundation language models, 2023. URL https://arxiv. wt//bbs22302.3377/
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cen- giz Pehlevan. Feature-learning networks are consistent across widths realistic scales, 2023.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533, PMLR, 2020.
Sho Yaida. Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pp. 165-192. PMLR, 2020.
Ge Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 14:17484-,7097, 2021.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11737. PMLR, 2021.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2003.03666, 2022.
Jacob Zavatone- Veth and Cengiz Pehlevan. Advances in Neural Information Processing Systems, 34:3364-3311 2021.
Jacob Zavatone- Veth, Abdulkadir Canatar, Ben Ruben, and Cengiz Pehlevan. Asymptotics of repre- sentation learning in finite bayesian neural networks. Advances in neural information processing systems, 34:24765-24777, 2021.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 12104-12113, 2022.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv: 09,321,,,, 2019.
13

---

Published as a conference paper at ICLR 2024
APPENDIX
A EXPERIMENTAL DETAILS
A.1 SIMPLIFIED RESNET
We design residual architecture following the footprint of Eq. and replacing the weight matrices with 1-strided convolutional layers with 3 kernel size. We max max-pool the feature map at three selected layers equally distributed across depth. For instance, for depth=12 model, we perform average pooling after the third, 6-th and 9-th layer. After each average pooling, we apply an extra convolutional layer to double the number of channels, which play the role of width in convolutional network. Hence, also the width doubled three times across the depth (one time per average pooling). The final feature map is then flattened and passed through final readout layer. The reported width in all the figures is the last width before the flattening operations. In the experiment with the 1 1/vL- scaling, we divide by the total number of layers instead of the number of residual blocks (i.e. = 3 # blocks). This extra factor does affect the scaling limit.
A.2 RESNET
Compared to the simplified ResNet architecture of the previous section, the ResNet family as pro- posed in He et al. (016)) augmented with two convolutional layers per residual block, each followed by batch normalization normalization. The subsampling is performed by 2-strided convo- lutional layers, for total of times across the depth. For further details, refer to He al. (2016) In some experiments (e.g. Fig. [B.4), we replace batch norm with LayerNorm applied across each two dimensional channel.
A.3 VISION TRANSFORMER
We use the standard VIT architecture as in Dosovitskiy et al. (2020), i.e. we tokenize the input images by splitting it into patches of fixed size, and embed it to N dimension using linear layer. Then we apply number of the usual Transformer blocks, each one composed of two residual blocks, Softmax-based self attention block, and MLP block with GeLU activation function. If used, we place LayerNorm at the beginning of each residual branch (Pre-LN configuration). Finally, we apply the following modification make compatible with μP Yang et al. 2021):
We re-scale the Softmax logits (i.c. the product of queries and keys) of each attention layer by 1/Ng. The usual convention would be to use instead //v Ng.
Initialize the queries to
Initialize all the other weights to have standard deviation of I/VN, where N here is the model dimension.
A.4 DETAILS
We always use data augmentation, in particular the following random transformations:
CIFAR-10: horizontal flips, 10-degree rotations, affine transformations, color jittering.
Tiny ImageNet and ImageNet: crops, horizontal flips.
Across all the experiments, we fix the batch size to 64. Unless stated otherwise, all the hyperpa- rameters that are not tuned are set to default value: 3. og 1. For ViTs, use patch size We use the μP implementation of the μ- u-Readout layer, and optimizers (SGD, and Adam under pP parametrizstion) in the released jP package
14

---

Published as a conference paper at ICLR 2024
B ADDITIONAL EXPERIMENTS
B.1 DEPTH WIDTH TRADE-OFF
While tuning the architecture small widths/depths, is interesting study along which direction width or depth - the architecture should be scaled. To test this, we train the residual convolu- tional model (Sec. A.I) relatively large scale (up almost billion parameters), for 20 epochs on CIFAR-10 10 with fixed learning rate of 0.046 (using data augmentation, see Sec. Sec.[A). We report the re- sults in Fig. Fi.. where we plot the train loss as a function of parameters. we highlight models of same widths, while in Fig. B.I(a) models of equal depth. We find that the Pareto optimal curve coincides by the depth 12 model for the range of widths/depths used for the experiment. As shown in Fig. [B.1 we then fit power laws to the curves of constant depth, which seems to predict an optimal depth 30 at very large widths. Notice that we fix the number of epochs. We hypothesize the optimal curve shift to different trade-off's as the number of epochs increases. We also expect different trade-off's for different datasets and models. However, extensively studying width-depth trade-off is beyond the scope of this work, thus we leave it potential direction for future work.
particular, in Fig. B..
9 90-1
Width
10-4
depth
8 x0-0-
6
7 x
256
Train Loss
512
384
Laa Oss
6 6 10
5 10-1
10
1024
1152
1280
40-1
1920
4 x x10-
2560
3200
103
10f
10?
10a
109
10.0
10a
10"
107
10a
10a
1010
Parameters
n_params
(a) Envelope of Depth-Width Width Trade-off
(b) Optimal Intermediate Depths
Figure B.1: Since all models have hyperparameter transfer across depths and widths, can explore trade-off between depth and width at the optimal learning rate. We illustrate this idea with simple experiment CIFAR-10. (a) We plot the final train loss function of parameters for different widths. (b) intermediate depth of ~ 9 12 preferred at fixed parameter count.
1.1
Depth
1.4
1.0
Train Loss
Train Loss
1.2
0.8
1.0
0.7
512
0.6
0.6
10-3
10-2
10-2 Learning Rate
Learning Rate
(a) ResNet, SP
(b) - ResNet SP
Figure B.2: Learning rates transfer across both depth and width in our proposed parameterization but not in standard parameterization (SP) or μP. Loss is plotted after 20 epochs on CIFAR-10. All the missing datapoints indicate that the corresponding run diverged. In particular, all depth 30 runs diverged when using pP without 1/ / scaling.
15

---

Published as a conference paper at ICLR 2024
B.2 How LONG SHOULD YOU TRAIN FOR ACCURATE TRANSFER?
An important question is after how many steps the optimal hyperparameters fixed and can hence be selected without training for longer. In Fig. [.3 we plot the training loss profiles for two different depths and 4 different learning rates for residual convolutional network of width 512 trained on CIFAR-10. Notice how after the first few epochs it is already possible to decide the optimal learning rate. We confirm this for ViT trained on Tiny ImageNet in Fig. bB3].-c), where the optimal learning rate at two selected epochs (3 and 9) is the same.
Ir
Depth
12
Trai Loos
0.004
Mani Aeccura
1.5
10
24
depth
mmm
10
12
0.5
10
15
20
10
10-1
10
10"
epochs
Learning Rate
Learning Rate
(a) LR dynamics
(b) Epoch 3
(c) Epach 9
Figure B.3: (a) dynamics of train loss for different depths and learning rates. The network is the residual convolutional network trained on CIFAR-10. (b-c) learning rate tranfer plots for the same VIT two different checkpoints (epoch and 5). Notice how the optimal learning rate is determined after the first few epochs of training.
B.3 OTHER EXPERIMENTS
In Fig. B.2 show how the learning rate does transfer under SP for the residual model con- sidered in this work, both with and without the 1/ L/VL-scaling In Fig. B.4 we train models from the ResNets family, replacing BatchNorm with LayerNorm, concluding that the version with 1/ L/VL scaling exhibits a better transfer. In Fig. B.6 we show our results on learning rate transfers for Ima- geNet, and in Fig. B.8 we show how also weight decay transfer under the proposed 1/ LL-salinng and has very consistent dynamics. The model used is the simplified residual convolutional model, trained for 20 epochs on CIFAR-10.
Depth
1.50
1.0
Train Loss
Train Loss
1.25
26
0.8
1.00
Width
512
0.6
0.75
1024
2048
0.4
0.50
10-3
10
10-1
10
10
10
Learning Rate
Learning Rate
(a) ResNet + μP, LayerNorm
(b) VT.
ResNet pP, LayerNorm
Figure B.4: Effect of normalization layers. The above examples are taken after epochs on CIFAR- 10. Normalization layers slightly improve consistency across depths in standard ResNets, but consistency is much more reliable with the 1/ VL scaling.
16

---

Published as a conference paper at ICLR 2024
Depth
0.40
0.7
18
Train Loss
0.35
Train Loss
0.6
34
Width
0.30
0.5
512
0.25
0.4
1024
2048
0.20
0.3
10
10
10
10
100
Learning Rate
Learning Rate
(a) ResNet jPP, BatchNorm
(b)
ResNet pP, BatchNorm
Figure B.5: Effect of normalization layers. The above examples are taken after 20 epochs on CIFAR- 10. Normalization layers can slightly improve consistency across depths in standard ResNets, but consistency is much more reliable with the 1/ //VL scaling.
depth=18
5.0
6.5
depth-26
depth-18
Train Loss
depth=34
depth=26
6.0
dopth-42
Train Loss
depth-34
4.6
depth=42
5.5
5.0
4.4
4.2
4.0
4.0
2
10
10
Epochs
Learning Rate
(a) Dynamics
(h) Learning Rate Transfer
Figure B.6: Learning rate depth sweep ImageNet for epochs with depth extrapolations of ResNet- 18 without any normalization layers. Training is performed with SGD with momentum 0.9 and batchsize 128. (a) The loss curves for different learning rates (each learning rate is distinct linestyle) and for different depths (colors). Dynamics are strikingly consistent across depths for each learning rate. (b) The optimal learning rate preserved across depths.
C A REVIEW OF VARIOUS WIDTH PARAMETERIZAATIOS
In this section, we discuss various parameterizations of neural networks with respect to the width of the model. We will discuss and compare their large width behaviors. First we will discuss the Neural-Tangent parameterization (NTP) which gives rise to kernel limit. Next, we will discuss the Mean Field /p parameterization which allows for feature learning at infinite width. In all models, focus
VNy w o(h²)
The choice of whether one NTP or μP depends on how yy is scaled with respect to N. When training with learning rate 17 = 7 no (to ensure 4 t=o (1)), we have the following scale of internal preactivation updates (Bordelon & Pehlevan 2022b
di h - 06N-112)
Depending on how is scaled with N, learning with respect to width N. In discrete time, a diverging parameterization is unstable (Yang & Hu 202T). Below discuss two of the popular stable scalings for analysis.
either have vanishing, constant, diverging feature
17

---

Published as a conference paper at ICLR 2024
10-4
10"4
10-4
10
1-
N 04
16
N:
N = 126
32
128
N 256
25.5
5112
10
10
10c
101
10
30'
10
(a) Fixed Convergence in N
(b) Average Square Error
(c) Error of Ensembled Logits
Figure B.7: Further numerical support of our finite N, L error decomposition in Appendix [ through ensembling logits over E = 15 random seeds in the same experimental setting as Figure (a) At each depth L, the finite width N network converges in square error as ON,L(N' independent of T.. (b) The full expected square enror from Figure 5 shows convergence rate of O(I-1) for sufficiently large N, consistent with theory. (c) The ensemble averaged logits (fL,N) converge the
infinite depth proxy at faster rate, which is consistent with the theoretical OLL 2) for sufficiently large N.
depth 3
10
2.0
2.0
Train Loss
Train Loss
depth 9
10-2
depth
1.5
10-1
1.5
Depth
12
0.5
0.5
10-5
10-4
10-2
10-2
10-1
10
20
Weight Decay
Epochs
(a) Weight Decay Dynamics
(b) Final Loss vs Depth
Figure B.8: Weight decay dynamics also transfer across depths.
C.I NEURAL TANGENT PARAMETERIZATION
In NTP, the parameter 7 = (1) 1) is a constant with respect to width N. As a consequence finite width N models only exhibit ONN11/2 movement of preactivations and O(N-1) movement of kernels. In this parameterization, models very large width well approximated by kernel method (Jacot et al. 2018 Lee et al. 2019).
C.2 MEAN FIELD/PP PARAMETERIZATIONS
On the other hand, i = 7o N for some constant 7o. then the network converges to feature learning infinite width limit. This is precisely the mean-field/uP parameterization (Yang & Huj 2021 Bordelon & Pehlevan] 2022b a). The work of Yang & Hu (002l) also gives precise sense in which this parameterization maximal: each layer learning features so that the preactivations in the sense that
VN w o(h') 0(1)
(10)
Thus, the dynamics for d h
h + VN (4 w') o(h') + VN W' ($0(h))) = (1).
de
de
Because the underlying feature learning updates within the network approximately constant as width N increased, this model exhibits more consistent dynamics and better hyperparameter
transfer across widths (Yang et al., 2022 Vyas al. 2023).
We note that the code for μP released in the paper by Yang et al. (2022) uses an equivalent version of this dynamical system where learning rates do not need be explictly scaled with N.
18

---

Published as a conference paper at ICLR 2024
D DERIVATION OF THE N - DO LIMIT AT FIXED L
In this section, we characterize the feature learning dynamics of the - 0o limit fixed depth L residual networks. Our derivation is extension of the approach presented in (Bordelon & Pehlevann (022b). Before taking the limit, we first identify our order parameters and weight space dynamics. Then we take the N oo limit using a technique known Dynamical Mean Field Theory (DMFT).
D.1 FINITE N, FUNCTION AND WEIGHT DYNAMICS
We start by computing the neural tangent kernel (NTK). For the sake of brevity, we will first write the equations where the read-in 8 weights are held fixed and will add their contribution to the dynamics in Appendix K.3j Excluding these weight matrices from the present discussion does not alter any of the scalings with width N or depth L. The dynamics of the prediction ʃ under gradient flow take the form
[(; = no En K(x, x' :1)((a)))
(11)
dt
where A(x;t) t) aC(((a()) is an error signal and K is the NTK of this model, which takes the (f(m.t)
form
0f(x) 0f(x')
K(0, = NO
(12)
oW oW
Explicitly calculating the gradients of the output with respect to the weights W gives
0f(x) 1 g (z))(''()))))
0f(x)
Nvo OW
of Noo @hf(z)
~ ON.((1)
(13)
VNL
Using this above formula, find that the NTK is O(1) with respect to N, L
K(x, x²)
(xt
(14)
L 8=1
where we introduced the O(1) feature kernel & and gradient kernels G which have the form
Φ'xx xx) The base cases for the feature and gradient kernels are given by
N (hh (x:) ((h()=))) (2,
N (xx) gf ga())
(15)
1 Th (x, x.) = D x', Gz. (r, (x) =
(16)
We next start by computing the dynamics of the weights under gradient flow
d
W
700700 En A(a; ()g + ())( (x)
(17)
dt
VLN
Integrating this equation from time 0 to time we get the following
h 8- (x,
(2,
(),,t)=h's (x( + VNT.
(0))((h))
7oTo L
[ dsE A(x s s)g (x' s) (x, t,)
(18)
~ (x t)) ((h (x'ss) the double time-index generalization of the feature kernel. Similarly, we obtain backward pass relation of the form
where x';,tt,,,t)) =
1
ge (x, - 8+1 (xx + VNL.
Sh' (x, 1) T
+ MoYo 3(t (x,
(1,()x))) t)) 3 S dslE Alw s) s) (x, itt,))))))(( (19)
19

---

Published as a conference paper at ICLR 2024
These equations define update rules for the preactivations h' and the gradients g" in terms of the feature and gradient kernels {Ø", G*). However the above equations explicitly depend on the ran- dom initial weights (0) through which time-varying signals propagate. Our ultimate goal is to gain insight into the effective dynamics over random initializations in the large limit. To do that, will invoke dynamical mean field theory. In the next section, characterize how these terms behave at large N using a saddle point argument.
D.2 INTUITIVE EXPLANATION OF LARGE WIDTH FEATURE LEARNING LIMIT
In the limit of - oo with To held fixed, the above equations can be simplified. The two key effects that occur at large width are
Kernels (feature kernels, gradient kernels, NTK) and output logits concentrate by a law of large numbers.
At each layer, each of the hidden neuron's 's activities becomes independent draws from single site stochastic process which is characterized by the kernels.
Once the single site density is known, the kernels can be computed averages the single site density. The next section derives this result formally from Martin-Siggia-Ross Path Integral derivation of the Dynamical Mean Field Theory (Martin et al. 1973).
D.3 PATH-INTEGRAL AND SADDLE POINT EQUATIONS IN N + 0o LIMIT
This computation sketches the DMFT derivation of the limiting large N process. Examples of detailed computations of the DMFT action S can found in the Appendix of Bordelon & Pehlevan 2022bba) a). Our notation is chosen to match the derivations found in their Appendix.
To characterize the effect of the random initial weights, we will attempt to characterize the distribu- tion of h' h'xx, (x,t), g (x, t)
by tracking the moment generating functional of these fields
Zl1i . = (()(())) exp Σ
S dæ / fdt (a,(tt h (x 1) + (y g(x D]]]
(20)
Moments of the fields can be obtained with differentiation near zero source, for example
lim (3% 0 oth,, (x (2) 85h.s, (xx. (x) ZKi .Jgh1 = (n (x1, (x1, t (axa ta)).
5
(21)
Now, following Bordelon & Pehlevan (022b), we introduce the following useful auxiliary fields
1 w'(0) 4+ (x, (0)) 0) E' 1) = VN
++ (x,
(22)
VN
The choice to introduce these auxiliary fields is apparent once one realizes that the h and fields be regarded as functions of x and the kernels Φ G. We enforce the constraints defining these fields with Dirac-delta functions by repeatedly multiplying by unity
1= / dx 1+1 (x, (2, 1) (x, VN
dx
(()))((h((,tt)
dx (x, t)xx))))+))))))) (2,, (zx
exp (x, (x))(2)))
w (0)o(h (0))hh((
(2T) N
NN
(23)
for all É, t. And similarly for the gf fields. After this insertion, the averages over {W* can be performed with Gaussian integration. Concretely, obtain the following average over Gaussian
20

---

Published as a conference paper at ICLR 2024
initialization (0)
In exp (- ( (0) / dx f /at (m, t)) + (x, (x,
VN
dtdsdæda' x (x, t) x s) x'.
{ dtdsdrdz' (xx t) (x' (x,
dtdsdada' x xt gf + (al .) A ax,
(24)
where we introduced the following order parameters
(x, - N Oh (z, ()) (h' 8)) ('(2, = 1 ( (x, a' (t,s) 6hh (a, (xx t)) (a',
(25)
(26)
As we did with the X, fields, can enforce the definition of the above order parameters with integral representations of Dirac-delta functions. For example, the kernels can be enforced by inserting
1= =f d (x, s)ts)ddb tt;
exp N (x, x' s), (xx ttstts))
2viN
cxp (- a': sh (z, t)) o(h (a),))))) (27)
where the O integral is performed along the imaginary axis (-ioc, in the complex plane. The same trick is performed for G and A which have conjugate variables G B respectively, After inserting these new variables, we find that the moment generating function can be expressed as an integral over a collection of order parameters 4.
q Vec (oc( (z. x' (t, s), &6 (2, (2, (z,, G A (1.. (28)
x' 8)},
which vectorized over all layers time points t, and samples x, w'. The moment generating function takes the form
Zlihnigg11 - ff dq exp (NSlg.
(29)
where S is the (1(1) DMFT action which has the following form
S >
/ dxdxx dtds (2, (f:t,) x':t.s) + G a ((,, st,s)
N
- -E dxdda' dtds A' a', (,,) s,t) +
2 jn. Jal
(30)
N
The functions Z, are single-site moment generating functions (MGFs) which express the distribution of h4, for a given value of the order parameters and sources jn Jg.
Z, = / II (x t)d (x, exp -
(31)
t.
21

---

Published as a conference paper at ICLR 2024
where Hf are the single-site Hamiltonians
H
1
dwda''
dids dids ( a'; 8)x (x) (x'
2
/ dxdx'd dids
+
i (2,
(1)
1) (z), ,)
+ i f dxda'dtds dtds [A
(e' (x o' s) dxdt (xx + EL (z. t)5 az, t) + fn t)h (a: t) + 5g a (x)
s))
(x. t)
(32)
In the above formulas, the h. g fields should be regarded as functions of the X, E fields.
Since the moment generating function has the form Z I dq exp (NS(q)), the N - oo limit can thus be characterized by the steepest-descen method ssaddle-point integration over q). The result is that all of the order parameters take on definite values (concentrate) at infinite width. The saddle point equations as
0 state that each of the kernels takes on definite values
3q
'' (2,
(x, s)))
N iml
Σ
(o (x',
1'(2,2':1,8)
(t. (a))
N
8
(g - (x, (x'
N
('(2,2';t,5)
txx
2N
(2
(((a,))((('ss)))
(33)
2N
where (, denotes averaging over a single site stochastic process defined by the moment generating function 2,. At zero source j -> all of the single site MGF functions Z, are identical so all averages (-), are also identical. At zero source, the kernels have the simple expressions such as (((x,, ot; 1,,ss) (0(h((a) (a, ())(()((')) +))).
D.3.1 HUBBARD- ITREAADISSTRATOOVIICII TRANSFORM
Now that we know the kernels Tº G take on definite values at N - co, we can perform Hubbard- Stratonovich transformation to introduce Gaussian sources
exp -
dadæ' dtds x (x, (x',, (or,
exp (-1 (-Ifaadee daodt 80 t)u (2. ((-SP(0,0r-))
(34)
I dadx dtds (x 1) (z,,
exp
(-i) dxdd E (le,,trr( (t,)
exp
(35)
/ re0gP(0,0(11)
22

---

Published as a conference paper at ICLR 2024
After introducing these Gaussian sources u, rf, we have linearized the single site Hamiltonians H. with respect to X and ç. Integration over these variables gives the following Dirac-Delta function constraints
(x, u (x, + dads A x' at s)
5x t) = r (z, + dxds B' a'; 5) s) s))
(36)
Using similar manipulations, can simplify the expressions for A B", 1 & zero source
A st, = (x,h"(x,
(r"(x',
(x. t)
B'(x,x';t,s) 6ul ,)
o ( 6 x'; 8)= == =0
(37)
Plugging the dynamics for X,E into the formulas for h², gf, we obtain the following final equations for the stochastic process of h. .
1
VT. / da' ds A ()) (x',
1
h (x, (xx,
VI. " (x,
7o70 Ew / ds dss s) 1000((x (z, x'it,) gf (xt) ot GP(O, x(-1)
L
((xx,) t) = - (x, + 1 V (z) (x)
(xxt) = (x,,) + dx' ds B' a' s) s0(( (z' re GP(O, G
70 Ex S ds A(x' s) d'; s,)0 VT.
(38a)
where the order parameters satisfy the the following saddle point equations
Φ (2, = (o(h' ()) s))) (2, = (g'(0, t)g (x'(
12 xx;,z';,,)) (z) B² (2 óg (x,
8r'(z',s)
b (x',s)
df (x, t) noB K(2 x' () A(z,t) af(xt
OL
dt
K(x, x' ,t)
1 L
Σ
cl (2, (t (x,
(39a)
Up to the residual skip connections and the factors of J this agrees with the N + limit DMFT derived in Bordelon & Pehlevan (022b). At fixed this above equation can be solved with a self- consistent Monte Carlo sampling procedure. For deep linear networks, the equations close since all h, fields are Gaussian. the next section, we analyze the behavior of the above equations at large depth L, so that we can take large depth limit.
E THE L + Oo LIMIT FROM THE DMFT SADDLE POINT EQUATIONS
In this section, we consider how the infinite width dynamics behaves at large depth. We will show that the limiting dynamics takes the form of system of ODEs for the kernels in terms of layer time
23

---

Published as a conference paper at ICLR 2024
E/LLL. First, will show that the response functions need to be rescaled with respect to L. Then we will take a continuum limit over layers to arrive at the final ODEs for the & and G correlation functions.
E.1 UNFOLDING THE LAYER RECURSION
useful first step which we will use the following analysis is "unfolding" the recurrence to get explicit formulas for he in terms of all of the Gaussian sources and all of the gradient fields g².
1(x,1) (x, +
A (xx 1)
VL Σ
da' ds A a'; ; (,8gg (a',
VL k=0
70
B / ds A(x' 8)
Σ
(x, (x, x', ,t, (
(40)
k=0
and similarly for the backward pass, we have
(x (0 +
1
Σ (a, (x,1)
VL
1
((h((z)))) / da'ds ds (x st) x'
km
Tl 70
((h ())Ex d' 8) Σ x' d.)
(41)
L
A very useful fact that can be immediately gleaned from the above equations is that
Yk L, (h(( 1) (x, t) Wk 2 L, (x,
VT.
(x,
(42)
VL
These facts will be utilized to obtain the scale of the response functions in the large depth limit.
E.2 CHARACTERIZING RESPONSE FUNCTIONS LARGE DEPTH
First, we will study the response functions A B We note that from the saddle point equations,
(85(1 (ax, ('))
1x( x1;1,8)
(r' (x'
(43)
Since
86(h)
5(h) öl it suffices consider the scale of the derivative 5r br
8h (m,t") or(a(ss)
A (2,x''tt,)) = VL
1
Σ
/ dæ" dt' AMxx, tt; (0l18(o, t)) 5g (
(r(()0))))
k=1
5g 50"( Enp (' di' (1''(,', (2) (0(1)(o, (x,, ())
noYo
1))
(44)
8r (2')))
Next, we use the fact that (0"("" (LL-1/2). Therefore, the second sum (the one which (ar.))
depends T.-I) scales as I.-0%o Solving the A recursion iteratively from the first layer, we arrive at the conclusion that
( x' st, O
(45)
24

---

Published as a conference paper at ICLR 2024
Indeed, sanity check, plugging this in gives consistent solution to the above equation. Re- peating an identical argument for the backward pass reveals an identical scaling for B with depth L
B' (xx ;t,
:t,s) s) ~ O
(46)
We therefore need to introduce rescaled version of the response functions
VL
A A 21,,)
A (1 x' ,t,)
(47)
110
VL
(2 x ,1,8) t
B' (a, x' ,,,))
(48)
Yo'ra
After this rescaling, we have
VL To
80(h (x,t))
A axx ,t,,'t,t t,,)
(49)
5r (x',
To calculate this rescaled response function, we can derive closed set of equations for all of the cross layer correlators
70 8r (x²,8)
ōh (h t)
Σ
f.
/
dæ" ds (2e, x", 11, 8) 8r"
(x",))
L
k=
(x,
1
O E)8(x =z)) s)e(h (x,
8r (a',s)
VL
(ar,
+ VI
(ax
8r (x',
5(1 (cc, / die" do do x" ,,,,,, (0""""""" dr (x',
Yo
+
T,
(a) 8r (x''
/ do (2, x" s))(("
Yo
(50)
where we introduced the shorthand
(0, = A (x, s, + o- (0, 1x, at,,,,,ts s) s)p(x') (x, = (x, x' s, G (x, x' t, 8) A)) ss)((a))
(51)
Similar equations exist for all derivatives with respect to v' 1). Once the above equations are solved, then one can compute the response functions A Be by averaging the necessary field deriva- tives.
E.3 CONTINUUM LAYER-TIME LIMIT AT INFINITE DEPTII
In this section, we argue the infinite depth limit can be viewed as continuum limiting process defined by ODEs for the kernels and SDEs for the preactivation and gradient fields for layertime L/L. We start by examining the asymptotic structure of preactivation he dynamics
h' (x;tt) ~ h (x; +
u (x;t)
VI.
ds Eg O (xx w': 1, )g ) ())''s)
(52)
L
The Gaussian source variables = are all uncorrelated across layers. So the first sum is also a mean zero Gaussian with variance
(
wt (a' ,s)
Σ k=
(2, O(1)
(53)
k=1
25

---

Published as a conference paper at ICLR 2024
At large depth L, this term will behave integrated Brownian motion since each layer's contribution is independent and order vdr VI. Next, reason about the scale of feature learning update ! ds' S. S.'ds ds [s Ew A(x' ( ()
2 No 10 L
ΣΣЛ,
K=
(', OK (a, t'it,) (gut) ()g ta $1 (()
(54)
The above double sum is (1) because we have ( (1) for all k, k'. This fact arises from the skip connections, which generate long-range correlations across layers. This behavior is very different from non-residual networks where of and ge are independent for all o e (Bordelon & Pehlevan 2022b).
To take the continuum limit, define functions of continuous layer time variable 7 E [0, 1],
h(T, x1,, h ha(x, e=rr g(T, xx,) g t)l(=rb ( o0; tt, = (x x'; s)k=r G(T; x' t';
(55)
The key idea when taking the continuum limit is the following relations based notion in- finitesimal layer increment dr ~ E. Let F be deterministic function of the layer index and let F(T) be its continuous analogue which satisfies F(T) = F leart. Then,
lim Σ F ( dd'F(())
t=1
lim I. [ F
F
F(r).
(56)
We will use these relations when deriving continuum limits of the DMFT process. We note that stan- dard error analysis for Fuler discretizations can be used to argue that finite depth averages converge to the integral at rate
f ('d'F(')~'L ~ O(L
(57)
E=1
This is merely the approximation еттоr of Riemann sum which is proportional to the step size [Rudin 1953), which this case is 1/L. This idea can be used generate a convergence rate of the NTK and the network output predictions at large depth as we discuss in
E.4 PREACTIVATION DYNAMICS
As discussed in the main text, the large limit for the dynamics of ht can be viewed as a stochastic integral over and over Brownian motion for the n sum
h(T; x. f) 2, + [' du(t', xx,))
+ nono / dx' ( ds ( di' [0(T'' :a x. x' s))a + ",',0,,'1 g(r x',s)
(58)
where du Brownian motion term with covariance structure (du(T, f)du(r', x', s)) 5(r - 2', 1,
(59)
Similarly, the gradient fields have an analogous structure
g(T; x. t) = g(T; (1; x) 1) + [ dr(r', (,)(()
+ Bar ( do [: "
no7o / dxl L: ds L d''((h(T' (,
(60)
26

---

Published as a conference paper at ICLR 2024
where dr is Brownian motion with covariance structure
(dr(T, æ, ()dr(r', x', 8)) = 6(T - G') n. ,, 1,8)dT.
(61)
We note that, due to the nonlincarities. the h.g fields are non- Gaussinn at finite To as in the case of mean field networks Bordelon & Pchlevan (2022b). We note that an alternative the above integral formulation is a stochastic differential equations
dh, a0, 2t,
+ dr 7071 / dx' ( ds [0(r),,,';,, s) (s)((a)))))))) + A(T;,, (),,, x' x'
(62)
The % - 0 limit of the / equation at 8 0 coincides with the SDE derived at initialization by
Hayou (2023): : Hayou Yang (2023). However, we now see that feature learning term (the
which depends on 7o)) adds a drift term to the layerwise dynamics compared to the diffusion term from the initial weights.
E.5 KERNEL LAYERWISE ODES IN T FEATURE LEARNING REGIME
We can start from the stochastic difference equation
dh(T, xz = du(T, x', + dr 7o f do' S: ds Ca(T, x,, x' ,t, s)g(T,
Cn x, x', ,,, = A(T, x. x', s) + Φ x, x', ,t, t) ss)((x'.
(63)
and derive ODE for the feature kernel at layer-time Applying Ito's lemma obtain the following feature kernel dynamics
Itō
1944
2,, t))
θr
8²
dēdā' dt' ds ( a. W., Öh(T, ā, ()0h(T, a¹, (o
+ +1000 / dx" de x, x", ( (9(T, ('))",,)T)"""eee ())
/ dæ"ds Cn(T x x.,5)
+ 10700
8') (g(T, ""' s')0( ())((((,,z' ())
(64)
The first term on the right hand side comes from signal propagation, which persists even in the lazy limit. The second set of terms come from feature learning corrections. This equation should be integrated from 5 = to T Similarly, there is backward pass ODE for the G kernel
0r G(T, x. xt,, 1, - (,) ,i)) G(T,2, xt,,8)
- no00 / dx" I' dt',(T, x, x,",,,,"",, t,t) (((h(T t)) x' ()) x",,,)))))
7o%o / dxd 1 ds' ds'C x" 1,t)) (1) ()) x² ,'))
CT x,, x',,),, B(T,
G(T, x. x' s)A(2' s)p(x')
(65)
which should be integrated from T = 1 to = 0.
We note again, that convergence to the ODE will occur at rate O(I-11) in absolute error.
E.6 RESPONSE FUNCTION EQUATIONS IN THE CONTINUUM LIMIT
The response functions can also be determined in the + oc limit using the following relations
B(T, 2(,, no <
ög(T, a.
A(T, xr, x' 1)
(66)
7o7o 8 dr(T,
du(T, x',
27

---

Published as a conference paper at ICLR 2024
The reason that the variation is computed with respect to the differential Brownian motion du, dr is to match the scale of the increment vI which appears in the finite L response functions A
VI
bd(h5) 5r ) Using our continuum limit equations, we obtain the following field derivatives 8h(TT x.1)
nono
OT - --0)6(0 x' 2')0(1- s)
odu(r',
7070 ʃ dr" dx" I Cn("""" 2,," ,t, ōg(T", x. ,t) 11) ōdul(' 1 x',
Sdu(r', ,0', 71070 dr dr (0M"" ,e,1)) fde" da" C far x² ,,,1', 66 (6(h("" ,,",,),")))
6g(T,
8h(T, xx,1)
710 / dt" f da" f df ",," xx,x", ,t,t) x,, odr(T', 22,9)
"g(T","",1"
,dr(r', a.,s))
8g(T, 22,
( -T)0(x -z')0(t - 5) o(h((T,z,
5dr(r', (',)
+ 7o / dr"'(h(("" (0,1(("",,,)))))) da" / df x, ".,",,tt) t,t) Ct'""""
80(1(T",2"" x²,,
Ödr(T', x',
(67)
Once these recursions are solved, the response functions can be computed by averaging the neces- sary field derivatives at each layer time 5
(o(h(T,,,t)) A(T, x. xt, = 6dr(T,2',s)
6g(T, xt,) ōdu(T,
B(T, x, x',,,)
(68)
E.7 FINAL RESULT: INFINITE WIDTH AND DEPTH LIMIT
In this section we combine the results obtained in the preceeding sections to state our main theo- retical result which is the infinite width and depth limit the learning dynamics in our ResNet. Below state the result for gradient flow, but minor modifications also characterize discrete time SGD, momentum, weight decay etc as we discuss in Appendices|K K
Proposition 3 Consider the N, - oo limit of the ResNet pP- uP--/. ./. model presented in the main text trained with gradient flow on data distribution p(x). The preuctivation hi(T, x., 1) and grudient g(T, xx 1) stochastic fields obey the equations
a(, h,0,,,t) + S[¹ du(t',z,t) + 700 Sf dr' fdaa da' S: ds x',,,,,, x',
g(r,x,t) = x,,+ + [ 5(h(T' x (()r(,r))))d(ffccc x.t)
+ [ dr' (((((( de' S" ds Cg(T', a' t)))((((T' ())((((t(,
(69)
where the zero mean Brownian motions du(r, x, t), dr(T, x, t) have the following covariance struc- ture
(du(r, a. f)du(r' ,a", {)) (T,, x. a' 5)5(T +')dr
(dr x, ())rr(rr ,',,,'',,s)) = G(T, x. , s) 7')dF, the deterministic kernels C₁ and Cg have the form (T, x x',, Cy 21, t, G(T, (, x' t, 8) Aa) ssp('') B(T, x. x'. t,s)
(70)
(T, x' 8) A(x s))('')) + A(T, x, x',t,,)
(71)
where the kernels $ G and response funncioss A, B are computed as averages over the h 9 stochas- tic process
( x' s) (0)(((,,,))))((((7223)))) G(T, x' s) (g(T, x, t)gg((]TT) x'
x', s))
1
Ög(T, 11
A(T, n, a x1,,,)== 7oo
B(T, ,t,,) 7o70 ('du(T,x',6 (72)
5dr(T,
28

---

Published as a conference paper at ICLR 2024
From the feature and gradient kernels, compute the dynamical NTK and the dynamics of network predictions
K( x. ,1) = ( dr G(T, a.,,,,,,1111)) ,1,1)(T, ,t,1
d fx, =100) da' dl')) (') t), A(x,t) Oc
(73)
Oxxx)t))
This is closed system of equations relating preactiation and gradient marginals, the kernel dynam- ics and the network prediction dynamics. In this limit, the kernels, response functions and network outputs are all deterministic.
F INFINITE DEPTH THEN INFINITE WIDTH LIMIT
In this section, we compute the limiting process in the other direction, by directly working with the distribution of order parameters q at finite N. We show that at large depth I. the DMFT action converges to a limiting object Sax (q). We the show that the large N limit the distribution recovers the same saddle point derived in Appendix [ for the large N then large L limit. Thus, this section constitutes a theoretical argument that the width and depth limits commute for the whole course of training dynamics. The derivations are performed at the level of rigor of physics, and we leave open for future work rigorous result about commutativity limits throughout training.
F.1 LARGE T. LIMIT FIXED N
We derived the finite N, L distribution of order parameters in Appendix Section D The result was that the log-density of the order parameters q over random initialization of network weights was given by the DMFT action In p(q) x NS) Suppressing unnecessary (for the purposes of this section) sample and time indices, the action at depth Sr had the form
Sr. Σ T +
A²B8 3(q
(74)
where the single site density Z has the form
2 = / II dh dñ dg
Be exp
D Re 11t1 - n
u
1
Ti -T (
VL"
VL
exp D ge 00 65--tth -g---
- )
1
1
5(h o(h))
no7o 2(1 ((h))
VI
VL
L
E (0h1)
- Eg
exp
(75)
where the averages over wf, of which are Gaussians with covariances Φ--1 and G respec- tively. In the above equation we work with h, directly rather than working with the intermediate fields X. which only depend on the initialization.
To obtain a proper large depth limit, we rescale the conjugate order parameters as
:'
L Th. G + 6e A
1o 4 VL
B
7o D VI
(76)
L
We note that such rescaling will not alter the distribution of real network observables { G, f.... (Bordelon Pehlevan [022b 2023). After this rescaling, we take L - oo which gives continuum limit of the resulting layer sums, obtain the following simplified infinite depth
29

---

Published as a conference paper at ICLR 2024
action in terms of the kernels and features at layer time [0.1]
S00 = / -
DiDDDDDDg Bu(r))((r) exp (
( -u() (o:o)))))oCC
((()) dr()))))())((( -odo---dddddduutt--oooo-oo----o-
exp
(-fdrti ((
exp
(77)
where Cb, Cg are given in the main text and the du( dr(r) are Brownian motion processes over layer time We therefore find that the distribution of order parameters at finite width N and infinite depth + 00 is
lim (q(g) Poo (q) x exp (NS (q))
(78)
L-000
This constitutes taking the infinite depth limit L - co at fixed N. this stage, the order parameters do not obey a deterministic set of equations, but are samples from the above density. In the next section, we study infinite width of this density and find that we recover the same result that we derived previously in the other direction.
F.2 SADDLE POINT EQUATIONS ON INFINITE DEPTH ACTION
This section considers the large limit the infinite depth stochastic process. In particular, we evaluate the saddle point of S.c and show that it recovers the same limit obtained by first taking a saddle point and then taking the large L limit. This shows that the limits commute.
68.0
((T) (0(()))))((())))
(T((T)
68.0
G(r) - (o(())(())
8G(T)
8800
-1 - inooo (h(+))(7)) =0
8A(T)
6S,0
Tá A To (117)((((11)()66)))
(79)
8B(T)
For the last two expressions, recognizing that
6
(-i h(r))u(()) = -ih(() exp (-i fite)au(t))
ddu((")) exp
(-i (-i (()))((())))))()))
-ig(T))(((T)) exp (0))0(()))a)
(
8dr(7')
exp
(80)
We can thus express the saddle point equations in terms of the left hand side derivatives with re- spect to Gaussian sources. Performing integration by parts, arrive at the usual response function formulas
< ddr(T)) B(T) no 'o-'
/66(((T))
8g(T)
A(r) = 'oo:oo
(81)
ddu(T)
Lastly, we can integrate over i((T) and ğ(T) which give Dirac-Delta functions that define the h and g stochastic processes
dh(T) du(r) + (oT0CL(T))(T)) dg(T) -dr((r)(((())) - -mxo((())))))
(82)
The rest of the indices (over samples and training times) can be incorporated into the above argument and the result is easily verified to match that derived in Appendix E from the large width then large depth limit. This exercise thus verifies that the large width and large depth limit continue to commute throughout training, extending the commutativity at initialization proven by Hayou & Yang 20233
30

---

Published as a conference paper at ICLR 2024
G EXACTLY SOLVEABLE CASES
In this section, give exact solutions to the dynamics which do require any non-Gausd Gaussian integrals. The two cases where this can be achieved are
1. the lazy limit + 0 of network training where the features are frozen networks with linear activations, which preserve Gaussianity of h, g any To.
We discuss these two limits in great detail below.
G.1 LAZY LIMIT
In the lazy limit To O, the flow for the preactivations İiT,, ær, h(T, ax) and the gradient fields g xr, () = g (T, az) are constant over training time. Similarly, the kernels & G are also constant throughout the dynamies. We thus have the following forward and backward dynamics
h(T, -(0,0)) + S du(r',
(83)
0(r,z)=g((1,2) + [. ò(h(T', (o)d,()))r((TF...
(84)
We note that the above equations indicate that 11(T, æ)) are all jointly Gaussian with zero mean. Since is Gaussian, it makes sense to characterize its covariance HT,, æ²) (M(T, ()) at layer-time т. This object satisfies the differential equation
H(T, (,,, $(T, x²)
ÖT
$ (T, 20, x (6(1(1,=)))0(ht;; h(r,0,x') (N(OH
(85)
Since h(T, æ) are zero-mean jointly Gaussian random variables, the average to compute & can be computed in closed form for polynomials and several practically used nonlinearities (ReLU, ERF, etc) (Cho & Saul 2009). Now, to compute the backward pass, we must
i θT G(T,z, = (
,a))) = G(T,,xxx))
(86)
Again, the bivariate Gaussian integral ( (δ(k(T, ())((,))))(((((+ x')) has closed form expression for many regularly used activation functions d. Given that we have solved for 4(T) and G(T), the dynamics of the predictor is
d
Oc
á f(x, 1) = =no (a,x')A(at A(('' f), A(x,t) 0f(x,t)
K K(x, (xxx0)=11 = S ('7,,,
(87)
This is an exact description of the network predictor dynamics the N 0o, L. - xo and to - limit. However, the 7o 0 limit lacks fundamental phenomenon of deep neural networks, namely feature learning. To gain analytical insight into the feature learning dynamics, we will study linear networks in the next section.
G.2 INFINITE DEPTH RESNETS WITH LINEAR ACTIVATIONS
In the case where the activations linear also close the equations at the level of the kemels (since activations remain Gaussian even after feature learning). To provide an exactly solveable model of the dynamics, consider computing the kernels after one step of feature learning.
31

---

Published as a conference paper at ICLR 2024
G.2.1 FULL DMFT EQUATIONS IN DEEP LINEAR NETWORKS
The DMFT equations for the preactivation processes have the form
" h(T, x 100 x1 + L du(r', xx,) + 7070 [ dr' dæ' f CnlT',x, 2t, ())((('', x²,s)
g(T, (0, 21,, + f dr(f', xx, 1070 [ dr' f fda' da' I Cglt's nr of ,1, 1)((((' x'
(88)
where Ch. Cg have the usual expressions. The field derivatives (to get the response functions) have the form
8h(T, xt.) O(T - -7)8(x 2')0(t s) odu(T', x),s)
1700 [" dr" I dx" Cn''" x, x" ,,,')
8g(T" x²" ('
nn
ōdu(t's x', 8)
_du('', 770 dr f do" d' a0, t,
(g xt.)
8h .t)
odu(t' xt' 5g( ,²) ddr (ir((T,,0 x.,
10760 I'" dr" {ad" dz" f CN (,, 5dr(T', x',s))
8h(T,
6g(T, m,,)
0(T T)8(x x' 1')0(t-
ödr(t' x),s)
+ +100 J: dr" du")
J Ct"'"""",,,, ,,,,,,, ,tt,, 8h(T" x²,1)
J
(89)
8dr(T', x¹,
Since the activations are linear, these derivatives are no longer stochastic and we have the equal layer time 7 = r' derivatives the response functions
A(T, x. x', 8) 8h(T, 1)
5g(T,
( B(r,z,0',t, s) 5du(T, x),
(90)
Next, we compute the kernels as simple Gaussian averages
H x st, (h(T, x, t))(T, x, t)} G(T, ,t,,, s, = (g(T, x, (Tg(T, xx, In the next section, we show a simple example where these equations become simple linear differ- ential equations.
(91)
G.2.2 ONE STEP ONE SAMPLE FEATURE KERNEL EQUATIONS
Before training, the initial conditions for the kernels Ho(T) and Go(T) are
θ, Ho(T) = Ho((), G, Go(T) Ho(T) = e Go(T) = + -
(92)
After a step of gradient descent, the preactivation at layer time T has the form
L. du(r') + 7o%o [ dr'(A(r') - Ho(T')ylon(**
h(T) (()) +
(93)
We let HT)), G(T) represent the kernels after one step. The response function after one step has satisfies the differential equation
A+A(T) A(7) Ho(T)y, A(O) A(7) Te
Te" y
(94)
The kernel H(T) after one step therefore satisfies
H(T) H(0) + '' dz'H(7')
+ '0 f dr'
+ Ho(+)))(((") + Ho(r")))
=((0) ("
dr' 2
[[ [
+r" 1|[7" + +1]. (95)
32

---

Published as a conference paper at ICLR 2024
where T') (90(r)90(7))- Differentiating both sides with respect to gives the integro- differential equation
0, H, (T) 275 1 I' dr dr'ul T') (F'+1)
(96)
Note that GT, 7') = Go(T) = 6 +1 c- forr < T. So we can pull this out of the integral and we are left with
L d'''''l'' [ =1]=re" TE
(97)
We are finally left with the ODE
d,H(T)= = H[) + 2110 Toy 2 le + +1 - +T]
20
(98)
Integrating gives
H(r) = c H()) 2ng
7
3 (1+e-)) Te" + - -1)
(99)
3
This solution gives the profile of the feature kernels after single step of gradient descent. This is verified numerically in Figure At small T this goes as HT) Ho(() oh + (1 +1+ ]]
H FINITE WIDTH CONVERGENCE: FLUCTUATION DO NOT SCALE WITH L
At fixed L. one compute systematic asymptotic corrections in Ox-NNI) to the statistical proper- ties of q. This follows closely from the finite width analysis around mean field dynamics performed by Bordelon & Pchlevan )023). The idea of their expansion is to extract from S information about finite size deviations from the infinite width DMFT process. The key fact which appears in the next to leading order terms is that at finite width N, the order parameters q become random with covariance
Σd + +NN(N 3), Σ I-=I-v'S(( Cov(g) ~ ~N=
(100)
The matrix N T gives the leading order (in powers of 1/N) covariance of the order parameters over random initialization at width N depth L. The goal of this section to show that the entries of this covariance matrix Σ that correspond to the NTK and (consequently the predictor f) do not diverge with L for the architecture of this paper (residual networks with I/VL branch scaling), but rather approach a well defined limit at rate O(L We again stress that this is unique to this architecture. Non-residual deep networks networks have propagator entries which scale linearly with L since variances accumulate rather than average out over layers (Bordelon & Pehlevan, 2022a 2023). We note that the predictor can be included in the set of order parameters q depth L, so that the finite width deviation scales as
((FxLL (oo)))
~ N 27
(101)
where averages are taken over random initialization.
H.1 COMPONENTS OF THE HESSIAN
Since we are primarily interested in the scaling of the finite width effects with respect to N and L, we will suppress the sample and time indices of the kernels and instead focus on the dependence of the fluctuations on layer index f. This will not change the analysis of the scaling but will save us from tracking a tremendous amount of indices. We will let &' to represent the feature kernel at layer& and G the gradient kernel at layer é with no time sample indices. Under this simplified notation, the Hessian has entries of the form
a²s
(o(H)) (0h"))())))0 (o(n (00))))))))) "TTT.""
"
0PS
(6(1))((9)) 50.4 a
a
1 D
0
I
OPS
(102)
33

---

Published as a conference paper at ICLR 2024
First, note that since the h fields have long-range correlations, the sff is non-vanishing for all E' We next note that D (11 since
80 ~ O(L and by Price's theorem,
0h Ou On 10 jo(()))) + 2L (0(h))0()
2²
D&E
O(1)
(103)
O5e
Similarly, we also need to calculate the terms involving G
0S
0 g0) -
G² KgG
E,E
aGEaG
0SS
DGB (g' òe 1 D
D
OG
L
0SS
0
(104)
DGE
Lastly, we must consider the cross terms involving both { and {G*)
0S
(0(1))()))) sf) O'
lil'
0S
dG ((())) 1
D
OD
SS
Oge (g'5)) = - D
O
8
0
(105)
The tensor represents instantancous source of variance for the kernels, while the D tensor represents the sensitivity of the kernel at layer to a perturbation in the kernel at layer All entries of each D matrix are (1). Now, compute the propagator, construct block matrices
K.
K
Do
Do.G
K
KG.
KGG
=
Da.
DG
(106)
Now, obtain the propagator E, we must compute the inverse of the Hessian V
(I- T-11D D T-
(I (I- D) 1i (1- DD)
- (I DD-
1 D
-T- -
0
(107)
The observables of interest are the top-left block which gives the joint covariance over all { G*) variables. We introduce the following x 2 block matrices
E BaG
Dil' bg N
Due D
Rip
D&E
(108)
De
KG
KG
The above equation gives get the following 2 2 matrix equations
E 1 E
DeugEe - I
D N +
1 L2
D
kk. D eil (109)
We see that the solutions E6." the above equation will be Oy This precisely due to the factors of I/L. which appear in the above sums. One way to see this is to consider a large L. > limit where the above layer sums converge integrals over
E(T,T') / dT, DT, E(T"" / f (r"' ((')(""""
/ D(T D(T',, x(T,T')
(110)
So find that the covariance of kernels each pair of layers E(T, is Oz()).
34

---

Published as a conference paper at ICLR 2024
H.2 VARIANCE OF THE NTK AND PREDICTOR
Using the above fact that DeEe Or,(1), we can reason about the variance of the neural tangent kernel K
Var(K) =
Σ
G
G
Σ E
NTP
NT?
EL'
Σ 2
G
OLN((N
(111)
NT2
This demonstrates that the NTK will have initialization variance that scales only with the width N but not the depth L. Using the fact that St = K A, we see that f will inherit fluctuations on the same scale of ON-L(N-1/2), The depth L NTK and predictor covariance converge to the infinite depth propagator with rate (I.-I).
I FINITE WIDTH AND DEPTH ERROR ANALYSIS
In this section, we combine the results of the previous sections. We consider the following error decomposition where we introduce the norm ||=|| = (=²)) and () denotes an average over random
initializations of the network. Applying the triangle inequality this norm, we have
|l - fs0,0ol || S || (fN,L) || (Finite Width Variance)
+ || (fN,L) - 1.
(Mean Predictor Finite Width Error)
+ Ilfoc.L - foc.soll (Finite Depth Error at Infinite Width)
(112)
From the preeceeding section the first term has asymptotic behavior
(IfNLL - (fx,,)))) x 0(N-3)
1 1 1
Eア +0
1
(113)
NL
N2
where the last term comes from the fact that Σ approximates the infinite depth propagator Σ> with error L/L. Next, can apply results developed from Bordelon & Pehlevan (2023) which show that the mean predictor satisfies (JN,L.) = Soc, + O(N-)) 1 to find that mean finite width error scales as
((f,L) - f ~ O(N-2)
(114)
Lastly, the finite depth error in the infinite width DMFT arises from the 1/L. discretization effect. Thus the square error goes
1 fo0,00 2)
(115)
which follows from the discretization error of the limiting large depth process for the NTK and thus the predictor. Altogether, we have
Ex o (O NL +o (
+0 (2))
ll foc,ooll
(116)
We see that at large width and depth this dominated by the first term which comes from the finite width fluctuations. However, if the model is ensembled, then we expect faster rate of convergence from the last two sources of error. Indeed, we provide an experiment verifying that ensembling improves the rate of convergence in Figure B.7
J WHY PARAMETERIZATION INFLUENCES HYPERPARAMETER TRANSFER AT FINITE WIDTHS AND DEPTHS
One may wonder why other parameterizations which also admit a scaling limit (such as Neural Tangent Parameterization) do not transfer well as networks in μP why proposal transfers over depth but others do not. To achieve successful transfer across finite widths, it is insufficient to reason about the infinite width/depth scaling limit, rather needs to identify parameterization
35

---

Published as a conference paper at ICLR 2024
that minimizes the approximation error between finite networks and the scaling limit (of course subject to the obvious constraints training in finite time, learning features at the desired rate etc).
Alternative parameterizations that do not keep feature updates constant, be thought of as choos- ing a function To (NN, L) which describes how much the features evolve with N, L. Suppose further that this has limiting value L "ro L.) % For example, neural tangent parameter- ization in our VL ResNet corresponds to "Yo VN with % 0. We let the logits for finite
width/depth model in this parameterization be JN. (N.L)
represent a finite width and depth net- work predictions with feature learning scale (Yo N, L). The approximation error this (N, L) network
and another (N', L') network has two potential sources of error
(N.2) (NO L') N.L JN'L
NL 500.00
LL) (No(N.
11 $00,00 2''1I
(117)
Heuristically, eliminate the final approximation error by choosing parameterization where (N,, L) = Yo is wutthddepth-independem constant. Then one is left with finite approximation errors to infinite model with the same rate of feature learning.
K ARCHITECTURAL + OPTIMIZATION EXTENSIONS
In this section, explore various extensions of the model written in the main text. Most of these extensions do not modify the scaling rule or the procedural ideas which we used to characterize the limit. We will explore convolutions, multiple layers per block, training read-in and read-out matrices, discrete time training dynamics (finite step size SGD), momentum, and LayerNorm.
K.1 CONVOLUTIONAL MODELS
A DMFT for convolutional ResNets can be easily obtained by introducing spatial coordinates b of each layer's 's representation (Bordelon & Pehlevan, 2022b). The preactivation vectors h (x t) R represent the activity of all N channels in layer at spatial position a. The trainable weights take the form W R N'N describe the filter at spatial displacement from center. The recursion of the ResNet
1 VNL
h h (a. () (((
(118)
where the over b runs over all valid spatial positions in the filter. The relevant order parameters are still feature-feature correlations, but now with additional spatial indices
V (2, (2 t)) ((h((a'
(119)
A similar gradient-gradient kernel is also crucial G (x, x' sts, From these objects, one can construct limiting DMFT at N 00
K.2 MULTIPLE LAYERS PER RESIDUAL. BLOCK
Many residual neural networks, including the popular ResNet-18 and ResNet-50 models, there more than a single convolution layer on a residual branch. In this section, we show that this does not alter the scaling rule the convergence the number branches goes to infinity with number of layers per block fixed. We therefore work with the case of K O(1) hidden layers each branch and still take the branch number
well defined large width and depth limit provided
36

---

Published as a conference paper at ICLR 2024
to infinity. The recurrence of interest is
f NoNN wt k(h))
h1+1 = h +
h ,EE (1,,...,b-11
VT.
hc4+1 W k fI...., 1}
VN
h W o(h')
VN
h
W
(120)
VD
The goal is to verify that gradient based learning on width and depth limit N and L - The weight dynamics are
gives a well defined infinite
d Wik = 7070 Ex A(x, t) g (x, t) h (2c, k (t. K 1}
dt d W
VNL
7o En A(x, 1) g e+ (x, t)o (x,
(121)
di
VNT.
where we have defined the gradient fields
(x, NYo I 0f(x,t) O1)(())
(h.*(xxt)
0f ON.((1)
g' (x, Nyo (h'
(122)
We next examine the scale of the updates to preactivation features he. and h
he (x, = 1 w &l (O))(())))))))) + 70 Ew S. ds A(x', s) (x, (a, x'; s)
VN
VL
h (x,,) = h (2r, 1)
WeE (0)
VLN
110 700 Ban / ds Aax', 8) s) (x, x', at,,,,ta ())
(123)
L
We that the h. features will be uncorrelated across k since there skip connections within a block. However, the h' exhibit long range correlations as before. To understand the cumulative impact of the feature learning corrections to the preactivations, must reason about how correlated the gradient signals g and g" l. are across different layers. We
g (xxt)) t)
(h l.k (x, (e, @ w T 1
(x 1)
= = -6 (x,)) + VNI o('' (z. f)) @ ((
(124)
We see that the skip connections keep long range correlations in the g signals, however, the g are uncorrelated with one another, in standard feedforward network.
K.2.1 THE LIMITING BEHAVIOR OF MULTI-LAYER PER BLOCK MODEL
Based on the correlation structures, of the of and ge gradient featuers, we have the following DMFT equations at infinite width and large depth
2+1
e+1 (x,t) = h1'xx, t)
Σ
(x. t)
VL
8+11
'' ds A(w' "0x',s)00 -I,
Exg
( (, (a' 111-1132 (125)
37

---

Published as a conference paper at ICLR 2024
where the Gaussian source has the form
w (x,t)~ ~
(126)
where 4AK K. determined through the lazy within-block recursion
ph. x', 1) (o(u fM (z, ()) ()r,t)))all (1)
E (I,..., 1}. (127)
Though all the kernels are dynamical, it as if the weights W are static for k {I, K K- 1}.
This is a consequence of the decorrelation of gf across k. We leave as a future problem whether it is possible to reparameterize the intermediate weights so that each of the internal block weight's dynamics move by (1) and contribute additional feature learning terms in the dynamics.
K.3 EFFECT OF TRAINING READ-IN WEIGHTS
in addition to training the weights of the ResNet "body" we also train the readout and readin weights, they will also contribute to the total NTK
K(x, x²) -
Σ t=1
G (a, (2) + ( (x)
(128)
where (a x') D D x' is the input kernel. We that the addition of these layers does not change the fact that the NTK (1) I) in this parameterization. Now, we need to verify that the first layer weights move by the appropriate scale
d (t)
W
700 Ea,(z, t)9. t)æ, ~ (1.
(129)
dt
vD
These equations lead to the following base case equations for the fields (xx t) and h (x,
h (ax t) = u (xx,) + No '/" ds $)g (
s)K (x,
(x, (x, ò(h² (x, t))r") (x, t) + moro (x, t)) I. ds
(130)
These fields will therefore evolve by (1) in this parameterization.
K.4 DISCRETE TIME UPDATES
DMFT for feedforward networks also be modified for discrete time (Bordelon & Pehlevan 20226 2023 ). A straightforward extension of these results will give discrete time dynamics for the hield dynamics h h'.g" of for our ResNet. The one subtlety about discrete time is that the dynamical neural tangent kernel (NTK) no longer governs the evolution of the predictor f. Instead,
EE A(w', s) x xt,,,s) + no / fdx dx'' Σ (z, Set
f(x,t) 170
(131)
sct
where B, is the minibatch at timestep and ( 11,)) VI. (80th"(a.tt)) is the final layer (r"(s''s))
y0770
response function at depth L. These are now evaluated for t, integer rather than real numbers. In the infinite width and depth limit N, L oo, we have the following discrete time update equations for the preactivations
N(T, = L'' du(r', ax,) + nono S Idx' ddd''e
Cn('', 2, o' s))(") xs,ss
1, = A(T, x. x' t,s) + |Bg 2,
Σ 5( x.)a(a' 8) 5) x, (132)
Notice that this is almost identical the dynamics for gradient flow except that the integrals time become sums over discrete time and the population density p(æ) is approximated at each step by uniform density minibatch B, An analogous discrete time equation holds for g(7).
38

---

Published as a conference paper at ICLR 2024
K.5 MOMENTUM
Momentum can be easily handled within our framework. Let 0 = [ W w',.., W w²} represent the concatenation of all trainable network parameters. The momentum update controlled by a momentum parameter a which controls the exponential moving average of gradients
0 + (() + (1)
v(1) - a(1 - - (1 ao VoE,.(/(,,
(133)
Writing this out for a particular hidden layer yields
(t (t + 1) w (1) + 7o70 (1)
V*(1) = (1) a1)
1) VLN E. Fr, () et (ar, ()r ())
(134)
We can unfold the recurrence for V get the expression
v²(t) (1 (1-a) Σ a' E.A(E, (z0
(135)
VLN
Plugging this into the update equations for (() we get
w (t+ + 1) (t) + 11 ro a))
LN Σ a' B, + s)o(h (a (a, s))
(136)
From this equation, we can again unfold the recurrence for W for h g¹, verifying that we get feature learning updates of the correct scale
W (t) and write the recurrence equations
1
h (x, h (xx, t)+ +
(x,t)
VT.
770 70 (1 a)
ΣΣ
D' 8), ,8) t
(137)
L
e'tecr
where x (xx VN (O) We that this update has the correct scaling struc- ture. the random (uncorrelated across layers) field has scale of v and will behave Brownian motion and the feature learning update with the gradient field (strongly correlated across layers) has the correct J. scaling. The normal DMFT procedure can now be carried out on the x fields and the corresponding fields for the backward pass.
L WEIGHT DECAY
Weight decay dynamics can also be analyzed with DMFT. The gradient flow dynamics with weight decay with parameter À have the form
w (t) = Yo dt
E. A(x, () e+1 (x, (h"(())())))))))))))))))) tx) (
(138)
VLN
The dependence of (t) on the initial condition can be isolated with a simple integrating factor
w'(1)= = no w (0) +
7o
/
ds A0 (ax, 8)g (ax, (1 (1 8))
(139)
VNL
Computing the forward pass recursion we find
h c+ (x, t) = (x, t) + -oodt x
VT.
7070 / ds (o Ala', 8)g e
(140)
This clearly will admit infinite width and depth limit following the arguments in Appendices D & E The difference is the presence factors of e noht which suppress the influence of the initial condition late time.
39

---

Published as a conference paper at ICLR 2024
M OTHER LARGE DEPTH LIMITS
is possible to construct other large depth limits residual networks. general we can consider
h +++ h + Bg T. W/2 "''("")))) ((h))
(141)
Well-behaved signal propagation at initialization and throughout feature learning are preserved for any a 2 1/2 provided the learning rate is scaled as = 77 2 2 This ensures that d h' (1) and _ f (11) in the large width and depth limit N, + co. The DMFT approach can also be used to analyze these limits, though we leave this to future work.
We show an example of a E {1,2) parameterizations in Figure M. In that model, we train ResNets with K 4 layers per residual block and start with zero initialization of the readouts from each block so that the Brownian motion terms vanish for both parameterizations. This is similar to the Res-Zero initialization (Bachlechner al. (022), however we note that with zero initialization of the branch readout, the parameterization (the depth-scaling exponent a) still makes a difference the dynamics. We find that both parameterizations show approximate transfer (runs diverge for same learning rates at large enough depth). However, the ODE scaling a = shows monotone improvement with depth while the SDE scaling a 1 shows the opposite at early time.
Depth
2.0
2.00
1.75
16
32
64
Train
2.0
1.00
0.75
128
5000
10000 Steps
15000
20000
10-2
Learning Rate
(a) ODE 1 Dynamies
(b) ODE 1 Transfer
2.00
2.00
1.75
3.75
1.50
3.50
1.25
Depth
Train
2.00
1.00
0.75
0.5
0.50
32
5000
10000
15000
20200
Steps
Learning Rate
(c) SDE
Dynamics
(d) SDE
Transfer
Figure M.I: Hyperparameters transfer for both SDE and ODE scalings for multi-layer-prr-blok model K 4 models. Depth L varied for both models. Increasing L leads to monotone improve- ments in the ODE but not SDE when initializing with zero readout weights for each block. In (d) all runs diverge at the next biggest learning rate.
A more in-depth characterization of the differences in transferability and convergence rates of mod- els in each of these parameterizations will be left future work.
40


Reviews:
[{"review_id": "1pQKDMJC6F", "reviewer": "Reviewer_vjdW", "summary": "The paper presents a parameterization for Resnet neural networks which is claimed to enable the transfer of hyperparameters such as the learning rate over a large range of network widths and depths. The motivation for this work is to enable engineers to tune hyperpaprameters on small models and then apply them to large models, avoiding a costly fine tuning on large models.\n\nThe method extends the muP parameterization which, according to the authors, enables hyperparameter transfers across different widths but not different depths. The proposed parameterization consists of zero-mean unit-variance parameter initialization and scaling factors for the residual branches, output pre-activations and learning rate that depend on both the width and depth of the model, and no normalization layers.\n\nThe proposed parameterization is derived by theoretically analyzing the behavior of a (simplified) neural network in the limit of  infinite width and infinite depth, and choosing a parameterization which satisfy certain desiderata. Experiments on a small dataset CFAR-10 are used to support the analysis on realistic neural architectures.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The method tackles an important practical problem: the efficient tuning of hyperparameters.\n- Good presentation", "weaknesses": "- The experimental section is quite limited: only a single task and a very simple dataset are considered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "F90IQ3ykP0", "reviewer": "Reviewer_fkzv", "summary": "Paper proposes a novel parameterisation of deep residual networks that tackles width- and depth-dependent cost of hyper-parameter tuning, an issue, exacerbated by recent increase in SOTA models' sizes. A novel $1/L$ extension of $\\mu P$ parameterization of residual networks is proposed. Paper argues that hyperparameter transfer is consistent across both width and depth for a variety of architectures, hyperparameters, and datasets. \nThe work is primarily empirical whose experimental evidence is supported by statistical physics and NTK theory. In particular, the paper advances a width-invariant feature updates present in $\\mu P$ parameterisation and in the same spirit derives depth scaling to ensure feature updates invariance over depth using dynamical mean field theory (DMFT). \n\nExact solutions of the DMFT dynamics in the rich (i.e.non-kernel) regime for a simple deep linear network are provided. The suggested parameterisation is verified empirically by well selected range of experiments, mostly presented in appendices and  summarised in the Section 3. The exposition of the challenges, method and main results is well and accessibly written rounded by limitations and future directions.", "questions": "See section Weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ Rounded, well written and easy to follow timely paper presenting valuable empirical evidence put into NTK theory context.\n+ Future directions section proposes the method as suitable to study depth-width trade-offs while scaling neural networks since hyperparameter dependencies are consistent across widths/depths. In my opinion this can be very valuable addition to community even if the method does not turn out practically useful and adds positively to my overall rating.\n+ Convincing experiments (with a caveat, that I hope will be addressed in the rebuttal, see Weaknesses, ad 1)) cover well selected range of settings and architectures", "weaknesses": "1. Fig. 1 Loss levels reached seem to be higher in case of proposed parameterization compared to alternatives. Could authors add some comments on the topic? Especially, is a proposed parameterization capable of reaching an optimal value of hyper parameters in scope, e.g., learning rate in practical settings? I rate paper 'accepted' conditioned on this issue will be alleviated in the camera-ready version.\n\n2. Limited applicability(?) - Could authors argue otherwise? How could it be improved in the paper? Could authors elaborate on computational costs of proposed method vs. alternatives?\n\n3. Paper heavily depends on $\\mu P$ parameterisation properties derived in Yang & Hu (2021). In my view a short primer on $\\mu P$ parameterisation (even in Appendix) would improve self-consistency and readability of the paper.\n\n4. Proposition 2, Assumption $\\gamma_0 \\rightarrow 0$ renders ODE from Proposition 1 solvable due to time invariant kernels (and thus given by Gaussian initialization) and the same goes for the second example where linearity of network makes time invariance explicit. To what extend are conclusion transferable to any realistic non-linear scenario?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "ApiuYrFzpt", "reviewer": "Reviewer_SFFE", "summary": "The authors' contribution in this work lies in their introduction of a straightforward parameterization for residual networks. Through empirical observations, they have demonstrated the remarkable consistency of hyperparameter transfer across various dimensions of network architecture, such as width and depth, as well as across diverse hyperparameters and datasets. This empirical evidence is further substantiated by theoretical analysis, which reveals that the behaviors of hidden layers remain intricate, unaffected by changes in network dimensions, and do not approach vanishing limits.", "questions": "Please see  weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper demonstrates that a straightforward adaptation of the $\\mu P$  parameterization enables the transfer of learning rates across both depth and width in residual networks, incorporating $1/\\sqrt{depth}$ scaling on the residual branches. The authors conduct extensive experiments. The theory is also nice.\n\nI think these findings have the potential to significantly reduce the computational expenses associated with hyperparameter tuning, thereby enabling practitioners to train large, deep, and wide models just once while achieving near-optimal hyperparameters.", "weaknesses": "1. Since that the $\\mu P$ parameterization is not a commonly used method, the authors should provide a more comprehensive introduction.\n\n2. The $1/\\sqrt{depth}$ trick has been extensively investigated in previous works [1-3]. The novelty of the proposed method appears to be limited. Please clarify the differences. Moreover, many recent works propose initializing ResNet using \"rezero\"-like methods, which achieve similar performance to the $1/\\sqrt{depth}$ trick. The authors should provide some comparisons. In particular, I would like to see some discussions on the scaling parameter. Why use $1/\\sqrt{depth}$? What would happen if the scaling parameter is set as $1/depth$ or 0?\n\n3. As the authors admitted in the \"Limitation\" paragraph, most of the experiments are confined to 10-20 epochs. This limitation makes the empirical evidence less convincing. In fact, I don't understand the difficulty in training a ResNet with more epochs if one can train it with 10-20 epochs.\n\n[1] Tarnowski W, Warchoł P, Jastrzȩbski S, et al. Dynamical isometry is achieved in residual networks in a universal way for any activation function. The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019: 2221-2230.\n\n[2] Yang G, Schoenholz S. Mean field residual networks: On the edge of chaos. Advances in neural information processing systems, 2017, 30.\n\n[3] Zhang H, Dauphin Y N, Ma T. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv:1901.09321, 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "hlzjH9rJXf", "reviewer": "Reviewer_J7oQ", "summary": "The authors study hyperparameter transfer between models at different scales with an emphasis on the hyperparameter transfer of ResNets (or ResNets sub-components) with depth. Following parts of the rational behind muP parameterization, they argue that good parameterizations are those having stable infinite width/depth limits. They then consider one such recent limit (\"1/\\sqrt{depth}\" scaling) and extend prior theoretical results on ResNets at initialization to ResNets trained with gradient flow. Specifically, they establish that network updates are also stable at this limit. The manuscript contains several experimental results on CIFAR-10,Tiny ImageNet, and ImageNet showing successful hyper-parameter transfer based on this \"1/\\sqrt{depth}\" scaling. As an additional theoretical side, the DMFT equations are solved for an infinite widh&depth linear network acting on a single data point.", "questions": "1. Can the authors rationalize why the above two gaps could be excused? \n\n2. Can the authors explain why for CIFAR-10 their parameterization seems non-optimal?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The work offers a good combination of theory and experiment.  \n\nIt expands a useful and timely avenue of research, hyperparameter transfer, where theory can actually aid practitioners.\n\nIt studies relevant complex architecture acting on realistic datasets as opposed to linear or shallow networks.", "weaknesses": "Several gaps exist between the stated goal of the theory and its practical implementation. One such gap, which the authors themselves cite, is that the theory predicts stable results following a scale-up. Naturally, however, spending more compute and getting equal performance, is of little practical value. In practice, one expects to get better results which seems to be the case at least for the larger datasets they consider. This means that in examples where HP-transfer is worthwhile, one is by definition well away from the scaling limit. \n\nAnother gap, of a similar kind, is that, as far as I could tell, their DMFT framework is based on gradient flow but it is then used to transfer the learning rate. While gradient flow may be a good approximation for SGD in the asymptotically low learning rate regime, typically the optimal learning rates lay far away from this regime. \n\nFor CIFAR-10, the 1/\\sqrt{depth} scaling leads to poorer performance compared to muP scaling. \n\nReference to current literature could be broadened and sharpened. For instance, muP parameterization is also a unique parameterization having maximal updates, not just a stable limit (This fuller characterization also partially escapes the first gap described above). Citing further alternatives to muP scaling and adaptive kernel approaches to feature learning is also desirable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)