# A Spectral Condition for Feature Learning

Greg Yang
xAI
&James B. Simon∗
UC Berkeley & Imbue &Jeremy Bernstein∗
MIT Equal contribution.

###### Abstract

The push to train ever larger neural networks has motivated the study of
initialization and training at large network width. A key challenge is to
scale training so that a network’s internal representations evolve
nontrivially at all widths, a process known as _feature learning_. Here, we
show that feature learning is achieved by scaling the spectral norm of weight
matrices and their updates like $\smash{\sqrt{\texttt{fan-out}/\texttt{fan-
in}}}$, in contrast to widely used but heuristic scalings based on Frobenius
norm and entry size. Our spectral scaling analysis also leads to an elementary
derivation of _maximal update parametrization_. All in all, we aim to provide
the reader with a solid conceptual understanding of feature learning in neural
networks.

##  1 Introduction

Recent years have seen an unprecedented push to train deep learning systems
with more and more parameters, leading to powerful models across domains and
the unlocking of qualitatively new capabilities (Silver et al., 2016, Brown et
al., 2020, Ramesh et al., 2022). This continuing trend, combined with the
technical challenges of training large models, has motivated much recent study
of the dynamics of neural networks at large width, and more generally the
study of how their dynamics scale as network width grows. This program has
yielded a cornucopia of theoretical insights (Lee et al., 2018, Jacot et al.,
2018, Arora et al., 2019, Canatar et al., 2021) and practical scaling
recommendations (Yang & Hu, 2021b, Yang et al., 2021, Dey et al., 2023).

A key challenge when training a network of large width is to ensure that
feature learning occurs at hidden layers. By this, we mean that the
hyperparameters of the network are scaled in a manner such that the hidden
representations of the network (as obtained by partial evaluation of the
network up to a certain layer) change substantially over the course of
training. Naïve hyperparameter scaling rules, including the well-studied
“neural tangent parametrization” (NTP), in fact lose feature learning at large
width (Lee et al., 2019, Sohl-Dickstein et al., 2020). But ample evidence
supports the conclusion that proper feature learning is necessary for
achieving optimal performance on many tasks (Lee et al., 2020, Fort et al.,
2020, Atanasov et al., 2022, Vyas et al., 2022). Furthermore, scaling training
correctly can lead to new functionality such as hyperparameter transfer. For
instance, the recently proposed maximal update parametrization (Yang & Hu,
2021b, Yang et al., 2021) allows for transferring hyperparameters from narrow
models to wide models, avoiding the cost of tuning the wide model directly.

Maximal update parametrization ($\mu$P) is derived by fairly involved “tensor
programs” arguments that track feature distributions analytically in the
infinite width limit. Anecdotally, the principles underlying $\mu$P are not
well understood by the community. In this paper, we provide a new perspective
on $\mu$P, showing that its scaling relations can be obtained by elementary
linear algebra arguments. In short, we show that $\mu$P is equivalent to
scaling the spectral norm of any weight matrix or update like
$\smash{\sqrt{\texttt{fan-out}/\texttt{fan-in}}}$. This simple condition has
various favorable numerical properties that contrast sharply with heuristic
optimization strategies based on controlling the Frobenius norm (You et al.,
2017) or entry size (Kingma & Ba, 2015) of updates. In the authors’
experience, the spectral scaling condition both simplifies the implementation
of $\mu$P in code, and is significantly easier to work with theoretically,
leading to further conceptual advances in our research (Bernstein et al.,
2023).

On a more fundamental level, an important step to solving many problems in
classical computer science is to write down a suitable distance function for
the problem at hand (Dhillon & Tropp, 2008). This idea is of particular
importance in the design of optimization algorithms, where a notion of
parameter distance is needed (Nemirovsky & Yudin, 1983, Amari, 1998). While it
can be tempting to use the Euclidean norm on parameter vectors to measure
distance, this naïve choice risks discarding the structure of the problem. For
example, neural networks involve compositions of linear operators, which we
refer to as their operator structure. Past efforts to metrize the space of
neural networks while accounting for their operator structure have included
using the Frobenius norm to measure distance between matrices (Bernstein et
al., 2020a), which motivates various optimization algorithms that make
Frobenius-normalized updates (You et al., 2020; 2017, Shazeer & Stern, 2018,
Liu et al., 2021). This paper shows that the spectral norm provides a better
notion of distance between operators in the context of deep learning.

###  1.1 Road map for the paper

To begin, we state the precise conditions on hidden features that we wish to
ensure. We will ask for two things: both the features and their updates upon a
step of gradient descent must be the proper size.

######  Desideratum 1 (Feature learning).

Let ${\bm{h}}_{\ell}({\bm{x}})\in\mathbb{R}^{0pt_{\ell}}$ denote the features
of input ${\bm{x}}$ at layer $\ell$ of a neural network, and let
$\Delta{\bm{h}}_{\ell}({\bm{x}})\in\mathbb{R}^{0pt_{\ell}}$ denote their
change after a gradient step. We desire that:
$\displaystyle\left|\\!\left|{\bm{h}}_{\ell}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})\quad\text{and}\quad\left|\\!\left|\Delta{\bm{h}}_{\ell}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}}),\quad\text{
at layers }\ell=1,...,L\\!-\\!1.$

Let us unpack 1. These conditions treat the $\ell^{2}$-norms of the feature
vectors ${\bm{h}}_{\ell}({\bm{x}})$ and $\Delta{\bm{h}}_{\ell}({\bm{x}})$, a
framing which will prove convenient. 1 amounts to asking that the “typical
element size” of vectors ${\bm{h}}_{\ell}({\bm{x}})$ and
$\Delta{\bm{h}}_{\ell}({\bm{x}})$ is $\Theta(1)$ with respect to width
$0pt_{\ell}$ (we give a review of big-$\Theta$ notation in Section 2).
Enforcing that hidden features have $\Theta(1)$ element size has long been a
principle of deep learning parametrization, motivated by the fact that
activation functions are designed to take order-one inputs and give order-one
outputs (LeCun et al., 2002). Our second requirement stipulates that feature
entries also undergo $\Theta(1)$ updates during training. Note that any larger
updates would blow up at large width, and any smaller updates would vanish at
large width. We take 1 as our definition of feature learning.111Our notion of
feature learning might also be called “nontrivial feature evolution.” While
other authors may prefer different notions of “feature learning”—for example,
the learning of interpretable, visualizable functions at hidden nodes (Zeiler
& Fergus, 2014, Olah et al., 2017)—nontrivial feature evolution in our sense
is necessary for any other reasonable definition of the term.

Our main message is that feature learning in the sense of 1 may be ensured by
the following spectral scaling condition on the weight matrices of a deep
network and their gradient updates:

######  Condition 1 (Spectral scaling).

Consider applying a gradient update
$\Delta\\!{\bm{W}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times 0pt_{\ell-1}}$ to the
$\ell$th weight matrix ${\bm{W}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times
0pt_{\ell-1}}$. The spectral norms of these matrices should satisfy:
$\displaystyle\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta\left(\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}\right)\
\ \ \text{and}\ \ \
\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta\left(\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}\right),\quad\text{
at layers }\ell=1,...,L.$

We review the spectral norm in Section 2. The spectral scaling condition has
two components which will serve to enforce the respective components of 1. The
first component mandates that each _weight matrix_ has a spectral norm of a
certain size, which will serve to enforce that the layer passes forward
features of the correct size. The second component mandates that each
_gradient update_ has a spectral norm of a certain size, which will ensure
that subsequent features undergo a change of the correct size. We have
implicitly assumed that the input has size
$\left|\\!\left|{\bm{x}}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{0}})$, which
is standard for image data. Language models are an important counterexample,
where embedding matrices take one-hot inputs and the $\sqrt{0pt_{0}}$ in 1
should be replaced by $1$. Appendix E provides a unifying treatment of these
cases.

To get some quick intuition for the origin of 1, observe that under the
forward propagation
${\bm{h}}_{\ell}({\bm{x}})={\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})$, if the
layer input ${\bm{h}}_{\ell-1}({\bm{x}})$ aligns with the top singular vector
of the weight matrix ${\bm{W}}_{\ell}$, then
$\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}$.
The requirement that
$\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})$
then follows from 1. The scaling of
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}$ can similarly
be obtained by writing
$\Delta{\bm{h}}_{\ell}({\bm{x}})=\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})+\ldots$
and applying the same argument. The key missing step is to justify that layer
inputs actually do line up with the top singular subspaces of both weight
matrices and weight updates. As the paper will show, gradient descent training
actually induces this form of alignment.

The bulk of this paper is dedicated to thoroughly demonstrating that training
in accordance with our spectral scaling condition satisfies 1 in MLPs. As an
accessible path to this conclusion, we begin in Section 3 with a simple
model—a deep linear MLP trained for one step on one example—and then
successively extend to multiple training steps, a nonlinear model, and
multiple inputs. In the process, we give a scaling analysis of the dynamics of
feature learning. We then explain how 1 may be achieved in a standard deep
learning setting and compare-and-contrast the resulting scaling prescription
with others in the literature. In particular, we recover the recent “maximal-
update parametrization” ($\mu$P) (Yang & Hu, 2021b).

###  1.2 Summary of contributions

Concretely, our contributions are as follows:

  * •

We propose the spectral scaling condition (1) and show that it suffices to
achieve feature learning in neural networks even at large width.

  * •

We show how 1 may be implemented: either via direct spectral normalization, or
by layerwise initialization scales $\\{\sigma_{\ell}\\}$ and learning rates
$\\{\eta_{\ell}\\}$ that recover maximal update parametrization.222In fact, 1
is the _unique_ scaling on spectral norm that is equivalent to $\mu$P in its
usual definition in terms of learning rate and initialization scaling (Theorem
1).

  * •

We show that other popular scaling rules, including so-called standard
parameterization and neural tangent parametrization, fail to satisfy 1.

In the main text, we focus on MLPs trained via ordinary gradient descent for
clarity. Our results may actually be extended to cover any architecture and
any adaptive optimizer (for a suitable definition of _any_ , c.f.
LABEL:{sec:formal_theory}). Therefore our spectral scaling condition provides
a unifying hyperparameter scaling rule that remains the same whether the
underlying optimizer is, say, SGD or Adam. We suggest that, when one wishes to
determine how the hyperparameters of a new deep learning system should scale
with width, one might turn to the satisfaction of 1 as an overarching
principle.

##  2 Preliminaries

Here we review standard notations which we use in our scaling analysis.

Scaling notation. We will use the usual big-$O$ notation and variants to make
statements about how various quantities scale with network width. Intuitively
speaking:

  * •

$f(0pt)=O(g(0pt))$ means that $f(0pt)$ “scales no faster than” $g(0pt)$,

  * •

$f(0pt)=\Theta(g(0pt))$ means that $f(0pt)$ “scales like” or “is order”
$g(0pt)$,

  * •

$f(0pt)=\Omega(g(0pt))$ means that $f(0pt)$ “scales at least as fast as”
$g(0pt)$.

Formally, $f(0pt)=\Theta(g(0pt))$ is equivalent to the statement that there
exist constants $c,C>0$ such that $c\cdot g(0pt)\leq f(0pt)\leq C\cdot g(0pt)$
for all sufficiently large $d$. The weaker statements $f(0pt)=O(g(0pt))$ and
$f(0pt)=\Omega(g(0pt))$ entail only the upper and lower bounds, respectively.

We will only be concerned with scaling with respect to layer widths in this
paper. Big-$O$ notation will hide any dependence on other factors — such as
depth, dataset size, learning rate schedule, a global learning rate prefactor
— and our statements purely concern how quantities will or should scale with
model width.

