# Tensor Programs VI:
Feature Learning in Infinite-Depth Neural Networks

Greg Yang
xAI
&Dingli Yu∗
Princeton Language
and Intelligence
&Chen Zhu
Nvidia
&Soufiane Hayou
Simons Institute
UC Berkeley
Equal contribution.Work partially done at the National University of
Singapore.

###### Abstract

By classifying infinite-width neural networks and identifying the _optimal_
limit, [23, 25] demonstrated a universal way, called $\mu$P, for _widthwise
hyperparameter transfer_ , i.e., predicting optimal hyperparameters of wide
neural networks from narrow ones. Here we investigate the analogous
classification for _depthwise parametrizations_ of deep residual networks
(resnets). We classify depthwise parametrizations of block multiplier and
learning rate by their infinite-width-then-depth limits. In resnets where each
block has only one layer, we identify a unique optimal parametrization, called
Depth-$\mu$P that extends $\mu$P and show empirically it admits depthwise
hyperparameter transfer. We identify _feature diversity_ as a crucial factor
in deep networks, and Depth-$\mu$P can be characterized as maximizing both
feature learning and feature diversity. Exploiting this, we find that absolute
value, among all homogeneous nonlinearities, maximizes feature diversity and
indeed empirically leads to significantly better performance. However, if each
block is deeper (such as modern transformers), then we find fundamental
limitations in all possible infinite-depth limits of such parametrizations,
which we illustrate both theoretically and empirically on simple networks as
well as Megatron transformer trained on Common Crawl.

##  1 Introduction

Deep neural networks have showcased remarkable performance across a broad
range of tasks, including image classification, game playing exemplified by
AlphaGo [17], and natural language processing demonstrated by GPT-4 [15]. 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 a fixed depth) have been extensively studied in
the literature: recent work by Yang et al. [25] identified the _Maximal Update
Parametrization_ ($\mu$P) that guarantees maximal feature learning in the
infinite width limit.111Here maximal feature learning refers to $\Theta(1)$
change in features in the infinite width limit. This should be contrasted with
the lazy training regime where the change in features is of order
$\Theta(n^{-1/2})$. Another benefit of $\mu$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 $\mu$P to include depth
scaling. We call the depth scaling Depth-$\mu$P.

The issue of depth scaling has persisted over time. A decade ago, deep neural
networks experienced significant degradation problems — having more than a few
dozen layers would increase the training error instead of improving the
model’s performance. This was partly due to the vanishing or exploding
gradient problem that affects the efficient propagation of information through
the network. The introduction of residual networks (ResNet) [8, 9, 18] 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 [19].

Specifically, in a residual architecture, the $l$-th residual block is
formulated as

$x^{l}=x^{l-1}+g^{l}(x^{l-1};W^{l}),$

where $x^{l-1}$ is the input, $x^{l}$ is the output, $W^{l}$ are the
parameters of the block, and $g^{l}$ (often called the _residual branch_) is a
mapping that defines the layer (e.g. a stack of convolutions in ResNet, or
SelfAttention and MLP in a Transformer). In this work, we focus on the case
where $g^{l}$ is a biasless perceptron with (or without) activation.

Table 1: Difference between standard depth scaling and Depth-$\mu$P. The constants $a$ and $\eta$ in Depth-$\mu$P are transferable across depth, i.e., one can tune a smaller network and use the same constants for deeper networks. On the other hand, the learning rate of standard depth scaling requires separate tuning for models of different depth.  | Branch Multiplier | Learning Rate
---|---|---
Standard | 1 | ? (tuned)
Depth-$\mu$P (SGD) | $a/\sqrt{\textrm{depth}}$ | $\eta$
Depth-$\mu$P (Adam) | $a/\sqrt{\textrm{depth}}$ | $\eta/\sqrt{\textrm{depth}}$

The stacking of many residual blocks causes an obvious issue even at the
initialization — the norm of $x^{l}$ grows with $l$, so the last layer
features do not have a stable norm when increasing the depth. Intuitively, one
can stabilize these features by scaling the residual branches with a 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.

When each block $g^{l}$ has only one layer (one matrix multiply), we identify
the parametrization we call Depth-$\mu$P as the optimal parametrization for
deep networks. It maximizes both _feature learning_ and _feature diversity_
222We give a formal definition of feature learning and feature diversity later
in the paper. among all possible parametrizations of block multiplier and
learning rate with depth. Our framework extends the previous results on $\mu$P
which deals with optimal width scaling [25]. It completes the width scaling
and hence provides a full width and depth scaling recipe that guarantees
maximal feature learning and hyperparameter transfer across width and depth.
Depth-$\mu$P contains the following modifications to the standard practice:

  1. 1.

There is a multiplier for each residual branch before adding to its input,
which is inversely proportional to the square root of $L$ (where $L$ is the
depth). Formally, with a constant $a$ independent from $L$,

$\displaystyle x^{l}$ $\displaystyle=x^{l-1}+\frac{a}{\sqrt{L}}\cdot
g^{l}(x^{l-1};W^{l}).$ (1)

  2. 2.

We set the learning rate of $W^{l}$ so that the update of $W^{l}$ during
training is proportional to $1/\sqrt{L}$. We derive different learning rate
schemes for different optimization algorithms based on this principle. For
Adam, because it is scale-invariant to the gradient, the learning rate of
$W^{l}$ is set to be $\eta/{\sqrt{L}}$. On the other hand, the learning rate
of $W^{l}$ for SGD is set as a constant $\eta$ because the gradient of $W^{l}$
is already of size $1/\sqrt{L}$ due to the multiplier.

In block depth $1$ (i.e., $g^{l}$ is a biasless perceptron, $W^{l}$ is a
single matrix), this scaling leads to the following properties:

  * •

At the initialization, each one of the $L$ residual blocks contributes
$\Theta(1/\sqrt{L})$ to the main branch. These $L$ contributions are
independent of each other, so the sum of them is of size $\Theta(1)$.

  * •

During training, the contribution of the update of each residual block is
$\Theta(1/L)$ due to the combining effect of the learning rate and multiplier.
The contributions of the updates are highly correlated, so they sum up to
$\Theta(1)$.

More detailed intuition of this scaling approach can be found in Section 3
where we provide a simple analysis with linear networks after one gradient
step. We give a complete classification of depthwise parametrizations in
section 7.

Figure 1: Behaviors of scaling strategies with a branch multiplier
$L^{-\alpha}$ and parameter update proportional to $L^{-\gamma}$.

###  1.1 Optimality of Depth-$\mu$P.

We thoroughly compare Depth-$\mu$P with other scaling strategies with a branch
multiplier $\propto L^{-\alpha}$ and parameter update $\propto
L^{-\gamma}$.333It implies that the effective learning rate is proportional to
$L^{-\gamma}$ for Adam and $L^{\alpha-\gamma}$ for SGD if the network is
stable at initialization. As shown in Figure 1, the space of $(\alpha,\gamma)$
is divided into several areas, each resulting in a different behavior when
$L\to\infty$:

  * •

Having $\alpha\geq 1/2$ is required to stabilize the network at
initialization. This ensures that he hidden activations and the network output
do not explode at initialization;

  * •

For any $\alpha+\gamma<1$, the network is unstable during training. The change
in hidden activations or the network output explodes with depth during
training;

  * •

For any $\alpha+\gamma>1$, training outcome is trivial. The change of the
network vanishes as depth increases;

  * •

For any $\alpha+\gamma=1$ with $\alpha>1$, the network is _unfaithful_ (a
formal definition will provided later in the paper). The hidden activations
explode during training as depth increases;

  * •

For any $\alpha+\gamma=1$ and $\alpha\in(1/2,1]$, we show that the network
converges to a _redundant_ limit that lacks _feature diversity_ , in that
close layers have similar outputs (in a neural ODE fashion).

  * •

The only choice of $\alpha$ and $\gamma$ left is $\alpha=\gamma=1/2$, which
corresponds to Depth-$\mu$P.

The rigorous definitions and proofs are presented in Section 7.

###  1.2 Hyperparameter Transfer for Depth.

The optimality of Depth-$\mu$P implies (under some assumptions) that the
optimal hyperparameters of the networks also converge as the depth ($L$)
increases. This convergence suggests that the optimal hyperparameters of
shallower networks are approximately equal to those of deeper networks. As a
direct implication, we can leverage this property to infer the hyperparameters
for deeper networks from the shallower ones, effectively reducing the cost
associated with hyperparameter tuning. With Depth-$\mu$P, we successfully
train networks comprising thousands of residual blocks, while also showcasing
the transferability of hyperparameters across depth.

###  1.3 Impossibility Results for Block Depth $\geq 2$

While the block depth 1 case admits a positive result, we show that the block
depth $\geq 2$ case does not and cannot (section 9). The basic issue is the
weights in different layers within a block is forced to interact additively
instead of multiplicatively when depth is large, if one wants to retain
diversity. This causes block depth $\geq 2$ to have worse performance than
block depth $1$ and for the optimal hyperparameters to shift with depth. We
demonstrate this pedagogically on resnet with MLP blocks but also on Megatron
transformer [16] trained on Common Crawl. These observations entail the need
to rethink the current approach to hyperparameter transfer.

##  2 Related Works

###  2.1 Width Scaling and $\mu$P

The infinite-width limit of neural networks has been a topic of extensive
research in the literature. Numerous studies have predominantly focused on
examining the behavior of various statistical quantities at initialization.
Some works have gone beyond the initialization stage to explore the dynamics
of feature learning in neural networks.

##### Lazy training.

With standard parametrization, a learning rate of order
$\mathcal{O}(n^{-1})$,444We also obtain the lazy infinite-width limit with the
NTK parametrization and a $\mathcal{O}(n^{-1/2})$ learning rate. $n$ being the
width, yields the so-called lazy training regime in the infinite-width limit,
where the features remain roughly constant throughout training [3, 25]. This
regime is also known as the Neural Tangent Kernel (NTK) regime and its
convergence properties have been extensively studied in the literature [10, 1,
2, 28].

##### Feature learning and $\mu$P.

Recent empirical studies (e.g. [25]) have provided compelling evidence that
feature learning plays a 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 a natural choice to enhance the performance of
such models.

In this context, $\mu$P (Maximal Update Parameterization), introduced in [25],
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, $\mu$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.

The derivation of $\mu$P leverages the Tensor Programs framework [22, 20, 21,
23, 25], which provides valuable tools for capturing the behavior of neural
networks in the infinite-width regime during the training process.

