Tensor Programs VI: Feature Learning in Infinite Depth Neural Networks
Authors: Greg Yang
Dingli Yu
Chen Zhu
Soufiane Hayou
Venue: ICLR 2024 poster (2024)
Keywords: Tensor Programs
mup
deep learning
optimization
optimal hyperparameter transfer
Primary area: optimization
TL;DR: We introduce Depth-$\mu$P, a principled approach for depth scaling, allowing for the training of arbitrarily deep networks while maximizing feature learning and feature diversity.

Abstract: Empirical studies have consistently demonstrated that increasing the size of neural networks often yields superior performance in practical applications. However, there is a lack of consensus regarding the appropriate scaling strategy, particularly when it comes to increasing the depth of neural networks. In practice, excessively large depths can lead to model performance degradation. In this paper, we introduce Depth-$\mu$P, a principled approach for depth scaling, allowing for the training of arbitrarily deep architectures while maximizing feature learning and diversity among nearby layers. Our method involves dividing the contribution of each residual block and the parameter update by the square root of the depth. Through the use of Tensor Programs, we rigorously establish the existence of a limit for infinitely deep neural networks under the proposed scaling scheme. This scaling strategy ensures more stable training for deep neural networks and guarantees the transferability of hyperparameters from shallow to deep models. To substantiate the efficacy of our scaling method, we conduct empirical validation on neural networks with depths up to $2^{10}$.

Published as a conference paper at ICLR 2024
TENSOR PROGRAMS VI: FEATURE LEARNING IN INFINITE DEPTH NEURAL NETWORKS
Greg Yang" xAI
Dingli Yu" Princeton Language and Intelligence. Princeton University
Chen Zhu Nvidia
Soufiane Hayou? Simons Institute, UC Berkeley
ABSTRACT
Empirical studies have consistently demonstrated that increasing the size of neural networks often yields superior performance in practical applications. However, there is lack of consensus regarding the appropriate scaling strategy. particularly when it comes to increasing the depth of neural networks. practice, excessively large depths can lead to model performance degradation. In this paper, we introduce Depth-uP, principled approach for depth scaling, allowing for the training of arbitrarily deep networks while maximizing feature learning and feature diversity among nearby layers. Our method involves dividing the contribution of each residual block and the parameter update by the square root the depth. Through the use of Tensor Programs, we rigorously establish the existence of limit for infinitely deep neural networks under the proposed scaling scheme. This scaling strategy ensures more stable training for deep neural networks and guarantees the transferability of hyperparameters from shallow deep models. To substantiate the efficacy of our scaling method, we conduct empirical validation on neural networks with depths up to 210
INTRODUCTION
Deep neural networks have showcased remarkable performance across broad range of tasks, including image classification, game playing exemplified by AlphaGo (Silver et al. (016), and natural language processing demonstrated by GPT-4 (OpenAl 2023). A prevailing trend in developing these networks is to increase their size and complexity, with empirical evidence indicating that using the same computation resources, models with more parameters tend to exhibit better performance. There are two ways to increase any network size: width and depth. The properties of the width (given fixed depth) have been extensively studied in the literature: recent work by Yang et al.] (2022) identified the Maximal Update Parametrization (nP) that guarantees maximal feature learning the infinite width limit Another benefit of μP is hyperparameter transfer which enables hyperparameter tuning on smaller models; the optimal hyperparameter choice for the smaller model remains optimal for larger models (i.e., models with larger width). However, despite the achievements of large-scale deep models and the theoretical understanding of scaling width, increasing the depth of neural networks still has both practical limitations and theoretical difficulties. In practice, increasing depth beyond some level often results in performance degradation and/or significant shifts in the optimal hyperparameters. In theory, unlike increasing width, increasing depth introduces new parameters that significantly change the training dynamics. In this paper, we aim to solve this problem by extending JPP to include depth scaling. We call the depth scaling Depth-ud..
The issue of depth scaling has persisted over time. A decade ago, deep neural networks experienced significant degradation problems having more than few dozen layers would increase the training error instead of improving the model's performance. This was partly due to the vanishing
Equual contribution.
Work partially done the National University of Singapore.
'erre maximal feature learning refers to @(1) change in features in the infinite width limit. This should be contrasted with the lazy training regime where the change features is of order (1n3)/

---