Vector and matrix norms. We will use the standard $\ell^{2}$-norm
$\left|\\!\left|\cdot\right|\\!\right|_{2}$ to assess the size of a vector.
For matrices, we will principally use the spectral norm
$\left|\\!\left|\cdot\right|\\!\right|_{*}$ (a.k.a. operator norm) defined as
follows:

######  Definition 1 (Spectral norm).

The spectral norm of a matrix ${\bm{A}}\in\mathbb{R}^{m\times n}$ is given by

$\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}\vcentcolon=\max_{{\bm{v}}\in\mathbb{R}^{n}\backslash\\{\mathbf{0}\\}}\frac{\left|\\!\left|{\bm{A}}{\bm{v}}\right|\\!\right|_{2}}{\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}}.$
(1)

That is, the spectral norm is the largest factor by which a matrix can
increase the norm of a vector on which it acts. The spectral norm of a matrix
is equal to its largest singular value. We will sometimes contrast the
spectral norm with the Frobenius norm
$\left|\\!\left|\cdot\right|\\!\right|_{F}$ given by
$\smash{\left|\\!\left|{\bm{A}}\right|\\!\right|_{F}^{2}=\sum_{ij}A_{ij}^{2}}$.

Properties of the spectral norm. Let ${\bm{A}},{\bm{B}}\in\mathbb{R}^{m\times
n}$ be arbitrary matrices and ${\bm{v}}\in\mathbb{R}^{n}$ be an arbitrary
vector. As with all norms, the spectral norm is subadditive, meaning that it
obeys the triangle inequality
$\left|\\!\left|{\bm{A}}+{\bm{B}}\right|\\!\right|_{*}\leq\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}+\left|\\!\left|{\bm{B}}\right|\\!\right|_{*}$.
The spectral norm is also submultiplicative in the sense that
$\left|\\!\left|{\bm{A}}{\bm{v}}\right|\\!\right|_{*}\leq\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}$
and
$\left|\\!\left|{\bm{A}}{\bm{B}}\right|\\!\right|_{*}\leq\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{B}}\right|\\!\right|_{*}$.
If we interpret a vector ${\bm{v}}\in\mathbb{R}^{n}$ as a $1\times n$ matrix,
then the $\ell^{2}$, spectral and Frobenius norms are equivalent:
$\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}=\left|\\!\left|{\bm{v}}\right|\\!\right|_{*}=\left|\\!\left|{\bm{v}}\right|\\!\right|_{F}$.

Special cases of the spectral norm. For a rank-one matrix ${\bm{A}}$, which
can be written as an outer-product ${\bm{A}}={\bm{u}}{\bm{v}}^{\top}$, it
holds that
$\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}=\left|\\!\left|{\bm{A}}\right|\\!\right|_{F}=\left|\\!\left|{\bm{u}}\right|\\!\right|_{2}\cdot\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}$.
A matrix ${\bm{B}}\in\mathbb{R}^{m\times n}$ is semi-orthogonal if either
${\bm{B}}^{\top}{\bm{B}}={\bm{I}}_{n}$ or
${\bm{B}}{\bm{B}}^{\top}={\bm{I}}_{m}$. A semi-orthogonal matrix has unit
spectral norm: $\left|\\!\left|{\bm{B}}\right|\\!\right|_{*}=1$.

##  3 The spectral scaling condition induces feature learning

In this section, we show that our spectral scaling condition (1) achieves
feature evolution of the correct scale in multilayer perceptrons (MLPs). We
begin with a toy example which conveys key intuitions and then give a series
of extensions which recover a much more general case.

###  3.1 Warmup: deep linear MLP, one step of SGD, on a single example

We begin with a simple model: a deep linear MLP trained for one step on a
single input. While elementary, this example will be sufficient to capture the
intuition for a much more general case.

Model definition. We will study an MLP composed of $L$ successive linear
transformations. We consider a single input ${\bm{x}}\in\mathbb{R}^{0pt_{0}}$
with norm
$\left|\\!\left|{\bm{x}}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{0}})$. The
first hidden representation is ${\bm{h}}_{1}({\bm{x}})={\bm{W}}_{1}{\bm{x}}$,
with the rest of the network following recursively as:

${\bm{h}}_{\ell}({\bm{x}})={\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\qquad\text{for
}\ell=2,\ldots,L.$ (2)

We let ${\bm{h}}_{L}({\bm{x}})\in\mathbb{R}^{0pt_{L}}$ be the network output.
We will keep the input dimension $0pt_{0}$ and output dimension $0pt_{L}$
fixed and consider scaling with respect to the hidden dimensions
$0pt_{1},\ldots,0pt_{L\\!-\\!1}$.

We let the global loss be $\mathcal{L}=g({\bm{h}}_{L}({\bm{x}}),{\bm{y}})$
where $g$ and ${\bm{y}}$ are a loss function and target vector, respectively.
During training, we will take gradient steps at each layer as
$\Delta\\!{\bm{W}}_{\ell}=-\eta_{\ell}\cdot\nabla_{{\bm{W}}_{\ell}}\mathcal{L}$,
where $\eta_{\ell}$ is a layerwise learning rate. We will ultimately solve for
the scale of $\eta_{\ell}$, but for now we will be content to discuss the
perturbation $\Delta\\!{\bm{W}}_{\ell}$ directly. By Equation 2, hidden vector
updates at subsequent layers are related by:

${\bm{h}}_{\ell}({\bm{x}})+\Delta{\bm{h}}_{\ell}({\bm{x}})=({\bm{W}}_{\ell}+\Delta\\!{\bm{W}}_{\ell})({\bm{h}}_{\ell-1}({\bm{x}})+\Delta{\bm{h}}_{\ell-1}({\bm{x}})).$
(3)

Hidden vector sizes. To reiterate 1, we wish for features at the $\ell$th
layer to have a norm which scales as
$\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})$.
Upon a gradient update, this feature vector should undergo an update of size
$\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})$.
We will show that 1 is sufficient to achieve these aims at all layers.

Plan of attack. For simplicity, we will first focus on the first step of
gradient descent after random initialization. We will argue recursively in
depth, showing that if the features at layer $\ell-1$ and their updates
satisfy 1, then so will those at layer $\ell$. In order to verify the desired
scalings, we will show upper and lower scaling bounds separately: we will
first show that the features and their updates are not larger than asked by 1,
and then show that they are in fact also not smaller than asked by 1.

Hidden vector updates. By the subadditivity and submultiplicativity of the
spectral norm, Equations 2 and 3 imply that:

$\displaystyle\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}$
$\displaystyle\leq\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}});$
(4)
$\displaystyle\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}$
$\displaystyle\leq\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}+\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|\Delta{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}+\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|\Delta{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}}),$
(5)

where on the right hand sides of the inequality we have inserted 1 and 1. The
spectral scaling condition thus gives features and feature updates obeying the
correct upper bounds, and we need merely show comparable lower bounds.

Tightness of bounds via matrix-vector alignment. The upper bound in the
submultiplicativity property
$\left|\\!\left|{\bm{A}}{\bm{v}}\right|\\!\right|_{2}\leq\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}$
can be very loose—in particular, this is the case when the vector only
interacts with the small singular values in the matrix. We will now show that
this is not the case in deep network training, and that these upper bounds
provide a fairly accurate description of the way things scale. We make two
observations regarding random weight matrices and gradient updates:

######  Claim 1 (Alignment of initial weight matrices).

Fix a feature vector
${\bm{h}}_{\ell-1}({\bm{x}})\in\mathbb{R}^{0pt_{\ell-1}}$. Assume
${\bm{W}}_{\ell}$ in $\smash{\mathbb{R}^{0pt_{\ell}\times 0pt_{\ell-1}}}$ is
sampled using a common weight initialization strategy (e.g., Gaussian or semi-
orthogonal init). Provided that fan-out is no less than fan-in
($0pt_{\ell}\geq 0pt_{\ell-1}$), then with high probability:

$\left|\\!\left|{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\Theta(\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}).$

######  Claim 2 (Alignment of updates).

For an update $\Delta\\!{\bm{W}}_{\ell}$ given by gradient descent with batch
size 1,

$\left|\\!\left|\Delta{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}.$

In words, 1 states that random hidden weight matrices scale incoming vectors
by factors commensurate to their spectral norms, so long as their fan-out is
not smaller than their fan-in, which we will assume is the case for all but
the final layer of the network. 2 states the same of weight updates, but the
proportionality constant is precisely one and requires no condition on
dimensionality.333Note that ${\bm{x}}$ in 2 is the same input that induced the
gradient $\Delta{\bm{W}}_{\ell}$ in the previous step. 2 is generally not an
equality if this not true.

We now justify these claims in turn. For 1, first suppose that
${\bm{W}}_{\ell}$ is a random semi-orthogonal matrix as is a popular
initialization strategy. Then all singular values of ${\bm{W}}_{\ell}$ are one
and, since fan-out exceeds fan-in, the null-space of ${\bm{W}}_{\ell}$ is
empty. Taken together, these observations imply the equality:
$\left|\\!\left|{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}$.
Fortunately, if the elements of ${\bm{W}}_{\ell}$ are instead sampled i.i.d.
from a centered Gaussian distribution with standard deviation $\sigma_{\ell}$,
then the situation is similar. It is easily shown by the law of large numbers
that
$\left|\\!\left|{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}\approx\sigma_{\ell}\sqrt{0pt_{\ell}}\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}$,
and it is a standard result in random matrix theory that
$\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\approx\sigma_{\ell}(\sqrt{0pt_{\ell-1}}+\sqrt{0pt_{\ell}})$
(Rudelson & Vershynin, 2010, Vershynin, 2018). 1 for Gaussian intialization
follows by combining these results.

2 is easily verified as follows. Observe that we can write the update at layer
$\ell$ as the outer-product:

$\Delta\\!{\bm{W}}_{\ell}=-\eta_{\ell}\cdot\nabla_{{\bm{h}}_{\ell}({\bm{x}})}\mathcal{L}\cdot{\bm{h}}_{\ell-1}({\bm{x}})^{\top}.$
(6)

So the update $\Delta\\!{\bm{W}}_{\ell}$ is rank-one with right singular
vector ${\bm{h}}_{\ell-1}({\bm{x}})$. To verify the claim, observe that:

$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\eta_{\ell}\cdot\left|\\!\left|\nabla_{{\bm{h}}_{\ell}({\bm{x}})}\mathcal{L}\right|\\!\right|_{2}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}^{2}=\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}.$
(7)

With the claims established, we can now get lower bounds on hidden vector size
which serve to verify 1. The features at initialization scale correctly as:

$\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta\left(\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}\right)=\Theta(\sqrt{0pt_{\ell}}),$
(8)

where we have first used 1 and then inserted 1. To bound the size of
$\Delta{\bm{h}}_{\ell}({\bm{x}})$, let us first observe from Equation 3 that

$\Delta{\bm{h}}_{\ell}({\bm{x}})=\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})+{\bm{W}}_{\ell}\Delta{\bm{h}}_{\ell-1}({\bm{x}})+\Delta\\!{\bm{W}}_{\ell}\Delta{\bm{h}}_{\ell-1}({\bm{x}}).$
(9)

So long as the first term
$\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})$ does not perfectly
cancel with the latter two, we have that:

$\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Omega(\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell}\right|\\!\right|_{2})=\Omega(\sqrt{0pt_{\ell}}),$
(10)

where in the last step we have inserted 1. Combining Equation 10 with our
matching upper bound from Equation 5, we conclude that
$\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})$
as desired.

We have achieved both clauses of Desideratum 1 at layer $\ell$, completing a
recursive step from layer $\ell-1$. The norm of the input ${\bm{x}}$ is by
assumption the correct size to serve as a base case, and thus we recursively
have the correct feature scaling at all layers.

####  3.1.1 Key intuitions

We now pause to discuss key intuitions from the above argument which will
carry through to the general case.