###  2.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-art models use deep architectures.
The introduction of skip connections [8, 9] 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 a
challenging task [12]. Moreover, tuning hyperparameters for large depth
networks is a time-and-resource-consuming 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 or in
the kernel regime [7, 4, 26, 13, 5, 6, 14, 27]. However, these works lack
insights into the dynamics with feature learning. For instance, one might
argue that features can still experience explosive growth if the learning rate
is not properly chosen. Therefore, an 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-$\mu$P, which offers a
comprehensive framework for depth scaling. Depth-$\mu$P encompasses block
multipliers and learning rate scaling, providing a complete recipe for
training deep networks. In the case of Multi-Layer Perceptrons (MLPs) (no skip
connections), Jelassi et al. [11] showed that a 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.

##  3 Warm-Up: An Intuitive Explanation with Linear Networks

Let us begin with a simple example that provides the necessary intuition
underpinning our depth scaling strategy. Given a depth $L$, width $n$,
consider a linear residual network of the form

$\displaystyle x^{0}$ $\displaystyle=U\xi,$ $\displaystyle\forall
l\in[L],\quad x^{l}$ $\displaystyle=x^{l-1}+\frac{1}{\sqrt{L}}W^{l}x^{l-1},$
$\displaystyle f$ $\displaystyle=V^{\top}x^{L},$

where the weight matrices $W^{l}\in\mathbb{R}^{n\times n}$ and $U,V$ are input
and output weight matrices that we assume to be fixed during training.

###  3.1 Optimal Scaling of the Learning Rate

To simplify the analysis, we consider gradient updates based on a single
datapoint. The first gradient step is given by

$W^{l}_{1}=W^{l}_{0}-\eta G^{l}_{0},$

where $\eta$ is the learning rate, and $G^{l}_{0}$ is a matrix with update
directions. For instance, we have the following expressions for $G^{l}_{0}$
with SGD and Adam:

  * •

SGD: $G^{l}_{0}=\frac{1}{\sqrt{L}}\delta x^{l}\otimes x^{l-1}$, where $\delta
x^{l}\mathrel{\raisebox{-1.29167pt}{$\mathbin{\overset{\text{{def}}}{=}}$}}\frac{\partial\ell}{\partial
x^{l}}$ for some loss function $\ell$.555We use $\delta$ for gradient because
we want to distinguish from $d$ in depth differential equations that appear
later in the paper.

  * •

Adam666For the sake of simplification, we consider SignSGD in this section,
which can be seen as a memory-less version of Adam. The analysis is valid for
any training algorithm that gives $\Theta(1)$ gradients.:
$G^{l}_{0}=\textrm{sign}\left(\frac{1}{\sqrt{L}}\delta x^{l}\otimes
x^{l-1}\right)$.

In both cases, $\delta x^{l}$ and $x^{l-1}$ are computed for a single training
datapoint $\xi_{0}$. The last layer features $x^{L}$ (for some input $\xi$)
are given by
$x^{L}=\prod_{l=1}^{L}\left(I+\frac{1}{\sqrt{L}}W^{l}\right)x^{0}$.777 To
avoid any confusion, here we define the matrix product by
$\prod_{l=1}^{L}A_{l}=A_{L}\times A_{L-1}\dots\times A_{1}$. We use the
subscript $t$ to refer to training step. After the first gradient step, we
have the following

$x^{L}_{1}=\prod_{l=1}^{L}\left(I+\frac{1}{\sqrt{L}}W^{l}_{1}\right)x^{0}=x^{L}_{0}-\frac{\eta}{\sqrt{L}}A_{L}+\mathcal{O}(\eta^{2}),$
(2)

where
$A_{L}=\sum_{l=1}^{L}\left[\prod_{k>l}\left(I+\frac{1}{\sqrt{L}}W^{k}_{0}\right)\right]G^{l}_{0}\left[\prod_{k<l}\left(I+\frac{1}{\sqrt{L}}W^{k}_{0}\right)\right]x^{0}$.
We argue that $A_{L}$ behaves as $\Theta(L)$ (in $L_{2}$ norm). This is the
due to the $1/\sqrt{L}$ scaling factor. To see this, we further simplify the
analysis by considering the case $d_{in}=n=d_{out}=1$ (single neuron per
layer) and the squared loss. In this case, the term $A_{L}$ simplifies to

$A_{L}=\sum_{l=1}^{L}\prod_{k\neq
l}\left(1+\frac{1}{\sqrt{L}}W^{k}_{0}\right)G^{l}_{0}x_{0}.$

##### Scaling for SGD.

With SGD, we have that $G^{l}_{0}=\frac{1}{\sqrt{L}}\prod_{k\neq
l}\left(1+\frac{1}{\sqrt{L}}W^{k}_{0}\right)x_{0}\delta x^{L}$, where $\delta
x^{L}=(Vx^{L}-y(\xi_{0}))$ and $y(\xi_{0})$ is the target output. Therefore,
it is easy to see that

$\displaystyle\mathbb{E}A_{l}^{2}=\frac{1}{L}\mathbb{E}\left(\sum_{l=1}^{L}\prod_{k\neq
l}\left(1+\frac{1}{\sqrt{L}}W^{k}_{0}\right)^{2}\delta
x^{L}x_{0}^{2}\right)^{2}=\Theta\left(\frac{1}{L}L^{2}\right)=\Theta(L),$

where we have used the fact that
$\mathbb{E}\left(1+\frac{1}{\sqrt{L}}W^{k}_{0}\right)^{2p}=1+\Theta(L^{-1})$,
for any positive integer $p$.

Hence, the magnitude of the first order term in eq. 2 is given by

$\mathbb{E}\left[\left(\frac{\eta}{\sqrt{L}}A_{l}\right)^{2}\right]=\Theta(\eta^{2}),$