Published as a conference paper at ICLR 2024
exploding gradient problem that affects the efficient propagation of information through the network. The introduction of residual networks (ResNet) (He et al. 2016 bj Srivastava et al, 2013) has partially resolved this issue, allowing for the training of deeper networks with improved performance. ResNet is constructed by layering residual blocks, which are composed of a series of convolutional layers and then an element-wise addition with the input. This element-wise addition (commonly referred to as skip connection) is a significant innovation of ResNet and remains an important ingredient in modern architectures including Transformers (Vaswani et al. (017))
Specifically, in a residual architecture, the /-th residual block is formulated as
x! = x + g'xx'~,;
where x/-1 is the input, of is the output, W are the parameters of the block, and g' (often called the residual branch) is mapping that defines the layer (e.g. a stack of convolutions in ResNet, or SelfAttention and MIP in a Transformer). In this work, we focus on the case where g' is a biasless perceptron with (or without) activation.
The stacking of many residual blocks causes obvious issue even at the initialization the norm of x' grows with I, so the last layer features do not have stable norm when increasing the depth. Intuitively, can stabilize these features by scaling the residual branches with depth-dependent constant. However, scaling the residual branches with arbitrarily small constants might result in no feature learning in the large depth limit since the gradients will also be multiplied with the scaling factor. In this paper, we propose Depth-juP, a principled approach to scaling up the depth of residual networks, enabling the training of arbitrarily deep networks while achieving feature learning and maximizing feature diversity among nearby layers. Our framework extends the previous results on μP which deals with optimal width scaling (Yang et al. 2022). It completes the width scaling and hence provides full width and depth scaling recipe that guarantees maximal feature learning and hyperparameter transfer across width and depth. Depth-uP contains the following modifications to the standard practice:
1. There is a multiplier for each residual branch before adding its input, which is inversely proportional to the square root of L (where is the depth). Formally, with constant a independent from L. = x'-1
Trivial
VI" g'(''"t: (1)
P-y
Smooth
Unfatthiul
We set the learning rate of W' so that the update of W' during training is proportional to 1/VL. We derive different learning rate schemes for different optimization algorithms based on this principle. For Adam, because it scale-invariant to the gradient, the learning rate of W is set to be n// /.VL. On the other hand, the learning rate of W for SGD is set as a constant v because the gradient of W is already of size 1/VL due to the multiplier.
Unstable
0
1/2
1
Init stable
Figure 1: Behaviors of scaling strategies with branch multiplier -a and parameter update size proportional to T.-?.
Figure compares Depth-AP with general branch multiplier scaling L-a and update size scaling 1/= see Section 4 for more details), showing the uniqueness of Depth-uP. Specifically, limiting our analysis block depth 1 (i.e., g is biasless perceptron, W' is a single matrix we show that Depth-uP leads to the following properties: a) At the initialization, each one of the residual blocks contributes 0(1/VT.) to the main branch. These I. contributions independent of each other, the sum of them is of size (11): b) During training, the contribution of the update of each residual block is ((1/L) L) due to the combining effect of the learning rate and multiplier. The contributions of the updates are highly correlated, so they sum up O(1); e) More
L- (where Depth-juP has
'In Appendix K. we include a thorough discussion on the scaling strategy for block depth 2 and above. We find limitations all possible infinit-depphh limits of such parametrizations, which we illustrate simple networks as well as Megatron transformer trained on Common Crawl (see Appendix
2

---

Published as a conference paper at ICLR 2024
importantly, classify all depth limits and show Depth-uP yields the optimal limit. This implies that the optimal hyperparameters of the networks in Depth-juP are approximately invariant wrt depth. With Depth-ud,, successfully train networks comprising thousands of residual blocks, while also showcasing the transferability of hyperparameters across depth.
2 RELATED WORKS
Recently, /P (Maximal Update Parametrization), introduced in Yang al. (2022). has emerged as a promising approach for maximizing feature learning while simultancously preventing feature explosion as the network width increases, given fixed depth. Notably, μP facilitates hyperparameter transfer across varying network widths. This means that instead of tuning hyperparameters directly on large models, one can optimize them on smaller models and utilize the same set of hyperparameters for larger models.
While increasing the width can lead to improved performance, increasing the depth of the network also yields significant performance gains, and most state-of-the-art models use deep architectures. To address the challenges associated with training deep networks, several studies have proposed scaling the network blocks using depth-dependent scaler to ensure stability of features and gradients at initialization (Hanin & Rolnick] 2018 Zhang et al.] 2019 Hayou et al. 2021 Noci
al. 2022 Hayou 2023 Hayou & Yang 2023 Noci et al. 2023). For Multi-Layer Perceptrons
(MLPs) (no skip connections), Jelassi et al. (2023 showed that a learning rate scaling of depth-3/2
guarantees stability after the initial gradient step. However, remains unclear how learning rate should be adjusted beyond the first step, and this scaling is not suitable for architectures with residual connections. We refer the reader to Appendix A for a more comprehensive literature review.
WARM-UP: DEEP LINEAR RESIDUAL NETWORKS
Let us begin with simple example that provides the necessary intuition underpinning our depth scaling strategy. Given depth L, width n, consider a linear residual network of the form
o = x' -- + W
for I = 1,..., L,
VT.
where _0 Ug the input layer, f(E) V is the output layer, W E R are the weight matrices, and U, V are input and output weight matrices assumed to be fixed during training.
3.1 OPTIMAL SCALING OF THE LEARNING RATR
To simplify the analysis, we consider gradient updates based on single datapoint. The first gradient step is given by W| = w - n/G, where is the learning rate, and Go is a matrix with update directions. For instance, we have the following expressions for G with SGD and Adam:
SGD: Go J dx @ xl-1, where dæl del at for some loss function f. del
Adam: G/ = sign
J dr' @ 1 where is the Kronecker product
In both cases, del and - are computed for a single training datapoint Eo- The last layer features 1.w)
a (for some input {) are given by xー = I (1+ training step. After the first gradient step, we have the following
W
2 4 We use the subscript to refer to
of = II
I l=l
VI
1
wi)
20 at
JT O(7)),
1
where At = Σtn Mest (T+ ++++ Go [Tec (I+ +++ W/ 20. We argue that Ar behaves as (L)) (in second norm). This is the due the I/VL scaling factor. To see this, we
'The step Adam can be viewed as the SignSGD update. Our analysis Adam is valid any training algorithm that gives @(1) gradients. "To avoid any confusion, here define the matrix product by IT A A AL-1... .... A
3

---

Published as a conference paper at ICLR 2024
further simplify the analysis by considering the case dnn dout 1 (single neuron per layer) and the squared loss.
Scaling for SGD. With SGD, we have that G/ VT. IR (1+ vT. W mndo where dx
(Va** 3/(E)) and y(Eo) the target output. Therefore, is easy to see that
EAA- ET
wx)
2
da' 20
do
+ (1 OLL)
VL
where we have used the fact that E (1 (1+
1
W
3p
1 O(L for any positive integer p.
VL
2 is given by E ((E @ which
Hence, the magnitude of the first order term in eq.
shows that the update is stable in depth as long 1 (1)) in depth. More precisely, this is the maximal choice of learning rate that does not lead to exploding features as depth increases.
Scaling for Adam. With Adam, we have Go =1 and therefore we obtain
EA E
ΣI
1+
With VL
1
20
(L3),
where we have used the same arguments as before. In this case, the first order term in eq. (2) is 3 4: O(r² L-), and therefore, the maximal learning rate that one choose
given by E
without exploding the features given by 1 = O(L-1/2).
Summary: By ensuring that parameter update is (1/vVT.) the features remain stable while feature update is (1). This ((1) update is due to the accumulation of (1/L)) correlated terms across depth. We refer the reader to Appendix BB for more in-depth analysis of this simple model.
3.2 SGD TRAINING DYNAMICS BEYOND ONE TRAINING STEP AND =
The intuitive example above assumes there is only one single neuron per layer (n = 1). However, rigorous extension the general finite width setting (n 1) is hard. This be solved by introducing Tensor Programs framework (Yang & Littwin 2023). which allows us to study the infinite-width network and be worry-free from noise caused by finite width. Tensor Programs also enable analysis of training dynamics beyond the first training step. Particularly, we rigorously derive the training dynamies of SGD for the linear residual network when the width and the depth sequentially go to infinity. Our analysis of the linear case reveals our main technical tools with minimal notations and prerequisite knowledge, so it is more friendly to readers who are interested in the technical details. Below provide our technical road map. The full version of the analysis can be found in Appendix
1. We first take the width of the network infinity by the Tensor Program framework. As result, instead of tracking vectors and matrices along the training trajectory. we track random variables that correspond the vectors, that is, for vector æ e R" that appears in the computation of the training, the coordinates of x can be viewed as iid copies of random variable aæ)) (called a ker) when n + co.
2.
Since the network is linear, every random variable can be written a linear combination of set of zero-mean "base" random variables by the Master Theorem of Tensor Programs (Yang & Littwin
track the random variables by analyzing the coefficients of their corresponding linear combinations, and the covariance between the "base" random variables.
2023). Then
We aggregate the coefficients of all random variables and represent them by a six-dimensional tensor r, where two of the dimensions have shape L because the number of random variables and the number of "base" random variables scale linearly with L. The next step is to find r converges when L grows. However, any convergence directly based on I is impossible because the shape of I is different when grows. We thus convert to a set of functions I' whose input domain is (0, 1] (0, 1]. where the two inputs to I corresponds to the fractional index of the two
4

---

Published as a conference paper at ICLR 2024
dimensions of I of shape L. Similarly, the covariance between the "base" random variables can represented by another set of functions C' whose imput domain is (0, 1]. Finally, we claim that the functions I' and C' converge when L - oo, and identify their limits as the solution of set of functional integrals Proposition C.4
In Appendix L.. we conduct thorough empirical verification of our theory in the linear case. The experiments clearly show the convergence of deep linear residual networks under Depth-udP.
4 PRELIMINARIES FOR THE GENERAL CASE
Notation Let L be the depth of the network, i.e., the number of residual blocks, and n. be the width of the network, i.c. the dimension of all hidden representations x" ,...,,".. x², Let Rd be the input of the network, U E g=xdm be the input layer, and V E gu*x be the output layer, so that x0 = US and the model output w.r.t. & is (() A VT Let & be the loss function absorbing the label, and dr' be the gradient of x' w.r.t. the loss. We denote variables at f-th training step by adding t as a subscript, c.g., the input at step t is & the hidden representation of /-th layer at step t is ø and the model output step is Jt. Let Th the number of training steps. Let I be the identity matrix, and 1 be the full one vector. For m E N+ let [m] [1,]= m...
4.1 UNIFIED SCALING FOR SGD AND ADAM
We extend the definition of entrywise update (Yang & Littwin 2023 for depth scaling, allowing us to study the unified depth scaling for SGD, Adam, and other optimization algorithms that perform only entrywise operations.
Definition 4.1. A gradient- based update of parameter w with both width and depth scaling is defined by set of functions Q =Q, 1- R]220 and c, d, 8, 7, The update at time t of the optimization
w + w 1pm 1pp==1 -. (n go...., ...,, 2 (.
where Jo. s O, .....,te t, are the gradients of at time s.
For SGD, 1 go. .o...., g..g.....,n,,d.gg)
1 1,5 gt, and the "true" learning rate vmn -m-ddr For
1- Ei 1.
Adam, (n go, .0,...,, Tn.
and the "true" learning rate is 1-8 El " g.
just 1p -
The purpose of multiplying the gradients nd 1 before Q to make sure the inputs to Qr are ((1) w.r.t. and 1 otherwise, the update might be trivial when n and L become large. For example, gradients are o(1) entrywise, then, in Adam, directly feeding gradients to Qt will always give an output of 0 because of the constant >> 0. In this paper, we will only consider d, 8 such that n L ( As a result, the output of Qe is also ((1) in general. Therefore, n n~ decides the seale of the update and should be our focus. We call mp Ln the effective learning rate.
4.2 SETUP
We consider L-hidden-layer residual network with biasless perceptron blocks: [L]. h' w I' / + - MS(o(h))).
where = = UG, the network output / = VT vf & the nonlinearity'activtionn function and MS refers to Mean Substraction and is given by MS(x) We follow μP YYngg 2020b) for the widthwise scaling, i.c., the initialization of U, V. w are i.d. zero-mean Gaussian with variance 1, -² vyn ym--L--" for U, V and W are 0, 1, respectively, i.e., the learning rate of W is ryn pm-1-_-y LL-
(z, 1n = Gz with G = II 11 /n.
respectively, and the in the effective learning rate
In sum, the (a, 7) pair decides the depthwise scaling of whole setup.
'Here, the input is used to perform one gradient step at training step t. We will see later that our claims should in principle hold for batched versions of the training algorithm.
It is called faithfulness in Yang Littwin (2023).
Note c, d, 8, 17 in Definition 4.1 can be different for parameters, so it is possible to make every parameter satisfy the condition. In this paper, we focus on matrices W which share the same of c, 7.
5

---

Published as a conference paper at ICLR 2024
Maximal update parametrization (nP) Yang
2020b) considers the change of initialization and learning rate of each weight matrix in the network when width scales up It provides a unique initialization and learning rate of each weight matrix a function of width that makes the update of each weight matrix maximal (up to a constant factor). The benefit of μP is not only the theoretical guarantee but also the hyperparameter stability when scaling up the width (Yang & Hu (02T).
R
Mean Substraction (MS). The use of mean substraction is key in our analysis. First of all, it acts in a similar fashion to LayerNorm without the normalization part. We believe that mean substraction in LayerNorm plays a more important role than the division by the norm. This is because with the right depth scaling, it is generally the case that features will remain stable as depth grows (no exploding/vanishion
however, the correlation will induce non-zero means, hence the need for mean substraction to avoid exploding behaviour. Secondly, substracting the mean only decreases the dimension of the layers by 1, which is insignificant in the large width limit. Lastly, one can think of generalized parametrization in which means of the activations are computed and added back after scaling the centered block. This will result in similar scaling patterns to the ones we propose in our setup. For the sake of simplification, we do not consider this general setup in this paper.
5 CLASSIFICATION OF DEPTHWISE PARAMETRIZATIONS
In this section, we provide a comprehensive description of the impact of depth parametrization on stability and update size. For this purpose, we only have two scalings keep track of: the branch multiplier and the learning rate scaling because the initialization scale is fixed by the faithfulness property (defined below). Requiring that the features don't blow up at initialization means that the branch multipliers must be at most O(1/vL). Assuming the updates are faithful (i.e., input to gradient processing functions are (1) entrywise), the update size can be at most L/L for the hidden layers, by an (Jacobian) operator corattooooooomooooo argument, but potentially much less. Naively speaking, there can be a trade-off between update size and initialization: if initialization is large, then the update may need to be small so as not to blow up the other parts of the network; likewise if the initialization is small, then the update size can be larger. But one may be surprised that careful calculation shows that there is no trade-off: we can maximize both initialization and update size at the same time.
Before delving into the details, let first define the notions of training routine, stability, faithfulness, and non- con-triviality. Hereafter, all the aymptotic notations such as O, Q and o should be understood in the limit "n 20, then L -> oc". For random variables, such notations should be understood in the sense of weak convergence (convergence in distribution). When we use the notation x = (() for some vector
(x1,,- F.,) R", it should understood in the sense that for all i E [n], an; (1)) Lastly, we will use bold characters (e.g. h instead of h) to denote 'batched' versions of the quantities. This just to emphasize that the following claims should hold for batched quantities as well.
Remark: in this section, we state the results as "claims" instead of theorems. In Appendix G.4 we provide "heuristic" proofs that can be made rigorous under non non-trivial technical conditions. We believe this additional layer of complexity is unneeded and do not serve the purpose of this paper. To showcase the correctness, we provide rigorous proof for all the claims in the linear case in Appendix H which also serves as intuitions for the general case.
Definition 5.1 (Training routine). A training routine is the package of 17, Q, and the input batches.
Definition 5.2 (Stability). We say a parametrization is
1. stable at initialization if h/g ah = O(1), E [4], and fo = O(1).
2. stable during training if for any training routine, any time 2 0, 1 E [T.], we have Ahi, Aa! (1)), vI [Z., and Af:- = O(1, where the symbol 'A' refers to the change after one gradient step.
We say the parametrization is stable if it is stable both at initialization and during training. Definition 5.3 (Faithful). We say parametrization is faithful at step t if h¹ O() for all I [4]. We say the parametrization is faithful if it is faithful for all t. We also say it is faithful at initialization (resp. faithful during training) if this is true at t 0 (resp. for t 2 1).
"Reparametrization is also included in the original jPP but it is not necessary for the purpose of this paper.

---

Published as a conference paper at ICLR 2024
Note faithfulness here refers "faithfulness to o", meaning the input to á is @(1). This is different from the definition of faithfulness in Yang & Littwin (202)), where faithfulness refers to "faithfulness to Q" meaning the imput to Q is ((1). "faithfulness to Q" is already assumed in this work as mentioned in Section 4.1
Definition 5.4 (Nontriviality). We say parametrization is trivial if for every training routine and any time 2 I, ft fo a 0 in the limit "n + oc, then + oo" (i.e., the function does not evolve in the infinite-wiin somnitewwidth-theee width- -het- limit). We say the parametrization is nontrivial otherwise.
Definition 5.5 (Feature Learning). We say a parametrization induces feature learning in the limit "n + 0o, then I + oo", if there exist training routine, and 2 1, and any X > 0, we have Aht [ML] (1).
5.1 MAIN CLAIMS
We are now ready state the main results. The next claim provides a necessary and sufficient condition under which a parametrization is stable at initialization.
Claim 5.1. A parametrization stable at initialization iff a 1/2.
Claim 5.1 is not new and similar results were reported by Hayou et al. (202T). However, Hayou al. 2021) focuses on initialization and lacks a similar stability analysis during training. In the next result, we identify two different behaviours depending on the scaling of the learning rate.
Claim 5.2. Consider parametrization that is stable initialization. Then the following hold (separately from each other): It is stable during training as well iff a + 1; b) It is nontrivial iff a + 7 S 1. Therefore, is both stable and nontrivial iff a + =1
From Claim 5.1 and Claim 5.2 having a+ 1 and > 1/2 is necessary and sufficient condition for a parametrization be stable and nontrivial throughout training. In the next result, we therefore restrict our analysis to such parametrization and study their faithfulness.
Claim 5.3. Consider a stable and nontrivial parametrization. The following hold (separately from each other).
It is faithful at initialization iff a 2 1/2. As a result, a = 1/2 is the minimal choice of a that guarantees faithfulness.
It is faithful during training iff S1. 1.
Therefore, a stable and nontrivial parametrization is faithful iff a E [1/2, 1].
The first claim follows from well-known calculations of randomly initialized residual networks Hayou et al. (202T). For the second claim, the intuition here is just that if a + Y = 1 and a > then Y 0, i.e., the update size blows up with depth. This would then cause the input the nonlinearities to blow up with size.
stable, nontrivial, and faithful parametrizationss have narrowed down the space of (a, 7) to a ca---= and a E [1/2, 1]. However, in order to achieve optimal hyperparameter transfer, we need identify the optimal parametrization among them. (See Appendix for detailed explanation on what causes hyperparameter transfer.) To accomplish this, we first define the notion of feature diversity exponent, which relates to the similarity in the features of adjacent layers.
Now by focusing
Definition 5.6 (Feature Diversity Exponent). We say parametrization has feature diversity exponent x 2 if is the maximal value such that for all [o, 1] and sufficiently small 0, and all time
1
[A+0)4]
-a! AA =
26e 1=))
Vn
where 2(1) should be interpreted in the limit "n - oo, then L + oc, then é - 0". We say a parametrization is redundant if x O.
In other words, the feature diversity exponent & is a measure of how different the outputs are in layers that close to each other. With 0, the output of each layer is essentially the same the output of the previous layer in the sense that the rate of change from one layer to the next bounded (at least locally), and hence the network is intuitively "wasting" parameters.

---

Published as a conference paper at ICLR 2024
Claim 5.4. Consider stable and nontrivial parametrization that is furthermore faithful during training (but not necessarily at initialization). Then it is redundant if E (1/2,1].
To understand the intuition behind Claim 5.4 let us see what happens when a > 1/2. In this case, the randommess of the initialization weights will have no impact on training trajectory as depth increases. To see this, consider some layer index [AL]. The blocks are divided by L which is larger than the magnitude of accumulated randomness (of order (AL) This basically destroys all the randomness from initialization and therefore the randomness in the learned features will consist only of that coming from U and V (input and oulput matrices). When depth goes infinity, the contribution of the randomesss in two adjacent layers becomes less important, we end up with adjacent layers becoming very similar because the gradients to these layers are highly correlated.
In contrast, we have the following result, which defines Depth-uuP.
Claim 5.5 (Depth-uP). a = 7 = 1/2 is the unique parametrization that is stable, nontrivial, faithful, induces feature learning, and achieves maximal feature diversity with K 1/2.
Note feature learning is not the unique property of Depth--ul but in terms of feature diversity, phase transition phenomenon occurs when a 1/2. More precisely, for Depth-juP, we can show that xi ||\اا| (c2//2) while the same quantity is OEe) for all a E (1/2, 1]. which suggests that Depth-uP yields rough path for Tt. This allows the features to change significantly from one layer to the next, hence efficiently using the parameters. For readers who are familiar with rough path theory, the 1/2 continuity exponent is a result of Brownian increments in the path TO
n 1/2
[(++))L]
It
Moreover, with a 1, there a phenomenon of feature collapse in the sense that the features will be contained in the o-algebra generated by the input and output layers, but contains randomness from the hidden layers (see Appendix G.2). Intuitively, the case of a = 1 is analogous to width situation, where deep mean field collapses same). For depth, the features (layers) are still relatively different but the redundancy does not allow significant variety in these features.
single neuron (all neurons become essentially the
In Appendix J we provide more insight into feature diversity and shed light on its importance. We show that the choice of nonlinearity and placement of nonlinearity can affect feature diversity, and indeed the best architecture that maximizes feature diversity also performs better in the experiments. We also provide empirical verification of our claims about feature diversity in Appendix L.4
EXPERIMENTS
In this section, we provide empirical evidence to show the optimality of Depth-uP scaling and the transferability of hyperparameters across depth. We train vanilla residual network with block depth 1 (1 MLP layer in each residual block) on CIFARIO dataset using Adam optimizer, batch size 64, for 50 epochs. The parametrization of the network and the weight update rule are as follows,
x' _ + a x L ())'"W"''')),, W w n x n n- L- " go.................... .n....,,L' ,nL' ge):
where the learning rate 1) and the block multiplier a are the hyperparaneters The nonlinearity & is ReL.U. The values of a, y depend on the parametrization of choice. For Depth-puP, we have y 1/2, and for standard parametrization, we have 0,7
Learning rate transfer (1)). In Figure 2 we show the training loss versus learning rate for depths 2*. for k (3.4...,10)) For Depth-uP, convergence pattern can be observed for the optimal learning rate as depth grows. Optimal learning rates for small depths (c.g. L = 2²) exhibit a mild shift which should be expected, as our theory shows convergence in the large depth limit. However, starting from depth L = 26, the optimal learning rate is concentrated around 10-3 With standard
Stable, nontrivial, and faithful parametrizations all induce feature learning
The reader might ask whether we can obtain an expoment smaller than 1/2. This is indeed possible, but will entail using correlated weights. We leave this question for future work.
Note that y here is constant across depths and widths, and the effective learning rate given by 1yn -I In standard parametrization, there is generally no rule to scale the learning rate with depth, and the optimal learning rate is typically found by grid search. Here, we assume that in standard parametrization, the learning rate is scaled by T. to preserve faithfullness.
8

---

Published as a conference paper at ICLR 2024
Depth jaP
1t0p=2:1000 logu(a)=1.00
slsp=37000: 35000, 150
2
2
0.5
-10
Logloss
Loguose
-15
-20
0.52
-25
0.90
0.48
-35
Logsin/1e 11 3)
Logain 3)
SP (cr 0,7 I)
xtep=0-1000 1000 (agy(a)e-3
stop=57000:39000 Loggla)=-30
0.5
00.
Log.o05
-1.6
-25
0.5
-3.5
Log,in (1-33)
SP (cr =y == 0)
step-0:1000, Logglain-4. 4.0
dtep-3700036600, 40.
0.90
2',
L
-05
0.85
-66
080
20
0.55
-25
-0
-35
Loggln/Te-3)
Logadn (te-3)
Figure 2: Training loss versus lcarning rate at varying depths for Depth-juP (top row) and two Standard Parametrizations (boltom rows). The network consists of MLP blocks with width
256 and block depth 1. We fix its input/output layers and train its hidden layers for 50 epochs with batch size 64 on CIFARIO 10 using Adam. Each curve represent the average training loss over a time slice of 1000 steps for depths 2 for Confidence intervals are based on 5 seeds. The results show that Depth-udP preserves
k E f3, 4,..,, 10}. optimal learning rate while consistently improving the training loss depth increases.
parametrization with learning rate scaling (a- = 0,7 1), the optimal learning rate exhibits significant shift with depth and training loss degrades when the depth is too large. Here we have already set 2 making the curves look better compared to the standard practice with 1. For standard parametrization without any depth scaling (a = = 0). the optimal learning rate exhibits significant shift as depth grows (even with smaller = 2 4), suggesting that standard parametrization is not suitable for depth scaling. Additional figures with multiple time slices are provided in Appendix M
Block multiplier transfer (a). In Figure 14, we investigate the stability of the hyperparameter a in Depth-uP as depth increases. The results suggest that the optimal value of this constant converges as depth grows, which suggest transferability. Additional experiments with multiple time slices are provided in Appendix M
7 DISCUSSION
This research contributes a foundational understanding of the optimal depth scaling for neural networks, showcasing the effecttivness of Depth-uP through rigorous theoretical justifications and practical experiments. As the field of deep learning continues to evolve, our work serves as a valuable guidepost for developing more efficient and powerful neural network architectures while also paving the way for exciting future investigations into the scaling of advanced architectures.
While our current theory establishes the optimality of depth- in the case of MLP blocks, the applicability of this scaling principle to more advanced architectures, such as Transformers, should in-principle yield favorable outcomes. We leave this for future work.
9

---

Understanding the Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
We thank Huishuai Zhang, Jeremy Bernstein, Edward Hu, Michael Santacroce, Liyuan Liu for their helpful comments and discussion. D. Yu was supported by NSF and ONR. Part of this work was done during D. Yu's internship at Microsoft.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization, 2019.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport, 2018.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming, 2020.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture, 2018.
Soufiane Hayou. On the infinite-depth limit of finite-width neural networks, 2023.
Soufiane Hayou and Greg Yang. Width and depth limits commute in residual networks, 2023.
Soufiane Hayou, Eugenio Clerico, Bobby He. George Deligiannidis. Arnaud Doucet, and Judith Rousseau. Stable resnet. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference Artificial Infelligence and Statistics, volume 130 of Proceedings
of Machine Learning Research, pp. 13322. PMLR, 13-15 Apr 2021. URL https:// proceedings mmrr. www.sepeesssvvi. haa1olllaayuaaa....m html
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 630-645. Springer, 2016b.
Arthur Jacot. Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks, 2020.
Samy Jelassi, Boris Hanin, Ziwei Ji, Sashank J. Reddi, Srinadh Bhojanapalli, and Sanjiv Kumar. Depth dependence of μp learning rates in relu mlps, 2023.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. difficulty training transformers. arXiv preprint arXiv:2004.0822 2020.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse, 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-widtt limit, 2023.
OpenAl. Gpt-4 technical report, 2023.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, L. Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Mare Lanctot, Sander Dieleman, Dominik Grewe. John Nham. Nal Kalchbrenner, llya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484 489, 2016.
10

---

Published as a conference paper at ICLR 2024
Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks, 2015.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation, 2020a.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture, 2020b.
Greg Yang. Tensor programs i: Wide feedforward or recurrent neural networks any architecture gaussian processes, 2021.
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 and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit, 2023.
Greg Yang, Edward J Hu, Igor Babuschkin. Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2003.03466, 2022.
Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization, 2019.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks, 2018.

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
A.1 WIDTH SCALING AND JP
The infinite-width limit of neural networks has been topic of extensive research in the literature. Numerous studies have predominantly focused on examining the behavior of various statistical quantities initialization. Some works have gone beyond the initialization stage to explore the dynamics of feature learning in neural networks.
Lazy Training. With standard parametrizatiin, a learning rate of order O(n n. being the width, yields the so-called lazy training regime in the infinite- width limit, where the reatures remain roughly constant throughout training (Chizat et al. 20200 Yang et al. 2022). This regime is also known as the Neural Tangent Kernel (NTK) regime and its convergence properties have been extensively studied in the literature (Jacot et al. 2020 Allen-Zhu et al. 2019 Chizat & Bach 2018 Zou et al. 2018).
Feature learning and J.P. Recent empirical studies (e.g. Yang et al. (2022) have provided
compelling evidence that feature learning plays crucial role in the success of deep learning. It is widely acknowledged that the remarkable performance achieved by deep neural networks can be attributed to their ability to acquire meaningful representations through the process of
training. Consequently, scaling the network architecture emerges as natural choice to enhance the performance of such models.
In this context, μP (Maximal Update Parametrization), introduced in Yang et al. (2022). has emerged as a promising approach for maximizing feature learning while simultaneously preventing feature explosion as the network width increases, given a fixed depth. Notably. jaP facilitates hyperparameter transfer across varying network widths. This means that instead of tuning hyperparameters directly on large models, one can optimize them smaller models and utilize the same set of hyperparameters for larger models.
The derivation of μP leverages the Tensor Programs framework (Yngg) 2021 2020a b Yang & Hu
2021 Yang et al. 2022), which provides valuable tools for capturing the behavior of neural networks in the mnfinitewwitti regime during the training process.
A.2 DEPTH SCALING
While increasing the width of neural networks can lead to improved performance, increasing the depth of the network also yields significant performance gains, and most state-of-the-ar models use deep architectures. The introduction of skip connections (He et al.) 201 b) played a pivotal role in enabling the training of deep networks. However, it became apparent that even with skip connections and normalization layers, training deep networks remains challenging task (Liu et al.) 20220). Moreover, tuning hyperparameters for large depth networks is time-and-res ree--aandd-------oomeee consuumin task.
To address the challenges associated with training deep networks, several studies have proposed scaling the network blocks using a depth-dependent scaler to ensure stability of features and gradients at initialization (Hanin & Rolnick] 2018 /Zhng et al. 2019 HHaou et al., 2021 Noci cct al., 2022 Hayou 2023 Hayou & Yang 2023 Noci et al. (023). However, these works primarily Tocus stability at initialization and lack insights into the dynamics during the training process. For instance, one might argue that features can still experience explosive growth if the learning rate is not properly chosen. Therefore, effective depth scaling approach should not only ensure stability at initialization but also provide guidelines for scaling the learning rate.
This motivation underlies the development of Depth-uuP, which offers comprehensive framework for depth scaling. Depth-jP encompasses block multipliers and learning rate scaling, providing a complete recipe for training deep networks. In the case of Multi-Layer Perceptrons (MLPs) (no Jelasssi
skip connections), Jelassi et al. (2023) showed that learning rate scaling of depth -3/2 guarantees stability after the initial gradient step. However, it remains unclear how the learning rate should be adjusted beyond the first step, and this scaling is not suitable for architectures with residual connections.
We also ohtain the lazy infinitewwidth limit with the NTK paramterization and O(n -1/2) learning rate.
12

---

Published as a conference paper at ICLR 2024
B CONTINUATION OF THE INTUITIVE EXPLANATION WITH LINEAR NETWORKS
Recall the network model
VI [L],
d'=
W VL
where xº = US, resp. ((5) = VT is the input, resp. output, layer, and W E R are the weight matrices, and U, V are input and output weight matrices that we assume be fixed during training.
B.I CONVERGENCE WHEN DEPTH GOES TO DC
Let us look at of again in the simple case din = dout = n 1 and analyze its behaviour when I, + co. This paragraph is only intended to give an intuition for the convergence. A rigorous proof of such convergence will be later presented in the paper. Let us consider the case with SGD training with learning rate v and let Mr,1 II (1 J.. w.) and =(Vrg v(G)))= With this, we have the following
L
of
II 1=1
1 VL ETMM, x..
W
WLOG, let assume that =0 O. Then, with high probability (the event that W VL, for some notion of "<"" occurs with probability of at least 1 e- 4 for some > 0 have that of > O. We can therefore look at log(xt) which simplifies the task. Taking the log and using Taylor expansion under high probability event, obtain
M + El +O(L-1+0
1og(z//20)
W
VT.
I
L
1=1
It Et=(w)))
Σ
w
Txn
+
+ +LL-1t+))
VT. 1=1
L
W1
L
1=1
VI
for some c > 0. The first and third terms VT. ΣLe
Wi and (w)) converge (almost surely) to
standard Gaussian and 1, respectively. The second term also converges naturally, since af converges in Le to Log-Normal random variable (Hayou (2023) and with delicate treatment (involving high probability bounds), one can show that the term E I+ W converges (in La norm) at large
depth. This implies that one should expect o to have some notion of weak convergence as depth grows. Note that the same analysis becomes much more complicated for general width > To avoid dealing with high probability bounds, convenient method consists of taking the width to infinity first + co, then analyzing what happens as depth increases. We discuss this in the next section.
B.2 A DISCUSSION ON THE GENERAL CASE
Difficulty of generalizing to the nonlinear case. The extension to the general width scenario intricate treatment of the term A₁ to find optimal scaling rules, yet the
(n 1) necessitates proposed scaling remains optimal for general width. This preliminary analysis lays the groundwork for proposing a specific learning rate scaling scheme that maximizes feature learning. Moreover, demonstrating the optimality of this scaling strategy in the presence of non-linearities trivial task. The primary challenge stems from the correlation among the post-activations induced during the training process. Overcoming these challenges requires a rigorous framework capable of addressing the large depth limit of crucial quantities in the network.
non-
For this purpose, we employ the Tensor Program framework to investigate the behavior of essential network quantities in the infinitewwdtth siittt-wiith-haeee width-then--eeye then-depth limit. By leveraging this framework, our 14This follows from simple concentration incqualities for sub-exponential random variables.
13

---

Published as a conference paper at ICLR 2024
theoretical findings establish that the aforementioned scaling strategy remains optimal for general networks with skip connections. Our framework considers the setup where the width is taken to infinity first, followed by depth. This represents the case where < depth < width, which encompasses most practical settings (e.g. Large Language Models).
The critical role of Initialization. A naive approach to depth scaling can be as follows: since the weights W might become highly correlated during training, one has to scale the blocks with I/L. To understand this, let us assume block multiplier of L" and consider the scenario of perfect correlation where all weights are equal, i.e., W W for every k E 1,..., L. In this case, the last layer features can be expressed as x (I + L "-) zy. When 1/2. the features are likely to exhibit an explosive growth with increasing depth, while opting for a = is guaranteed to stabilize the features.
However, in this paper, we demonstrate that this intuition does not align with practical observations. Contrary to expectations, the features do not undergo an explosive growth as the depth increases when 1/2. This phenomenon is attributed to two key factors: random initialization and learning rate scaling with depth. These factors ensure that the weight matrices never become highly correlated in this particular fasshion during the training process.
In summary. while a naive depth scaling strategy based on scaling blocks might suggest the need for to stabilize the features,
findings reveal that in practice, this is not the case. The interplay of random initialization and learning rate scaling effectively prevents the features from experiencing explosive growth, even with the choice of 1/2.
C CONTINUATION OF SGD TRAINING DYNAMICS
In this section, continue to study the linear neural network with residual connections under Dupth-uuP. We derive the training dynamics of SGD for this linear network when the width and the depth sequentially go to infinity. The road map of our analysis consists the following three steps.
1. We first take the width of the network to infinity by the Tensor Program framework (Yang & Littwin 2023). As a result, instead of tracking vectors and matrices along the training trajectory, we track random variables that correspond to the vectors, that is, for vector x E IR=" that appears in the computation of the training, the coordinates of can be viewed iid copies of random variable [z) (called a ket) when n co. 15
2. Since the network is linear, every random variable can be written as linear combination of set of zero-mean "base" random variables by the Master Theorem of Tensor Programs (Yang & Littwin 2023). Therefore, we can track the random variables by analyzing the coefficients of their corresponding linear combinations, along with the covariance between the "base" random variables.
3. Since the number of random variables and the number of "base" random variables scale linearly with L, the coefficients of all random variables can be represented by six dimensional tensor, where two of the dimensions have shape L. We then map the tensor set of functions whose input domain is [0, 1] [0, 1). Finally, we claim that the functions converge when L + oo, and identify their limits as the solution of a set of integral functional.
Formally, under mild assumptions we first prove that for any depth I,, when n 00, the training dynamics of SGD can be characterized by two sets of functions (Tr,,.. : [0, 1] x (0, (o.....
(... T......-1).see (x...__.....)..e (0.1) and {Cr,s,a : (0, 1] + [R), -1.. a0 where T is the total number of training steps. The network output time t is I, (1, in the infinite width limit, and the formal definition of I' and C' can be found in Appendix CC
The definition of [æ]) requires the coordinates of is (1) w.r.t. n, and [ιz) is trivial if the coordinates of o(1) w.r.t. n. Therefore, for whose coordinates not @(1), we normalize by multiplying polynomial of the resulting vector has coordinates (1).
''tther than the assumptions of Tensor Programs (e.g. the loss function is well cell-conditioned), the rest assumptions (e.g., the input and output dimension is I; the batch size is 1) are only for the convenience of the notations and can be easily generalized within the Tensor Program framework.
14

---

Published as a conference paper at ICLR 2024
Then we take the depth L to infinity, and find I and C' is the solution of the following functional integral in the limit as in Proposition C.4
Assumptions and Notations
Recall the linear network is given by
x² = US. x'= a
AI E [L],
w
1 x'-1
VT. f=V
For convenience, assume any notation to denote the same object but at t-th training step, c.g., the input at step is a single datapoint &x, the hidden output of (-th layer at step t is x/,, and the model output at step is fe. Let T be the number of training steps. Let & be the loss function absorbing the label at time t, and X be the derivative of the loss at time f, i.e., Xt (fte. Let dx, de,/ /dz,, and its normalized version ozt = n dal.
the SGD learning rate of W is 1. We add t subscript to
The Tensor Program analysis heavily depends the scaling of initialization and learning rate of U, V, W w.r.t n. In this paper, we use JuP as the scaling w.r.f. n since maximizes feature learning in the large width limit (Yng & Hul 202I). Without loss of generality, we follow Yang & Huj (2021) and assume the input and output dimension is 1, i.e., & R, E R. For a clean presentation, we additionally assume U, are frozen during training in this section and each coordinate of W is initialized with i.i.d. Gaussian of variance 1/n.
C.I WIDTH LIMIT UNDER pP
As the first step, we take the width of the network to infinity using Tensor Programs (TP). As briefly mentioned in the road map of the section, the TP framework characterizes each vector involved in the training procedure by a random variable when n + oc. For a vector x IK" that has roughly iid coordinates, we write [x] E I (called a ket) to denote a random variable such that x's entries look like iid copies of [æ). Then for any two vector 2, y E R" that have roughly iid Elz) [v). which
coordinates, their limiting inner product by we write succinctly as (alv) Deep linear network with SGD is a simple example for this conversion from vectors to random variables. As shown in Program 1 we define series of scalars (f, and X:)
be written as lim =Iy
and random variables ([U"), [nV), [at), [8x(), a 1 W sx())) using the ket notations. For better understanding, we provide a brief introduction to TP below.
Tensor Programs (TP) in nutshell. When training neural network, one can think of this procedure as a process of successively creating new vectors and scalars from an initial set of random vectors and matrices (initialization weights), and some deterministic quantities (dataset in this case). In the first step, the forward propagation creates the features 2 where the subscript refers to initialization, and the scalar fo. which is the network output. In the first backward pass, the output derivative Xo is computed, then the gradients day are backpropagated. (Since the coordinates of dab vanish to when n - x0, TP instead tracks its normalized version Sarg def n da(-) New vectors are created and appended infinity, vectors of size in the TP (c.g., the features z/e and normalized gradients dx() see their coordinates converge to roughly iid random variables (e.g.. [ar)) and [8x.) in Program [b. and other scalar quantities (e.g., fo and X²) converge to deterministic values (c.g.. fc and X in Program I under proper parametrization (μPP). The Master Theorem (Yang et al., 2022) captures the behaviour of these quantities by characterizing the infinite-width limit of the training process. For more in- depth definitions and details about TP, refer the reader to Yang al. (2022).
the TP as training progresses. When the width goes to
Now when we look back to Program I the definitions scalars and random variables should be clear (except for W and [[ dz()). One can find straightforward correspondence between those and their finite counterpart, for example:
fe corresponds to fú. and Xt corresponds Xx:
(xf) corresponds to a and [[xx.] corresponds to 8z/. (Recall dx; = n -dd is the normalized version of dsr(.)
15

---

Published as a conference paper at ICLR 2024
Program 1: Random Variables induced from Tensor Program for the Linear Network with LR ? = and frozen U.V. V.
Initial random variables: [U), (nV) are independent standard Gaussian. for t = O....,
do
(x-] dd [,0';;
for = 1,..., do
(W/=t-1) an W 1 J Es=o (ots
VI
(at) = lot + S II 1
end
ft def (af(nV);
) de V)
for 1 L.....11 do
[W: 5xt) del Iwo )z - 1 (ox2_
VT.
1 = [62]] + Jr [ Sx!));
end
end
where
(W/rt and [Wo 6xc)) are defined in Definition C.1
By SGD, W/ = w - + de's x x-- which corresponds to [W/x[-7) =
[W VT. E.= (xt (11--111--11
Now we can dive into the definition of [W and [W Sxt)- Let W be the set of initial random matrices of size n X n, i.e., (W) (W)..... W.',, and wV del (WT, W E W]. Let Vy denote the set of all vectors in training of the form W for some Then for every W SUUT and Wy E Vw, can decompose [Wy) into the sum of [W y) and [[y)) where [yy) is a random variable that act as if W were independent of y. and [Wy) is the random variable capturing the correlation part between W and y. Specifically, let us briefly track what happens to Whr! during training. In the first step, we have W6ap-1
(1 which has roughly Gaussian coordinates (in the large width limit). In this case, we have [W/2/ 1-1 O. After the first backprop, we have day dzb + 1 W means that the update in W will contain a term of the form W 2 for some vector This implies that W will contain a term of the form W!W:T z' for some vector 2'. This term induces an additional correlation term that appears when we take the width infinity. The [ [W is defined by isolating this additional correlation term from WoW:"= z'. The remaining term Gaussian in the infinite-width limit, which defines the term [Wart-") Formally, we present the following definition.
VL dag. which
Definition C.I. We define [Wy) def [Ww) + [Wy) for every W E WUWT and Wy E Vw, where
|| y) is a Gaussian variable with zero mean. VW E WUWT, Wy, Wz e Vw.
Cov (IW D) (W=] set (y1z).
W, W" E WUWT, Wy E Vw, W's E Vw [Wy) and [W'2] are independent
W 4 W. WW is also independent from [U] and [nV].
[Wy) defined to be linear combination of {[z) WT E VwT) ). Then we can unwind
any ([)) inductively linear combination of [..000 [U) and [nV), which allows fully
16

---

Published as a conference paper at ICLR 2024
define
Qly) O1
[Wy) def
Σ WT WT WT
C.2 DEPTHWISE SCALING OF RANDOM VARIABLES
As mentioned in Definition C.1 both (xf) and 18x)-1) can be written as linear combination of
"base" random variables: [Wo"""" (0s ( Sxx"))
) (ee((.....)...m.e [U) and
[nV). Moreover, the coefficients of the linear combinations can be calculated in a recursive way: by expanding [1 using Definition
C.1
we have
[x/) = -
J VL
W
1 VL
(0a")
a|c
1
(at 1
1ot
Till dz!))
VL
The recursive formula for [ōx{) is similar.
Using this induction, we claim in the linear combinations, the coefficient of every |.) O(I/VL). and the coefficient of (U) and (mV) is O(1). We also claim the covariance between any pairs of random variables in the form of [xx] and [ is (1).
Proposition C.2. Vt, Vs S t, Vl, m, vly) (0x4)
oly)
Olu)
0lv) 0IU) 0n
= (1), ayy)
= 0
alW m
Vt. s,l, m, Vly) flaf), [6x))} }, v/2) E fI=T), |,00">>
(1 = (().
The reasoning of Proposition is provided in Appendix D Note the computation of covariance also be written a recursive formula. The reasoning relies essentially on inductive argument.
C.3 INFINITE DEPTH LIMIT
Now formalize our argument above and obtain the formula describing the dynamics of the network when - co. We first write the coefficients of the linear combinations as a six dimensional tensor T b.s.,a,b.I,. where i, 8 E {0....,
1), a, (0, 1,,,, I, m [T.]. Specifically, I D,b,,,, and [W?" 6æ.). Here, we use 0 to
represents the derivative of [x=]) and [8x{) w.r.t. [W x.
denote kets appears in the forward pass ([=))) and [Wa" x am 1), and to denote kets in the backward pass (60z4) and [Wam (x Formally, Dr,,,0,0.1.m
01
8(a{)
W -11 D1..0011.... Ow 6xm
0(62()
T f...1..... ajWr="" 2 I.S.I.I.. = alwa 6xm")"
However, is hard to describe the limit I because its size increases along with L. Therefore, we define the following set of functions (T (Tr, : [0, 1] (0, 1] + R)r 111,,. be, (0. For 2 0,
Te (p. 4) = VT.. T.,.,, b.....b(L.)..(
For s = 11 1,,01,00 (p. = 0)o(4o)
012[4pl, knv) P T., =
aiu; 11, (p.4) =
01 blly TT,11.1. (v. q) = alvV) 0
Here 1, normalized to [o, 1] the input domain of I's are identical for different L; I Ti,,,a,bll. multiplied by VL because Ir,. (L1 (VLL/L) by Proposition C.2 and the extra -1 case helps us also capture the derivative w.r.t. [U') and [nV).
Similarly, we can also define another set of function {Cr,s,a (0, 1] + [........... describe the covariance between the "base" random variables: Vp (0, 1) let / = =[Tp],
to
17

---

Published as a conference paper at ICLR 2024
Cr,s,0 (p) def Cor([W/a} Cor([ 1. (a) 1
Ct,,,1 (p) Cov([W 5x() TWo dx!)) = (dx//6x²))
For = -1, C-1,-1.0 (p) (U)( [U)) = 1, and C-1,-1.1 (p) t Co (nv)) = 1, By
Definition C.T] the "base" random variables of different "groups" independent, we only tracks the covariance listed above.
Using this definition I and C, convenient to write their recursive formula in the following lemma.
Lemma C.3 (Finite depth recursive formula for I' and C' (Informal version of Lemma D.I. I and C can be computed recursively as follows:
De,r0.b (2:99
1 (-1 (((t=:/0--0(==1111
= b.,
9
9
L
t-1
[TeaOs 1-1 ( Cr (t))
+
r ambb
(
L ==0
1-1
1
Tir111b
L .9
Tim.b
(.9)
L . ("
(Tisio ( (T.T.) Cr,s.1 (
T
1
[ Te aa q)dg.
Cr, Σ Σ Σ
1'=-10== 10-10'==1be0000 10.1)
where 1 = [[p] - a 0, 0, and Lp1 1f
The proof of Lemma C is straightforward from Program In Appendix D we also give formal proof that I and converge when L grows to infinity, the case where is powers of 2. The restriction on L being powers of is imposed for the convenience of the proof, and the convergence of T and is true the general case. Moreover, we derive the infinite depth behavior based on the recursion of T and C' in Lemma CK
Proposition C.4 (Infinite depth limit of T and C' (Informal version of Proposition D.2 ) In the limit + 0o, we have
Te..o,(p..) g) =I(=-n) +
Σ T T...((1'.9)) (aaaao//))-((T.a...p.. (Tr.... Ce,,,((p)))d )dp':
Pe.1b(p. g) = [((t=))((==))((ps))) +
E Temonll' 2) (Trssolp' .')
(p) = Σ ΣΣ> Σ
/ Ter o
(Q)) a.st aa g)dg.
(=-13'==1bE(0.. (0.
The proof of Proposition C.4 follows from Lemma A rigorous proof requires first showing the existence of solution of the integral functional satisfied by the couple (T, The solution typically fixed point of the integral functional in Proposition C.4 After showing the existence, one needs to show that (T) converges to this limit. This typically requires controlling the
difference between finite-depth and infinite-depth solutions and involves obtaining upper-bounds on error propagation. The existence is guaranteed under mild conditions on the integral functional. We omit here the full proof for existence and assume that the functional is sufficiently well-behaved for this convergence result to hold. The formal proof of the convergence I and C for L = 2 (k N) in Appendix D is showcase of the correctness of the proposition.
18

---

Published as a conference paper at ICLR 2024
D DETAILS OF THE LINEAR CASE
D.1 PROOF SKETCH OF PROPOSITION
Here we provide proof sketch of Proposition C.2 I and C' in the infinite depth limit.
the formal prove is implied by the existence of
Proof sketch. The claims can be reasoned by induction on and Let us take [af) as an example, since [ is symmetrie with [zr]]- By expanding the definition of [x+]), we have
[x/) = |x 1
VL 0
5+
1
VL (ot 1ol
VL
Note by induction, (c4 = (1) and
Bm
OL1/VD), so
W bx!))
(xd) -1o +
1
W W
) + +0 (7))
Σ (00!))
VL
W at + (
= E.10> +
Σ Σ
16
VL
m'=1s'=1
0 (+z).
Then by unwinding 1oxy' and noting that by induction, Vs t.
x² 15
)
8 88877'2 oye
O (1), 01227' JnaV)
(1), we have
B[W )
a|
O 1 (1z)
0(x))
O
0
JnrV 0
amWa damr)
VL
0
Also by unwinding, v(v) E {1-))) {6x;")))
Oly)
(w/+)) = ΣΣΣ
fam
(amm-11
(am''1laa"" 1omm'" -1)
Ol xe
ml x
m' o'
alr!)
aiwan dangy (Boo (oxng")
0|+1)
aly)
+ 0|a)) aly)
+ oule Uile VllV) 0lr)
=0(1).
D.2 FORMAL RECURSIVE FORMULA OF T AND C
By the same way of expanding [Ir]] and (wl=?), we formally derive the recursive formula for I' and C' below.
Lemma D.1 (Finite depth recursive formula for I and C). I' can be computed recursively as follows:
Fort = O,...,T 1,
V (0, 1T..11.0g Er.
19

---

Published as a conference paper at ICLR 2024
For 1 I,....,,. L, Vr I Vp E (77 L. Vq E (0, vb {0.1].
1-1
CLa,o(p)= =
Ter.o. 1 Cu ( 3,66
L .9
dg;
===1 8 ='=- be (0.1)
1-1 I .9
Penos (p,g) = Temob
+ M(0=r)A00=000011110
T.... (1.9) (Fesos
(
(1-1 1) ( Cr,s.0 (1)
L L
==0
f. (1,11,
X ((f...
Vg E (1.1,,,,,1,,,11000) =
For L....,,, vr st. E (Z. Vg E (0,1), Ab e {0,1}.
Cra.1 pt
1 1)
Σ Σ Σ
I Ternusl/T. 2.00// g)dq
p==1 8' 16(00,1
Te (p, = -I 1,, (t.1a)
+ [((=nnA(=D))==TT0
('-1.9) (Tis.1
( Cr, (
+ Σ T..
The proof is straightforward from Program infinite-depth behavior.
| The recursive nature of I and C yields the following
Proposition D.2 (Infinite depth limit of I and C). In the limit - cc, we have Vp [0, 10,6 E 10,
Tr.-100.00.0 q) Eti
Temoo(p. 9) 110=::00=0)002000 + J Σ Tar.Lo ()) - (Tr,,,, (../))
t-1
Cr..o(p))
j. = (1, d. 6-=60 C(fo::
Tt.-1i1(1,g) = =Xi:
PeM.1,(p. g) I::=:)
/
Σ
Tano,ll' .)) (Tr..ool .F)) (p))) 4=0
(p) ΣΣΣ (==-1'==-16(00.1) be(0,1)
(
D.3 CONVERGENCE OF AND WHIEN = 2*
In this section, prove I and C will converge when L + For convenience, will only consider the case when L = 2* for some integer k. To distinguish I and C' corresponding to different L, we add the depth as the superscript, i.e., 1 and CL. Theorem D.3. VL T, < 1 E (0, (0, 13. Vp (0, 1),4 (0,
(p. kers a Cauchy sequence,
20

---

Published as a conference paper at ICLR 2024
( (P))kN is Cauchy sequence.
The proof is by induction on t. We will prove the following claims (A) (B) (C) (D) on t given they are satisfied for any < L. For = 0, (A) (B) (C) (D) are trivial.
Assumption on <t Assume 3e >>> such that VT. >> I' and I. = =2 f vr <x..
(A) vp E (0, 1. 1......1..99 (0.
|T T (p.9) T a 4)) S se/L,
CC - C C.r (0, a)l S Sel.
(B) I s.. 4) S IC (p) cc
(C) C (pra(p) (p) is c-Lipschitz w.r.t. and I' I (p, q) is c-Lipschitz w.r.t.
(D) I s. (p- pt +) - o 2. S c/L, I 1, - too p- t)1 S c/L.
Remark (A) indicates that (Tat..,pl and (Cr h converge. We only care about r < s because
C 8,3,a will never be used, and r a, is known: for {0, ......1}.
A) pz. (e.,) I[(a - A A 2 10- II(a=
Proof for /-th step (the forward pass) In following subsections, we will prove inductively on increasing order of all L L' and L 2h, and increasing order of {0. 1/.1/ 1) that Vs < t,
(DO) Ih o (p.p+ + ) I 0, (p.p
+ + + expl(cpp)/
(CO) For s < l, [ t, (p.)) r (p t. 4)1 l00g explei f))/L:
(BO) T (p.. 9) S c2 exp(e ));
(A0) 1//2 (p, 9) 12 (p. 4) s c3 explc( (p)
(CI) OP ) - of (p) s c102 exp(ci +))/
(B1) Cl + } < c2 exp(cip);
(A1) 10,00 ) C. o (p + / S c5c2
where 02 ma [&l] exp(c: (c2/222/11 22)), c3
3ct, Ca 4t(t 1)c² + 2tc, cs catla
tat(Act + 2c, + 29) + tc(3c, + 14) + c(2c + 2c).
Proof for t-th step (the backward pass) induction on decreasing order of p.
Similar bounds also apply to De.. and Ct by
Conclusion Combining both backward pass and forward pass at time shows (C(()) also hold for s with larger (but constant) Thus, (A)())C(() hold for any constant by induction on training steps.
D.3.1
T (g, IN FORWARD PASS (PROOF FOR DO, CO, BO, A0)
We first consider
(p- p ( 9 + [[(t=r)A(b=1)) (Lp = [Lg]]]
1
r Th.. (p.9) = PL
1-1
(Ptho1 (p-
-Ct (0))
+ Σ
L. (p,9)
P) (p-
==0
21

---

Published as a conference paper at ICLR 2024
(D0) Difference between I (p. + 3) and o... + 2) Assume p 1/L (p 0 is
trivial), let q = ++ 1/L, q' = p 2/L, note that r s.. 1 s. 9') since p + 1/ v q S q', so for < t, r
Ti, (p, - -r, (v.4))
p I - Th,,,006 (p-t.1)
1
$ [T.
I 1, (p.g) - -I (p.0') [
p
1 T.P)
-C, (p)
+
1
1 7
S 02 expl(c
11/1
1. c/ 202 explei (p-
I
1
(1+ 2ct
1/L5
cxpl(cp))
as C1 2 2ct.
(CO) Lipschitz w.r.t. p For r <t,
1
(roo((p.g) Tiroslp I 7.911
= - Σ
p B.i (p.0) (Tisno (p-1,10
1 -) - Cheo (p)
1
(x2xxp(116 + + e2 exp(cile +
$ I
e
- 7
ctc exp(ci
(BO) Bounded Again assume 2 1/L (p 0 is trivial because 2 1 exp(er22)))) since IT. 1 t. S 02 exp(ci +)), we can bound I (p.
(t..)))) 2) S <c2exp(< 1 dez dl2 ---00
c2 cxp(ec(p - L )) + ct/L)
1
$ c2exp(e1(p 2L ))
as long C1 2 2ct.
(A0) Difference between L and L/2 bounded When = 0, it is trivial. When p = 1/L, it is also trivial by Lipschitz w.r.t. p. which results
Teros (p. q) 1 o.b. (p. s 33te2/L S c302 c302 exp(cx/2L)/L
When p 2 L/L, since
1/22 1./2 tr,,,,6
Gr,0,6 (p.g) =
(p
p I
2
+
2
Σ
1/2 1,
(p.9)
(Tr..0. (p
2 p) I P c0)
L
s=0
we compare it with L'
trr.ub (p. g) expanded based on its previous two steps
Th. (v.4) = = 1r..0
p
2
T,
+ I
Σ r
Am (1,0)
(Thaui
1 ( (p-
T. (p-T.P)
-t.. (p)
s=0
L
1
rb
1,b
p 1
(Thsoo
p-
--1---
-
2 1
1)
L
Cra0 (p p-
==0
22

---

Published as a conference paper at ICLR 2024
In order to bridge the two above, namely matching the inputs for I and C, need a middle term
FL GroD (p. -TT= 1,0000 (p-2) 1.
(0-
2
1-1
T I...I. (0,
(Tiao: (p-
2 p. - Ci,a, (p)
L ==0
Now bound [T, 0,,000 (p.9) - Yiros (p. 9) and [T (0,9) - 1L/2 - 1r00000 (p. q)) separately, which add up to be the bound for 1r,,0 (p.g) r (p.9))-
I Lro (p, 4) T2 1,.,
1 r 1. (p. 9)) 1 (D-
2 p- (D- T
- 1,s0111
rhe smb p --
(ri (Nisou)
2
1
p L. )
1
/ -=:P--
C,.000
2
Chi2, (0) 60)
-r sm1. (p.9)
p-
I I:P)
2
2t e/ -2 exp(ci (p - )
< et ctc, exp(c1(p
11/4 +
L
L
L
1
2 ct c102 exp(ci(p 7 ))
+
at 02 exp(ci1p
+
L
L
2²2
3ct + cacl
L C2 exexppe
- 1).
and
1//// tr0,b (p.q) - P'L
t,,0, (p.9))
$ 112/22 tr,06b
(p)
2
- I tr,0,bb
p (p-2 2
(p-
p [.9)
L22
(p)
2
-Iis.0011 pp- (p-2 (.P)
c.. (p) + Cle. (p)
1,,,0,1
p
(.P)
L
2
+ ++C (( (1))
[1..00011
p
TIP)
(c302 exp(ci - 1 ct(es + es)o exp(ci 11/ L 2t i 02exp(ci(p- - 111)
Ga + clcs + cs + +2/LL ex2exp(c 1
L
In sum, as ei 2 2lestestattarts)
3
L122 tro. (p.g) Ti (p.0) C3 + es
+ +o c2 exp(cc )
L
1
S <cpexp(ci0p
2T 1)/L.
D.3.2 Ct.,o(p + 1)) IN FORWARD PASS (PROOF FOR C1, B1, A1)
Now consider Cla By expanding
Cha,ol++
3 f==-1zz-16e bE (0.1)
[
T Tt. (4)) & a, (p, g)dg.
23

---

Published as a conference paper at ICLR 2024
will have
[-1
Chaoll+ 1
E
E
E bE(0,1)
ae q)dq
Σ
Cl
(00) s. pp g)dg.
s'=0
(CI) Lipschitz
Since Car,,
and l%
,,, are bounded and Lipschitz,
Ct + + Ckx,ol))
1 1.9)
Σ
Σ
9) e'dg
(E(0)))
Σ
Iri 1 1. c. dq
1 ..91.c. L
(bE()
Σ L
10
+ Σ L 10 ddg.
#=0
$ 1//L (20(s 1) dee explarlp 1) 6² + 21(x 1)ez exple(( - + à
+ 12.22explell (v + c 8- c2expcex 1
(41(s + 1)c² + 2sc) IL expla +
I cxce (p - + 1//
(B1) Bounded Since ICt,,, S c2 explci +)), we bound Ch + +) as:
CC + + S exple 1) (1 ca/L) c2 exp(cip),
as long as Cj 2 C4.
(A1) Difference between and L/2 bounded It easy to see that for p
we will prove that for p Ch. + +
CL./2 -C,./2 + } =0 0,
1C (p+ } C//22 (p } s 02 expl(xpp))
1
Then by (CI), for p
1/T),
ICL,,(pp + I
1 }
c (p= + I
) < exp(c( I )) 1
er c2 explarl expaaa 1/ < (ca + l)c2 exp(c1p)/L
=302exp(ccp)
Suppose p E (2/L,4/T, 1}, we compare Ct, +) Cha,olP 2) and C/ - CLI2 t). Intuitively, both of them are (IIIT.,, and their difference is O(1/T.2.. In particular,
24

---

Published as a conference paper at ICLR 2024
both of them can be written into four parts:
Ct + 1
ChsolP -
[-1
2 2:00)
Σ Σ
Σ
Ct Ct s0, qdq
(E))
bE(0,1)
[ Theossl 2
2 L. 0t (Fa.r.oll.9)) -r. ,o L. 2.00 dq (E2)
Σ
Σ
bE(U.1)
Cl a, o)
(E))
Σ lo
2 L 9))dq
C ( 5 (.....)) - 0,,0p
(g
and CL (p -)- C/ 2)= = + E 1/122 + E//² + & E/// where E/ defined in the
same way E but with C and 1/2 instead of C and r Next we bound | E one by one:
1. The only hard part to bound in |8! - c///311 is
(..., - F o. L - (T...o0(p.)))
2
(T
T...0,o0e L 9)))
LL22
By almost the same proof of (A0).
I. ②) rt. op
L 9) (T (.
1//2
10000 L 9))
2
dles + + at + ct + 5)
C2 exp(ci(p 71)
122
Then we have
18 - c4/2 / 11)
cl(cs + cs + dl + +5)
c2 expexp(1( 1 c.c
I
1
+ 4eto2 exp(c(
c/L.c.
L
+ 4etoz exp(c1(p
- I)/L.
c's(cs + cs + ct + + 13)
exp(c((p }
L
2. Bounding E
E2
cL/2
is similar
1 cL/2 where first bound
Ih o. (p.9) PL
2
-r (T (p.9) o :. -.9))
14/2
2 ()) S 9c2t/12
Then we have
1 cL/2 <L/21/(2(($ + 1))
< c30 63c2 exp(ci - /)/L c. 2c/L
+ c2 exp(ci
c/L 20/L
2
ex2exp(c1(p
96 1/21/12
( +2 + 9ct)
C2 cxp(c1(p 2 3)
<
L
25

---

Published as a conference paper at ICLR 2024
3. For E8 -//2,, we first simplify
cL/2
2
22/2
(p)) 2.000,
3
L
Le
5
and
Ege
1 L
Σ
oo(p.)) dk o. + olp
1 I
1
w o..... -/
Again, introduce intermediate term
E/ = L
Σ s'=0
Cla ,.0
Then we can bound
Ej
Eg 21//2
$ [ Eh
+ 13
Eg c/
2
[)/T c + 2)1 c/T.)
(cr02 exp(ci
L
21
7 }} /) +2 e2ep((c ) c/L) 1
+ T (c502 expeeexpci10 (p
te(cs + 1 + 20s +
02 exp(ci 1
1,2
4. For E - -ET/3/ we use
Ih. Tp 2.
2 I. - -T (.,.,,)))))
1/2
2 2,00,,,6
I I.9)) s 9//1/T?
which is used in 1 1/ Finally,
1 - -
2 20/L
2402 expe((p
[
L//LLLLLL
2
cxpcxp(e1
)) 9c
L
c(2c) + 9ct)
(2 })
02
L
In sum,
Chaole+ 1 Ct. Cp 1p 1 CL/2
1 C./2 }
(p
1
< e'((cc + 2c, + 14) + c2 + 15ct) + te(3c, + 5) + c(2c, + 9ct).
explei (p - ))
12
t'A(Act + 2c, + 29) + tc(3c1 + 14) + c(2c, + 2c)
02 exp(ci 1) ))
L22
Therefore, since ei @.(Acl 2cs + 29) 1o(3cs + 14) + c(2cs + 2c).
1 ) - CL + 1.
CCt.ooD+ +
1
C//2 - 7,11 + c1/LL2222 02 expllc )
< <ICt.ole
7
1 11
< (1+ + ci/ L/))ez2xp(cci
c exp(err)/L
26

---

Published as a conference paper at ICLR 2024
E NOTATIONS FOR THE GENERAL CASE
This section provides an introduction to the new TP notations from Yang & Littwin (023). We only require the definition of the inner and outer products in this paper.
Averaging over When x E R", we always use greek subscript a, .,... E [n] to index its entries. Then (Ta)a a denotes average entry. This notation will only be used to average over n-dimensions, but not over constant dimensions.
E.1 THE TENSOR PROGRAM ANSATZ: REPRESENTING VECTORS VIA RANDOM VARIABLES
From the Tensor Programs framework YYng al. 2022), we know that as width becomes large, the entries of the (pre-)actee Jactivattia vectors and their gradients will become roughly iid, both at initialization and training. Hence any such vector's behavior can be tracked via a random variable that reflects the distribution of its entries. While call this the "Tensor Program Ansatz", it is completely rigorous calculus.
E.1.1 KET NOTATION
Concretely, if æ E IR" one such vector, then we write [x) E (called a ket) for such a random variable, such that I's entries look like iid samples from ]rr any two such vectors y E RR. (Fa: Va) R2 for each a will look like iid samples from the random vector (Jx),, [u)), such that, for
example, lim z_y E [>] [1)), which write succinctly just (alv)- Here (xrl is called n bra, interpreted as a sort of "transpose" to [ar). In our convention, (ar) is always a random variable independent of and always has (11) typical entry size
This notation can be generalized to the case where x E RRnk y E R In this case, we can think of (ar y) as the j matrix given by ((zalv)) (1lo))))sss ISaSk.
1865)
Because we will often need to multiply ket with a diagonal matrix, we introduce shorthand:
[ = DiagX))
(4)
if a is n x l and x is a k-dimensional vector.
E.1.2 OUTER PRODUCT
Likewise, if both æ and have shape n x k the expression
[x] represents the limit of xy E =2
More formally, [x] is defined as operator that takes ket [z] E R and return the ket
([x) (= = (x(l)) e R
i... it returns the random vector [x) e I² multiplied by the deterministic matrix (y/[²) kg on the right. This corresponds to the limit of az 2 z/n. Likewise, Iæ=) (y] acts on bra (uzl E IR by
(w](()) (vl) = ((w)l=)) E R
which corresponds to the limit of Lw T T. T This definition of [=) (y] makes the expressions
n
[y])), (w]lx)(() (w|>><[<>)
unambiguous (since any way of ordering the operations give the same answer).
Remark E.1 (Potential Confusion). One should not interpret [xx) (y]l the scalar random variable
[x) 1y) = Ei [=>))l")>> which would act on a ket [x] to produce ((zl (wD)> = E([x) [[x)
Iz), which is deterministic. On the other hand, [x)) (yl 2²) is always linear combination of [x).
nondeterministic random variable in general. In particular, any correlation between [.)) and (ly] does
not directly play role in their outer product [æ) (yil: we always have [>) (u]=) [ (1)
where (Iz) [1 [ m) is an iid copy of ([v). [z)) independent from [æx).
1i.c., || / /n (1)) as
27

---

Published as a conference paper at ICLR 2024
Outer Product with Diagonal Inserted Finally, if RR is deterministic, then (consistent with eq. we define [x) x (yll as the operator that acts on kets [2] E RV by
R.. Morally, (x) is just a shorter way of writing (z) (z)g and represents the limit of xDiag(X)y In particular, [z) (y) [
([a)) (v))))
x1]=) [=>
E.1.3 NONLINEAR OUTER PRODUCT
xyy Rax= is the (linear) outer product of two vectors a E R" and y E R", then day ), the entrywise application of nonlinear IR + R to xy ", is kind of nonlinear outer product Passing to the ket notation, in general we define Φ (1ll) as the operator that acts on kets as
① ((=)
|=) def E db
Σ Io |
DZ)
where (Iv))_,... ..... | | an iid copy of (Iv"),... ly ). [[]) independent from [æ=) and the expectation is taken only over the former. This is just like, in the finite case,
(ay') z/n Φb Σ s'y' z/n.
Moreover, (w) E ER, R2,12) E R then
(wlo (1x) (vl)) [z] = (wll (z0 |z) Rix
Σ | |
(Iw)) @ l=)"
= E +
where & denotes outer product of vectors and expectation is taken over everything.
More generally, if Φ RR' -> R. then @ (ua)) (.,,) defined by
an operator taking kets to kets,
( (1x1)) (BIl..... [xx)) (xee (z)))
Σ MoiM
Σ (of) Iet)
Remark E.2 (Potential Confusion). Note <0[ (v1) is not the image of the operator (ax) (yl under ø in the continuous function calculus of operators, but rather a "coordinatewise application" of For example, if ((t) = t², then ([[)) is not [x)) (yl (y], the latter being what typically "squaring
an operator" means, but rather ( (yl2 x 2) (y vl.
1.1.4 COMPARISON WITH PREVIOUS Z NOTATION
For readers familiar with the Tensor Programs papers, this new "bra-ket" notation (aka Dirac notation) relates to the old z notation by
(x> Z", (xl1)) EZ
The new notation's suceinetness of expectation inner product should already be apparent. Furthermore, the old notation is not very compatible with multi-vecton whereas [æ)) makes it clear that) represents the constant dimension side. Consequently, (nonlinear) outer product awkward to express in it, especially when its contraction with random variables requires an explicit expectation symbol E.
F INFINITE- WITTH LIMIT WITH THE BRA-KET NOTATION
As before, when the width n of the program goes infinity, infer how the program behaves via a calculus of random variables. We define them below via the new ket notation instead of the earlier Z notation.
28

---

Published as a conference paper at ICLR 2024
Ket Construction. We recursively define the random variable [a)) (called a ker) for each vector and deterministic number for each scalar 0 in the program. For a vector Wæ in the program, we also define random variables [W z) and [Wx) (called hat-ket and dot-ker respectively) such that [W>) (Wx) [Wax). These the same as Z and Z in the old TP notation YYang et al. 2022 and they satisfy
Hat All hat-kets are jointly Gaussian with zero-mean and covariance 18
(ov(Ww)]\Uy)) I(W (>
(5)
Doz Every dot-ket is linear combination of previous kets, expressed by the following equation
0|x))
(Wa)
E
W
yEu
eq. 66) is the same equation as in( Yang et al.
2022
Zdot) but formulated much more succinetly in
the bra-ket notation:
Σ Z Oz azw y
Yang et al. 2022 Zdot),
zwz
y
There is an alternative notion for [W 2) in Yang & Littwin (2023) that write
[ a) = [W (WTx)x).
This is more convenient to write as we introduce the operator view.
We can see the ket (Wx) as the result of the action of an operator on the ket [z].
Definition F.1. Let W he an initial matrix a Tensor Program. We define [W]. [Wî, [wi to be the linear operators on kets that act by
(W]x) [Wx)
(W]x) de (wx)
[W]x) d [w( iwijx)
Any linear operator that is equal to [W] for some initial matrix W is called initial operator.
We also define the adjoint relations between the operators:
jwi' = iw
jwi' TW
[w]t |.
Parameter Update In the SGD case, the parameter update of W is simple. With the operator notation and outer product notation, we can write
[W = [w/] nloh() xx (x/-11.
In eq. 5) I(W = U) is the deterministic number that is I iff W and U are the same matrix (as symbols in the program) and 0 otherwise. This should nor be interpreted random variable that is precisely when W and U take the same values.
19 To be rigorous, we need to specify the "Hilbert space" of kets. This is somewhat pedantic and not crucial to the key points of this paper, but the Hilbert space can be constructed as follows: Let (T)) be the r-algebra generated by the kets of the program Let E(m) def U+ a(T) be the union (more precisely, the direct limit) of (xx') over all programs x' extending Then the Hilbert space in question is the I,2 space of random variables over the E of program.
29

---

Published as a conference paper at ICLR 2024
In this work, A denotes change for step, i.e.,
|AW 4+1 = --n]o)))) (1
Å denotes total change, i.e.,
1-1
|AW! = =- - n x fa-
which we write succinetly [&w - - -n x ( (Compared to Yang & Littwin (023). A
and Á are changed from and A because we want to use ó for gradients instead of d, which is now used for depth differentiation).
Note in the general case,
AAW = -nlöh' <1 xsx (x'sz')
where
Jöh st)
(a'st 1-1 (1((h) xa 1 16hh)) -
Xsi
So
2-1
= 17 Σ
löh'
xet (0 D.
(7)
7=0
For the rest of the paper, we write [W! = -n/5hh ax) (xc 2t
for convenience. The generalization
eq. 7 follows Yang & Littwin 2023
G HEURISTICS FOR THE PROOFS IN THE GENERAL CASE
The notation in this section is mostly defined in appendix The complete notation is defined in Yang & Littwin (003)).
G.1 DEPTH-UP
Let MS(x) = x - (x, 1)/n = Gr where G - - /n, where R", Recall the definition of the network and the normalized gradients
x UE
x a + VI Go(h')
f(E) VT at
6x nV
5h o' (1h) @ (Goz')
8z dx' + VT.
WiT 8h
where V (1//n) coordinatewise, dx (11) coordinatewise and W 0 1 coordinate- wise.
We also abuse the notation of G and use it as operator on kets: Glax) def (xr) R|xx).
30

---

Published as a conference paper at ICLR 2024
Forward.
Similar to the linear case, show that under technical conditions (mostly on the activation function) that the infinite-depth limit of the TP follows the dynamics
d|t)) = vaAGo (IWalu)) + vaNjaWMa?)
vaxGe
(1W)) dAGe' (IW6l=))) AW?
=VAGG (WWTTa)) + (Wo + dAGo' (1W6)=)1) (AW
= =VDAGo (Iwo)Ta1) dAGe' (1W6)11)) (iw3lat) + (AWM(N?)
where [o. 1] refers to the fractional layer index (A represents layer index [XL] as + oo), refers to the training step. [W&] the matrix operator (defined in Appendix (F), and the tilde symbol refers to the "normalized" version the object, i.c., multiply the ket with (dX) for some c such that the multiplication (normalized ket) is ((1) w.r.t. I, and same for the normalized operators. The first term represents a Gaussian noise.
In the linear case, we have
dl0;)= Vdd (IWoTa?)) + +dd (iwdla?)+1AW 13w>)u?))
Note
(Wbllt) vdx 88221 WAT ( Vdxioni.) on
Using multi-vector notation, we write
(áWMa))> = 50
(ar (x< |x,)) = 11
(8h2) X. (x/12?)
$<1
(AW = -nlonn
|1,) = -1
Σ
1811)) X. (1))
Backward.
Similar to the forward prop, we obtain the following dynamics for the infinite-depth
TP
-d(dx=)) = vaxiwA I @ (Gox;))
Vax (IWO i VANAAWAT) O (INGLEt) (INGLES)+ VaNIIN + vaNjaWAl>) ) (GBa?)
-TaiiWAT +dijwati o" (Iw {
+ vaxiwaT (d((W + +aNJaaA! ( (d([W>2;))
LAWAL [GBxt)
Here the (dx)3/2 term got dropped. The individual terms can simplified as follows
|AWill (d([W)=))]] -nn (8h
Gor?>> -nl (
iwaTi [<(W3>+)]] = ||
06
2| (GGa;,)
= - E
8' d'(
| R
2 W3
=(VdX)
where the other terms from the product rule drops because
O[Gox;)
0
0 0W
31

---

Published as a conference paper at ICLR 2024
G.2 1/L BRANCHES
G.2.1 FORWARD:
d|a;)= = dDG 0 (1W212)) +1AW) 11, vo)]
dAGB @ (IWTTa)) + -(WW><))) 1100. vo)
where the equality follows because [ is contained the o-algebra of [Uo, so) jwa(x?) 0. Since (WWAl ((Un, (o)),, Vo)), |&WM)a?) E a(U6, Vo)), and the expectation is really just over [Wo (=))-
G.2.2 BACKWARD
<lox;) = dAE [ (o("x>)) @ (G8z;)) [U, vo)l = dE E [A AT 1o' (W @ (G8x;)) 11V.. vo]]
Here the [W and [ drop out because the former is zero-mean and indepenent from [Uo, Vo) and the latter drops out because ( is contained the o-algebra of [U, Vo).
G.3 1/L BRANCHES, (1/2, 1]
G.3.1 FORWARD
dlt;) = (d)) @ (a))))))++ 11o. vo)
= dAE of (IWaT GLAW>1
because the reason as above.
G.3.2 BACKWARD
d (or7) = dAE |
10 (WAx?) o (Gox7)) 10 Vo)
dAE [[ [ [[W>"T (Gox²)) I1U, Vo]]
= dAE ['([W6 x21) (
Here,
1o' ( @ (G8x²)) = [o'( Glba2) + (dx)+-a
[ ((([WW
GLāW? (6022))
G.4 JUSTIFICATIONS OF THE CLAIMS
Claim 5.2. Stability during training when + is straightforward (some technical conditions on the activation function are required). This is because the weight updates are of order L-a-7 and feature updates involve no more than L terms of size L - (plus higher order terms that do not contribute to the update in the large depth limit). When > 1, the contribution of sum of at most I. terms of order T-a-- will decrease zero, the network output fa will converge to fo in this case, yiclding trivial limit. However, when 7 1, the updates remain important in the infinite depth limit, yielding a non-trivial limit.
Consider a stable and nontrivial parametrization (i.e. a ++ = 1). Faithfullness at initialization is achieved only when > 1/2. This was proven in Hayou et al. (2021) in a more general setup. Faithfullness during training is ensured as long as a because feature updates are always ((1) in depth. With 7 < 0 and the weight updates explode with depth in this case,
Claim 533
32

---

Published as a conference paper at ICLR 2024
which yield exploding behaviour for h.
Claim 5.4
When E (1/ obtain smooth limiting dynamics when L - demonstrated This limiting process is smooth process (no Brownian jumps) that satisfies the
in Appendix G.3
required definition of redundancy.
Claim 5.5 dynamics this case (Appendix G.I) With Depth-uP, the randomness of the initialization in the hidden layer remains present throughout training, inducing a Brownian-like term that breaks redundancy.
It remains to prove that Depth-puP is non-redundant. This is a result of the limiting
G.S SUBLETY: LAYERWISE (LOCAL) LINEARIZATION BUT NOT GLOBAL LINEARIZATION
Definition G.I. We say parametrization induces layerwise linearization iff each layer can be linearized without changing the network output when > oo, that is, vI E [L].
L (o(W(a) ( o(W(a(")) - - <'(W(=t-") @ ((W)- - w))= = o(L=
Claim G.1. A stable and nontrivial parametrization induces layerwise linearization iff E (1/2,1)
In Depth-uP, W W is ((1/VL) which is much smaller than W Therefore, @W!x
6( @Wijat o' ( to the depth-uP case, for a [1/2, 1), the activation in the forward pass can be linearized which indicates layerwise linearization when + + 1.
-1) ) @ ((W) Wo) o(1/ L/V,), thus satisfies Definition G.1 Similar
However, note that this does not imply the entire network is linearized (w.rt. all the parameters in sense Neural Tangent Kernel). our setup, where the widthwise scaling follows jP and the input/output layers are initialized at constant scale (w.r.t. L), it is actually not possible to have kernel limit. Even in our linear case Section 3.2 one can see the learned model is not linear w.rt. the parameters.
the initialization the output layer is L times larger than our setup (assuming L < n so the widthwise scaling still follows /P)) it may induce parametrization that linearize the entire network. In that situation, the learning rate has to be L times smaller than Depth-uP to obtain stability during training, so the change of parameters is also L times smaller, which can lead to the linearization of the entire network. Since we focus on maximal feature learning, the rigorous argument is beyond the scope of this paper.
33

---

Published as a conference paper at ICLR 2024
H CLASSIFICATION OF DEPTHWISE PARAMETRIZATIOOS IN LINEAR CASE
We discuss the classification results on the linear residual networks with SGD training and give rigorous proof's for the claims in this simplified setting. Recall the linear residual networks:
Al [L], = x'-1 aL aL=""
where h' = W and the effective learning rate of W is 1-- Without loss of generality, we assume 1) = 1.
H.1 INITIALIZATION
At initialization, we have
(m/)) 1 + L -hh))
where
(1b) = [W 1 = 11 11
Since | ) is independent from [Wi'bot) 57, we have
(=b(=))) (ab | (NGllh)) = (0b- 100-1 1000 1 -2a (1b 1ab = (1+1- (1+1-20)(00 (a00-11000) 1
Using this recursion, we can write
(=bla)) +1+ L 2a (=g (1)
Therefore, (xb1=;) = (11) iff a 2 1/2, otherwise (1 + L 22 as م explodes with large L.
A similar argument stands for 1ig and fo. Therefore, have proved Claim 5.1
Similarly, we can get the stability of the first backward pass, i.e., 8z ((1) for a 2 1/2. Given a 2 1/2, we can also settle the size of 5hg that
T = O(L
which implies
AW 1+ -116 38 a//1
H.2 AFTER THE FIRST STEP OF GRADIENT UPDATE
Now we look at the second forward pass, and assume the input is the same, i.e., (xf) = [=])) we have
1o+) = 1- + L (1 (Wo +
where |AWl| = - 1ohh (zb-)))))))) - -- (02g) (ab and (ōh _ I<15h is the normalized
version of (8h(), which happens to equal (ōog))- By the definition of [W and [Wor get similar formula to the Depth-uuP case:
b+t) = 12| [ E L 1002))
L ( (20-l=1-1)
Now write b' = Lx a1
and ạt - 10n- 1) then
o/le T 6m6
[++) = | + L "[W"(a -
17 + L (b
(b' + ct (ozt))
By expanding |ō2! ( - ( 8ate =1ōat) D L L- dam") we have
1.
(af) = 1a- + L W a 17 + a (b' + +c) | + Σ
L L Zmge)
m=l+1
= -1=1)
W 1 Σ --a-"(60 cth (oxbt)
min mmn 1.1)
(6' + d'
(8)
m=
34

---

Published as a conference paper at ICLR 2024
Note the four terms in eq. (8) are independent of each other.
Now it is easy to compute c' because only the first two terms in eq. have correlation with zb:
d = I, 20) = (1))
with a 2 1/2. For bf, we have the following recursive formula:
8+ = L 20x Σ
= -(1-1-20)))
mL
Stable during training and nontrivial. Finally, can reason about the fi (note fo 0, so Ai fi), which indicates whether the parametrization is stable during the first step and whether the parametrization is nontrivial for the first step:
20
I. i = (nViat) Σ L L +cc O(L
Therefore, we have proved Claim 5.2 that the parametrization is stable during training iff + 2 1, and is nontrivial iff a + S 1.
Faithfulness. Although there is no activation in the linear case, we still prove Claim 5.3 to enlighten the proof of the general case.
At the initialization, h/ and ar/-1 have the same size, therefore, faithfulness is equivalent stability, which means it happens iff a 2 1/2.
During training, we can expand [k1] in similar way to eq.
[Mt) 1 L I (b' + cf)
(00%)
Σ
W Bagg) = L
L
m=l+1
Therefore, it is faithful iff y 2 0. It is equivalent to < because we have a: +7=1.
Feature diversity exponent. first expand x dteL of
To simplify the analysis, we assume that eL always an integer. We
LeL.
[attcl) - (at)
Σ m=l+1
L L
m-
Σ
L-a + (oatt)
m=l+1
min Im l.+c))
Σ
L-a-y
Σ
(b" + +c I 1--0oo ōzgg).
--iin(m-1.1+++
With + 1, it is clear that the first term is O(L VeLL - the second term O -a+12), Therefore, there are only two cases here: if a 1/2, the overall size O ((2222 +c) = (c2/2)) if 1/2, the first and the third term vanish as L - oo, so the overall size is (c). In sum, we have proved Claims 5.4 and 5.5
has size (e)), and the third term has size O(VI.
Layerwise linearization. Claim G.1 is trivial in this simplified setting because layerwise linearization is always true for linear nets. To enlighten the proof of the general case, we recap that [ [AW|I=/->> O. If there were an activation function, the linearization would bring an error of o(L-") in n/, which means an error o(L-7-9) = o(L-1) to x/.
-=)e((02)) = (L-")) which is much smaller than [W +(1) when
We need AT and Ah for stability, but they similar Af
35

---

Published as a conference paper at ICLR 2024
H.3 BEYOND ONE STEP
The argument above is in general tracking the derivatives and covariance, in other words, I and C in the Depth-uP case.
Now we generalize Lemma
C.3 and obtain the following recursion for I' and C'
T1,r,0,6 (1.9)
1 + L 11/2
= Ti.no
9
L
(t.9) (T7-11/T 1.0.11
1 1-1 1 Cr,s.0 (
+ 1
L L.
==0
1-1
1tr11
L L .9)
= =Irrb
6.9
+ L
("z+
(1-11/10s.110
(TTT) Cr. )
+ L -La07
Temoo
#=0
(
Cr,s,a(p== = Σ Σ Σ Σ
1=-1s'=-1bE(0.1) (1caa
where / = [Lp] - 1f = 0, and = Lpp l
Then all the claims can be reasoned by tracking the order of I and C.
Distinguish parametrizations with + and 1/2 S a S 1. The parametrizations with a + = 1 and 1/2 < a < are all nontrivial, stable, and faithful. However, there is a large gap between = 1/2 (Depth-juP) and > 1/2 in terms of the difficulty of tracking I and C. For > 1/2, we can see that Ct,.,a = ((1. It. = @(1) and Tr,..a.b o(1) for 8 2 0. In this case, simplify the recursion by ignoring I T. with 0:
11,-1,0.b (
( Cr, (t)
= 1,0,
L
L.1.b
8=0
1-1 L
(1
(1
Cr, (i)
11,1
11,-1,1
Σ
8; 1,0,b
L
L
==0
Cr,
Σ (1-),a
bE10,1}
where / [Tp] -1 if 0, and 1 [Tp]] a= 1. Note Tr,-1,a, q) is simplified function that only depends on because r,, q) is constant when fixing p.
This simplification means the randomness in any W/J does not have an effect on the dynamics in the infinite depth limit the complicated functional integrals for a 1/2 in Proposition C.4 simplified to be ODEs when a > 1/2. This ODE dynamic also directly implies that the feature diversity exponent is 0 for a >> 1/2.
WHAT CAUSES HYPERPARAMETER TRANSFER?
In a popular misconception, hyperparameter transfer is implied by the existence of a limit. For example, the fact that μP transfers hyperparameters, in this misconception, is because of the existence of the feature learning limit (aka the limit), the limit μP as width goes to infinity. However, this is not the case. Indeed, there are plethora of affinttewwitth limits, such the NTK limit, but there can only be one way how the optimal hyperparameters scale, so existence cannot imply transfer. In stronger version of this misconception, transfer is implied by the existence of
36

---

Published as a conference paper at ICLR 2024
"feature learning" limit. But again, this is False, because there are infinite number of feature learning limits (where the y limit is the unique maximal one).
Instead, what is true is that the optimal limit implies the transfer of optimal hyperparameters. For example, the width limit case, μP is the unique paramctrization that yields a maximal feature learning limit. Compared all other limits, this is obviously the optimal one. Hence jP can transfer hyperparameters across width.
So far, there is priori definition for the "optimality" of limit: One can only tell by classifying all possible limits; it turns out only a small number of different behavior can occur in the limit, and thus one manually inspect for which limit is the optimal one.
Similarly, in this work, to derive a depthwise scaling that allows transfer, we need to classify all possible infinite depth limits and Depth-uP will turn later in the paper[ More interestingly than the width case, here we have multiple modes of feature learning when taking the depth limit and is important to discern which mode of feature learning is optimal. Thus, again, it is insufficient to derive any one limit, even with feature learning, and be able to infer it yields HP transfer.
be optimal in sense that we define
In appendix L. we provide experiments with 1/T. block scaling (a. = (1, 0). aka ODE scaling, which provably induces feature learning in the infinite-depth limit, but sub-optimal. Our results show a significant shift in the optimal learning rate with this parametrization.
J FEATURE DIVERSITY
In this section, we show that the choice of nonlinearity and placement of nonlinearities can affect feature diversity greatly.
J.1 GRADIENT DIVERSITY
Gradient diversity is an important factor toward feature diversity. Observe that the gradient &a at x continuous in 1 in the limit I. - cc. In linear model (or the pre-nonlin model, where nonlinearity is put before the weights), this causes 5/' = L-"5x' to be very similar between neighboring blocks. As a result (because the weights W receives an update proportional to öh @ x x²-1), in the next forward pass, neighboring blocks contribute very similarly to the main branch z'. This leads to waste of model capacity.
J.2 PRE-NONLIN LEADS TO POOR PERFORMANCE
For example, in Figure 3 for a relu pre-nonlin resnet (i.e. blocks are given by (' instead of () we see that although Depth-uP indeed transfers hyperparameters (as predicted by theory), the performance is dramatically worse than the post-nonlin resnet in Figure [T and depth gives no performance gains beyond 8 layers. Specifically, it is because 8h = L- oz like the linear case, and (áx!-1) is also similar between neighboring blocks. As result, the gradient of the weights W', proportional to Sh @ (z has little diversity compared to nearby blocks.
J.3 MAXIMIZING FEATURE DIVERSITY WITH ABSOLUTE VALUE NONLINEARITY
In a nonlinear model, we have 5h 8x @ (H+'.. Because / is almost independent from all other At * in the Depth-uP limit, of (ht' can serve to decorrelate the 8h', depending on what @ is. For example, if é is relu, then ø' is the step function. h' is approximately a cero-mean Gaussian in the Depth μP limit, that ø' o'(A') ) is approximately 0 or I with half probability each. This decorrelates 8h' much better than the linear case. But of course, this line of reasoning naturally leads to the conclusion that ø' sign would be the best decorrelator of 5h' and the maximizer of feature diversity (with & among the class of positively 1 I-homogeneous functions) - then 8h and 5hm are completely decorrelated for o m.
There are important nuances here that will be spelled out an upcoming paper. For example, if space of hyperparameters is not chosen correctly, then it could appear that no limit is optimal in any manner. For example, if one in (widthwise) SP. one only thinks about the ID space of the global learning rate, then all infinite-wddth limits are defective and indeed there is no hyperparameter transfer where the bigger always does hetter.
37

---

Published as a conference paper at ICLR 2024
(o. yl 12/2, 2/2), tttp://0000900001ggg.
slop=37000:39000
peroonn
-05
-1.0
-0.5
-10
-25
-90
-35
-9
10
Loggle/1o-81
Logg
Figure 3: Pre-Nonlin Leads to Poor Performance Although Depth-uu for prenonlin resnet indeed transfers hyperparameters (Left), depth gives no performance gains beyond 8 layers and the performance is dramatically worse than the post-nonlin resnet (Right). In right plot, the "Min LogLoss" is minimal log loss over all block multiplier and learning rate. Networks are trained on CIFAR- 10 with Adam. See Figure for more details about the setup.
5500-55000000000 logu
stop=37000:3900
(=22, km
30
7.0
4.0
8.0
50
9.
60
-2.0
-25
-30
-2
10
(oggle//m-3)
Logg1
Figure 4: Improving performance with absolute value non-linearity, which maximizes feature diversity. (Networks are trained CIFAR-10 with Adam.). See Figure [ for more details about the setup.
Indeed, as shown in Figure swapping in absolute value for & dramatically improves the training performance of deep (block depth 1) resnets.
In general, in lieu of absolute value, any even nonlinearity would suffice.
J.4 FEATURE DIVERSITY IS IN TENSION WITH LAYERWISE LINEARIZATION
The reason that d' (h')) can decorrelate 5h is very much related to layerwise linearization. Recall that Depth-uP, 1/' can decomposed to a zero-mean Gaussian part T/' of size @(1) and a correction term it of size ((L-1/3) (corresponding to the decomposition [k+) [h [ Te' is independent from hm for m # I but it' can be very strongly correlated all other htm Thus, d'(h') can decorrelate 5h' precisely because N dominates h², and this also precisely the reason we have layerwise linearization.
In the 1/L scaling (a, 7) (), Tr' is the same order ht and layerwise linearization does occur, but also (h') can no longer effectively decorrelated oh!.
Once again, remind the reader that layerwise linearization in this case is not detrimental (in this block depth 1 case) because N in fact accumulate contributions from the learned features of all previous blocks and thus strongly depends the learning trajectory (in contrast to the (widthwise) NTK case where TV is already determined at initialization).
38

---

Published as a conference paper at ICLR 2024
step-1 0000 11000, rMiL res LN
-1.0
5tag-37000 38000.
0.00
block d:pth
depth
-1.5
-0.7
Log.ogg..oo
1.00
-25
-9.0
-200
-35
10
Logy(Totalayers)
LayyyTosa
Figure 5: Block Depth 2 performs worse than Block Depth 1, Relu. In relu resnet with no LN, block depth does worse than block depth I when matching total number of layers (and thus parameter count). However, training longer (38000 steps, Right) helps it catch up (compared to 11000 steps, Left). The y-axis is minimal log loss over all block multiplier and learning rate
s0ep-10000: 11000, abs, with
slsp=37000:88000 abs, with
Nork depth
-1.
blo0k depth
0.5
-1.0
-25
-15
-30
-35
-25
-1.
-30
-4.
Logi(Tocak (oyyaaaakaayers)
Log, (Todal ayeers)
Figure 6: Block Depth 2 performs worse than Block Depth 1, Abs. In abs resnet with LN, block depth 2 does significantly worse than block depth 1 when matching total number of layers (and thus parameter count). Training longer (38000 steps, Right) does not close the performance gap (compared to 11000 steps, Left). The y-axis is minimal log loss over all block multiplier and learning rate
K BLOCK DEPTH 2 AND ABOVE
Remark on notation: Here and in the next section, all big-O notation is in L only; the scaling in width is assumed to be in p.P.
In most of this work, we have considered depth-1 MLP for g' in eq. ( it's straightforward to derive and classify the infinite-widh width-then wtwte-infnniteeeeppeee limits for larger depths in each block. In particular, the following I/VL scaling still makes sense in this more general setting with block depth & and leads to a well defined limit:
x' x.
(x(-'11, Wi ...., W ((1) initialization scale, 0(1/VL) learning rate VL
(9)
This is what we call Depth-uPii in the block depth I case, but we shall not use this name in the general block depth case because this parametrization is no longer optimal"
K.1 BLOCK DEPTH 2 2 IS DEFECTIVE
very clear symptom of this is that the performance of block-depth-2 resnets is worse than that of block-depth-I networks, when matching parameter count, although they can (but not always) catch up after training for a long time (figs. 5 and 6). Simultaneously, we are seeing nontrivial or even significant hyperparameter shifts as the total number of blocks increases (fig. 7).
What exactly mean by optimal will be explained below.
39

---

Published as a conference paper at ICLR 2024
yf (1/2.1/2) stop-31002:88000 log 2.0. rol,,
yl -0pp--57000::00000 logy
-05
-10
-15
-20
2
-25
-30
Logulg
Logg
Figure 7: Block Depth 2 Hyperparameter Shift in relu resnet with no LN (Left) and abs resnet with LN (Right).
K.2 DEFECT OF 1/VT. SCALING IN BLOCK DEPTH 2
The reason that the I/VL scaling is no longer fine the block depth 2 2 case is the linearization of the multiplicative interaction between the layers in the block. Indeed, just like the block depth case, the I.vT. scaling forces the weight updates AW of each weight matrix to be O ((VL) smaller than the initialization Wo. Thus, within the block, the training dynamics when depth L is large is in the kernel regime, where the contribution to the block output g(æ; W*) is only a summation, instead of product, of individual contributions from each layer's 's weights updates.
When aggregated over all L blocks, the result is that there is only multiplicative interaction of AW across blocks but not within layers. In other words, the network output is dominated, for example in the linear case, by the contributions of the form M4... ... M where each M' can be one of T, 2 W/2Awi or AW22W11 but NOT AW'2 All other contributions (which all involve within-block interactions like 2 AW") are subleading. In the general nonlinear case, replacing the block
with the linearized version
<(h')) + d' (1(h)) o JAW ((h)
+ + ('((()()))) (hh,) @ ( ④ (AW"-'-11
will achieve the same performance as depth L + oo, where h/, = W! (1²,) and h4 W 2'1 When block depth & (our main subject of study in this work), all interactions are included but this is no longer true when k > 1.
In fig. 8 the heatmap of loss function of block multiplier and learning rate demonstrates this vividly for block depth 2.
Small depth The optimal sublevel set of (learning rate, block multiplier) has slope R -2 when the number of blocks is 2'. In other words, around the optimum, double the learning rate while dividing the block multiplier by has similar performance. This is because AW and AW interact mulripticatively, so that doubling their sizes leads to quadrupling their contribution the block output. The simultaneous decrease of block multiplier by 4 then roughly keep their contribution invariant in size.
Large depth On the other hand, the optimal sublevel set has slope x -1 when the depth is 200. Doubling the learning rate while halving the block multiplier has similar performance. This reflects the fact that AW and A 1W12 now interact additively.
Intermediate depths interpolate this phenomenon, as seen in the plot for depth 25.
In the same heatmaps, one can see the optimal (learning rate, block multiplier) (in the 1/vL parametrization) shifts from the middle of the grid to the upper left as depth goes from 2 to 200 demonstrating the lack of hyperparameter transfer.
This change in slope is seen in relu networks well, with without layernorm.
40

---

Published as a conference paper at ICLR 2024
lop
12.
11.
12.
9.
05
73
03
2
53
43
34
24
13
03
8
8
8
Figure 8: The "slope" of the optimal sublevel set in the (learning rate, block multiplier) space changes from -2 to 1 as depth goes from 2 to 2 Here we use absolute value nonlinearity with layer normalization, block depth 2, and networks are trained for 50 epochs with Adam CIFAR- 10.
Finally, note that the 1/VT. scaling still yields I. -> limit where the network still learns features as a whole, even though within each block this is no longer true. Thus, this another reminder that mere "feature learning" does not imply "hyperparameter transfer""
K.3 CLASSIFICATION OF PARAMETRIZATIONS
These heatmaps already demonstrate that no parametrization of (global learning rate block multiplier) can transfer hyperparameters robustly, because any such parametrization can only shift the heatmaps but not stretch them, so one cannot "transfer" a sublevel set of one slope into sublevel set of another slope.
But even if we allow learning rate to vary between layers in block, no stable, faithful, nontrivial parametrization can avoid the linearization problem described above.
For simplicity, fix of hyperparameters consisting of the learning rate for each of the layers in block, as well as the block multiplier (one for each block); WLOG all weights are initialized (1)) This yields space of dimension blockdepth
nonlinearity and block depth 2
23 We consider the space
35
Indeed, for this to happen, the weight update AW must be at least of order 2(11) (size of initialization) for some i. But this would contribute a drift term to the block output gf = ga(a!=1; that is large the noise term. This then implies that either the parametrization unstable (if the block multiplier I- is ())/L)) or lacks feature diversity (if the block multiplier T. is O()T.)).
For example, in a linear model,
L"lg') = [ = - (WW -WW2
[W -/ is independent and zero-mean across / (the noise term), while [W/ 2/111 + (A is correlated across (the drift term). [Wj' is always @(1) because the W!? W are. If !w is I211,, then [AW'
2(1) well, making the drift term as large as the noise term. If AW"ll (2)), then W AW 221). causing WW'
[WPWWi--/-) 2 ( A =111 to be 2(1) 26
The same argument be straightforwardly adapted to nonlinear MLPs (with subtraction) and arbitrary block depth 2 2, and as well to general nonlinearities that are not necessarily positive- homogencous, with hyperparameter space enlarged include initialization.
meaning, the lcarning tied across all layers in block 24 but our arguments generalize trivially to arbitrary block depth > 2
25hi WL.G because the nonlincarities are homogeneous
9((1), then by symmetry the backward pass suffers the same problem. But for general block depth, this argument does not say anything about the middle layers, while the argument presented above implies that A t cannot he 2(1) for any
26 2'nne can also observe that if Aw"ll
41

---

Published as a conference paper at ICLR 2024
neddl
roul agar
hidcen Jepwou15
tidsin Jayens
Jayers
fadsin ltyesudd
lsyor
optimasr
ogatint N
211
width
14lth
660th
Figure 9: Trained linear network converges to its infinite width limit which is computed recursively based I and C. Depth fixed at 64, width varies between 2 2 ..... 213 Networks trained with SGD for 10 steps. The root mean square statistics (y-axis) at 1st, 5th and 10th steps are plotted using solid lines where the z-axis the width. The root mean square values are computed on the outputs of some of the layers (including the input layer, output layer. and hidden layers at each quarter). The corresponding value for the infinite width indicated with dashed lines.
Jayer
Jayerss
Neyees
loyer,
Iedown 1epex.1.0
sudpal luper
againt disj
Figure 10: Under Depth-uP, infinite wide linear network training converges when increasing the depth. Infinite wide linear networks of depth 2 ...,,..,................. 2 are computed recursively based on I' and C. The root mean square statisties (y-axis) 1st, 5th and 10th steps are plotted across the depth (x-axis).
K.4 So WHAT IS THE OPTIMAL PARAMETRIZATION
of the above considerations suggest that we are missing crucial hyperparameters in our consideration when increasing the complexity of each block. Our study right now is akin to the naive study of the I-dimensional hyperparameter space of the global learning rate in SP. Discovering these missing hyperparameters will be an important question for future work.
L FULL EXPERIMENTS
L. VERIFYING THE THEORY IN THE LINEAR CASE
showed that complete description of the training dynamics of linear networks can be formulated in terms of and C. this section, we provide empirical results supporting our theoretical findings. We first verify the finite-depth recursive formula for I in Lemma C.3 is the correct limit when the width goes to infinity, then proceed to show that the infinite-depth limit the correct one.
In Section 3.2
Infinite-width limit. In Figure 9 we train a series of 64-layer linear networks of width ,r,28..... 2 with 5, 10 steps on MNIST, and plot the root mean square") of the layer outputs using solid lines. We also compute the infinite width limit of the corresponding statistics using the recursive formula for and plot them as dashed horizontal lines. For clarity of the figure, we only
The root mean square of a vector x = (II,...,,5, F.. is
which denoted as "12" in Figures
T0
42

---

Published as a conference paper at ICLR 2024
plot the statistics of the imput layer, output layer, and hidden layers of index 16, 32, 48, and 64. It is clear that as the width grows, the solid lines converge to the dashed lines consistently across the training steps. It indicates that our computation of the infinite width limit is correct.
Infinite-depth limit. We verify that the infinite width limit above converges when the depth grows. We consider linear networks of the same architecture but vary the depth from 2 2 We again compute the root mean square values of the layer outputs using the recursive formula for I, and plot them in Figure I0 with depth being z-axis. For clarity of the figure, we only plot the statistics of the input layer, output layer, and hidden layers of index T./4, 1./2, 3T//4, and I.. One can observe that the statistics of the layer outputs converge quickly when the depth grows from 24 to 29, which verifies our convergence result.
1.2 HYPERPARAMETE TRANSFER
In this section, we provide empirical evidence to show the optimality of Depth-uP scaling and the transferability of some quantities across depth. We train vanilla residual network with block depth 1 (1 MLP layer in each residual block) on CIFAR-10 dataset using Adam optimizer, batch size 64, for 50 epochs (input and output layers are fixed). The network parameterized as follows
s' 2' + a x L (W
and the weights are trained with the rule
W - W x L. go..... nL gt),
where the learning rate y and the block multiplier a are the hyperparameters 28 The values of a; Y depend the parametrization of choice. For Depth-ud,, we have = 1/2, and for standard In our experiments, we assume base depth 8, meaning
paramctrization, we have a = 7 that we replace I. by L/8 in the parametrization above.
Learning rate transfer (n). In Figure we show the training loss versus learning rate for depths 2 for k (3.4... 10). For Depth-uuP, a convergence pattern can be observed for the optimal learning rate as depth grows. Optimal learning rates for small depths (e.g. I, = 29) exhibit a mild shift which should be expected, as our theory shows convergence However, starting from depth L = 26, the optimal learning rate is concentrated around 10 For parametrization that only scales the multiplier but not LR (a = 1/2, 7 = 0), we observe the optimal learning rate shifts significantly. For standard parametrization without any depth scaling (a = Y = 0), the optimal learning rate exhibits a more significant shift as depth grows Moreover, even if one picks the optimal learning rate for each depth, the performance still degrades when the depth is very large, suggesting that standard parametrization is not suitable for depth scaling. Additional figures with multiple time slices are provided in Appendix M
the large depth limit.
Is feature learning sufficient for HP transfer? In Appendix !| we explained when and why hyperparameter transfer occurs. Precisely, to obtain HP transfer, one needs to classify all feature learning limits and choose the optimal one. We introduced the notion of feature diversity and showed that Depth-uuP is optimal in the sense that it maximizes feature diversity. To show that optimality is needed for HP transfer,
train a resnet with (a, = (1,0) which is also a feature learning limit. Figure 12 shows that in this case the learning rate exhibits a significant shift with depth. Interestingly, the constant 1) in this case seems to increase with depth, suggesting that the network is trying to break from the ODE limit, which is sub-optimal. Note that in Figure 11] with Depth-puP obtain better training loss compared the ODE parametrization in Figure 12
Do we still have transfer with NoyerNorm (LN)) Our theory considers only Mean Substraction (MS), and Figure [T shows the results with MS. To see wether LN affects HP transfer, we train resnets with the same setup Figure I1 with absolute value non-linearity and LN applied to r
28 Note that y) here is the constant, and the effective learning rate is given by 1pn -
In standard parametrization, there is generally no rule to scale the learning rate with depth, and the optimal learning rate is typically found by grid search. Here, we assume that in standard parametrization, the learning rate scaled by T.-1 to preserve faithfulness.
43

---

Published as a conference paper at ICLR 2024
y) (1,2, 1/2), stop=0:1000 log, (a)=1,0
(c. y) (1/2, 1/2). dlep-37000030000
logz (e)-1.0
0.80
1-22,,
0.0
-05
-10
0.59
10
Logk
0.54
-25
0.50
-30
-35
Log2/n/1e-3
Logaln (te-3)
(a.10 10 (1/2, D), 0tsp=0:1000 Logx(a)=1.0
(a. y) (3/2. 0), stop=37000-9000 Loogy(a)=1.
0.0
2 k=
0.58
1.0
0.50
10
-1.6
0.54
-2.0
0.52
-25
9.0
0.48
-8
2
Logaln (10-3)
Logsin (1e.3)
yl 0), 10ep=0;1000,
logy 1a)=1.0
[a, 1/ - (C. 36p==37000:8000 loge (a)=1.0
1.0
1.0
0.5
Log 20.
0.8
0.7
-1.0
10
-18
0.6
-20
0.5
-2
3
22
Logato (te33)
Loguón (ees3)
Figure 11: Train logloss versus learning rate for width 256 and varying depths. The network consists of MLP blocks (with block depth 1), trained for 50 epochs on CIFARIO dataset using Adam. The batch size is fixed to 64. We tune the depth 2 network to obtain the optimal (a),, (n/
3)) = (1, 0). and scale all deeper networks using 23 as base depth. The reader check that the /. = 2 curves in each columns
the same. We show the logloss versus the learning rate of the hidden layers (input/output layers fixed) for three parametrizations: Depth-uP (Top), Scaling only the blocks (no LR scaling), i.e. & 0 (Middle), and Standard Parametrization without any scaling (a 7 0) (Bottom). Each curve represents the average training loss over time slice of 1000 steps for depths 2² for k 21, 2..,, 10). Confidence intervals are based on 5 seeds. The results show that Depth-uP preserves the optimal learning rate while consistently improving the training loss as depth increases. If we only scale the blocks without scaling the LR (a = 1/2, = 0) when training with Adam, the optimal learning rate shifts significantly with depth. With standard parametrization without any depth scaling (common practice), the results show significant shift in the optimal learning rate well. For SP, cap the log loss at 1, which is why for depth 2,,"" 10 we have black horizontal line at Log 1.
before matrix multiplication with W (preLN), We keep MS after non non-linearity although it can be removed since LN is applied in the next layer. Our results, reported in Figure suggest that Depth- μP guarantees learning rate transfer with LN as well.
Block multiplier transfer (a). In Figure [4] we investigate the stability the hyperparameter in Depth-pP depth increases. The results suggest that the optimal value of this constant converges
44

---

Published as a conference paper at ICLR 2024
vi 11. stspm0:1000. Loggteim0.5
(a: y) (1,0), stip=37000:380000 Logg(a)=-0.5
2
1.0
Logk 005
0.076
-15
0,550
200
25
LogyCn/te-3)
Logutn (1o-33
Figure 12: Same setup as fig.
11
for the parametrization
(a, = (1, 0) (the ODE limit).
y) (1/2, 5tep-37000.
logg(a)~330
y)
shxp=57000.58000 loggla)>-3.0
LocLoss
-10
Ligs(g/1e-3)
LogaIn/1e-3)
Figure 13: Same setup as Figure 11 with Abs non non-lineariit instead of ReLU and LayerNorm applied to ./-1 before matrix multiplication with W We show the logloss versus the learning of the hidden layers (input/output layers fixed) for two parametrizations: Depth-uP (Leff) and scaling only the blocks without LR scaling ((a,7) (1/2, 0)) (Right). The results show that Depth-uP preserves the optimal learning rate while consistently improving the training loss as depth increases. we only scale the blocks without scaling the LR (a 1/2, 7 = when training with Adam, the optimal learning rate shifts significantly with depth.
depth grows, which suggest transferability. Additional experiments with multiple time slices provided in Appendix M
stop=0:1000, Logg(11to333000
stxp=37000.38000 www.:/o....g.oee..oo 1og.lo"t--3pa.
1.0
0.5
0.9
-0.5
LogLose
0.8
1.0
2
-20
-30
0.5
-6
-2 Logola)
Logula)
Figure 14: Train logloss versus block multiplier a for varying depths. Same training setup in fig. The results suggest that Depth-uP stabilizes the hyperparameter as depth increases.
45

---

Published as a conference paper at ICLR 2024
batoos-22448.5570666
2:16
2.14
L k= 1
2.08
208
Logg(a)
Figure 15: Modern transformers are insensitive to block multiplier
beltchw327088 65536
2.373
23300
2,9055
2:08
2:08
22900
2:04
2 2701
2.2
2200
2:00
2075
30
35
4.0
4.5
5.0
3.5
0.0
85
7.0
30
35
4.0
5.0 Logult)
5.5
6.0
0.5
7.0
2ogg(4)
Figure 16: In (Megatron) Transformer trained on Common Crawl, deeper does worse initially (Lefft) but eventually does better (Right).
L.3 WHAT HAPPENS IN TRANSFORMER?
Because transformers have block depth 2, as discussed in appendix [.44 suspect that parametrization of (learning rate, block multiplier) will be able to robustly transfer hyperparameters across depth for transformers.
we have plenty of reasons to
Here do large scale experiment using Megatron trained Common Crawl and catalogue
our observations 30 In summary, in our particular setup (which should be close to most large language model pretraining), we see that the 1/vL scaling seems to transfer hyperparameters at the end of training (Figure 17 Right)). However, we also see that deeper does worse initial training (Figure [6][eft)) and 2) optimal hyperparameters scale like O(1) in the middle of training (Figure [eft)))) Combined with the theoretical insights of Appendix this leads us to conclude that while the 1/ 1/VL scaling can potentially be practically useful in transformer training, is likely to be brittle to architectural and algorithmic changes, or even simple things like training time.
In fact, we observe that transformers are insensitive the block multiplier (Figure 15), so that the only relevant hyperparameter is really just learning rate. Thus, empirically measuring the scaling trend of the optimal learning rate, as done modern large scale pretraining, can be a practically more robust way to transfer hyperparameters.
Here L the number of transformer layers, each of which consists of an attention layer and an MLP layer (each of which has depth 2).
L.4 FEATURE DIVERSITY
In this section, we empirically verify claims about feature diversity exponent (Claims 5.4 and 5.5). We use the same setup as in Appendix L.2] i.c., we train deep residual networks of width 256 on CIFAR-10 dataset with Adam and batch size 64. In Figure [8] we compare
We train the models for 3900 steps, using cosine decay schedule with 500 warmup steps. We use sequence length of 4096, hatch size 256, resulting in approximately 4B tokens per training run.
46

---

Published as a conference paper at ICLR 2024
2445
2.50
2.40
L L=2".k km
a.
4.0
6.0
6.0
7.0
2.
1=2%,
4.0 7.0
3.0
60
2.30
235
2:15
2.10
Logdln M1
Logséqii (1-4))
Figure 17: In the middle of (Megatron) transformer training, optimal learning rate is approximately invariant (Left), while at the end of training, it approximately scales like 1/VL. However, the I scaling transfers the maximum viable learning rate better in either case.
two parametrizations, Depth- uP (a 7 1/2) and the ODE parametrization (a, (1, We measure |z [(A++))) _!A d(e) at = 1000 for the two parametrizations and varying depth.
For each parametrization and depth L, we rescale function by multiplying a constant e such that for clean presentation. One can observe
d(1/256) 1. and then plot the rescaled function clearly that Depth-juP has feature diversity exponent (almost) 1/2 for any I,, while the curves for ODE parametrization move from cl/2 to e when grows. This exactly fits our theory that Depth- μP maximizes the feature diversity, while other parametrizations (even with feature learning) have smaller feature diversity exponents that should go to O in the infinite depth limit.
Growth along with L and (. In Figure 19 it by dividing additional tu and a constant such that d(1/256)
100, 500, 1000, and rescale t.p 1, and then plot the rescaled function d/(c....5 for clean comparison between and 0.5 We observe that for both Depth- uP and ODE parametrization, the slopes of the curves grow along with I, and t. The growth along t can be explained by the cumulative correlation between layers. The growth along for ODE parametrization is because the independent components between nearby layers decrease when grows. We do not have a clear understanding for the growth along L for Depth-uP and we leave future work.
measure d(e)
Absolute value activation increases feature diversity. In Figure 20 we plot the same curves as in Figure 19 but comparing ReLU activation and absolute value activation under Depth-ud.. We observe that the slope of the curves for absolute value activation is smaller than ReLU activation. It matches our theory in Appendix that that absolute value activation increases feature diversity.
M ADDITIONAL FIGURES FOR HYPERPARAMETER TRANSFER EXPERIMENTS
47

---

Published as a conference paper at ICLR 2024
a= 1/2
a=1
15.0
@(E)
15.0
8
12.5
12.5
9
10.0
10.0
10
11
Scaled ||
7.5
5.0
5.0
2.5
(
0.0
0.0
0.00
0.01
0.02
0.03
0.04
0.05
0.06
0.00
0.01
0.02
0.03
0.04
0.05 0.06
E
Figure 18: Difference between feature at layer [M.] and feature at layer L(A + <)T.J as a curve of for width n 256 and varying depths. For a clean presentation, each curve scaled by a constant so it always passes (1/256, 1). The feature diversity exponent depends on the growth of the curve when L + oo. For Depth-uP (left), the curve is always close to e//2, meaning 1/2. For ODE parametrization (right), the curve shifts from 1/2 to when L grows, indicating its i. goes to 0 in the infinite depth limit.
100
t=500
t= 1000
logit)
2.5
2.5
2.5
Scated Norm
2.0
12
1.5
1.5
155
13
1.0
0.5
0.5
0.5
0.01
0.52
0.03
0.05
0.01
0.02
0.03
0.05
0.08
0.01
0.82
0.25 0.06
Figure 19: Same setup as Figure 18 but at step / 100, 500, 1000, and each curve is scaled by dividing constant and additional 2/2 so it always passes (1/256, 1). The curve indicating feature diversity exponent exactly 1/2 should be horizontal line at 1. For Depth-uP (a 0.5), the curves are almost horizontal. For ODE parametrization larger L and larger L.
1), slopes of the curves are larger with
100
500
1000
1.125
1.25
logiui
1.100
1.005
1.10
1.30
1.005
1.05
1.00
Abs
BelU
0.08
0.00
0.02
Figure 20: Same setup Figure 19 but comparing Depth-uP with ReLU activation and absolute value activation. Each curve is scaled by dividing a constant and 2//2 so it always passes (1/256, 1). The curve indicating feature diversity exponent exactly 1/2 should be horizontal line at For both activations, slopes of curves are small, but growing along with L and f. The slopes with absolute value activation (ő Abs) are slower than the slopes with ReLU activation (ò ReLU), indicating feature diversity is higher with absolute value activation.
48

---

Published as a conference paper at ICLR 2024
step=0
step=1000
step=2000
step=3000
step=4000
0.60
LogLoss
0.55
0.50
step=5000
6000=6000
step=7000
step=8000
step=9000
LogLoss
-1
step= 10000
step= 11000
step=12000
step= 13000
step= 14000
LogLoss
-2
step= 15000
step=16000 16000
step=17000
step= 18000
step= 19000
LogLoss
-2
step=20000
step=21000
step=22000
step=23000
step=24000 24000
LogLoss
step= 25000
step=26000
step=27000
step=28000
step=29000 29000
LogLoss
-2
30000
step=31000
step=32000
step=33000
step=34000
LogLoss
-2
0.0
2.5
0.0 2.5
0.0
2.5
2.5 Log(n /1e-3)
0.0
2.5 Log(n /1e-3)
0.0
Log(n/1e-3)
Log(n/1e-3)
/1e-3)
Log(n /1e-3)
Figure 21:
Same as fig.
11 (Up. Depth-uP) with multiple time slices.
49

---

Published as a conference paper at ICLR 2024
step=0
step=1000
step=2000
step=3000
step=4000
LogLoss
0.8
0.6
step=5000
step=6000
step=7000
step=8000
step=9000
LogLoss
step=10000
step=11000
step=12000
step= 13000
step=14000
LogLoss
-2
step=15000
step=16000
step=17000
step= 18000
step=19000
LogLoss
-2
step=20000
step=21000
step=22000
step=23000
step=24000
LogLoss
22
step=25000
step=26000
step=27000
step=28000
step=29000
LogLoss
22
step= 30000
step=31000
step=32000
step=33000
step=34000
LogLoss
-2
0
10
10
10
10
10
Log(n x L /1e-3)
Log(n x /1e-3)
Log(n x /1e-3)
Log(n x L /1e-3)
Log(n L /1e-3)
Figure 22: Same as fig.
11 [Middle, Standard Parametrization with 1) with multiple time slices.
50

---

Published as a conference paper at ICLR 2024
step=0
step=1000
step=2000
step=3000
step=4000
0.8
LogLoss
0.6
step=5000
step=
step=7000
step=8000
step=9000
LogLoss
step=10000
step=11000
step=
step=
step= 14000
LogLoss
-2
step=15000
step= 16000
step=17000 17000
step=18000
1900=10000
LogLoss
-2
step=20000
step=21000
step=22000
step=23000
step=24000
LogLoss
12
step=25000
step=26000
step=27000
step=28000
step=29000
LogLoss
22
step=30000
slep=31000
step=32000
step=33000
step=34000
LogLoss
-2
-5
-5
0
-5
0
-5 0
-5 0
Log(n /e-3)
Log(n /1e-3)
Log(n /1e-3)
Log(n /1e-3)
Log(n /1e-3)
Figure 23: Same as fig. multiple time slices.
(Bottom, Standard Prameterization with no scaling, a = 0, = 0) with
51

---

Published as a conference paper at ICLR 2024
step=0
step=1000
step=2000
step=3000
step=4000
1.0
LogLoss
0.8
step=5000
step=6000
step=7000
step=8000
step=9000
LogLoss
0.5
0.0
-0.5
step=10000
step=11000
step=12000
step= 13000
step= 14000
LogLoss
step=15000
step=1600
step=17000
step= 18000
step= 19000
LogLoss
step=20000
step=21000
step=22000
step=23000
step=24000
LogLoss
step=25000
step=26000
step=27000
step=28000
step=29000
LogLoss
step=30000
step=31000
step=32000
step=33000
34000
LogLoss
5
0
5
Log(a)
Log(a)
Log(a)
Log(a)
Log(a)
Figure 24: Same as fig.
14 with multiple time slices.
52


Reviews:
[{"review_id": "Sdk95OH2e7", "reviewer": "Reviewer_AGJB", "summary": "The paper studies feature learning parameterizations of deep residual networks when depth goes to infinity (after taking infinite width in a muP regime), and studies the impact of two scaling parameters: a multiplier $L^{-\\alpha}$ of each block before adding on the residual stream, and a multiplier $L^{-\\gamma}$ on the learning rate, where $L$ is the depth.\n\nThe authors find that the best regime, termed Depth-muP, is for $\\alpha = \\gamma = 1/2$, leading to various desirable properties. Of particular interest and novelty is the \"feature diversity\" property, which leads to diverse features across neighboring layers and is only achieved for $\\alpha = 1/2$.", "questions": "Here are a few points that should be addressed in order to strengthen the paper.\n\n* feature diversity: it'd be helpful to provide additional experiments on this point, especially for $\\alpha > 1/2$, something which seems to be missing in the current draft. It'd be helpful to provide more intuition about the benefits of feature diversity (which seem closely related to benefits of depth), perhaps with concrete examples. Can the quantity in Def 6.6 be plotted during training to see the effect of $\\alpha$?\n\n* presentation: while the warmup section 3 is quite insightful, I found section 4 to be much more obscure, and perhaps not essential for the later discussions. I encourage the authors to either provide some more intuition on the \"physical\" meaning of the different equations, or to shorten the section and defer details to the appendix. Perhaps an informal version of the actual tensor program would be more insightful? Also, how does this related to the AMP and DMFT literature? some more comparison to related literature would be helpful (in particular, this concurrent [paper](https://arxiv.org/abs/2309.16620) seems highly related and should be cited and compared to in the final version)\n\n* mean subtraction: could you clarify a bit more the role of this? Is MS specifically needed because of the activation? Would adding an MLP output layer as in transformers drop this requirement? Note that in practice LLM practitioners often drop the mean subtraction part of layer-norm, using RMSnorm instead (e.g. in LLaMa, likely for computational reasons, but it seems to convey that MS isn't really needed?)\n\n* scaling: do you expect any different behaviors in different limits, e.g. in the proportional deep-wide limit studied e.g. by Hanin, or what do you expect would change if the number of GD steps can grow with the number of layers? any comparison or intuition here would be useful.\n\nOther comments:\n- on feature diversity: how does the quantity in Def 6.6 behave at initialization vs later in training? what do you mean by \"maximal\" in claim 6.5? any more intuition beyond the footnote on how to get smaller exponents than 1/2?\n- \"analogous to width situation where deep mean field collapses to a single neuron\": can you elaborate? what is deep mean field, and how are these related?\n- Figures: any observations on what happens when changing width and depth together, as opposed to fixing the width?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The findings in the paper are of significant interest for scaling of neural networks to large depths, both for theoreticians and practitioners.\nWhile the multiplier scalings had been studied before, especially in kernel regimes, its extension to feature learning is crucially important, and provides a meaningful criterion for the benefits of $\\alpha = 1/2$, namely feature diversity (while higher $\\alpha$ leads to redundancy across layers).", "weaknesses": "see questions", "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": "ImIzsogQ3E", "reviewer": "Reviewer_Mvpr", "summary": "This work proposes a desirable scaling of hyper-parameters to achieve feature learning in networks with skip connections, infinite width, and infinite depth. In other words, it is about $\\mu$P for extremely deep networks with skip connections. First, the authors introduce two scaling factors depending on the depth — $\\alpha$ for the forward propagation and $\\gamma$ for the gradient. These factors are key to enabling feature learning in such deep networks. Second, the authors focus on a deep linear network and obtain an analytical expression for feature learning under the proposed scaling factors. Finally, the effectiveness of these scaling factors is empirically tested through some experiments,  particularly those examining the transferability of learning rates.", "questions": "(i) Where in the paper do the authors utilize the tensor program formulation to derive the equation $\\alpha+\\gamma=1$? It appears that the dynamics of linear networks discussed in Sections C & D presuppose $\\alpha=\\gamma=1/2$. Furthermore, the general case presented in Section E merely offers a tensor program representation of the training process without shedding light on how we should determine ($\\alpha$, $\\gamma$). \n\nI hypothesize that the assessment of $A_l^2$ in Section 3 leads to the determination of $\\alpha = \\gamma = 1/2$. However, I am unable to find this evaluation in Sections C, D and E.  Could you please elucidate precisely where in these sections the authors derive both $\\alpha + \\gamma = 1$ and $\\alpha = \\gamma = 1/2$, and explain how the tensor program formulation is applied to achieve these conclusions?\n\n(ii) Feature learning is defined as $\\Delta \\boldsymbol{h}_t^{\\lfloor\\lambda L\\rfloor}=\\Theta(1)$ (Definition 6.5). Contrarily, the authors derive the unique parameterization for feature learning (Claim 6.5) from the non-redundant exponent $\\kappa=1/2$. However, the manuscript does not clearly show the connection between definition 6.5 and $\\kappa$. Could you write down explicitly how these two quantities are equivalent to each other? Furthermore, why and in what manner is the case of $\\kappa=1/2$ *unique* for enabling feature learning outlined in Definition 6.5?\n\n(iii) What is the activation function utilized in Figure 2? Additionally, what logarithm base is employed for the log loss depicted on the vertical axis? These may appear as minor details, but I believe they are crucial.  Since the analysis in Section 4 and Figures 3 & 4 focus on the linear network, it seems no guarantee that the Depth-$\\mu$P works for general nonlinear activation functions.  For instance, could you empirically verify whether the results obtained are applicable to Tanh and ReLU activation functions?\n\nRegarding the scale of the vertical axis, the issue lies with the right side of Figure 2.  If the base is $e$,  the model exhibits Loss ~ exp(-3.5) ~ 0.03. This implies that the training is not finished. Consequently, it becomes difficult to judge whether the learning rate transfer is successful in the models that have been completely trained.\n\n(iv) Difference from Jelassi et al. (2023) \n\n> Jelassi et al. (2023) showed that a learning rate scaling of depth−3/2 guarantees stability after the initial gradient step\n\nThe distinction between the current work and Jelassi et al. (2023) should be more explicitly stated.  I think that $\\mu$P in the current work is also consistent with the initial gradient step. I mean, the $\\mu$P derived from the initial gradient step would remain consistent across general t steps. Are you suggesting that your Depth-$\\mu$P differs between the initial step and subsequent steps where $t > 1$?\n\n(v) In section G.5., the authors say \n>In our setup, where the input and output layers are initialized at a constant scale (w.r.t. L), it is actually not possible to have a kernel limit. Even in our linear case in Section 4, one can see the learned model is not linear.\n\nI am confused about this explanation and believe a more detailed and comprehensive elucidation is necessary.  Firstly, are you implying that the Depth-$\\mu$P is the unique parameterization that is stable, nontrivial, and faithful? Specifically, in very deep networks with skip connections, are you suggesting that the kernel regime is absent?  As far as I understand, the original work on $\\mu$P [Yang & Hu, 2021] characterizes feature learning as the boundaries of the kernel regime. Thus, it would be surprising (and interesting, if it is true) that the kernel regime disappears, leaving only feature learning as a stable state of training.  \nSecondly, in what sense, is the learned linear network not linear? By definition, the linear network is a linear mapping of the input x.  Could you clarify in what aspect the learned linear network deviates from this linear characteristic?\n\nDue to these ambiguous points, I have currently assigned a lower score. However, I am open to increasing this score if they are clarified or if any misunderstandings on my part are resolved.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "-  While some previous work proposed to use $\\alpha=1/2$ for the skip connections,  the scale of gradient updates for such networks has not been investigated. This study addresses this gap by demonstrating that $\\gamma=1/2$ is appropriate for feature learning in networks with infinite depth. \n\n- This study keeps up with the latest research on muP, particularly highlighting Yang & Littwin 2023 for entry-wise adaptive gradients, and further increases the usefulness of muP.", "weaknesses": "-**Not a few ambiguous points**\n\nAs I have noted in questions (i)-(v), there are numerous unclear aspects in this paper that are critical to understanding its main claim. In particular, the dependencies among theoretical claims are ambiguous, making it uncertain how Main Claim 6.5 is derived, as highlighted in my Questions (i) and (ii).\n\n-**Empirical verification is limited**\n\nFigure 2 is the only evidence that the proposed Depth-$\\mu$P works for non-linear networks because Figures 3 & 4 suppose linear networks and Figures 5-9 are part of Figure 2. However, the activation function used in Figure 2 is not mentioned.  In addition, it is also unclear whether the training is finished (Question (iii)). Thus, it is hard to judge whether the Depth-$\\mu$P works in realistic situations. Because a large part of this work focuses on the linear network, there is a concern that it may only function effectively in models close to the linear network.\n\nOne more concern is that the width is not so large (n=256) compared to the depth.  The theoretical framework of this study considers a scenario where depth << width. However, the experimental conditions seem closer to a proportional limit where depth ~ width. It is surprising that the theory and experiments align despite this discrepancy. It would be advisable to validate this with several types of widths, activation functions and datasets (or tasks).", "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": "OSZeFFq2mn", "reviewer": "Reviewer_f9Ax", "summary": "The paper addresses the problem of scaling the depth of neural networks and its impact on model performance especially the norm of the feature will blow up as depth grows. The authors introduce a principled approach called Depth-µP, which allows for training of arbitrarily deep networks while maximizing feature learning and diversity among layers. They propose dividing the contribution of each residual block and parameter update by the square root of the depth, ensuring more stable training and hyperparameter transferability.", "questions": "1. It would be better if the author could provide some experimental results on the performance of the modified structure. \nOtherwise, the analysis will be too specific but not related to any real applications.\n\n2. When analyzing the infinite-depth effect, I think it will be better the separate the infinite-width effect. \nHowever the current analysis seems to combine these two cases, and difficult to tell the contribution of the analysis of the infinite depth.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper introduces a novel depth scaling strategy called Depth-µP, which addresses the challenges of increasing the depth of neural networks. This method is designed for a specific network architecture where the feature output of each residual block is divided by a hyperparameter L to control the norm, not blowing up.\n2. The authors establish a theoretical foundation for Depth-µP by using Tensor Programs to rigorously analyze the behavior of infinitely deep neural networks under the proposed scaling scheme.", "weaknesses": "1. The motivation described in the introduction `The stacking of many residual blocks causes an obvious issue even at the initialization — the norm of $x^l$ grows with $l$,...` does not match the real design of the neural network. As there will be a normalization layer after or before the activation layers in common resnet or transformers.\n\n2. The Depth-$\\mu$P algorithm is designed on a modified structure of the residual network, but there is not any performance comparison of this structure and commonly used structure.\nEven though with such a modified structure we can train networks with an arbitrary depth, it does not indicate the performance is better.\n\n3. It is not so easy for me to follow the theoretical analysis. The introduction of the $\\Gamma$ and $C$ functions is not smooth. And why the analysis needs the width of the network to go to infinity is still not clear.", "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": "t0a5cBvcfm", "reviewer": "Reviewer_TGji", "summary": "This paper investigates how to transfer optimal hyperparameters for a network with different depths. The authors use the tensor program framework to analyze the covariance of a linear resnet. The authors verified their theory on toy resnets.", "questions": "See above", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper presents a solid theory for an important problem. Although from the reviewer's viewpoint, the question is far from being answered this paper makes a solid concrete first-step contribution.", "weaknesses": "- The resnet archtiecture is too simplifeid. The muP initialization can't work when the residual block is two-layer. How about the theory for linear resnet when the residual block is two-layer?  \n- the tensor program framework can deal with the activation function and normalization layer but this paper dropped the part there. Can even just analyze one step of gd？ (or linear net+noramlization layer)\n- the convergence in the tensor program is just loss function convergence and can't obtain the convergence of optimal hyperparameters. Can the author show that the optimal hyperparameter have faster convergence speed?", "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": "yzU9emh5XG", "reviewer": "Reviewer_5jxB", "summary": "The paper studies the scaling limits of residual models in the infinite depth limit $L \\to \\infty$ and their training dynamics. In particular, they study a parametrization where the residual branches are scaled by $1/\\sqrt{L}$ (previously studied at initialization by several prior works). In order to achieve a feature learning limit as both the width and depth tend (sequentially) to infinity, the authors adopt the $\\mu P$ parametrization that guarantees maximal feature learning in the infinite width limit ($N \\to \\infty$). For the proposed model, the authors devise equations describing the training dynamics using the Tensor Programs framework. The authors also introduce the concept of feature diversity, defined as the norm of the difference of the features of nearby layers. They show that among the class of admissible parameterizations in the limit $N,L \\to \\infty$, the proposed parametrization maximizes feature diversity. The authors apply their theory to hyperparameter transfer, showing that the optimal hyperparameters transfer as the model size increases (in terms of depth).", "questions": "1. I would like the authors to comment on why they think hyperparameters should transfer in the proposed feature learning limit (as opposed to the NTK limit for instance).  \n2. Training time. How much do the optimal hyperparameters (e.g. learning rate) shift as a function of training time t? It would be nice to comment on this in the main paper, as practitioners might be interested in tuning the architecture with minimal amount of training steps. \n3. Why use the mean-centered version of Resnet? For instance, in https://arxiv.org/abs/2302.00453 the weights are applied after the nonlinearity, thus achieving centering. \n3. Validity of infinite depth limit under the tensor program. It is my understanding of tensor programs that a finite number of tensor variables (i.e. finite number of layers) are required to apply the tensor program framework. How is the infinite depth limit handled? At initialization, the commutativity of the limits was proven in https://arxiv.org/abs/2302.00453. Is the order of the limits still commutative during training? \n\nOverall, I am in favor of acceptance of this paper, although I have some reserves that I hope will be addressed during the rebuttal period.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "The main strengths of the paper lie in the fact of contributing to multiple research areas, and in striking a combination of important theoretical results and practical consequences. In particular:\n1. The authors devise a parametrization that guarantees feature learning in the infinite width-and-depth limit, nicely extending existing works on feature learning in the mean-field (infinite width) limit. \n2. The authors have extensive theory describing the training dynamics in the limit (Section 4, Appendix).\n3. The concept of feature diversity is novel, and allows the authors to nicely classify the infinite depth parameterizations according to this measure. \n4. The problem of hyperparameter transfer is of utmost importance for practitioners. To my knowledge, this problem was previously addressed when scaling the width to infinity, but not the depth. Hence, this work allows the practitioners to tune the architecture at relatively small widths and depths, and then scale up the architecture with a rough guarantee that the hyperparameters also transfer.  \n5. Experiments, although they do not seem to test more widely used architectures such as convolutions and attention models, test very deep networks, which makes me confident about the validity of most of the proposed claims.", "weaknesses": "The main weaknesses lie around a few conceptual leap of faiths that in my view are not entirely justified either by theory or experimentally. \n1. **How does maximizing feature diversity result in the \"optimal\" parametrization?**: The authors identify a class of parameterizations that satisfies a number of desiderata (stable, nontrivial, faithful, feature learning). Among these, the authors argue at multiple points in the paper that Depth-$\\mu$P is the one maximizing feature diversity, and hence \"optimal\" (E.g. at the end of the introduction, the beginning of Experiment Section). In which sense is maximizing feature diversity optimal? Suppose it is in terms of convergence to low training loss. In that case, the authors should have included more experiments testing more models in the class of parameterizations that satisfy stability, nontriviality, faithfulness, and feature learning, but are not maximally diversifying features. However, all the baselines are for $\\alpha=0$, which are not in the regime mentioned above ($\\alpha \\geq 1/2$). \n2. **Why does hyperparameter transfer happen?**. While it is understandable that hyperparameters should not transfer under the SP baselines (initialization is unstable under $\\alpha=0$), the authors make little effort to justify (with a heuristic argument or experiments) why hyperparameters should transfer under the proposed feature learning regime. I guess in principle, the constant learning rate $\\eta$ could still shift constantly with increasing model size. Hence even an experiment testing (e.g. showing that other admissible parameterizations do not exhibit transfer) would have made the claims in the paper more supported by evidence. \n3. **Missing citations**. There are a number of works studying feature learning in infinite-width networks. I think some of these works should have been cited. Also, most of the literature on Gaussian process behavior in neural networks is not cited (e.g. see below). \n\nMinor:\n4. \"In Appendix G.4, we provide “heuristic” proofs that can be made rigorous under non-trivial technical conditions\". It is slightly confusing. I thought these claims were the result of the theory developed in Section 5. While I believe it is perfectly fine for a subset of the proofs to be non-rigorous, I am confused as to what are the additional assumptions beyond the ones in the rigorous theory of Section 5. \n\nSong Mei, Theodor Misiakiewicz, and Andrea Montanari. \"Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.\"\n\nB Bordelon, C Pehlevan, \"Self-consistent dynamical field theory of kernel evolution in wide neural networks\"\n\nLenaic Chizat and Francis Bach. \"Implicit bias of gradient descent for wide two-layer neural networks\ntrained with the logistic loss\"\n\nLee at al. \"Deep Neural Networks as Gaussian Processes\"\n\nLee et al. \"Wide neural networks of any depth evolve as linear models under gradient descent\"", "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)