Weight updates are low-rank and aligned. An important observation is that
weight updates are highly structured: they have low rank and align to incoming
vectors. This motivates the spectral norm (which is the degree by which a
matrix scales a “perfectly aligned” vector) as the correct measure of size.

Spectral variables enable simpler scaling analysis. Most prior work on
hyperparameter scaling (Yang & Hu, 2021b, Yaida, 2022) discusses layerwise
initialization scales $\\{\sigma_{\ell}\\}$ and learning rates
$\\{\eta_{\ell}\\}$ as the primary variables—although there are exceptions
(Bernstein et al., 2020a). By contrast, we work directly in terms of
quantities that these hyperparameters regulate: the spectral norms of
${\bm{W}}_{\ell}$ and $\Delta\\!{\bm{W}}_{\ell}$. This enables us to determine
the sizes of hidden vectors and their updates quite easily, whereas the same
calculation in terms of $\sigma_{\ell}$ and $\eta_{\ell}$ is more involved.
Section 4 shows how to recover $\sigma_{\ell}$ and $\eta_{\ell}$ from our
spectral scaling condition.

###  3.2 Extensions: additional gradient steps, nonlinearities, and multiple
examples

We now extend our warmup example to successively more complex settings,
ultimately recovering the general case. As we add back complexity, our
spectral scaling condition will remain sufficient to achieve feature evolution
of the proper size, and key intuitions from our warmup will continue to hold
up to minor modifications. Each extension requires making a natural
assumption. We empirically verify these assumptions for a deep MLP in Appendix
C.

####  3.2.1 Additional gradient steps

Our warmup argument relied on two properties of ${\bm{W}}_{\ell}$ at
initialization: its spectral norm being the correct size (1) and it passing
forward features ${\bm{h}}_{\ell}({\bm{x}})$ of the correct size (2).
Fortunately, it is easily seen that these two properties remain true of
${\bm{W}}_{\ell}$ after a gradient step, and we can therefore treat the second
(and later) steps exactly as we did the first. This follows quickly given the
following assumption.

######  Assumption 1.

Updates do not perfectly cancel initial quantities. That is:

$\displaystyle\left|\\!\left|{\bm{W}}_{\ell}+\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}$
$\displaystyle=\Theta\left(\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}+\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\right)$
(11)
$\displaystyle\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})+\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}$
$\displaystyle=\Theta(\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}+\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}).$
(12)

The sort of perfect cancellation required to violate this assumption will be
rare in practice (and adding a small amount of randomness to the learning rate
$\eta_{\ell}$ will fix any occurrence with high probability). It follows that
$\left|\\!\left|{\bm{W}}_{\ell}+\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})$
and
$\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})+\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}}).$
With these facts in place, the same argument we used in Section 3.1 for the
first step ensures that 1 also holds at later steps.

####  3.2.2 Nonlinearities

We now add a nonlinearity $\phi$ to each layer of our MLP. The modified
forward recursion relation is:

${\bm{h}}_{\ell}({\bm{x}})={\bm{W}}_{\ell}{\bm{h}}^{\prime}_{\ell-1}({\bm{x}}),\qquad{\bm{h}}^{\prime}_{\ell}({\bm{x}})=\phi({\bm{h}}_{\ell}({\bm{x}}));\qquad\text{for
}\ell=2,\ldots,L\\!-\\!1,$ (13)

with base case ${\bm{h}}_{1}({\bm{x}})={\bm{W}}_{1}{\bm{x}}$ and output
${\bm{h}}_{L}({\bm{x}})={\bm{W}}_{L}{\bm{h}}^{\prime}_{L-1}({\bm{x}})$. We
assume that the hidden features before and after the application of the
nonlinearity are of the same scale:

######  Assumption 2.

$\left|\\!\left|{\bm{h}}^{\prime}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta\left(\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}\right).$

This is the expected behavior for most activation functions (which are
designed to take in order-one inputs and return outputs which neither explode
nor uniformly vanish) and seems like a reasonable assumption. As in the linear
case, $\Delta\\!{\bm{W}}_{\ell}$ will be rank-one and align to incoming signal
as:

$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}=\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}$
(14)

at each step. All our scaling arguments from the linear case therefore carry
through: the term
$\Delta\\!{\bm{W}}_{\ell}{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})$ is sufficient
to induce a change $\Delta{\bm{h}}_{\ell}({\bm{x}})$ satisfying 1. (The other
terms which depend on $\Delta{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})$ will be no
larger.) 1 therefore still achieves correctly-scaled feature evolution.

####  3.2.3 Batch size greater than one

When training on a minibatch
$\mathcal{D}=\\{({\bm{x}}_{i},{\bm{y}}_{i})\\}_{i=1}^{B}$ with size $B>1$,
each gradient step is simply an average of the steps on each example:

$\Delta\\!{\bm{W}}_{\ell}=\frac{1}{B}\sum_{i=1}^{B}\Delta\\!{\bm{W}}_{\ell}^{(i)},$
(15)

where the sum runs over the minibatch index and
$\Delta\\!{\bm{W}}_{\ell}^{(i)}$ denotes the update that would result from a
step on the single example $({\bm{x}}_{i},{\bm{y}}_{i})$. While
$\Delta\\!{\bm{W}}_{\ell}$ is generally no longer rank one and cannot
perfectly align to all $B$ incoming vectors, Equation 15 makes it clear that
at least one summand will align to each ${\bm{h}}_{\ell}({\bm{x}}_{i})$ from
the batch. We first assume that this term is not perfectly cancelled by the
others:

######  Assumption 3.

$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell}({\bm{x}}_{i})\right|\\!\right|_{2}=\Theta(|\\!|\frac{1}{B}\Delta\\!{\bm{W}}_{\ell}^{(i)}{\bm{h}}_{\ell}({\bm{x}}_{i})|\\!|_{2})$.

We additionally make the assumption that the batch size is fixed and
independent of width:

######  Assumption 4.

The batch size is width-independent: $B=\Theta(1)$.

Combining these assumptions, we find that
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell}({\bm{x}}_{i})\right|\\!\right|_{2}=\Theta\left(\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}}_{i})\right|\\!\right|_{2}\right)$:
the batch update is aligned with incoming signal in a scaling sense. Our
previous scaling arguments in fact needed only alignment in this
“big-$\Theta$” sense (as opposed to the perfect rank-one sense of 2), and thus
they still carry through: our spectral scaling condition continues to suffice
to achieve proper feature learning as per 1.

Empirical observation: low-rank structure remains at large batch size.
Surprisingly, we observe numerically that MLP updates remain low (effective)
rank and aligned with incoming vectors even at large batch size $B$. This is
demonstrated in Figure 1.

Figure 1:  Gradient updates have low effective rank and high-alignment with
incoming hidden vectors even at large batch size in MLPs on CIFAR-10. We
randomly initialize MLPs with depth $L=3$, hidden widths
$0pt_{1}=0pt_{2}=300$, and ReLU and tanh activation functions. We then compute
gradient updates $\Delta\\!{\bm{W}}_{\ell}$ for layer $\ell=2$ on randomly-
sampled size-$B$ subsets of CIFAR-10. Left. As a measure of effective rank, we
report the stable rank
$\textrm{srank}(\Delta\\!{\bm{W}}_{\ell})\vcentcolon=\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}^{2}/\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}^{2}$
of gradient updates. The stable rank remains less than 10 even when $B$ is
large, which is much less than its maximal possible value of
$\min(0pt_{1},0pt_{2})=300$. Right. We report the average alignment
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}/\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\left|\\!\left|{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})\right|\\!\right|_{2}$
of the weight update $\Delta\\!{\bm{W}}_{\ell}$ to the incoming vector
${\bm{h}}^{\prime}_{\ell-1}({\bm{x}})$, averaged over ${\bm{x}}$ from the
batch. Observe that alignment does not decay substantially with large batch
size. Dashed lines on both axes in both subplots show the network width.
Shaded regions denote one standard deviation of variation over random
initializations and batches. Curves look similar after training. See Appendix
A for experimental details.

###  3.3 Adam and other adaptive optimizers

Adam (like most adaptive optimizers in deep learning) processes gradients into
updates via an entrywise function. For example, the momentum-less version of
Adam is SignSGD (Bernstein et al., 2018) which just applies the sign function
to each gradient entry. Via less elementary arguments detailed in Appendix B,
all the above discussion holds for $\Delta{\bm{W}}_{\ell}$ calculated from
these optimizers. In a gist, the main ingredients are 1) the nontrivial
insight from _Tensor Programs_ that gradients look like outer products of iid
vectors, and 2) the fact that entrywise processing preserves the Frobenius
norm of such a matrix up to multiplicative constants when width is large.

###  3.4 Uniqueness of spectral scaling condition

Technically, 1 can be satisfied by scalings other than 1. For example, if one
implements 1 at the first layer but sets
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=0$ for layers
$\ell=2,\ldots,L$, then 1 still holds. However, 1 is the unique _maximal_
scaling: if any of
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}$ or
$\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}$ exceeds 1, then training
will blow up as width is increased. See Appendix B for a proof.

Matrix | Shape | Stable Rank | Spectral Norm | Frobenius Norm | RMS Entry Size
---|---|---|---|---|---
${\bm{W}}_{\ell}$ | $0pt\times 0pt$ | $0pt$ | $1$ | $\sqrt{0pt}$ | $\sqrt{1/0pt}$
$\Delta\\!{\bm{W}}_{\ell}$ | $0pt\times 0pt$ | $1$ | $1$ | $1$ | $1/0pt$
${\bm{W}}_{\ell}$ | $1\times 0pt$ | $1$ | $\sqrt{1/0pt}$ | $\sqrt{1/0pt}$ | $1/0pt$
$\Delta\\!{\bm{W}}_{\ell}$ | $1\times 0pt$ | $1$ | $\sqrt{1/0pt}$ | $\sqrt{1/0pt}$ | $1/0pt$
${\bm{W}}_{\ell}$ | $0pt_{\ell}\times 0pt_{\ell-1}$ | $\min(0pt_{\ell},0pt_{\ell-1})$ | $\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}$ | $\sqrt{\min(0pt_{\ell},0pt_{\ell-1})\cdot\frac{0pt_{\ell}}{0pt_{\ell-1}}}$ | $\sqrt{\frac{\min(0pt_{\ell},0pt_{\ell-1})}{0pt_{\ell}\times 0pt_{\ell-1}}\cdot\frac{0pt_{\ell}}{0pt_{\ell-1}}}$
$\Delta\\!{\bm{W}}_{\ell}$ | $0pt_{\ell}\times 0pt_{\ell-1}$ | $1$ | $\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}$ | $\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}$ | $\sqrt{\frac{1}{0pt_{\ell}\times 0pt_{\ell-1}}\cdot\frac{0pt_{\ell}}{0pt_{\ell-1}}}$
Table 1: Conversion between matrix norms for initial weights ${\bm{W}}_{\ell}$
and updates $\Delta\\!{\bm{W}}_{\ell}$. Results are first given for square
matrices and vectors, then finally generalized to matrices of all shapes.
Table entries denote the $\Theta$-scaling of each quantity. To obtain the
Frobenius norm, one multiplies the spectral norm by the square root of the
stable rank. To obtain the root-mean-square entry size, one divides the
Frobenius norm by the square root of the number of entries. Notice that a
$n\times 0pt$ initial matrix ${\bm{W}}_{\ell}$ with $\Theta(1)$ spectral norm
has $\Theta(\sqrt{1/0pt})$ entry size, whereas a $n\times 0pt$ update
$\Delta\\!{\bm{W}}_{\ell}$ with $\Theta(1)$ spectral norm has only
$\Theta(1/0pt)$ entry size, which is a factor $\sqrt{0pt}$ smaller! The truth
of the statement “the weights move only a negligible amount in wide neural
networks” thus depends on whether this change is measured in entry size or
spectral norm.