which shows that the update is stable in depth as long as $\eta=\Theta(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 $G^{l}_{0}=\pm 1$, and therefore we obtain

$\displaystyle\mathbb{E}A_{l}^{2}=\mathbb{E}\left(\sum_{l=1}^{L}\prod_{k\neq
l}\left(1+\frac{1}{\sqrt{L}}W^{k}_{0}\right)x_{0}\right)^{2}=\Theta\left(L^{2}\right),$

where we have used the same arguments as before. In this case, the first order
term in eq. 2 is given by

$\mathbb{E}\left[\left(\frac{\eta}{\sqrt{L}}A_{l}\right)^{2}\right]=\Theta(\eta^{2}L^{-1}).$

Therefore, the maximal learning rate that one can choose without exploding the
features is given by $\eta=\Theta(L^{-1/2})$.

Summary: By ensuring that parameter update is $\Theta(1/\sqrt{L})$, the
features remain stable while feature update is $\Theta(1)$. This $\Theta(1)$
update is due to the accumulation of $\Theta(1/L)$ correlated terms across
depth.

###  3.2 Convergence when Depth goes to $\infty$

Let us look at $x^{L}_{1}$ again in the simple case $d_{in}=d_{out}=n=1$ and
analyze its behaviour when $L\to\infty$. 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 $\eta=1$ and let $M_{L,l}=\prod_{k\neq
l}\left(1+\frac{1}{\sqrt{L}}W_{0}^{k}\right)$ and
$\tau=(Vx^{L}_{0}-y(\xi_{0}))x^{0}$. With this, we have the following

$x^{L}_{1}=\prod_{l=1}^{L}\left(1+\frac{1}{\sqrt{L}}W^{l}_{0}-\frac{1}{L}\tau
M_{L,l}\right)x^{0}.$ (3)

WLOG, let us assume that $x^{0}_{0}>0$. Then, with high probability (the event
that $W^{l}_{0}\ll\sqrt{L}$, for some notion of “$\ll$”, occurs with a
probability of at least $1-e^{-L^{\alpha}}$ for some $\alpha>0$)888This
follows from simple concentration inequalities for sub-exponential random
variables., we have that $x^{L}_{1}>0$. We can therefore look at
$\log(x^{L}_{1})$ which simplifies the task. Taking the log and using Taylor
expansion under a high probability event, we obtain

$\displaystyle\log(x^{L}_{1}/x^{0})$
$\displaystyle=\frac{1}{\sqrt{L}}\sum_{l=1}^{L}W^{l}_{0}-\frac{1}{L}\sum_{l=1}^{L}\tau
M_{L,l}+\frac{\sum_{l=1}^{L}(W^{l}_{0})^{2}}{L}+\mathcal{O}(L^{-1+\epsilon})$
$\displaystyle=\frac{1}{\sqrt{L}}\sum_{l=1}^{L}W^{l}_{0}-\tau
x^{L}_{0}\frac{1}{L}\sum_{l=1}^{L}\frac{1}{1+\frac{1}{\sqrt{L}}W^{l}_{0}}+\frac{\sum_{l=1}^{L}(W^{l}_{0})^{2}}{L}+\mathcal{O}(L^{-1+\epsilon}),$

for some $\epsilon>0$. The first and third terms
$\frac{1}{\sqrt{L}}\sum_{l=1}^{L}W^{l}_{0}$ and
$\frac{\sum_{l=1}^{L}(W^{l}_{0})^{2}}{L}$ converge (almost surely) to a
standard Gaussian and $1$, respectively. The second term also converges
naturally, since $x^{L}_{0}$ converges in $L_{2}$ to a Log-Normal random
variable ([5]) and with a delicate treatment (involving high probability
bounds), one can show that the term
$\frac{1}{L}\sum_{l=1}^{L}\frac{1}{1+\frac{1}{\sqrt{L}}W^{l}_{0}}$ converges
(in $L_{2}$ norm) at large depth. This implies that one should expect
$x^{L}_{1}$ to have some notion of weak convergence as depth grows. Note that
the same analysis becomes much more complicated for general width $n>0$. To
avoid dealing with high probability bounds, a convenient method consists of
taking the width to infinity first $n\to\infty$, then analyzing what happens
as depth increases. We discuss this in the next section.

###  3.3 A Discussion on the General Case

##### Difficulty of generalizing to the nonlinear case.

The extension to the general width scenario ($n>1$) necessitates a more
intricate treatment of the term $A_{l}$ to find optimal scaling rules, yet the
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 is a non-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.

For this purpose, we employ the Tensor Program framework to investigate the
behavior of essential network quantities in the infinite-width-then-depth
limit. By leveraging this framework, our 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 $1\ll
depth\ll 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^{k}_{t}$ might become highly correlated during training, one has to scale
the blocks with $1/L$. To understand this, let us assume a block multiplier of
$L^{-\alpha}$ and consider the scenario of perfect correlation where all
weights are equal, i.e., $W^{k}_{t}=W$ for every $k\in{1,\dots,L}$. In this
case, the last layer features can be expressed as
$x^{L}=\left(I+L^{-\alpha}W\right)^{L}x_{0}$. When $\alpha=1/2$, the features
are likely to exhibit an explosive growth with increasing depth, while opting
for $\alpha=1$ 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 $\alpha=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 fashion during the training process.

In summary, while a naive depth scaling strategy based on scaling blocks might
suggest the need for $\alpha=1$ to stabilize the features, our 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 $\alpha=1/2$.

##  4 SGD Training Dynamics of Infinitely Deep Linear Networks

In this section, we continue to study the linear neural network with residual
connections under Depth-$\mu$P. Using the Tensor Program framework [24], we
rigorously derive the training dynamics of SGD for the linear residual network
when the width and the depth sequentially go to infinity. The road map of our
analysis consists the following three steps.

  1. 1.

We first take the width of the network to infinity by the Tensor Program
framework [24]. 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 a vector $x\in\mathbb{R}^{n}$ that appears in the
computation of the training, the coordinates of $x$ can be viewed as iid
copies of random variable
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle$
(called a _ket_) when $n\to\infty$. 999The definition of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle$
requires the coordinates of $x$ is $\mathcal{O}(1)$ w.r.t. $n$, and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle$
is trivial if the coordinates of $x$ is $o(1)$ w.r.t. $n$. Therefore, for $x$
whose coordinates are not $\Theta(1)$, we normalize $x$ by multiplying
polynomial of $n$ so the resulting vector has coordinates $\Theta(1)$.

  2. 2.

Since the network is linear, every random variable can be written as a linear
combination of a set of zero-mean “base” random variables by the Master
Theorem of Tensor Programs [24]. 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. 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 a six-dimensional tensor, where two of the dimensions have
shape $L$. We then map the tensor to a set of functions whose input domain is
$[0,1]\times[0,1]$. Finally, we claim that the functions converge when
$L\to\infty$, and identify their limits as the solution of a set of functional
integrals.

In Section 10.1, we conduct a thorough empirical verification of our theory in
the linear case. The experiments clearly show the convergence of deep linear
residual networks under Depth-$\mu$P.

##### Assumptions and Notations

Recall the linear network is given by

$\displaystyle x^{0}$ $\displaystyle=U\xi,$ $\displaystyle\forall
l\in[L],\quad x^{l}$ $\displaystyle=\frac{a}{\sqrt{L}}W^{l}x^{l-1}+x^{l-1},$
$\displaystyle f$ $\displaystyle=V^{\top}x^{L}.$

For convenience, we assume $a=1$, the SGD learning rate of $W^{l}$ is $1$. We
add $t$ as a subscript to any notation to denote the same object but at $t$-th
training step, e.g., the input at step $t$ is a single datapoint $\xi_{t}$,
the hidden output of $l$-th layer at step $t$ is $x_{t}^{l}$, and the model
output at step $t$ is $f_{t}$. Let $T$ be the number of training steps. Let
$\ell_{t}$ be the loss function absorbing the label at time $t$, and
$\chi_{t}$ be the derivative of the loss at time $t$, i.e.,
$\chi_{t}=\ell^{\prime}_{t}(f_{t})$. Let $\delta
x_{t}^{l}=\partial\ell_{t}/\partial x_{t}^{l}$, and
$\widetilde{\delta}x_{t}^{l}=n\delta x_{t}^{l}$ is the normalized version of
$\delta x_{t}^{l}$.

The Tensor Program analysis heavily depends on the scaling of initialization
and learning rate of $U,V,W$ w.r.t $n$. In this paper, we use $\mu$P as the
scaling w.r.t. $n$ since it maximizes feature learning in the large width
limit [23]. Without loss of generality, we follow [23] and assume the input
and output dimension is $1$, i.e., $\xi\in\bm{\mathbb{\mathbb{R}}}$,
$f\in\bm{\mathbb{\mathbb{R}}}$. For a clean presentation, we additionally
assume $U,V$ are frozen during training in this section and each coordinate of
$W$ is initialized with i.i.d. Gaussian of variance $1/n$.

###  4.1 Width Limit under $\mu$P

As the first step, we take the width of the network $n$ 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\to\infty$. For a vector $x\in\mathbb{R}^{n}$ that has
roughly iid coordinates, we write
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle\in\mathbb{R}$
(called a _ket_) to denote a random variable such that $x$’s entries look like
iid copies of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle$.
Then for any two vector $x,y\in\mathbb{R}^{n}$ that have roughly iid
coordinates, their limiting inner product by $n$ can be written as
$\lim_{n\to\infty}\frac{x^{\top}y}{n}=\mathbb{E}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x\rangle\cdot\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle$,
which we write succinctly as $\langle
x\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle$.
Deep linear network with SGD is a simple example for this conversion from
vectors to random variables. As shown in Program 1, we define a series of
scalars ($\mathring{f}_{t}$ and $\mathring{\chi}_{t}$) and random variables
($\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\delta
x_{t}^{l}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l}x_{t}^{l-1}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l\top}\delta
x_{t}^{l}\rangle$) using the ket notations. For better understanding, we
provide a brief introduction to TP below.

Initial random variables:
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$
are independent standard Gaussian.

for  _$t=0,\ldots,T-1$_ do

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{0}\rangle\mathbin{\overset{\text{{def}}}{=}}\xi_{t}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$;

for  _$l=1,\ldots,L$_ do

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l}x_{t}^{l-1}\rangle\mathbin{\overset{\text{{def}}}{=}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle-\frac{1}{\sqrt{L}}\sum_{s=0}^{t-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{s}^{l}\rangle\langle
x_{s}^{l-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle$;

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle\mathbin{\overset{\text{{def}}}{=}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle+\frac{1}{\sqrt{L}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l}x_{t}^{l-1}\rangle$;

end for

$\mathring{f}_{t}\mathbin{\overset{\text{{def}}}{=}}\langle
x_{t}^{L}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$;

$\mathring{\chi}_{t}\mathbin{\overset{\text{{def}}}{=}}\ell_{t}^{\prime}(\mathring{f}_{t})$;

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\delta
x_{t}^{L}\rangle\mathbin{\overset{\text{{def}}}{=}}\mathring{\chi}_{t}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$;

for  _$l=L,\ldots,1$_ do

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle\mathbin{\overset{\text{{def}}}{=}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle-\frac{1}{\sqrt{L}}\sum_{s=0}^{t-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{s}^{l-1}\rangle\langle\widetilde{\delta}x_{s}^{l}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle$;

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l-1}\rangle\mathbin{\overset{\text{{def}}}{=}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle+\frac{1}{\sqrt{L}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle$;

end for

end for

where
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle$
are defined in 4.1.

Program 1 Random Variables induced from Tensor Program for the Linear Network
with LR $\eta=1$ and frozen $U,V$.

##### Tensor Programs (TP) in a nutshell.

When training a 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 $x^{l}_{0}$ where the subscript $0$ refers to
initialization, and the scalar $f_{0}$, which is the network output. In the
first backward pass, the output derivative $\chi_{0}$ is computed, then the
gradients $\delta x^{l}_{0}$ are backpropagated. (Since the coordinates of
$\delta x_{0}^{l}$ vanish to 0 when $n\to\infty$, TP instead tracks its
normalized version
$\widetilde{\delta}x_{0}^{l}\mathrel{\raisebox{-1.29167pt}{$\mathbin{\overset{\text{{def}}}{=}}$}}n\cdot\delta
x_{0}^{l}$.) New vectors are created and appended to the TP as training
progresses. When the width $n$ goes to infinity, vectors of size $n$ in the TP
(e.g., the features $x^{l}_{t}$, and normalized gradients
$\widetilde{\delta}x^{l}_{t}$) see their coordinates converge to roughly iid
random variables (e.g.,
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x^{l}_{t}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x^{l}_{t}\rangle$
in Program 1), and other scalar quantities (e.g., $f_{t}$ and $\chi_{t}$)
converge to deterministic values (e.g., $\mathring{f}_{t}$ and
$\mathring{\chi}_{t}$ in Program 1) under proper parametrization ($\mu$P). The
Master Theorem [25] 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, we refer the reader to [25].

Now when we look back to Program 1, the definitions of scalars and random
variables should be clear (except for
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle$).
One can find straightforward correspondence between those and their finite
counterpart, for example:

  * •

$\mathring{f}_{t}$ corresponds to $f_{t}$, and $\mathring{\chi}_{t}$
corresponds to $\chi_{t}$;

  * •

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle$
corresponds to $x_{t}^{l}$ and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle$
corresponds to $\widetilde{\delta}x_{t}^{l}$. (Recall
$\widetilde{\delta}x_{t}^{l}=n\cdot\delta x_{t}^{l}$ is the normalized version
of $\delta x_{t}^{l}$.)

  * •

By SGD, $W_{t}^{l}=W_{0}^{l}-\frac{1}{\sqrt{L}}\sum_{s=0}^{t-1}\delta
x_{s}^{l}\otimes x_{s}^{l-1}$, which corresponds to
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{t}^{l}x_{t}^{l-1}\rangle=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle-\frac{1}{\sqrt{L}}\sum_{s=0}^{t-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{s}^{l}\rangle\langle
x_{s}^{l-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle$.

Now we can dive into the definition of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{t}^{l}\rangle$.
Let $\mathcal{W}$ be the set of initial random matrices of size $n\times n$,
i.e., $\\{W_{0}^{1},\ldots,W_{0}^{L}\\}$, and
$\mathcal{W}^{\top}\mathbin{\overset{\text{{def}}}{=}}\\{W^{\top}:W\in\mathcal{W}\\}$.
Let $\mathcal{V}_{W}$ denote the set of all vectors in training of the form
$Wy$ for some $y$. Then for every $W\in\mathcal{W}\cup\mathcal{W}^{\top}$, and
$Wy\in\mathcal{V}_{W}$, we can decompose
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\rangle$
into the sum of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$, where
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is a random variable that act as if $W$ were independent of $y$, and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$ is the random variable capturing the correlation part between $W$ and
$y$. Specifically, let us briefly track what happens to $W^{l}_{0}x^{l-1}_{t}$
during training. In the first step, we have $W^{l}_{0}x^{l-1}_{0}$ which has
roughly Gaussian coordinates (in the large width limit). In this case, we have
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{l}_{0}x^{l-1}_{0}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu=0$. After the first backprop, we have $\delta x^{l-1}_{0}=\delta
x^{l}_{0}+\frac{1}{\sqrt{L}}W^{l\top}_{0}\delta x^{l}_{0}$, which means that
the update in $W^{l-1}$ will contain a term of the form $W^{l\top}_{0}z$ for
some vector $z$. This implies that $W^{l}_{0}x^{l-1}_{1}$ will contain a term
of the form $W^{l}_{0}W^{l\top}_{0}z^{\prime}$ for some vector $z^{\prime}$.
This term induces an additional correlation term that appears when we take the
width to infinity. The
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{l}_{0}x^{l-1}_{1}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$ is defined by isolating this additional correlation term from
$W^{l}_{0}W^{l\top}_{0}z^{\prime}$. The remaining term is Gaussian in the
infinite-width limit, which defines the term
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{l}_{0}x^{l-1}_{1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$. Formally, we present the following definition.

######  Definition 4.1.

We define
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\rangle\mathbin{\overset{\text{{def}}}{=}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$ for every $W\in\mathcal{W}\cup\mathcal{W}^{\top}$ and
$Wy\in\mathcal{V}_{W}$, where

  * •

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is a Gaussian variable with zero mean. $\forall
W\in\mathcal{W}\cup\mathcal{W}^{\top},Wy,Wz\in\mathcal{V}_{W}$,

$\mathrm{Cov}\left(\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wz\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu\right)\mathbin{\overset{\text{{def}}}{=}}\langle
y\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}z\rangle.$

$\forall
W,W^{\prime}\in\mathcal{W}\cup\mathcal{W}^{\top},Wy\in\mathcal{V}_{W},W^{\prime}z\in\mathcal{V}_{W^{\prime}}$,
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{\prime}z\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ are independent if $W\neq W^{\prime}$.
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is also independent from
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$.

  * •

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$ is defined to be a linear combination of
$\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}z\rangle:W^{\top}z\in\mathcal{V}_{W^{\top}}\\}$.
Then we can unwind any
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle$
inductively as a linear combination of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\bullet\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$,
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$,
which allows us to fully define

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}Wy\dot{\rangle\mkern-3.0mu}\mkern
3.0mu\mathbin{\overset{\text{{def}}}{=}}\sum_{W^{\top}z\in\mathcal{V}_{W^{\top}}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}z\rangle\cdot\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{\top}z\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}.$

###  4.2 Depthwise Scaling of Random Variables

As mentioned in 4.1, both
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l-1}\rangle$
can be written as linear combination of “base” random variables:
$\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu\\}_{s\in\\{0,\ldots,t\\},m\in[L]},\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu\\}_{s\in\\{0,\ldots,t\\},m\in[L]},\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$.
Moreover, the coefficients of the linear combinations can be calculated in a
recursive way: by expanding
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\rangle$
using 4.1, we have

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle+\frac{1}{\sqrt{L}}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu+\frac{1}{\sqrt{L}}\sum_{s=1}^{t-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{s}^{l}\rangle\left(\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{s}^{l}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}-\frac{1}{\sqrt{L}}\langle
x_{s}^{l-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l-1}\rangle\right).$

The recursive formula for
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle$
is similar.

Using this induction, we claim in the linear combinations, the coefficient of
every
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\bullet\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is $\mathcal{O}(1/\sqrt{L})$, and the coefficient of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$
is $\mathcal{O}(1)$. We also claim the covariance between any pairs of random
variables in the form of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l-1}\rangle$
is $\mathcal{O}(1)$.

######  Proposition 4.2.

$\forall t,\forall s\leq t,\forall l,m$,
$\forall\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle\in\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle\\}$,

$\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}=\mathcal{O}\left(\frac{1}{\sqrt{L}}\right),\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}=\mathcal{O}\left(\frac{1}{\sqrt{L}}\right),\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle}=\mathcal{O}\left(1\right),\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle}=\mathcal{O}\left(1\right).$