##  4 Efficient implementation of the spectral scaling condition

We have thus far proposed the spectral scaling condition (1) and argued in
Section 3 that this condition induces feature learning. We now explain how
this condition can be implemented in practice. Several implementation
strategies are possible because relations exist between different notions of
matrix norm (Table 1 provides a summary). In this section, we first discuss an
implementation via direct spectral normalization, and ultimately recover a
more standard setup in which each layer’s weight matrix is controlled by an
initialization scale $\sigma_{\ell}$ and a (SGD) learning rate $\eta_{\ell}$
as follows:

######  Parametrization 1 (Spectral parametrization).

We claim that the spectral scaling condition (1) is satisfied and feature
learning is achieved (as per 1) if the initialization scale and learning rate
of each layer $\ell$ are chosen according to:
$\sigma_{\ell}=\Theta\left(\frac{1}{\sqrt{0pt_{\ell-1}}}\min\left\\{1,\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}\right\\}\right);\qquad\qquad\eta_{\ell}=\Theta\left(\frac{0pt_{\ell}}{0pt_{\ell-1}}\right).$

Naïve method: direct spectral normalization. The most naïve way to impose 1 is
to directly normalize the relevant quantities by their spectral norm. For
instance, to initialize the weight matrix
$\smash{{\bm{W}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times 0pt_{\ell-1}}}$ at the
$\ell$th layer to have spectral norm $\smash{\sqrt{0pt_{\ell}/0pt_{\ell-1}}}$,
one could sample a temporary matrix ${\bm{W}}^{\prime}_{\ell}$ using any
standard initializer and then re-normalize according to:

$\displaystyle{\bm{W}}_{\ell}=\sigma\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}\times\frac{{\bm{W}}^{\prime}_{\ell}}{\left|\\!\left|{\bm{W}}^{\prime}_{\ell}\right|\\!\right|_{*}},$
(16)

where $\sigma=\Theta(1)$ is a width-independent prefactor. Similarly, to
ensure that the gradient step $\Delta\\!{\bm{W}}_{\ell}$ also has a spectral
norm of the proper size, one might spectrally normalize the gradient according
to:

$\Delta\\!{\bm{W}}_{\ell}=-\eta\sqrt{\frac{0pt_{\ell}}{0pt_{\ell-1}}}\times\frac{\nabla_{{\bm{W}}_{\ell}}\mathcal{L}}{\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}},$
(17)

where $\eta=\Theta(1)$ is a width-independent prefactor.

It is quick to check that normalizing as in Equations 16 and 17 will satisfy
1, but computing spectral norms is expensive and (as we will now show) can be
avoided entirely by working out how
$\left|\\!\left|{\bm{W}}_{\ell}^{\prime}\right|\\!\right|_{*}$ and
$\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}$ will
scale and dividing by the appropriate factor.

Random initialization. Let us suppose that, as is common practice,
${\bm{W}}_{\ell}$ is initialized as
${\bm{W}}_{\ell}=\sigma_{\ell}\cdot{\bm{W}}^{\prime}_{\ell}$, where all
elements of ${\bm{W}}_{\ell}^{\prime}$ are initialized i.i.d. from a normal
distribution with mean zero and unit variance. The spectral norm of a matrix
thus constructed is roughly
$\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}\approx\sigma_{\ell}\cdot(\sqrt{0pt_{\ell}}+\sqrt{0pt_{\ell-1}})$
(Rudelson & Vershynin, 2010, Vershynin, 2018). To get the desired scaling
$\smash{\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})}$,
we need merely choose
$\sigma_{\ell}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}}\cdot(\sqrt{0pt_{\ell}}+\sqrt{0pt_{\ell-1}})^{-1})$.
Simplifying within the $\Theta(\cdot)$, we arrive at $\sigma_{\ell}$ scaled as
in the spectral parametrization (1). Initializing weights with a prefactor
$\sigma_{\ell}$ scaling in this manner achieves the correct spectral norm of
${\bm{W}}_{\ell}$. We note that the constant factor suppressed by the
$\Theta(\cdot)$ here will usually be small—for example, a prefactor of
$\sqrt{2}$ agrees with typical practice for ReLU networks at most layers. If
${\bm{W}}_{\ell}^{\prime}$ is instead a random semi-orthogonal matrix, then we
can simply use a prefactor
$\sigma_{\ell}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})$.

Gradient updates. Here we give two methods for obtaining weight updates
$\Delta\\!{\bm{W}}_{\ell}$ with the correct spectral norm. The first method is
to note that, for a matrix with low stable rank, the Frobenius norm scales
like the spectral norm (and is cheap to compute), so we may simply approximate
$\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}\approx\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{F}$
and use Equation 17 directly. This approach is useful if one wants to avoid
worrying about width-scaling pre-factors entirely. The second method—on which
we will spend more time—is to make standard updates

$\Delta\\!{\bm{W}}_{\ell}=-\eta_{\ell}\nabla_{{\bm{W}}_{\ell}}\mathcal{L}$
(18)

and find a scaling of $\eta_{\ell}$ such that
$\smash{\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})}$.

The main challenge lies in finding the scaling of the gradient
$\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}$.
Note that we expect each gradient update $\Delta\\!{\bm{W}}_{\ell}$ to induce
a change
$\left|\\!\left|\Delta{\bm{h}}_{L}({\bm{x}})\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{L}})$
in the output which induces a change $\Delta\mathcal{L}=\Theta(1)$ for common
loss functions $\mathcal{L}$. Taylor expanding the loss to first order, we
also expect that

$\Delta\mathcal{L}=\Theta(\langle\Delta\\!{\bm{W}}_{\ell},\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\rangle)=\Theta(\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}\cdot\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{F})=\Theta\left(\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\cdot\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}\right),$
(19)

where $\langle\cdot,\cdot\rangle$ denotes the trace inner product and we have
used the facts that the two arguments of the inner product are (a)
proportional to each other and (b) low-rank. Inserting
$\Delta\mathcal{L}=\Theta(1)$ and the spectral scaling condition
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})$,
we can conclude that

$\left|\\!\left|\nabla_{{\bm{W}}_{\ell}}\mathcal{L}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell-1}/0pt_{\ell}}).$
(20)

This result may also be reached via direct layerwise-recursive analysis of the
size of the gradient. Returning to LABEL:eqn:gra\width_update, we now see that
we achieve a properly-scaled update
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(\sqrt{0pt_{\ell}/0pt_{\ell-1}})$
if we take $\eta_{\ell}=\Theta\left(0pt_{\ell}/0pt_{\ell-1}\right)$ as
prescribed by 1.

In summary: training with layerwise initialization $\sigma_{\ell}$ and
learning rate $\eta_{\ell}$ scaled as in our 1 will implement the spectral
scaling condition and give features and feature evolution of the correct size.

##  5 Comparisons to existing parametrizations

We have shown that training in accordance with our spectral scaling condition
(1) at every layer suffices to achieve correctly-scaled feature evolution, and
we have given width scalings for layerwise hyperparameters $\sigma_{\ell}$ and
$\eta_{\ell}$ that suffice to put this condition into action in a standard
deep learning paradigm (1). Here we compare this “spectral parametrization”
with popular parametrizations. We find that our parametrization recovers the
“maximal update parametrization” ($\mu$P) at all layers and is different to
other parametrizations.

###  5.1 Comparison with “maximal update parametrization”

Maximal update parametrization ($\mu$P) was recently proposed as a scaling
rule that retains feature learning even at infinite width. $\mu$P as given in
Table 3 of Yang et al. (2021) may be recovered from our 1 by setting
$0pt_{0}=0pt_{L}=1$ and $0pt_{1}=0pt_{2}=...=0pt_{L-1}$. Our 1 actually
streamlines and generalizes $\mu$P: we provide a unifying treatment for any
rectangular matrix, rather than treating input, hidden and output layers
separately. In other words, from a spectral point of view, no layer is
special. Our parametrization also includes scaling with respect to the input
dimension $0pt_{0}$ and output dimension $0pt_{L}$ (as opposed to neglecting
them as agnostic $\Theta(1)$ quantities) and treats hidden widths of unequal
dimension ($0pt_{\ell-1}\neq 0pt_{\ell}$).

###  5.2 Contrast to “standard parametrization”

At present, the vast majority of deep learning systems use either “Kaiming,”
“Xavier,” or “LeCun” initialization (He et al., 2015, Glorot & Bengio, 2010,
LeCun et al., 2002) with layer-independent learning rates. Generically, we
refer to this as “standard parametrization” (SP), where layerwise
initialization and learning rates scale as:

$\sigma_{\ell}=\Theta(1/\sqrt{0pt_{\ell-1}})\qquad\text{and}\qquad\eta_{\ell}=\Theta(1).$
(21)

Notice that SP initialization exceeds 1 in any layer with fan-out smaller than
fan-in. This includes the final layer in sufficiently wide networks. So, while
1 implies that weight matrices have spectral norm $\Theta(\sqrt{\texttt{fan-
out}/\texttt{fan-in}})$ at initialization, under SP the spectral norms of
certain layers are initialized larger than this. This means that, under SP,
network outputs can blow up if training aligns the layer inputs with the top
singular subspaces (and in fact this alignment generally occurs).

###  5.3 Contrast to “neural tangent parametrization”

The neural tangent parametrization (NTP) of Jacot et al. (2018) parameterizes
a weight matrix as ${\bm{W}}_{\ell}/\sqrt{0pt_{\ell-1}}$ where the entries of
${\bm{W}}_{\ell}$ are sampled iid standard normal, and applies gradient
descent with a layer-independent step-size. Since dividing a weight matrix by
$\sqrt{0pt_{\ell-1}}$ also divides the gradient of that layer by the same
factor, NTP is equivalent to training in our setup with a layer-wise standard
deviation and learning rate:

$\sigma_{\ell}=\Theta(1/\sqrt{0pt_{\ell-1}})\qquad\text{and}\qquad\eta_{\ell}=\Theta(1/0pt_{\ell-1}).$
(22)

Comparing Equations 21 and 22, we see that NTP shares the same initialization
scaling as SP but uses a smaller step-size. To see the deficiency of NTP,
notice the output layer $\sigma_{L}$ is $\sqrt{0pt_{L-1}}$ larger than 1, so
that by the linearity of backpropagation, the gradient to any middle layer
${\bm{W}}_{\ell}$ is also $\sqrt{0pt_{L-1}}$ larger. Then the $1/0pt_{\ell-1}$
learning rate in NTP induces a change $\Delta{\bm{W}}_{\ell}$ that is
$\sqrt{0pt_{L-1}}/0pt_{\ell-1}$ (smaller) compared to 1, which prescribes a
learning rate of $\eta_{\ell}=1$. Because 1 guarantees
$\left|\\!\left|\Delta{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(1)$, NTP
causes $\Delta{\bm{W}}_{\ell}$ to vanish in spectral norm as hidden widths
$0pt_{\ell-1},0pt_{L}\to\infty$.

It bears noting that an MLP parameterized with the NTP can be made to undergo
feature evolution by simply rescaling the network output (and appropriately
scaling down the global learning rate) (Chizat et al., 2019, Bordelon &
Pehlevan, 2022). This operation transforms the NTP into $\mu$P.

###  5.4 Contrast to “Frobenius-normalized updates”

1 mandates that the spectral norm of the update at layer $\ell$ be
proportional to the spectral norm of the weight matrix to which it is applied:
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\propto\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}$.
This contrasts with a body of optimization work (You et al., 2017, Shazeer &
Stern, 2018, Bernstein et al., 2020a; b, Liu et al., 2021) that has suggested
instead that the Frobenius norm of the update at layer $\ell$ be proportional
to the Frobenius norm of the weight matrix to which it is applied:
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}\propto\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{F}$.
For instance, Bernstein et al. (2020a) analysed the operator structure of deep
neural networks and wrote down perturbation bounds on network activations in
terms of perturbations to the weight matrices at each layer, and used these
perturbation bounds to motivate making updates that are small in Frobenius
norm. The flaw in that analysis is the assumption that weight matrices and
their perturbations have identical conditioning structure (Bernstein et al.,
2020a, Condition 2 of Theorem 1), when in reality weight matrices have high
stable rank and updates have $\Theta(1)$ stable rank. In light of this fact,
the Frobenius-normalized proportionality rule should be modified to
$\smash{\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}\propto\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{F}/\sqrt{\min(0pt_{\ell},0pt_{\ell-1})}}$
in order to see proper feature evolution.

##  6 Demonstration: $\bm{\mu}$P versus NTP

Here we discuss a simple, illustrative experiment in which we directly verify
that $\mu$P obeys our spectral scaling condition (1) and achieves leading-
order feature evolution (1) and the NTP does not. We do so via direct
measurement of spectral quantities in MLPs of varying width trained on the
same task. Appendix A provides full experimental details, and results are
plotted Figure 2.

Model and data. We train MLPs with $L=3$ linear layers, no biases, and ReLU
activations. For demonstration purposes, we train on a small subset of $B=200$
examples from a two-class subset of CIFAR-10. The model has input dimension
$0pt_{0}=3072$, output dimension $0pt_{3}=1$, and uniform hidden dimension
$0pt_{1}=0pt_{2}=0pt$, with $0pt$ varied between training runs. We initialize
and train each MLP twice with hyperparameters obtained using $\mu$P and NTP
scalings, respectively. We train networks of widths $0pt\in[16,4096]$ to near-
zero training loss. After training, we compute various spectral quantities
which we now discuss.

Figure 2:  Spectral quantities are $\bm{\Theta(1)}$ under $\mu$P but decay
with width under NTP. We train multilayer perceptrons of varying width, and
plot the following quantities computed between the initial and final network:
(A) Average relative change in features. (B) Relative change in weights in
spectral norm. (C) Final layer alignment with incoming vectors (Equation 24).
(D) Relative change in weights in Frobenius norm. Labeled triangles show
predicted powerlaw slopes. Shaded regions show one standard deviation over
random data and initialization.

Norms of feature updates. We measure the average relative change in features
over training:

$\mathbb{E}_{{\bm{x}}}\\!\left[\frac{\left|\\!\left|{\bm{h}}_{2}({\bm{x}})-{\bm{h}}_{2}^{0}({\bm{x}})\right|\\!\right|_{2}}{\left|\\!\left|{\bm{h}}_{2}^{0}({\bm{x}})\right|\\!\right|_{2}}\right],$
(23)

where ${\bm{h}}_{2}^{0}({\bm{x}})$ and ${\bm{h}}_{2}({\bm{x}})$ are the second
(preactivation) hidden vector at initialization and after training
respectively, and the expectation is over samples ${\bm{x}}$ from the batch.
As shown in Figure 2A, this feature evolution ratio remains roughly fixed as
width $d$ grows when using $\mu$P, in satisfaction of 1. By contrast, it
decays as $1/\sqrt{0pt}$ for the NTP as predicted by Lee et al. (2019).

Spectral norms of weight updates. We measure the relative change in weights in
spectral norm:
$\left|\\!\left|{\bm{W}}_{2}-{\bm{W}}_{2}^{0}\right|\\!\right|_{*}/\left|\\!\left|{\bm{W}}_{2}^{0}\right|\\!\right|_{*}$,
where ${\bm{W}}_{2}^{0}$ and ${\bm{W}}_{2}$ are the second weight matrix
before and after training, respectively. As shown in Figure 2B, this ratio
remains roughly fixed with width in the case of $\mu$P, in accordance with 1.
By contrast, it decays as $1/\sqrt{0pt}$ for the NTP.

Final-layer alignment. We measure the alignment of the final layer to incoming
vectors as follows:

$\textrm{Final-layer
alignment}\vcentcolon=\mathbb{E}_{x}\\!\left[\frac{\left|\\!\left|{\bm{W}}_{3}{\bm{h}}^{\prime}_{2}({\bm{x}})\right|\\!\right|_{2}}{\left|\\!\left|{\bm{W}}_{3}\right|\\!\right|_{*}\cdot\left|\\!\left|{\bm{h}}^{\prime}_{2}({\bm{x}})\right|\\!\right|_{2}}\right].$
(24)

This quantity is $\Theta(1/\sqrt{0pt})$ at initialization. As shown in Figure
2C, it grows to $\Theta(1)$ when using $\mu$P but remains
$\Theta(1/\sqrt{0pt})$ when using the NTP.

Frobenius norms of weight updates. One often hears the claim that “the weights
don’t move” when training very wide neural networks. Here we show that the
validity of this claim crucially depends on the choice of metric. In Figure
2D, we show that the Frobenius norm is deceptive: the net relative change
$\left|\\!\left|{\bm{W}}_{2}-{\bm{W}}_{2}^{0}\right|\\!\right|_{F}/\left|\\!\left|{\bm{W}}_{2}^{0}\right|\\!\right|_{F}$
can decay with width even when the relative change in spectral norm is
constant. This provides crucial context for interpreting existing results in
the literature (Lee et al., 2019, Figure 1).

##  7 Related work

$\mu$P was derived heuristically from spectral norm considerations in talks
given by the first author in 2021 (Yang & Hu, 2021a). Earlier work (Bernstein
et al., 2020a) derived a spectral analysis of feature learning based on
perturbation bounds, but that work obtained the wrong scaling relation with
network width due to a flawed conditioning assumption on gradients. Below, we
review various strands of related work on feature learning and training
strategies.

Parametrizations for wide neural networks. Much work has examined the scaling
behavior of training dynamics of networks at large width. Together, works on
the “neural tangent kernel” (NTK) limit (Jacot et al., 2018, Lee et al.,
2019), the “mean-field” limit (Rotskoff & Vanden-Eijnden, 2022, Mei et al.,
2019, Sirignano & Spiliopoulos, 2022), and the related “feature learning”
limit (Geiger et al., 2020, Yaida, 2022, Yang & Hu, 2021b) (i.e. the $\mu$P
limit), paint a rich picture of a family of possible infinite-width scalings.
After healthy debate regarding the relative empirical performance of the (more
analytically tractable) NTK limit and the feature learning limit, a recent
consensus holds that learning features is usually beneficial in practical
large-scale deep learning settings (Chizat et al., 2019, Fort et al., 2020,
Vyas et al., 2022). Our spectral scaling analysis recovers the feature
learning limit in a simpler manner than previous analyses.

Spectral normalization. Spectral normalization emerged as a form of weight
normalization in the generative adversarial network literature (Miyato et al.,
2018). This form of normalization acts on the weight matrices, and is used
analogously to other normalization schemes such as batchnorm (Ioffe & Szegedy,
2015) and layernorm (Ba et al., 2016). In contrast to our spectral scaling
condition (1), this method treats only the weights ${\bm{W}}_{\ell}$, not the
updates $\Delta\\!{\bm{W}}_{\ell}$. Furthermore, spectral normalization
implementations typically set the spectral norm of weight matrices either to
one or to a tunable hyperparameter (Farnia et al., 2019), and do not include
the key factor of $\sqrt{0pt_{\ell}/0pt_{\ell-1}}$ in our 1.

Operator theory of neural networks. Neural networks are constructed by
composing linear operators with elementwise nonlinearities. One line of work
studies this operator structure and how it behaves under perturbation to
understand how step sizes should be set in gradient descent. For instance,
Bernstein et al. (2020a) derive perturbation bounds on the maximum amount of
feature change that can be induced by a gradient step in terms of the operator
properties of the weight matrices. Meanwhile, Yang & Hu (2021b) study the
operator structure of neural networks in the limit that width is taken to
infinity, proposing a parametrization that obtains feature learning in this
limit.

##### Optimization theory.

A body of literature studies optimization algorithms for deep networks that
take steps whose size is set relative to the weights to which they are applied
(You et al., 2017, Bernstein et al., 2020a; b, Liu et al., 2021, Carbonnelle &
Vleeschouwer, 2019, Shazeer & Stern, 2018). A particular focus has been placed
on setting the Frobenius norm of update steps to be small relative to the
Frobenius norm of the weight matrices (You et al., 2017, Bernstein et al.,
2020a; b, Liu et al., 2021). A main practical takeaway of this paper is that
the Frobenius norm should be replaced by the spectral norm to get proper width
scaling. The source of the difference between Frobenius and spectral norm is
that gradient updates tend to have low stable rank, as shown in Figure 1,
while the weights themselves tend to have high stable rank.

##  8 Conclusion

We have presented an analysis of the dynamics of feature learning in deep
neural networks, beginning with desired conditions on feature evolution (1)
and culminating in the demonstration that these conditions may be achieved by
simple scaling rules (1 and 1) applied uniformly to each layer. Our analysis
recovers and generalizes practically-important “feature-learning
parametrizations” and provides a simple, unifying perspective on the question
of parametrization in wide neural networks. For comparison, formal results
derived under the tensor programs framework are given in Appendix B.

Our discussion has focused principally on MLPs for clarity, but our feature
learning desideratum and spectral scaling condition can be directly applied to
structured architectures. The spectral scaling condition may be applied to
multi-index tensors as appear in convolutional architectures by applying the
condition to appropriate “slices” of the full tensor. Simple application of
our spectral scaling condition recovers $\mu$P scalings reported for these
model classes (see e.g. Yang et al. (2021), Table 8 and Section J.2). We also
give the hyperparameter scalings for biases (which are easily derived but
omitted in the main text for clarity) in Appendix D. This architectural
universality is also proven rigorously in Appendix B.

Finally, we note that under a natural redefinition of the $\ell^{2}$-norm of a
vector ${\bm{v}}\in\mathbb{R}^{0}pt$ incorporating a normalization prefactor,
all vector and matrix norms used in our spectral analysis become $\Theta(1)$,
permitting an elegant summary of our conclusions. We state and discuss this
nondimensionalization procedure in Appendix E. This generalization permits the
extension of our results to e.g. the input layer in language modeling, in
which one-hot embeddings violate our assumption that
$\left|\\!\left|{\bm{x}}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{0}})$.

## Acknowledgements

The authors thank Josh Albrecht, Blake Bordelon, Alex Wei, Nikhil Ghosh, and
Dhruva Karkada for useful discussions and comments on the manuscript. JS
gratefully acknowledges support from the National Science Foundation Graduate
Fellow Research Program (NSF-GRFP) under grant DGE 1752814.

## Author Contributions

GY developed our core insight regarding the utility of the spectral norm,
produced our tensor programs theory (Appendix B), and aided in refinement of
the paper. JS spearheaded the writing of the paper, led iteration towards
simple analysis which communicates our spectral picture, and ran experiments.
JB developed an early incarnation of the spectral picture (Bernstein et al.,
2020a), contributed key insights simplifying our exposition including unifying
all layers under single formulae, aided in writing the paper, and ran
experiments.