$\forall t,s,l,m$,
$\forall\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}y\rangle\in\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle\\}$,
$\forall\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}z\rangle\in\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{s}^{m}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{s}^{m}\rangle\\}$,

$\langle
y\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}z\rangle=\mathcal{O}(1).$

The reasoning of 4.2 is provided in Appendix C. Note the computation of
covariance can also be written as a recursive formula. The reasoning relies
essentially on an inductive argument.

###  4.3 Infinite Depth Limit

Now we formalize our argument above and obtain the formula describing the
dynamics of the network when $L\to\infty$. We first write the coefficients of
the linear combinations as a six dimensional tensor
$\mathbf{\Gamma}_{t,s,a,b,l,m}$, where
$t,s\in\\{0,\ldots,T-1\\},a,b\in\\{0,1\\},l,m\in[L]$. Specifically,
$\mathbf{\Gamma}_{t,s,a,b,l,m}$ represents the derivative of
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle$
w.r.t.
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$. Here, we use $0$ to denote kets appears in the forward pass
($\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$), and $1$ to denote kets in the backward pass
($\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$). Formally,
$\mathbf{\Gamma}_{t,s,0,0,l,m}=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}$,
$\mathbf{\Gamma}_{t,s,0,1,l,m}=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{l}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}$,
$\mathbf{\Gamma}_{t,s,1,0,l,m}=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m}x_{s}^{m-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}$,
$\mathbf{\Gamma}_{t,s,1,1,l,m}=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{l}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{m\top}\widetilde{\delta}x_{s}^{m}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu}$.

However, it is hard to describe the limit of $\mathbf{\Gamma}$ because its
size increases along with $L$. Therefore, we define the following set of
functions
$\\{\Gamma_{t,s,a,b}:[0,1]\times(0,1]\to\mathbb{R}\\}_{t\in\\{0,\ldots,T-1\\},s\in\\{-1,\ldots,t\\},a,b\in\\{0,1\\}}$:
For $s\geq 0$,

$\Gamma_{t,s,a,b}\left(p,q\right)=\sqrt{L}\cdot\mathbf{\Gamma}_{t,s,a,b,\lceil
Lp\rceil,\lceil Lq\rceil}$

For $s=-1$,
$\Gamma_{t,-1,0,0}\left(p,q\right)=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{\lceil
Lp\rceil}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle},\Gamma_{t,-1,0,1}\left(p,q\right)=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{t}^{\lceil
Lp\rceil}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle},\Gamma_{t,-1,1,0}\left(p,q\right)=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{\lceil
Lp\rceil}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle},\Gamma_{t,-1,1,1}\left(p,q\right)=\frac{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{t}^{\lceil
Lp\rceil}\rangle}{\partial\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle}.$

Here $l,m$ are normalized to $[0,1]$ so the input domain of $\Gamma$s are
identical for different $L$; $\mathbf{\Gamma}_{t,s,a,b,l,m}$ is multiplied by
$\sqrt{L}$ because $\mathbf{\Gamma}_{t,s,a,b,l,m}=\mathcal{O}(1/\sqrt{L})$ by
4.2; and the extra $s=-1$ case helps us also capture the derivative w.r.t.
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle$
and
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle$.

Similarly, we can also define another set of function
$\\{C_{t,s,a}:(0,1]\to\mathbb{R}\\}_{t,s\in\\{-1,\ldots,T-1\\},a\in\\{0,1\\}}$
to describe the covariance between the “base” random variables: $\forall
p\in(0,1]$, let $l=\lceil Lp\rceil$,

  * •

$C_{t,s,0}\left(p\right)\mathbin{\overset{\text{{def}}}{=}}\mathrm{Cov}(\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{t}^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l}x_{s}^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu)=\langle
x_{t}^{l-1}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x_{s}^{l-1}\rangle$,

  * •

$C_{t,s,1}\left(p\right)\mathbin{\overset{\text{{def}}}{=}}\mathrm{Cov}(\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{t}^{l}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l\top}\widetilde{\delta}x_{s}^{l}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu)=\langle\widetilde{\delta}x_{t}^{l}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\widetilde{\delta}x_{s}^{l}\rangle$,

For $t=-1$,
$C_{-1,-1,0}\left(p\right)\mathbin{\overset{\text{{def}}}{=}}\mathrm{Cov}(\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}U\rangle)=1$,
and
$C_{-1,-1,1}\left(p\right)\mathbin{\overset{\text{{def}}}{=}}\mathrm{Cov}(\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}nV\rangle)=1$,
By 4.1, the “base” random variables of different “groups” are independent, so
we only tracks the covariance listed above.

Using this definition of $\Gamma$ and $C$, it is convenient to write their
recursive formula in the following lemma.

######  Lemma 4.3 (Finite depth recursive formula for $\Gamma$ and $C$
(Informal version of Lemma C.1)).

$\Gamma$ and $C$ can be computed recursively as follows:

$\displaystyle\Gamma_{t,r,0,b}\left(\frac{l}{L},q\right)=$
$\displaystyle\leavevmode\nobreak\
\Gamma_{t,r,0,b}\left(\frac{l-1}{L},q\right)+\mathbb{I}_{[(t=r)\wedge(b=0)\wedge(l=\lceil
Lq\rceil)]}$
$\displaystyle+\frac{1}{L}\sum_{s=0}^{t-1}\Gamma_{s,r,1,b}\left(\frac{l}{L},q\right)\left(\Gamma_{t,s,0,1}\left(\frac{l-1}{L},\frac{l}{L}\right)-C_{t,s,0}\left(\frac{l}{L}\right)\right).$
$\displaystyle\Gamma_{t,r,1,b}\left(\frac{l-1}{L},q\right)=$
$\displaystyle\leavevmode\nobreak\
\Gamma_{t,r,1,b}\left(\frac{l}{L},q\right)+\mathbb{I}_{[(t=r)\wedge(b=1)\wedge(l=\lceil
Lq\rceil)]}$
$\displaystyle+\frac{1}{L}\sum_{s=0}^{t-1}\Gamma_{s,r,0,b}\left(\frac{l-1}{L},q\right)\left(\Gamma_{t,s,1,0}\left(\frac{l}{L},\frac{l}{L}\right)-C_{t,s,1}\left(\frac{l}{L}\right)\right).$
$C_{t,s,a}(p)=\sum_{t^{\prime}=-1}^{t}\sum_{s^{\prime}=-1}^{s}\sum_{b\in\\{0,1\\}}\int_{0}^{1}\Gamma_{t,t^{\prime},a,b}(l/L,q)C_{t^{\prime},s^{\prime},b}(q)\Gamma_{s,s^{\prime},a,b}(l/L,q)\mathrm{d}q,$