## References

  * Amari (1998) Shun-ichi Amari.  Natural gradient works efficiently in learning.  _Neural Computation_ , 1998.
  * Arora et al. (2019) Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  In _Neural Information Processing Systems_ , 2019.
  * Atanasov et al. (2022) Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan.  Neural networks as kernel learners: the silent alignment effect.  In _International Conference on Learning Representations_ , 2022.
  * Ba et al. (2016) Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton.  Layer normalization.  _arXiv:1607.06450_ , 2016.
  * Bernstein et al. (2018) Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar.  signSGD: Compressed Optimisation for Non-Convex Problems, February 2018.  URL https://arxiv.org/abs/1802.04434v3.
  * Bernstein et al. (2020a) Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming-Yu Liu.  On the distance between two neural networks and the stability of learning.  In _Neural Information Processing Systems_ , 2020a.
  * Bernstein et al. (2020b) Jeremy Bernstein, Jiawei Zhao, Markus Meister, Ming-Yu Liu, Anima Anandkumar, and Yisong Yue.  Learning compositional functions via multiplicative weight updates.  In _Neural Information Processing Systems_ , 2020b.
  * Bernstein et al. (2023) Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue.  Automatic Gradient Descent: Deep Learning without Hyperparameters.  _arXiv:2304.05187_ , 2023.
  * Bordelon & Pehlevan (2022) Blake Bordelon and Cengiz Pehlevan.  Self-consistent dynamical field theory of kernel evolution in wide neural networks.  In _Neural Information Processing Systems_ , 2022.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  _Neural Information Processing Systems_ , 2020.
  * Canatar et al. (2021) Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan.  Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks.  _Nature Communications_ , 2021.
  * Carbonnelle & Vleeschouwer (2019) Simon Carbonnelle and Christophe De Vleeschouwer.  Layer rotation: A surprisingly simple indicator of generalization in deep networks?  In _ICML Workshop on Identifying and Understanding Deep Learning Phenomena_ , 2019.
  * Chizat et al. (2019) Lénaïc Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  _Neural Information Processing Systems_ , 2019.
  * Dey et al. (2023) Nolan Dey, Gurpreet Gosal, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, Joel Hestness, et al.  Cerebras-GPT: Open compute-optimal language models trained on the Cerebras wafer-scale cluster.  _arXiv:2304.03208_ , 2023.
  * Dhillon & Tropp (2008) Inderjit S. Dhillon and Joel A. Tropp.  Matrix nearness problems with Bregman divergences.  _SIAM Journal on Matrix Analysis and Applications_ , 2008.
  * Farnia et al. (2019) Farzan Farnia, Jesse Zhang, and David Tse.  Generalizable adversarial training via spectral normalization.  In _International Conference on Learning Representations_ , 2019.
  * Fort et al. (2020) Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli.  Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel.  _Neural Information Processing Systems_ , 2020.
  * Geiger et al. (2020) Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart.  Disentangling feature and lazy training in deep neural networks.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2020.
  * Glorot & Bengio (2010) Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _International Conference on Artificial Intelligence and Statistics_ , 2010.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification.  In _International Conference on Computer Vision_ , 2015.
  * Ioffe & Szegedy (2015) Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International Conference on Machine Learning_ , 2015.
  * Jacot et al. (2018) Arthur Jacot, Clément Hongler, and Franck Gabriel.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Neural Information Processing Systems_ , 2018.
  * Kingma & Ba (2015) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015.
  * Krizhevsky (2009) Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * LeCun et al. (2002) Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller.  Efficient backprop.  In _Neural Networks: Tricks of the Trade_. Springer, 2002.
  * Lee et al. (2018) Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S. Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Deep neural networks as Gaussian processes.  In _International Conference on Learning Representations_ , 2018.
  * Lee et al. (2019) Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  In _Neural Information Processing Systems_ , 2019.
  * Lee et al. (2020) Jaehoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein.  Finite versus infinite neural networks: an empirical study.  In _Neural Information Processing Systems_ , 2020.
  * Liu et al. (2021) Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue.  Learning by turning: Neural architecture aware optimisation.  In _International Conference on Machine Learning_ , 2021.
  * Mei et al. (2019) Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Mean-field theory of two-layers neural networks: Dimension-free bounds and kernel limit.  In _Conference on Learning Theory_ , 2019.
  * Miyato et al. (2018) Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for Generative Adversarial Networks.  In _International Conference on Learning Representations_ , 2018.
  * Nemirovsky & Yudin (1983) Arkady S. Nemirovsky and David B. Yudin.  _Problem complexity and method efficiency in optimization_.  Wiley, 1983.
  * Olah et al. (2017) Chris Olah, Alexander Mordvintsev, and Ludwig Schubert.  Feature visualization.  _Distill_ , 2017.
  * Poole et al. (2016) Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli.  Exponential expressivity in deep neural networks through transient chaos.  In _Neural Information Processing Systems_ , 2016.
  * Ramesh et al. (2022) Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.  Hierarchical text-conditional image generation with CLIP latents.  _arXiv:2204.06125_ , 2022.
  * Rotskoff & Vanden-Eijnden (2022) Grant M. Rotskoff and Eric Vanden-Eijnden.  Trainability and accuracy of artificial neural networks: An interacting particle system approach.  _Communications on Pure and Applied Mathematics_ , 2022.
  * Rudelson & Vershynin (2010) Mark Rudelson and Roman Vershynin.  Non-asymptotic theory of random matrices: Extreme singular values.  In _International Congress of Mathematicians_ , 2010.
  * Shazeer & Stern (2018) Noam Shazeer and Mitchell Stern.  Adafactor: Adaptive learning rates with sublinear memory cost.  In _International Conference on Machine Learning_ , 2018.
  * Silver et al. (2016) David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis.  Mastering the game of Go with deep neural networks and tree search.  _Nature_ , 2016.
  * Sirignano & Spiliopoulos (2022) Justin Sirignano and Konstantinos Spiliopoulos.  Mean field analysis of deep neural networks.  _Mathematics of Operations Research_ , 2022.
  * Sohl-Dickstein et al. (2020) Jascha Sohl-Dickstein, Roman Novak, Samuel S. Schoenholz, and Jaehoon Lee.  On the infinite width limit of neural networks with a standard parameterization.  _arXiv:2001.07301_ , 2020.
  * Vershynin (2018) Roman Vershynin.  _High-Dimensional Probability: An Introduction with Applications in Data Science_.  Cambridge University Press, 2018.
  * Vyas et al. (2022) Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran.  Limitations of the NTK for understanding generalization in deep learning.  _arXiv:2206.10012_ , 2022.
  * Yaida (2022) Sho Yaida.  Meta-principled family of hyperparameter scaling strategies.  _arXiv:2210.04909_ , 2022.
  * Yang & Hu (2021a) Greg Yang and Edward J. Hu.  Feature learning in infinite-width neural networks.  In _ICML 2021 Workshop on Over-parameterization: Pitfalls and Opportunities_ , 2021a.  https://slideslive.com/38963058/feature-learning-in-infinitewidth-neural-networks.
  * Yang & Hu (2021b) Greg Yang and Edward J. Hu.  Tensor Programs IV: Feature learning in infinite-width neural networks.  In _International Conference on Machine Learning_ , 2021b.
  * Yang & Littwin (2023) Greg Yang and Etai Littwin.  Tensor programs IVb: Adaptive optimization in the $\infty$-width limit.  _arXiv:2308.01814_ , 2023.
  * Yang et al. (2021) Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao.  Tensor Programs V: Tuning large neural networks via zero-shot hyperparameter transfer.  In _Neural Information Processing Systems_ , 2021.
  * You et al. (2017) Yang You, Igor Gitman, and Boris Ginsburg.  Scaling SGD batch size to 32K for ImageNet training.  Technical Report UCB/EECS-2017-156, University of California, Berkeley, 2017.
  * You et al. (2020) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training BERT in 76 minutes.  In _International Conference on Learning Representations_ , 2020.
  * Zeiler & Fergus (2014) Matthew D. Zeiler and Rob Fergus.  Visualizing and understanding convolutional networks.  In _European Conference on Computer Vision_ , 2014.

##  Appendix A Experimental details

Experimental details for Figure 1. We examine randomly-initialized MLPs with
depth $L=3$, widths $0pt_{0}=3072,0pt_{1}=0pt_{2}=d=300,0pt_{3}=10$, and ReLU
and tanh activation functions. We then pass single batches of CIFAR-10 data of
varying size $B$ through the model and compute a single gradient step
$\Delta\\!{\bm{W}}_{\ell}$ at layer $\ell=2$ with arbitrary learning rate. We
then compute the stable rank of $\Delta\\!{\bm{W}}_{\ell}$ and the alignment
metric

$\mathbb{E}_{{\bm{x}}}\\!\left[\frac{\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell-1}({\bm{x}}_{i})\right|\\!\right|_{2}}{\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\left|\\!\left|{\bm{h}}^{\prime}_{\ell-1}({\bm{x}}_{i})\right|\\!\right|_{2}}\right],$
(25)

where the expectation is taken over ${\bm{x}}$ from the batch. Shaded
envelopes in Figure 1 denote one standard deviation with respect to both
random network initialization and random batch selection over $10$ trials.

We conducted an additional experiment to try to ascertain the source of the
low effective rank of $\Delta\\!{\bm{W}}_{\ell}$. Let
${\bm{H}}^{\prime}_{\ell-1}=[{\bm{h}}^{\prime}_{1}({\bm{x}}_{1})\ldots{\bm{h}}^{\prime}_{1}({\bm{x}}_{B})]\in\mathbb{R}^{d\times
B}$ be the matrix formed from stacking the $\ell=1$ post-nonlinearity hidden
vectors from the full batch, and let
${\bm{G}}_{\ell}=[{\bm{g}}_{2}({\bm{x}}_{1})\ldots{\bm{g}}_{2}({\bm{x}}_{B})]\in\mathbb{R}^{d\times
B}$ with
${\bm{g}}_{\ell}({\bm{x}})=\nabla_{{\bm{h}}_{\ell}({\bm{x}})}\mathcal{L}$ be
the matrix formed from stacking the loss gradients at layer $\ell=2$ from the
full batch. It is the case that
$\Delta\\!{\bm{W}}_{\ell}\propto{\bm{G}}_{\ell}{{\bm{H}}^{\prime}_{\ell-1}}^{\top}$,
and if $\Delta\\!{\bm{W}}_{\ell}$ has low stable rank, it is a reasonable
guess that either ${\bm{H}}^{\prime}_{\ell-1}$ and ${\bm{G}}_{\ell}$ also has
low stable rank. In fact, we find that both ${\bm{H}}^{\prime}_{\ell-1}$ and
${\bm{G}}_{\ell}$ have low stable rank, as shown in
LABEL:fig:forwar\width_backwar\width_sranks.

Figure 3:  Stable ranks for the set of forward hidden vectors
${\bm{H}}^{\prime}_{\ell-1}$ and the set of backward hidden vectors
${\bm{G}}_{\ell}$ for $\ell=2$. Even at large batch size $B$, both are
consistently much lower than their maximal possible values of $d=300$. Dashed
lines on both axes in both subplots show the network width $d$. Note that the
stable rank of a matrix ${\bm{M}}$ is given by
$\smash{\left|\\!\left|{\bm{M}}\right|\\!\right|_{F}^{2}/\left|\\!\left|{\bm{M}}\right|\\!\right|_{*}^{2}}$.

Experimental details for Figure 2. We train MLPs with depth $L=3$, widths
$0pt_{0}=3072,0pt_{1}=0pt_{2}=d,0pt_{3}=1$, and ReLU activation functions. The
data consists of 200 samples from CIFAR-10 (Krizhevsky, 2009) from only the
classes airplane and automobile and uses $\pm 1$ targets.

We use two different hyperparameter schemes as follows. To implement $\mu$P,
we take

$\sigma_{1}=\sqrt{\frac{2}{0pt_{0}}};\quad\sigma_{2}=\sqrt{\frac{2}{d}};\quad\sigma_{3}=\frac{\sqrt{2}}{d};\quad\eta_{1}=\eta\cdot\frac{d}{0pt_{0}};\quad\eta_{2}=\eta;\quad\eta_{3}=\eta\cdot\frac{1}{d},$
($\mu$P)

with global learning rate $\eta=0.1$. To implement NTP, we follow Jacot et al.
(2018) and Lee et al. (2019):