where $l=\lceil Lp\rceil-1$ if $a=0$, and $l=\lceil Lp\rceil$ if $a=1$.

The proof of Lemma 4.3 is straightforward from Program 1. In Appendix C, we
also give a formal proof that $\Gamma$ and $C$ converge when $L$ grows to
infinity, in the case where $L$ is powers of $2$. The restriction on $L$ being
powers of $2$ is imposed for the convenience of the proof, and the convergence
of $\Gamma$ and $C$ is true in the general case. Moreover, we derive the
infinite depth behavior based on the recursion of $\Gamma$ and $C$ in Lemma
4.3.

######  Proposition 4.4 (Infinite depth limit of $\Gamma$ and $C$ (Informal
version of C.2)).

In the limit $L\to\infty$, we have

$\displaystyle\Gamma_{t,r,0,b}(p,q)=\mathbb{I}_{[(t=r)\wedge(b=0)\wedge(p\geq
q)]}+\int_{0}^{p}\sum_{s=0}^{t-1}\Gamma_{s,r,1,b}(p^{\prime},q)\cdot(\Gamma_{t,s,0,1}(p^{\prime},p^{\prime})-C_{t,s,0}(p^{\prime}))\mathrm{d}p^{\prime};$
$\displaystyle\Gamma_{t,r,1,b}(p,q)=\mathbb{I}_{[(t=r)\wedge(b=1)\wedge(p\leq
q)]}+\int_{p}^{1}\sum_{s=0}^{t-1}\Gamma_{s,r,0,b}(p^{\prime},q)\cdot(\Gamma_{t,s,1,0}(p^{\prime},p^{\prime})-C_{t,s,1}(p^{\prime}))\mathrm{d}p^{\prime};$
$\displaystyle
C_{t,s,a}(p)=\sum_{t^{\prime}=-1}^{t}\sum_{s^{\prime}=-1}^{s}\sum_{b\in\\{0,1\\}}\int_{0}^{1}\Gamma_{t,t^{\prime},a,b}(p,q)C_{t^{\prime},s^{\prime},b}(q)\Gamma_{s,s^{\prime},a,b}(p,q)\mathrm{d}q.$

The proof of 4.4 follows from Lemma 4.3. A rigorous proof requires first
showing the existence of a solution of the integral functional satisfied by
the couple $(\Gamma,C)$. The solution is typically a fixed point of the
integral functional in 4.4. After showing the existence, one needs to show
that $(\Gamma,C)$ 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 of $\Gamma$
and $C$ for $L=2^{k}$ ($k\in\mathbb{N}$) in Appendix C is a showcase of the
correctness of the proposition.

This gives a convergence in distribution:

######  Theorem 4.1.

In the $L\to\infty$ limit, the kets
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x^{L}_{s}\rangle,s=0,1,\ldots,$
converge in distribution as a zero-mean Gaussian process with kernel

$\langle
x^{L}_{s}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x^{L}_{t}\rangle=C_{t,s,1}(1).$

Thus, for each fixed neuron index $\alpha$, the collection $\\{x^{L}_{\alpha
s}\\}_{s\geq 0}$ converges in distribution to a zero-mean Gaussian process
with kernel $C_{t,s,1}(1)$ in the $n\to\infty$ then $L\to\infty$ limit.

For audience familiar with stochastic processes, we in fact have a weak
convergence of the entire continuous-depth-indexed process
$\\{\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}x^{p}_{s}\rangle,\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\delta
x^{p}_{s}\rangle\\}_{p\in[0,1],s\geq 0}$ in the Skorohod topology.

##  5 What Causes Hyperparameter Transfer?

In a popular misconception, hyperparameter transfer is implied by the
existence of a limit. For example, the fact that $\mu$P transfers
hyperparameters, in this misconception, is because of the existence of the
feature learning limit (aka the $\mu$ limit), the limit of $\mu$P as width
goes to infinity. However, this is not the case. Indeed, there are a plethora
of infinite-width limits, such as the NTK limit, but there can only be one way
how the optimal hyperparameters scale, so existence cannot imply transfer. In
a stronger version of this misconception, transfer is implied by the existence
of a “feature learning” limit. But again, this is False, because there are
infinite number of feature learning limits (where the $\mu$ limit is the
unique maximal one).

Instead, what is true is that the _optimal_ limit implies the transfer of
_optimal_ hyperparameters. For example, in the width limit case, $\mu$P is the
unique parametrization that yields a maximal feature learning limit. Compared
to all other limits, this is obviously the optimal one. Hence $\mu$P can
transfer hyperparameters across width.

So far, there is no _a priori_ definition for the “optimality” of a 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 can 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-$\mu$P
will turn out to be optimal in a sense that we define later in the
paper.101010There are important nuances here that will be spelled out in an
upcoming paper. For example, if the 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 1D space of the
global learning rate, then all infinite-width limits are defective — and
indeed there is no hyperparameter transfer where the bigger always does
better. More interestingly than the width case, here we have multiple modes of
feature learning when taking the depth limit and it 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.

In section 10, we provide experiments with $1/L$ block scaling
$(\alpha,\gamma)=(1,0)$, aka ODE scaling, which provably induces feature
learning in the infinite-depth limit, but is sub-optimal. Our results show a
significant shift in the optimal learning rate with this parametrization.

##  6 Preliminaries for the General Case

For the general case, we recall and extend the notation from the previous
sections and also define new ones.

##### 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.e. the dimension of all hidden
representations $x^{0},\ldots,x^{L}$. Let
$\xi\in\bm{\mathbb{\mathbb{R}}}^{d_{\mathrm{in}}}$ be the input of the
network, $U\in\bm{\mathbb{\mathbb{R}}}^{n\times{d_{\mathrm{in}}}}$ be the
input layer, and $V\in\bm{\mathbb{\mathbb{R}}}^{n\times e}$ be the output
layer, so that $x^{0}=U\xi$ and the model output w.r.t. $\xi$ is
$f(\xi)\triangleq V^{\top}x^{L}$. Let $\ell$ be the loss function absorbing
the label, and $\delta x^{l}$ be the gradient of $x^{l}$ w.r.t. the loss. We
denote variables at $t$-th training step by adding $t$ as a subscript, e.g.,
the input at step $t$ is $\xi_{t}$111111Here, 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., the hidden
representation of $l$-th layer at step $t$ is $x_{t}^{l}$, and the model
output at step $t$ is $f_{t}$. Let $T$ be the number of training steps.

###  6.1 Unified Scaling for SGD, Adam, and All Entrywise Optimizers

We extend the definition of entrywise update ([24]) for depth scaling,
allowing us to study the unified depth scaling for SGD, Adam, and other
optimization algorithms that perform only entrywise operations.

######  Definition 6.1.

A gradient-based update of parameter $w$ with both width and depth scaling is
defined by a set of functions
$\bm{Q}=\\{Q_{t}:\mathbb{R}^{t+1}\to\mathbb{R}\\}_{t\geq 0}$, and
$c,d,\delta,\gamma,\eta$. The update at time $t$ of the optimization is

$w\leftarrow w-\eta
n^{-c}L^{-\gamma}Q_{t}(n^{d}L^{\delta}g_{0},\ldots,n^{d}L^{\delta}g_{t}),$

where $g_{s},s=0,\ldots,t$, are the gradients of $w$ at time $s$.

For SGD,
$Q_{t}(n^{d}L^{\delta}g_{0},\ldots,n^{d}L^{\delta}g_{t})=n^{d}L^{\delta}g_{t}$,
and the “true” learning rate is $\eta n^{-c+d}L^{-\gamma+\delta}$. For Adam,

$Q_{t}(n^{d}L^{\delta}g_{0},\ldots,n^{d}L^{\delta}g_{t})=\frac{\frac{1-\beta_{1}}{1-\beta_{1}^{t+1}}\sum_{s=0}^{t}\beta_{1}^{t-s}n^{d}L^{\delta}g_{s}}{\sqrt{\frac{1-\beta_{2}}{1-\beta_{2}^{t+1}}\sum_{s=0}^{t}\beta_{2}^{t-s}(n^{d}L^{\delta}g_{s})^{2}+\epsilon}},$

and the “true” learning rate is $\eta n^{-c}L^{-\gamma}$.

The purpose of multiplying the gradients $n^{d}L^{\delta}$ before $Q_{t}$ is
to make sure the inputs to $Q_{t}$ are $\Theta(1)$ w.r.t. $n$ and $L$121212It
is called faithfulness in Yang and Littwin [24].; otherwise, the update might
be trivial when $n$ and $L$ become large. For example, if gradients are $o(1)$
entrywise, then, in Adam, directly feeding gradients to $Q_{t}$ will always
give an output of $0$ because of the constant $\epsilon>0$.

In this paper, we will only consider $d,\delta$ such that $n^{d}L^{\delta}g$
is $\Theta(1)$.131313Note $c,d,\delta,\gamma,\eta$ in Definition 6.1 can be
different for parameters, so it is possible to make every parameter to satisfy
the condition. As a result, the output of $Q_{t}$ is also $\Theta(1)$ in
general. Therefore, $n^{-c}L^{-\gamma}$ decides the scale of the update and
should be our focus. We call $\eta n^{-c}L^{-\gamma}$ the _effective learning
rate_.

###  6.2 $\mu$P and Widthwise Scaling

Maximal update parametrization ($\mu$P) [21] considers the change of
initialization and learning rate of each weight matrix in the network when
width scales up.141414Reparametrization is also included in the original
$\mu$P, but it is not necessary for the purpose of this paper.  It provides a
unique initialization and learning rate of each weight matrix as a function of
width $n$ that makes the update of each weight matrix maximal (up to a
constant factor). The benefit of $\mu$P is not only the theoretical guarantee
but also the hyperparameter stability when scaling up the width [23].

In this paper, we assume the widthwise scaling follows $\mu$P. That is, the
$c$ in the effective learning rate $\eta n^{-c}L^{-\gamma}$ and the
initialization variance of each weight matrix follows Table 2.

Table 2: Widthwise scaling of $\mu$P, where $c$ (defined in Definition 6.1) describes the widthwise scaling of the effective learning rate. | Input weights | Output weights | Hidden weights
---|---|---|---
Init. Var. | $1$ | $n^{-2}$ | $n^{-1}$
$c$ | $0$ | $1$ | $1$
|  |  |

###  6.3 Our Setup

We consider an $L$-hidden-layer residual network with biasless perceptron
blocks:

$\displaystyle x^{0}$ $\displaystyle=U\xi,$ $\displaystyle\forall
l\in[L],\quad x^{l}$
$\displaystyle=L^{-\alpha}\,\mathrm{MS}(\phi(h^{l}))+x^{l-1},\quad
h^{l}=W^{l}x^{l-1},$ $\displaystyle f$ $\displaystyle=V^{\top}x^{L}.$

where $\mathrm{MS}$ refers to Mean Subtraction and is given by
$\mathrm{MS}(x)=x-\langle x,1\rangle/n=Gx$ with $G=I-11^{\top}/n$, for any
$x\in\mathbb{R}^{n}$. The initialization and learning rate of $U,V$ follows
$\mu$P. The initialization of $W^{l}$ follows $\mu$P, and the learning rate of
$W^{l}$ is $\eta n^{-1}L^{-\gamma}$.

##### Mean Subtraction ($\mathrm{MS}$).

In general, without mean subtraction, the mean of $\phi$ will dominate the
depthwise dynamics. For example, when $\phi$ is relu, each layer will only add
nonnegative quantities to $x^{l}$ that on average is positive. Its
accumulation over depth either causes the network output to blow up if the
multiplier $L^{-\alpha}$ is too large, or lack feature diversity otherwise. As
we shall see, mean subtraction removes this failure mode and enable more
powerful infinite-depth limits.151515Note that using an _odd_ nonlinearity
will also achieve similar results because they have no mean under a
symmetrically distributed input, which is approximately the case for $h^{l}$
throughout training. This is the case for $\phi$ = identity that we discussed
earlier. But it turns out odd nonlinearities minimize feature diversity, so
mean subtraction is a much better solution.

######  Definition 6.2.

Fix a set of update functions
$\bm{Q}=\\{Q_{t}:\mathbb{R}^{t+1}\to\mathbb{R}\\}_{t\geq 0}$. A _depthwise
parametrization_ of the MLP residual network above is specified by a set of
numbers $\\{\alpha,\gamma,\delta\\}$ such that

  1. (a)

We independently initialize each entry of $W^{l}$ from $\mathcal{N}(0,n^{-1})$

  2. (b)

The gradients of $W^{l}$ are multiplied by $nL^{\delta}$ before being
processed by $Q_{t}$: i.e., the update at time $t$ is

$W^{l}\leftarrow W^{l}-\eta
n^{-1}L^{-\gamma}Q^{l}_{t}(nL^{\delta}g_{0},\ldots,nL^{\delta}g_{t})$ (4)

where $g_{s},s=0,\ldots,t$, are the gradients of $W^{l}$ at time $s$ and
$Q_{t}$ is applied entrywise.

##### Miscellaneous notations.

For a vector $x$, let $[x]_{i}$ be its $i$-th coordinate. For a matrix $M$,
let $[M]_{i}$ be its $i$-th row. Let $I$ be the identity matrix, and
${\bm{1}}$ be the full one vector. For $m\in\mathbb{N}^{+}$, let
$[m]=\\{1,\ldots,m\\}$. Let $\otimes$ be the Kronecker product.

##  7 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 to 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 $\Theta(1/\sqrt{L})$.
Assuming the updates are faithful (i.e., input to gradient processing
functions are $\Theta(1)$ entrywise), the update size can be at most $1/L$ for
the hidden layers, by an (Jacobian) operator-norm 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 a 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 us first define the notions of training
routine, stability, faithfulness, and non-triviality. Hereafter, all the
aymptotic notations such as $\mathcal{O}$, $\Omega$ and $o$ should be
understood in the limit “$n\to\infty,\textrm{ then }L\to\infty$”. For random
variables, such notations should be understood in the sense of weak
convergence (convergence in distribution). When we use the notation
$x=\mathcal{O}(1)$ for some vector $x=(x_{1},\dots,x_{n})\in\mathbb{R}^{n}$,
it should understood in the sense that for all $i\in[n],x_{i}=\mathcal{O}(1)$.
Lastly, we will use bold characters (e.g. $\bm{h}$ instead of $h$) to denote
‘batched’ versions of the quantities. This is 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 F.4, we provide “heuristic” proofs that can be made
rigorous under non-trivial technical conditions. We also showcase the
correctness of the claims by proving them rigorously in our linear setting in
Appendix D. We believe this additional layer of complexity is unneeded and
does not serve the purpose of this paper.

######  Definition 7.1 (Training routine).

A training routine is the package of $\eta$, $\bm{Q}$, and the input batches.

######  Definition 7.2 (Stability).

We say a parametrization is

  1. 1.

_stable at initialization_ if

$\bm{h}_{0}^{l},\bm{x}_{0}^{l}=\mathcal{O}(1),\forall
l\in[L],\quad\text{and}\quad\bm{f}_{0}=\mathcal{O}(1).$ (5)

  2. 2.

_stable during training_ if for any training routine, any time $t\geq 0$,
$l\in[L]$, we have

$\Delta\bm{h}_{t}^{l},\Delta\bm{x}_{t}^{l}=\mathcal{O}(1),\forall
l\in[L],\quad\text{and}\quad\Delta\bm{f}_{t}=\mathcal{O}(1),$

where the symbol ‘$\Delta$’ 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 7.3 (Faithful).

We say a parametrization is _faithful at step $t$_ if
$\bm{h}_{t}^{l}=\Theta(1)$ for all $l\in[L]$. 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\geq 1$).

Note faithfulness here refers to “faithfulness to $\phi$”, meaning the input
to $\phi$ is $\Theta(1)$. This is different from the definition of
faithfulness in Yang and Littwin [24], where faithfulness refers to
“faithfulness to $Q$” meaning the input to $Q$ is $\Theta(1)$. “faithfulness
to $Q$” is already assumed in this work as mentioned in Section 6.1.

######  Definition 7.4 (Nontriviality).

We say a parametrization is _trivial_ if for every training routine and any
time $t\geq 1$, $\bm{f}_{t}-\bm{f}_{0}\xrightarrow{\mathrm{\mathrm{a.s.}}}0$
in the limit “$n\to\infty,$ then $L\to\infty$” (i.e., the function does not
evolve in the infinite-width-then-depth limit). We say the parametrization is
_nontrivial_ otherwise.

######  Definition 7.5 (Feature Learning).

We say a parametrization induces _feature learning_ in the limit
“$n\to\infty$, then $L\to\infty$”, if there exist a training routine, and
$t\geq 1$, and any $\lambda>0$, we have $\Delta\bm{h}_{t}^{\lfloor\lambda
L\rfloor}=\Theta(1)$.

###  7.1 Main Claims

We are now ready to state the main results. The next claim provides a
necessary and sufficient condition under which a parametrization is stable at
initialization.

######  Claim 7.1.

A parametrization is stable at initialization iff $\alpha\geq 1/2$.

Claim 7.1 is not new and similar results were reported by Hayou et al. [7].
However, Hayou et al. [7] 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 7.2.

Consider a parametrization that is stable at initialization. Then the
following hold (separately from each other).

  * •

It is stable during training as well iff $\alpha+\gamma\geq 1$.

  * •

It is nontrivial iff $\alpha+\gamma\leq 1$.

Therefore, it is both stable and nontrivial iff $\alpha+\gamma=1$.

From Claim 7.1 and Claim 7.2, having $\alpha+\gamma=1$ and $\alpha\geq 1/2$ is
a necessary and sufficient condition for a parametrization to be stable and
nontrivial throughout training. In the next result, we therefore restrict our
analysis to such parametrizations and study their faithfulness.

######  Claim 7.3.

Consider a stable and nontrivial parametrization. The following hold
(separately from each other).

  * •

It is faithful at initialization iff $\alpha\geq 1/2$. As a result,
$\alpha=1/2$ is the minimal choice of $\alpha$ that guarantees faithfulness.

  * •

It is faithful during training iff $\alpha\leq 1$.

Therefore, a stable and nontrivial parametrization is faithful iff
$\alpha\in[1/2,1]$.

The first claim follows from well-known calculations of randomly initialized
residual networks [7]. For the second claim, the intuition here is just that
if $\alpha+\gamma=1$ and $\alpha>1$ then $\gamma<0$, i.e., the update size
blows up with depth. This would then cause the input to the nonlinearities to
blow up with size.

One might argue that faithfulness at initialization is not important (e.g.
features at initialization could converge to zero without any stability or
triviality issues) and what matters is faithfulness throughout training. It
turns out that faithfulness at initialization plays a crucial role in the
optimal use of network capacity. To see this, we first define the notion of
feature diversity exponent, which relates to the similarity in the features of
adjacent layers.

######  Definition 7.6 (Feature Diversity Exponent).

We say a parametrization has feature diversity exponent $\kappa\geq 0$ if
$\kappa$ is the maixmal value such that for all $\lambda\in[0,1]$ and
sufficiently small $\epsilon>0$, and all time $t$,

$\frac{1}{\sqrt{n}}\left\|\bm{x}^{\lfloor(\lambda+\epsilon)L\rfloor}_{t}-\bm{x}_{t}^{\lfloor\lambda
L\rfloor}\right\|=\Omega(\epsilon^{1-\kappa}),$

where $\Omega(1)$ should be interpreted in the limit “$n\to\infty$, then
$L\to\infty$, then $\epsilon\to 0$”. We say a parametrization is _redundant_
if $\kappa=0$.

In other words, the feature diversity exponent $\kappa$ is a measure of how
different the outputs are in layers that are close to each other. With
$\kappa=0$, the output of each layer is essentially the same as the output of
the previous layer in the sense that the rate of change from one layer to the
next is bounded (at least locally), and hence the network is intuitively
“wasting” parameters.

######  Claim 7.4.

Consider a stable and nontrivial parametrization that is furthermore faithful
during training (but not necessarily at initialization). Then it is redundant
if $\alpha\in(1/2,1]$.

To understand the intuition behind Claim 7.4, let us see what happens when
$\alpha>1/2$. In this case, the randomness of the initialization weights will
have no impact on training trajectory as depth increases. To see this,
consider some layer index $\lfloor\lambda L\rfloor$. The blocks are divided by
$L^{\alpha}$ which is larger than the magnitude of accumulated randomness (of
order $(\lambda L)^{1/2}$). 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 output matrices). When
depth goes to infinity, the contribution of the randomness 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-$\mu$P.

######  Claim 7.5 (Depth-$\mu$P).

$\alpha=\gamma=1/2$ is the unique parametrization that is stable, nontrivial,
faithful, induces feature learning, and achieves maximal feature diversity
with $\kappa=1/2$.