$\sigma_{\ell}=\sqrt{\frac{2}{0pt_{\ell-1}}};\qquad\eta_{\ell}=\eta\cdot\frac{1}{0pt_{\ell-1}}$
(NTP)

at all layers, again with $\eta=0.1$. These parameterizations are equivalent
at $d=1$, which lets us view each parameterization as a particular scaling
prescription applied to a narrow base network.

We train full-batch for $10^{4}$ steps, which is sufficient for all widths to
drop below $0.01$ training loss on average by the end of training. We do not
expect that training for many more steps would saliently change the resulting
plots. Shaded envelopes in Figure 2 denote one standard deviation with respect
to both random network initialization and random batch selection over $10$
experiment trials.

It is perhaps worth emphasizing that these experiments worked much better than
they had to. Our theory strictly applies only to the case of a small number of
gradient steps relative to network width, but the net updates shown in each
subplot of Figure 2 reflect the accumulation of thousands of gradient steps, a
number which is larger than network width in all cases. We were thus surprised
by the very clear agreement of this experiment with predicted power laws.

##  Appendix B Tensor programs theory

For simplicity, assume all hidden widths $0pt_{1}=\cdots=0pt_{L-1}$ are the
same. Borrowing from Yang & Hu (2021b), an _abc-parametrization_ is just a
recipe for scaling (as powers of width) the multiplier, initializer scale, and
learning rate of all parameter tensors of a neural network. SP, NTP, and
$\mu$P are all examples of abc-parametrizations. A _stable_ abc-
parametrization is one whose (pre)activations and output do not blow up with
width at any step of training. Then (under the generous 5) the following
theorem is our main result:

######  Theorem 1.

In $\mu$P, for almost every learning rate (in the measure-theoretic sense),
LABEL:{cond:scaling} is satisfied at any time during training for sufficiently
large width. $\mu$P is the unique stable abc-parametrization with this
property.

The main insight leading to Theorem 1 is that:
$\Delta\\!{\bm{W}}_{\ell}/\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}$
converges to a Hilbert-Schmidt operator (in an appropriate space) as width
goes to infinity. This in turn implies:

######  Lemma 1.

Unless $\Delta\\!{\bm{W}}_{\ell}=0$,
$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{F}/\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}=\Theta(1)$
as width goes to infinity.444 It’s possible to make precise quantitative
statements here using a quantitative version of the Master Theorem but we
won’t be concerned with it here.

These statements are universal: they hold for any architecture and any
adaptive optimizers representable by tensor programs (including convolutional
neural networks, residual networks, and transformers, etc., as well as
RMSProp, Adam, etc.), not just MLP and SGD.

###  B.1 Proofs

As we have explained the core intuitions of our spectral perspective of
feature learning in the main text, here we focus on proving the most general
result in the most concise way.

Our proofs will rely on the following notions defined in prior work:

  * •

representable architecture (Yang & Littwin, 2023, Defn 2.9.1)

  * •

matrix/vector/scalar parameters (Yang & Littwin, 2023, Defn 2.9.1)

  * •

abcd-parametrization for representable architectures (Yang & Littwin, 2023,
Defn 2.9.7)

  * •

entrywise optimizer (Yang & Littwin, 2023, Sec 2.1)

  * •

ket and iid-copy notation (Yang & Littwin, 2023, Sec 1.2)

Everything here follows under the following:

######  Assumption 5.

Assume our architecture is representable by a
$\textsc{Ne}{\otimes}\textsc{or}\top$ program with pseudo-Lipschitz
nonlinearities, trained by an entrywise optimizer with pseudo-Lipschitz update
functions.

#### Random Initialization

The following is our main proposition for initialization:

######  Proposition 1.

In any abcd-parametrization, any matrix parameter ${\bm{W}}$ at random
initialization has
$\left|\\!\left|{\bm{W}}\right|\\!\right|_{F}/\left|\\!\left|{\bm{W}}\right|\\!\right|_{*}=\Theta(\sqrt{0}pt)$,
but any vector or scalar parameter ${\bm{W}}$ has
$\left|\\!\left|{\bm{W}}\right|\\!\right|_{F}/\left|\\!\left|{\bm{W}}\right|\\!\right|_{*}=1$.

###### Proof.

This is a claim about random matrices and follows from classical random matrix
theory. Here’s a quick sketch: If $\sigma$ is the standard deviation of a
matrix entry, then
$\left|\\!\left|{\bm{W}}\right|\\!\right|_{F}\approx\sigma\cdot 0pt$ from the
central limit theorem, and
$\left|\\!\left|{\bm{W}}\right|\\!\right|_{*}\approx 2\sigma\cdot\sqrt{0pt}$
as stated in the main text, from which the first part of the proposition
follows. For vectors and scalars, the stated ratio is always $1$. ∎

#### Matrix Updates

In a tensor program, each vector $x$ of the program converges to a 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_) as width goes to infinity in the sense that the scaled inner
product $\langle x,y\rangle/0pt$ converges almost surely to $\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=\operatorname*{\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\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$.
The kets form a Hilbert space $\mathcal{Z}$. Then any weight parameter
converges to a linear operator from $\mathcal{Z}$ to $\mathcal{Z}$; any vector
parameter converges to a linear operator from $\mathcal{Z}$ to $\mathbb{R}$ or
$\mathbb{R}$ to $\mathcal{Z}$. (Any scalar parameter converges to a random
real, but that is not too important here).

######  Proposition 2.

Consider any abcd-parametrization. For any matrix or vector parameter
${\bm{W}}$, at any step of training,
$\Delta\\!{\bm{W}}/\left|\\!\left|\Delta\\!{\bm{W}}\right|\\!\right|_{F}$
converges to a Hilbert-Schmidt integral operator.

###### Proof.

This is trivial for vector parameters. For matrix parameters, observe
$\widetilde{\Delta}\\!{\bm{W}}=\Delta\\!{\bm{W}}/\left|\\!\left|\Delta\\!{\bm{W}}\right|\\!\right|_{F}$
is always a nonlinear outer product

$\widetilde{\Delta}\\!{\bm{W}}=Q(\bm{x};\bm{y};\bm{c})$

for multi-vectors $\bm{x}$ and $\bm{y}$ and multi-scalars $\bm{c}$ and some
nonlinearity $Q$. Then the limit
$\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}\\!{\bm{W}}\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$}}$
acts on a ket
$\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$
by

$\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}\\!{\bm{W}}\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=\operatorname*{\mathbb{E}}_{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 1$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
1$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 1$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 1$}\kern
0.5pt}}}Q(\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$}}\bm{x}\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$}}\bm{y}\rangle^{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 1$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
1$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 1$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 1$}\kern
0.5pt}}};\mathring{\bm{c}})\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^{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 1$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
1$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 1$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 1$}\kern 0.5pt}}}$

integrating over
$\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$}}\bm{y}\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$}}z\rangle$.
The Hilbert-Schmidt norm 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$}}\widetilde{\Delta}\\!{\bm{W}}\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$}}$
is

$\left|\\!\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$}}\widetilde{\Delta}\\!{\bm{W}}\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$}}\right|\\!\right|_{HS}=\sqrt{\operatorname*{\mathbb{E}}_{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 0$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
0$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 0$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 0$}\kern 0.5pt}},\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 1$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
1$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 1$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 1$}\kern
0.5pt}}}Q(\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$}}\bm{x}\rangle^{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 0$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
0$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 0$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 0$}\kern
0.5pt}}};\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$}}\bm{y}\rangle^{\mathord{\mathchoice{\kern
1.0pt\framebox{$\displaystyle 1$}\kern 1.0pt}{\kern 1.0pt\framebox{$\textstyle
1$}\kern 1.0pt}{\kern 0.7pt\framebox{$\scriptstyle 1$}\kern 0.7pt}{\kern
0.5pt\framebox{$\scriptscriptstyle 1$}\kern 0.5pt}}};\mathring{\bm{c}})^{2}}$

which is finite by the Master Theorem. Therefore
$\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}\\!{\bm{W}}\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$}}$
is Hilbert-Schmidt. ∎

######  Proposition 3.

In Proposition 2, for large enough width, unless $\Delta\\!{\bm{W}}=0$, both
spectral norm and Frobenius norm of
$\Delta\\!{\bm{W}}/\left|\\!\left|\Delta\\!{\bm{W}}\right|\\!\right|_{F}$ are
$\Theta(1)$.

This proposition implies Lemma 1.

###### Proof.

This is obvious for Frobenius norm since it converges to the Hilbert-Schmidt
norm of the operator limit.

However, the spectral norm cannot be expressed directly in such form. But, by
definition of spectral norm and of the the ket space $\mathcal{Z}$, one can
construct a nonzero vector $z$ in an extension of the program that defined
$\Delta\\!{\bm{W}}$, such that

$\left|\\!\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$}}\widetilde{\Delta}\\!{\bm{W}}\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\right|\\!\right|=\theta\left|\\!\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$}}\widetilde{\Delta}\\!{\bm{W}}\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$}}\right|\\!\right|_{*}\left|\\!\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$}}z\rangle\right|\\!\right|,$

for some $\theta\in(1/2,1]$.555In fact, we can pick $\theta\in[1-\epsilon,1]$
for any $\epsilon>0$. This implies, by the Master Theorem, that

$\frac{\theta}{2}\frac{\left|\\!\left|\widetilde{\Delta}\\!{\bm{W}}z\right|\\!\right|}{\left|\\!\left|z\right|\\!\right|}\leq\left|\\!\left|\widetilde{\Delta}\\!{\bm{W}}\right|\\!\right|_{*}$

for sufficiently large width. The LHS is $\Theta(1)$, so we are done.

∎

######  Proposition 4.

In $\mu$P, at any fixed time during training, as width $\to\infty$, any matrix
parameter ${\bm{W}}$ has spectral norm $\Theta(1)$ and Frobenius norm
$\Theta(\sqrt{0}pt)$.

###### Proof.

In $\mu$P, it’s trivial to see that
$\left|\\!\left|\Delta\\!{\bm{W}}\right|\\!\right|_{F}=\Theta(1)$. So any
nonzero $\Delta\\!{\bm{W}}$ (at any point of training) has both spectral norm
and Frobenius norm $\Theta(1)$ by the above.

Simple calculation then shows that ${\bm{W}}$ at any fixed time has
$\Theta(\sqrt{0}pt)$ since this is the case at initialization by Proposition
1. This means the quadratic mean of the singular values of ${\bm{W}}$ is
$\Theta(1)$, so its max singular value must be $\Omega(1)$. But it furthermore
must be $\Theta(1)$ because ${\bm{W}}$ at initialization and all of its
updates have $O(1)$ spectral norm. ∎

######  Theorem 2.

In $\mu$P, for all but a measure-zero set of learning rates,
LABEL:{cond:scaling} is satisfied at any time during training for sufficiently
large width. $\mu$P is the unique stable and faithful abcd-parametrization
with this property.

###### Proof.

In $\mu$P, by Proposition 4, all matrix parameters satisfy
LABEL:{cond:scaling} no matter what the learning rate is. However, for vector
parameters, it’s possible for some specific learning rate to cause the weights
to vanish after an update, but at most a (Lebesgue) measure-zero set of
learning rates will cause this to happen. Assuming this vanishing does not
happen, ${\bm{W}}$ has $\Theta(1)$ Frobenius norm at initialization and at all
times during training. By Proposition 3, ${\bm{W}}$ also has $\Theta(1)$
spectral norm, so LABEL:{cond:scaling} is satisfied. A similar but easier
argument applies to all scalar parameters. This shows $\mu$P satisfies
LABEL:{cond:scaling}.

Since any other stable and faithful parametrization essentially just rescales
the initialization and the update, we see that no other parametrization can
satisfy LABEL:{cond:scaling}.∎

For SGD, all abc-parametrizations are equivalent to a faithful abcd-
parametrization because $Q$ is identity, so Theorem 2 recovers Theorem 1.