In terms of feature diversity, a phase transition phenomenon occurs when
$\alpha=1/2$. More precisely, for Depth-$\mu$P, we can show that
$n^{-1/2}\left\|{\bm{x}^{\lfloor(\lambda+\epsilon)L\rfloor}_{t}-\bm{x}_{t}^{\lfloor\lambda
L\rfloor}}\right\|=\mathcal{O}(\epsilon^{1/2})$ while the same quantity is
$\mathcal{O}(\epsilon)$ for all $\alpha\in(1/2,1]$, which suggests that
Depth-$\mu$P yields _rough_ path for $\bm{x}_{t}$. 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.161616The
reader might ask whether we can obtain an exponent smaller than $1/2$. This is
indeed possible, but it will entail using correlated weights. We leave this
question for future work.

Moreover, with $\alpha=1$, there is a phenomenon of feature collapse in the
sense that the features will be contained in the $\sigma$-algebra generated by
the input and output layers, but contains no randomness from the hidden layers
(see Section F.2). Intuitively, the case of $\alpha=1$ is analogous to width
situation, where deep mean field collapses to a single neuron (all neurons
become essentially the same). For depth, the features (layers) are still
relatively different but the redundancy does not allow significant variety in
these features.

###  7.2 Sublety: Layerwise (local) linearization but not global linearization

######  Definition 7.7.

We say a parametrization induces layerwise linearization iff each layer can be
linearized without changing the network output when $L\to\infty$, that is,
$\forall l\in[L]$,

$L^{-\alpha}G\left(\phi(W^{l}_{t}\bm{x}_{t}^{l-1})-\phi(W^{l}_{0}\bm{x}_{t}^{l-1})-\phi^{\prime}(W_{0}^{l}\bm{x}_{t}^{l-1})\odot((W_{t}^{l}-W_{0}^{l})\bm{x}_{t}^{l-1})\right)=o(L^{-1})$

######  Claim 7.6.

A stable and nontrivial parametrization induces layerwise linearization iff
$\alpha\in[1/2,1)$.

However, note that this does not imply the entire network is linearized
(w.r.t. all the parameters in the sense of Neural Tangent Kernel). 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.

If the initialization of the output layer is $L$ times larger than our setup
(assuming $L\ll n$ so the widthwise scaling still follows $\mu$P), it may
induce a parametrization that can linearize the entire network. In that
situation, the learning rate has to be $L$ times smaller than Depth-$\mu$P 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.

##  8 Feature Diversity

In this section, we show that the choice of nonlinearity and placement of
nonlinearities can affect feature diversity greatly.

###  8.1 Gradient Diversity

_Gradient diversity_ is an important factor toward feature diversity. Observe
that the gradient $\delta x^{l}$ at $x^{l}$ is continuous in $l$ in the limit
$L\to\infty$. In a linear model (or the pre-nonlin model, where nonlinearity
is put before the weights), this causes $\delta h^{l}=L^{-\alpha}\delta x^{l}$
to be very similar between neighboring blocks. As a result (because the
weights $W^{l}$ receives an update proportional to $\delta h^{l}\otimes
x^{l-1}$), in the next forward pass, neighboring blocks contribute very
similarly to the main branch $x^{l}$. This leads to a waste of model capacity.

###  8.2 Pre-Nonlin Leads to Poor Performance

For example, in Figure 2, for a relu pre-nonlin resnet (i.e. blocks are given
by $W^{l}\phi(x^{l-1})$ instead of $\phi(W^{l}x^{l-1})$), we see that although
Depth-$\mu$P indeed transfers hyperparameters (as predicted by our theory),
the performance is dramatically worse than the post-nonlin resnet in Figure
10, and depth gives no performance gains beyond 8 layers. Specifically, it is
because $\delta h^{l}=L^{-\alpha}\delta x^{l}$ like the linear case, and
$\phi(x^{l-1})$ is also similar between neighboring blocks. As a result, the
gradient of the weights $W^{l}$, proportional to $\delta
h^{l}\otimes\phi(x^{l-1})$, has little diversity compared to nearby blocks.

Figure 2: Pre-Nonlin Leads to Poor Performance Although Depth-$\mu$P 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 10 for more details about the setup.

###  8.3 Maximizing Feature Diversity with Absolute Value Nonlinearity

In a nonlinear model, we have $\delta h^{l}=\delta
x^{l}\odot\phi^{\prime}(h^{l})$. Because $h^{l}$ is almost independent from
all other $h^{m},m\neq l$ in the Depth-$\mu$P limit, $\phi^{\prime}(h^{l})$
can serve to decorrelate the $\delta h^{l}$, depending on what $\phi$ is. For
example, if $\phi$ is relu, then $\phi^{\prime}$ is the step function. $h^{l}$
is approximately a zero-mean Gaussian in the Depth $\mu$P limit, so that
$\phi^{\prime}(h^{l})$ is approximately 0 or 1 with half probability each.
This decorrelates $\delta h^{l}$ much better than the linear case. But of
course, this line of reasoning naturally leads to the conclusion that
$\phi^{\prime}=\mathrm{sign}$ would be the best decorrelator of $\delta h^{l}$
and the maximizer of feature diversity (with $\phi$ among the class of
positively 1-homogeneous functions) — then $\delta h^{l}$ and $\delta h^{m}$
are completely decorrelated for $l\neq m$.

Indeed, as shown in Figure 3, swapping in absolute value for $\phi$
dramatically improves the training performance of deep (block depth 1)
resnets.

Figure 3: Improving performance with absolute value non-linearity, which
maximizes feature diversity. (Networks are trained on CIFAR-10 with Adam.).
See Figure 10 for more details about the setup.

In general, in lieu of absolute value, any even nonlinearity would suffice.

###  8.4 Feature Diversity is in Tension with Layerwise Linearization

The reason that $\phi^{\prime}(h^{l})$ can decorrelate $\delta h^{l}$ is very
much related to layerwise linearization. Recall that in Depth-$\mu$P, $h^{l}$
can be decomposed to a zero-mean Gaussian part $\widehat{h}^{l}$ of size
$\Theta(1)$ and a correction term $\dot{h}^{l}$ of size $\Theta(L^{-1/2})$
(corresponding to the decomposition
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}h^{l}\rangle=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}h^{l}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}h^{l}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu$). $\widehat{h}^{l}$ is independent from $\widehat{h}^{m}$ for $m\neq l$
but $\dot{h}^{l}$ can be very strongly correlated to all other $\dot{h}^{m}$.
Thus, $\phi^{\prime}(h^{l})$ can decorrelate $\delta h^{l}$ precisely because
$\widehat{h}^{l}$ dominates $\dot{h}^{l}$, and this is also precisely the
reason we have layerwise linearization.

In the $1/L$ scaling $(\alpha,\gamma)=(1,0)$, $\widehat{h}^{l}$ is on the same
order as $\dot{h}^{l}$ and layerwise linearization does not occur, but also
$\phi^{\prime}(h^{l})$ can no longer effectively decorrelated $\delta h^{l}$.

Once again, we remind the reader that layerwise linearization in this case is
not detrimental (in this block depth 1 case) because $\widehat{h}^{l}$ in fact
accumulate contributions from the learned features of all previous blocks and
thus strongly depends on the learning trajectory (in contrast to the
(widthwise) NTK case where $\widehat{h}^{l}$ is already determined at
initialization).

##  9 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 $\mu$P.

In most of this work, we have considered depth-1 MLP for $g^{l}$ in eq. 1,
it’s straightforward to derive and classify the infinite-width-then-infinite-
depth limits for larger depths in each block. In particular, the following
$1/\sqrt{L}$ scaling still makes sense in this more general setting with block
depth $k$ and leads to a well defined limit:

$\displaystyle x^{l}$ $\displaystyle=x^{l-1}+\frac{a}{\sqrt{L}}\cdot
g^{l}(x^{l-1};W^{l1},\ldots,W^{lk}),\quad\text{$\Theta(1)$ initialization
scale,}\quad\text{$\Theta(1/\sqrt{L})$ learning rate}$ (6)

This is what we call Depth-$\mu$P in the block depth 1 case, but we shall not
use this name in the general block depth case because _this parametrization is
no longer optimal_.171717What we exactly mean by _optimal_ will be explained
below.

###  9.1 Block Depth $\geq 2$ is Defective

A very clear symptom of this is that the _performance of block-depth-2 resnets
is worse than that of block-depth-1 networks_ , when matching parameter count,
although they can (but not always) catch up after training for a long time
(figs. 4 and 5).

Figure 4: Block Depth 2 < Block Depth 1, Relu. In relu resnet with no LN,
block depth 2 does worse than block depth 1 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

Figure 5: Block Depth 2 < 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

Simultaneously, we are seeing nontrivial or even significant hyperparameter
shifts as the total number of blocks increases (fig. 6).

Figure 6: Block Depth 2 Hyperparameter Shift in relu resnet with no LN (Left)
and abs resnet with LN (Right).

###  9.2 Defect of $1/\sqrt{L}$ Scaling in Block Depth 2

The reason that the $1/\sqrt{L}$ scaling is no longer fine in the block depth
$\geq 2$ case is the _linearization of the multiplicative interaction_ between
the layers in the block. Indeed, just like the block depth 1 case, the
$1/\sqrt{L}$ scaling forces the weight updates $\Delta W$ of each weight
matrix to be $\Theta(\sqrt{L})$ smaller than the initialization $W_{0}$. 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(x;W^{\bullet})$
is only a _summation_ , instead of _product_ , of individual contributions
from each layer’s weights updates.

When aggregated over all $L$ blocks, the result is that there is only
multiplicative interaction of $\Delta W$ 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 $M^{L}\cdots M^{1}$ where each $M^{l}$
can be one of $I,W_{0}^{l2}W_{0}^{l1},W_{0}^{l2}\Delta W^{l1},$ or $\Delta
W^{l2}W_{0}^{l1}$, but NOT $\Delta W^{l2}\Delta W^{l1}$. All other
contributions (which all involve within-block interactions like $\Delta
W^{l2}\Delta W^{l1}$) are subleading. In the general nonlinear case, replacing
the block

$\phi(W^{l2}\phi(W^{l1}x^{l-1}))$

with the linearized version

$\phi(h^{l}_{\wedge})+\phi^{\prime}(h^{l}_{\wedge})\odot[\Delta
W^{l2}\phi(h^{l}_{\vee})]+\phi^{\prime}(h^{l}_{\wedge})\odot[W_{0}^{l2}(\phi^{\prime}(h^{l}_{\vee})\odot[\Delta
W^{l1}x^{l-1}])]$

will achieve the same performance as depth $L\to\infty$, where
$h^{l}_{\wedge}=W_{0}^{l2}\phi(h^{l}_{\vee})$ and
$h^{l}_{\vee}=W_{0}^{l1}x^{l-1}$.

When block depth $k=1$ (our main subject of study in this work), _all_
interactions are included but this is no longer true when $k>1$.

In fig. 7, the heatmap of loss as a 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
$\approx-2$ when the number of blocks is $2^{1}$. In other words, around the
optimum, double the learning rate while dividing the block multiplier by 4 has
similar performance. This is because $\Delta W^{l1}$ and $\Delta W^{l2}$
interact _multiplicatively_ , so that doubling their sizes leads to
quadrupling their contribution to 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 $\approx-1$ when the
depth is $2^{10}$: Doubling the learning rate while halving the block
multiplier has similar performance. This reflects the fact that $\Delta
W^{l1}$ and $\Delta W^{l2}$ now interact _additively_.

Intermediate depths interpolate this phenomenon, as seen in the plot for depth
$2^{5}$.

Figure 7: The "slope" of the optimal sublevel set in the (learning rate, block
multiplier) space changes from $-2$ to $-1$ as depth goes from $2^{1}$ to
$2^{10}$. Here we use absolute value nonlinearity with layer normalization,
block depth 2, and networks are trained for 50 epochs with Adam on CIFAR-10.

In the same heatmaps, one can see the optimal (learning rate, block
multiplier) (in the $1/\sqrt{L}$ parametrization) shifts from the middle of
the grid to the upper left as depth goes from $2^{5}$ to $2^{10}$,
demonstrating the lack of hyperparameter transfer.

This change in slope is seen in relu networks as well, with or without
layernorm.

Finally, we note that the $1/\sqrt{L}$ scaling still yields a $L\to\infty$
limit where the network still learns features as a whole, even though within
each block this is no longer true. Thus, this is another reminder that mere
"feature learning" does not imply "hyperparameter transfer"!

###  9.3 Classification of Parametrizations

These heatmaps already demonstrate that no parametrization of (global learning
rate181818meaning, the learning tied across all layers in a block, 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 a sublevel set of another
slope.

But even if we allow learning rate to vary between layers in a block, no
stable, faithful, nontrivial parametrization can avoid the linearization
problem described above.

For simplicity, fix a positive-homogeneous nonlinearity and block depth
2.191919but our arguments generalize trivially to arbitrary block depth $\geq
2$ We consider the space of hyperparameters consisting of the learning rate
for each of the layers in a block, as well as the block multiplier (one for
each block); WLOG all weights are initialized $\Theta(1)$.202020This is WLOG
because the nonlinearities are homogeneous This yields a space of dimension
$\mathrm{blockdepth}+1=3$.

Indeed, for this to happen, the weight update $\Delta W^{li}$ must be at least
of order $\Omega(1)$ (size of initialization) for some $i$. But this would
contribute a drift term to the block output $g^{l}=g^{l}(x^{l-1};W^{\bullet})$
that is as large as the noise term. This then implies that either the
parametrization is unstable (if the block multiplier $L^{-\alpha}$ is
$\Omega(1/{L})$) or lacks feature diversity (if the block multiplier
$L^{-\alpha}$ is $O(1/L)$).

For example, in a linear model,

$L^{\alpha}\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}g^{l}\rangle=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W^{l2}W^{l1}x^{l-1}\rangle=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\Delta
W^{l2}W^{l1}x^{l-1}\rangle.$

$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is independent and zero-mean across $l$ (the noise term), while
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\Delta
W^{l2}W^{l1}x^{l-1}\rangle$ is correlated across $l$ (the drift term).
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\widehat{\rangle\mkern-3.0mu}\mkern
3.0mu$ is always $\Theta(1)$ because the $W^{l2}_{0},W^{l1}_{0}$ are. If
$\Delta W^{l2}$ is $\Omega(1)$, then
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}\Delta
W^{l2}W^{l1}x^{l-1}\rangle=\Omega(1)$ as well, making the drift term as large
as the noise term. If $\Delta W^{l1}$ is $\Omega(1)$, then
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}\Delta
W^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern 3.0mu=\Omega(1)$, causing
$\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu=\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}W_{0}^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern
3.0mu+\mathchoice{\scalebox{0.7}[1.0]{$\displaystyle\talloblong$}}{\scalebox{0.7}[1.0]{$\textstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptstyle\talloblong$}}{\scalebox{0.7}[1.0]{$\scriptscriptstyle\talloblong$}}W_{0}^{l2}\Delta
W^{l1}x^{l-1}\dot{\rangle\mkern-3.0mu}\mkern 3.0mu$ to be
$\Omega(1)$.212121One can also observe that if $\Delta W^{l1}=\Omega(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 $\Delta W^{li}$ cannot be $\Omega(1)$
for any $i$.

The same argument can be straightforwardly adapted to nonlinear MLPs (with
mean subtraction) and arbitrary block depth $\geq 2$, and as well to general
nonlinearities that are not necessarily positive-homogeneous, with
hyperparameter space enlarged to include initialization.

###  9.4 So What is the Optimal Parametrization?

All 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 1-dimensional
hyperparameter space of the global learning rate in SP. Discovering these
missing hyperparameters will be an important question for future work.

Figure 8: Trained linear network converges to its infinite width limit which
is computed recursively based on $\Gamma$ and $C$. Depth is fixed at 64,width
varies between $2^{7},2^{8},\ldots,2^{13}$. Networks are 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 $x$-axis is 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 is indicated with dashed lines.
Figure 9: Under Depth-$\mu$P, infinite wide linear network training converges
when increasing the depth. Infinite wide linear networks of depth
$2^{4},2^{5},\ldots,2^{9}$ are computed recursively based on $\Gamma$ and $C$.
The root mean square statistics ($y$-axis) at 1st, 5th and 10th steps are
plotted across the depth ($x$-axis).

##  10 Experiments

###  10.1 Verifying the Theory in the Linear Case

In Section 4, we showed that a complete description of the training dynamics
of linear networks can be formulated in terms of $\Gamma$ and $C$. In this
section, we provide empirical results supporting our theoretical findings. We
first verify the finite-depth recursive formula for $\Gamma$ in Lemma 4.3 is
the correct limit when the width goes to infinity, then proceed to show that
the infinite-depth limit is the correct one.

##### Infinite-width limit.

In Figure 8, we train a series of $64$-layer linear networks of width
$2^{7},2^{8},\ldots,2^{13}$ with $1,5,10$ steps on MNIST, and plot the root
mean square222222The root mean square of a vector $x=(x_{1},\ldots,x_{n})$ is
$\sqrt{\frac{\sum_{i=1}^{n}x_{i}^{2}}{n}}$, which is denoted as “l2” in
Figures 8 and 9. of the layer outputs using solid lines. We also compute the
infinite width limit of the corresponding statistics using the recursive
formula for $\Gamma$ and plot them as dashed horizontal lines. For clarity of
the figure, we only plot the statistics of the input 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^{4}$ to $2^{9}$. We again compute the root mean square values of the
layer outputs using the recursive formula for $\Gamma$, and plot them in
Figure 9 with depth being $x$-axis. For clarity of the figure, we only plot
the statistics of the input layer, output layer, and hidden layers of index
$L/4$, $L/2$, $3L/4$, and $L$. One can observe that the statistics of the
layer outputs converge quickly when the depth grows from $2^{4}$ to $2^{9}$,
which verifies our convergence result.

Figure 10: Train logloss versus learning rate for width $n=256$ and varying
depths. The network consists of MLP blocks (with block depth 1), trained for
50 epochs on CIFAR10 dataset using Adam. The batch size is fixed to $64$. We
tune the depth $2^{3}$ network to obtain the optimal
$(\log_{2}(a),\log_{2}(\eta/1e-3))=(1,0)$, and scale all deeper networks using
$2^{3}$ as base depth. The reader can check that the $L=2^{3}$ curves in each
columns are the same. We show the logloss versus the learning rate of the
hidden layers (input/output layers fixed) for three parametrizations:
Depth-$\mu$P (Top), Scaling only the blocks (no LR scaling), i.e. $\gamma=0$
(Middle), and Standard Parametrization without any scaling ($\alpha=\gamma=0$)
(Bottom). Each curve represents the average training loss over a time slice of
1000 steps for depths $2^{k}$ for $k\in\\{1,2,\dots,10\\}$. Confidence
intervals are based on 5 seeds. The results show that Depth-$\mu$P preserves
the optimal learning rate while consistently improving the training loss as
depth increases. If we only scale the blocks without scaling the LR
($\alpha=1/2,\gamma=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 a significant shift in the
optimal learning rate as well. For SP, we cap the log loss at 1, which is why
for depth $2^{9},2^{10}$, we have a black horizontal line at $LogLoss=1$.

###  10.2 Hyperparameter Transfer

In this section, we provide empirical evidence to show the optimality of
Depth-$\mu$P 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 is parameterized
as follows

$x^{l}=x^{l-1}+a\times L^{-\alpha}\textrm{MS}(\phi(W^{l}x^{l-1})),$

and the weights are trained with the rule

$W^{l}\leftarrow W^{l}-\eta\times
n^{-1}L^{-\gamma}Q^{l}_{t}(nL^{\delta}g_{0},\ldots,nL^{\delta}g_{t}),$

where the learning rate $\eta$ and the block multiplier $a$ are the
_hyperparameters_.232323Note that $\eta$ here is the constant, and the
effective learning rate is given by $\eta n^{-1}L^{-\gamma}$. The values of
$\alpha,\gamma$ depend on the parametrization of choice. For Depth-$\mu$P, we
have $\alpha=\gamma=1/2$, and for standard parametrization, we have
$\alpha=0,\gamma=1$.242424In 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 $L^{-1}$ to preserve
faithfulness. In our experiments, we assume base depth $8$, meaning that we
replace $L$ by $L/8$ in the parametrization above.

##### Learning rate transfer ($\eta$).

In Figure 10, we show the training loss versus learning rate for depths
$2^{k}$, for $k\in\\{3,4\dots,10\\}$. For Depth-$\mu$P, a convergence pattern
can be observed for the optimal learning rate as depth grows. Optimal learning
rates for small depths (e.g. $L=2^{3}$) exhibit a mild shift which should be
expected, as our theory shows convergence in the large depth limit. However,
starting from depth $L=2^{6}$, the optimal learning rate is concentrated
around $10^{-3}$. For parametrization that only scales the multiplier but not