##  Appendix C Empirical checks of 1, 2 and 3

In Section 3, we first illustrated how 1 satisfies 1 in a minimal model—a deep
linear network trained for one step on one sample—and then iteratively
extended our argument to multiple steps, nonlinearities, and multiple samples.
Each extension came with a mild assumption. These assumptions are intended to
be natural conditions one expects from generic network dynamics.

In this section, we restate each assumption, explain further why it is
expected to hold, and then present a validating experiment. All experiments
use the same setup as that of Figure 2. Plotted quantities depending on the
sample ${\bm{x}}$ are averaged over all ${\bm{x}}$ in the batch, with shaded
envelopes showing one standard deviation of this mean over five experiment
trials.

######  Assumption 1.

Updates do not perfectly cancel initial quantities. That is:

$\displaystyle\left|\\!\left|{\bm{W}}_{\ell}+\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}$
$\displaystyle=\Theta\left(\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{*}+\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{*}\right)$
(26)
$\displaystyle\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})+\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}$
$\displaystyle=\Theta(\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}+\left|\\!\left|\Delta{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}).$
(27)

Recall that the cancellation of two high-dimensional tensors (i.e. matrices or
vectors) tends to be unlikely as dimension grows (unless there is a good
reason to expect cancellation, which there is not here). For the small
learning rate used in our experiment, 1 is in fact true simply because small
updates are not big enough to cancel the initial quantities even if aligned.
In Figure 4, we verify that 1 also holds when $\Delta\\!{\bm{W}}_{\ell}$ and
$\Delta{\bm{h}}_{\ell}({\bm{x}})$ are respectively replaced by
${\bm{W}}_{\ell}-{\bm{W}}_{\ell}^{0}$ and
${\bm{h}}_{\ell}({\bm{x}})-{\bm{h}}^{0}_{\ell}({\bm{x}})$, the total updates
across all of training.

Figure 4:  Verification of 1. Subplots show the ratios between left- and
right-hand sides of conditions of 1 at various layers. Envelopes show
variation over trials. The ratios are $\Theta(1)$ (showing no obvious decay
with network width), verifying Assumption 1.

######  Assumption 2.

$\left|\\!\left|{\bm{h}}^{\prime}_{\ell}({\bm{x}})\right|\\!\right|_{2}=\Theta\left(\left|\\!\left|{\bm{h}}_{\ell}({\bm{x}})\right|\\!\right|_{2}\right).$

Recall that
${\bm{h}}^{\prime}_{\ell}({\bm{x}})=\phi({\bm{h}}_{\ell}({\bm{x}}))$, and that
${\bm{h}}_{\ell}({\bm{x}})$ contains elements of size $\Theta(1)$.
Satisfaction of 2 merely requires that $\phi$ maps a nonvanishing fraction of
preactivations to nonzero quantities. Its violation would require
preactivations to concentrate in regions of $\mathbb{R}$ which $\phi$ maps to
zero. This might occur, for example, in the unlikely scenario in which
$\phi=\textrm{ReLU}$ and almost all preactivations are negative.666It will
also occur if the activation function is $\phi(z)=0$, but we can neglect this
edge case. For activations that do not map nonzero inputs to zero (for
example, $\phi=\tanh$), 2 may be dropped altogether. Figure 5 verifies 2 in a
deep ReLU MLP. See Poole et al. (2016) for a theoretical framework sufficient
to check 2 at initialization.

Figure 5:  Verification of 2. At all layers, the norm of the activation vector
scales as the norm of the corresponding preactivation vector, both before and
after training. (Note that the ratio in this case is close to $1/\sqrt{2}$,
which is what one expects from an approximately-mean-zero random variable
passed through a ReLU nonlinearity.)

######  Assumption 3.

$\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell}({\bm{x}}_{i})\right|\\!\right|_{2}=\Theta(|\\!|\frac{1}{B}\Delta\\!{\bm{W}}_{\ell}^{(i)}{\bm{h}}_{\ell}({\bm{x}}_{i})|\\!|_{2})$.

Like 1, a violation of 3 would require a perfect cancellation of high-
dimensional matrices, which is unlikely. We verify 3 in Figure 6. 3 is also
verified implicitly by the right-hand subplot of Figure 1.

Figure 6:  Verification of 3. We compute
$\|\Delta\\!{\bm{W}}_{\ell}^{(i)}{\bm{h}}_{\ell}({\bm{x}}_{i})|_{2}/\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}{\bm{h}}_{\ell}({\bm{x}}_{i})\right|\\!\right|_{2}$
at layer $\ell=2$ for a single step both at initialization and after training.
As expected, this ratio remains $\Theta(1)$ as width grows — note the lack of
decay with width. We note as a curiosity that this quantity hovers around
$1/\sqrt{B}$ (dashed line), which one expects from the central limit theorem
if, for example, $\Delta\\!{\bm{W}}_{\ell}$ resembles a sum of $B$ terms
similar to $\Delta\\!{\bm{W}}_{\ell}^{(i)}$ but with random signs.

##  Appendix D Scalings for biases

Here we extend our spectral analysis to biases. Let
${\bm{b}}_{\ell}\in\mathbb{R}^{0pt_{\ell}}$ be a bias vector which enters
during forward propagation as
${\bm{h}}_{\ell}({\bm{x}})={\bm{W}}_{\ell}{\bm{h}}^{\prime}_{\ell-1}({\bm{x}})+{\bm{b}}_{\ell}$.
We may choose to view the bias vector as a weight matrix
${\bm{b}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times 1}$ connecting an auxiliary
layer with width $1$ and output $1$ to the $\ell$th hidden layer, after which
we may simply apply our scaling analysis for weight matrices. The spectral
scaling condition (1) prescribes that
$\left|\\!\left|{\bm{b}}_{\ell}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})$
and
$\left|\\!\left|\Delta{\bm{b}}_{\ell}\right|\\!\right|_{2}=\Theta(\sqrt{0pt_{\ell}})$,
and 1 prescribes that the initialization scale and learning rate should be
$\sigma^{b}_{\ell}=\Theta(1)$ and $\eta^{b}_{\ell}=\Theta(0pt_{\ell})$. In
practice, one may usually just take $\sigma^{b}_{\ell}=0$.

##  Appendix E Nondimensionalization and natural norms

While there are conventional notions of vector norms and matrix spectral
norms, one can equip any norm on the input and output spaces of a matrix,
thereby inducing a corresponding spectral norm on it. For example, the
conventional spectral norm of a matrix is defined assuming the typical
$\ell^{2}$ norm on the input and output spaces; if the input space is instead
given the $\ell^{1}$ norm, then the spectral norm is something very different.
As we explain below, the vectors encountered in deep neural networks naturally
come with a specific notion of norm, which we refer to as the “natural norm.”
When we equip vectors at various layers of the neural network with their
natural norms, we induce what we call a “natural spectral norm” on the weight
matrices. Our primary criterion 1 is then simplified to stipulating that these
_natural spectral norms_ should be $\Theta(1)$ in the limit of large network
width (2).

###  E.1 Dense and sparse vectors

Vectors in the context of deep learning typically fall into one of two
categories: dense and sparse. Dense vectors are characterized by having every
entry contribute a constant amount to the squared Euclidean norm of the
vector, while sparse vectors have only a constant number of entries that do
so. Here, the term “constant” means $\Theta(1)$ with respect to the size of
the vector.777Actually, a vector of $d$ iid samples from a unbounded
subgaussian or subexponential distribution will have its extreme elements
scale like $\mathrm{polylog}(d)$, so the really correct thing to say here is
$\tilde{\Theta}(1)$, but in the main text we convey the key intuition in the
main text without being too pedantic. This fine detail does not affect the
conclusion of this section.

For example: in the input layer, image data typically takes the form of dense
vectors, while the one-hot encoding in language models takes the form of
sparse vectors. The output vector of a network is typically a dense vector.
All hidden vectors (pre- or post-activation) are dense.

###  E.2 Defining natural norms

######  Definition 2 (Natural $\ell^{2}$-norm).

The _natural $\ell^{2}$-norm_ for a _dense_ vector ${\bm{v}}\in\mathbb{R}^{m}$
is the RMS norm

$\left|\\!\left|{\bm{v}}\right|\\!\right|_{\tilde{2}}\vcentcolon=\frac{1}{\sqrt{m}}\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}.$
(28)

The _natural $\ell^{2}$-norm_ for a _sparse_ vector ${\bm{v}}$ is simply the
usual $\ell^{2}$-norm:

$\left|\\!\left|{\bm{v}}\right|\\!\right|_{\tilde{2}}:=\left|\\!\left|{\bm{v}}\right|\\!\right|_{2}.$
(29)

While we say “dense” or “sparse” vector, these adjectives really apply to the
space that contains such a vector. For example, the set of one-hot encodings
form a “sparse” input space, while the pre-activations at layer $\ell$ forms a
“dense” hidden space.

######  Definition 3 (Natural spectral norm).

Given a parameter matrix ${\bm{A}}$, equip its input and output spaces with
their natural norms. Then the _natural spectral norm_
$\left|\\!\left|{\bm{A}}\right|\\!\right|_{\tilde{*}}$ of ${\bm{A}}$ is
defined as the induced spectral norm.

For example, if both the input and output spaces are dense, then the natural
spectral norm of ${\bm{A}}\in\mathbb{R}^{m\times n}$ is

$\left|\\!\left|{\bm{A}}\right|\\!\right|_{\tilde{*}}\vcentcolon=\frac{\sqrt{m}}{\sqrt{n}}\left|\\!\left|{\bm{A}}\right|\\!\right|_{*}.$
(30)

###  E.3 Spectral scaling in natural norms

By equipping vectors and matrices with their natural norms, we simplify the
conditions 1 and 1 from the main text:

######  Desideratum 2 (Feature learning, natural norms).

Let ${\bm{h}}_{\ell}({\bm{x}})\in\mathbb{R}^{0pt_{\ell}}$ denote the features
of input ${\bm{x}}$ at layer $\ell$ of a deep neural network, and let
$\Delta{\bm{h}}_{\ell}({\bm{x}})\in\mathbb{R}^{0pt_{\ell}}$ denote their
change after a gradient update. We desire that:
$\displaystyle\left|\\!\left|{\bm{h}}_{\ell}\right|\\!\right|_{\tilde{2}}=\Theta(1)\quad\text{and}\quad\left|\\!\left|\Delta{\bm{h}}_{\ell}\right|\\!\right|_{\tilde{2}}=\Theta(1),\quad\text{
at layers }\ell=1,...,L\\!-\\!1.$

Note the norm here is just RMS norm because we only talk about hidden vectors
(which are always dense).

######  Condition 2 (Spectral scaling, natural norms).

Consider applying a gradient update
$\Delta\\!{\bm{W}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times 0pt_{\ell-1}}$ to the
$\ell$th weight matrix ${\bm{W}}_{\ell}\in\mathbb{R}^{0pt_{\ell}\times
0pt_{\ell-1}}$. The spectral norms of these matrices should satisfy:
$\displaystyle\left|\\!\left|{\bm{W}}_{\ell}\right|\\!\right|_{\tilde{*}}=\Theta\left(1\right)\
\ \ \text{and}\ \ \
\left|\\!\left|\Delta\\!{\bm{W}}_{\ell}\right|\\!\right|_{\tilde{*}}=\Theta\left(1\right),\quad\text{
at layers }\ell=1,...,L.$

In summary, using these rescaled norms (that we call “natural norms”), our
problem is nondimensionalized: feature vectors, feature vector updates, weight
matrices, and weight matrix updates are $\Theta(1)$ in norm. These natural
norms provide a universal framework that covers specific cases, such as one-
hot embeddings in language models, that are not handled directly by 1.