Spectral Scaling for Feature Learning
Spectral Scaling for Feature Learning
Abstract
arXiv:2310.17813v1 [[Link]] 26 Oct 2023
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 fan-out/fan-in, in contrast to widely used but heuristic
p
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 (Brown et al.,
2020; Ramesh et al., 2022; Silver et al., 2016). 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 (Arora et al., 2019; Canatar et al., 2021; Jacot et al., 2018; Lee et al.,
2018) and practical scaling recommendations (Dey et al., 2023; Yang & Hu, 2021b; Yang et al., 2021).
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 (Atanasov et al., 2022; Fort et al., 2020; Lee
et al., 2020; 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 (µP) is derived by fairly involved “tensor programs” arguments that track
feature distributions analytically in the infinite width limit. Anecdotally, the principles underlying µP are
not well understood by the community. In this paper, we provide a new perspective on µP, showing that
its scaling relations can be obtained by elementary linear algebra arguments. Inpshort, we show that µP
is equivalent to scaling the spectral norm of any weight matrix or update like fan-out/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 µP
in code, and is significantly easier to work with theoretically, leading to further conceptual advances in our
research (Bernstein et al., 2023).
∗ Equal contribution.
1
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 (Amari, 1998; Nemirovsky & Yudin, 1983). 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 (Liu et al., 2021;
Shazeer & Stern, 2018; You et al., 2017; 2020). This paper shows that the spectral norm provides a better
notion of distance between operators in the context of deep learning.
Desideratum 1 (Feature learning). Let hℓ (x) ∈ Rnℓ denote the features of input x at layer ℓ of a
neural network, and let ∆hℓ (x) ∈ Rnℓ denote their change after a gradient step. We desire that:
√ √
||hℓ ||2 = Θ( nℓ ) and ||∆hℓ ||2 = Θ( nℓ ), at layers ℓ = 1, ..., L−1.
Let us unpack Desideratum 1. These conditions treat the ℓ2 -norms of the feature vectors hℓ (x) and ∆hℓ (x),
a framing which will prove convenient. Desideratum 1 amounts to asking that the “typical element size”
of vectors hℓ (x) and ∆hℓ (x) is Θ(1) with respect to width nℓ (we give a review of big-Θ notation in
Section 2). Enforcing that hidden features have Θ(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 Θ(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 Desideratum 1 as our definition of feature learning.1
Our main message is that feature learning in the sense of Desideratum 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 ∆Wℓ ∈ Rnℓ ×nℓ−1 to the ℓth
weight matrix Wℓ ∈ Rnℓ ×nℓ−1 . The spectral norms of these matrices should satisfy:
r r
nℓ nℓ
||Wℓ ||∗ = Θ and ||∆Wℓ ||∗ = Θ , at layers ℓ = 1, ..., L.
nℓ−1 nℓ−1
We review the spectral norm in Section 2. The spectral scaling condition has two components which will
serve to enforce the respective components of Desideratum 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 ||x||2 = Θ( n0 ), which is standard for image data. Language
√
models are an important counterexample, where embedding matrices take one-hot inputs and the n0 in
Condition 1 should be replaced by 1. Appendix E provides a unifying treatment of these cases.
To get some quick intuition for the origin of Condition 1, observe that under the forward propagation
hℓ (x) = Wℓ hℓ−1 (x), if the layer input hℓ−1 (x) aligns with the top singularpvector of the weight matrix
Wℓ , then ||hℓ (x)||2 = ||Wℓ ||∗ · ||hℓ−1 (x)||2 . The requirement that ||Wℓ ||∗ = Θ( nℓ /nℓ−1 ) then follows from
1 Our 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 (Olah et al., 2017;
Zeiler & Fergus, 2014)—nontrivial feature evolution in our sense is necessary for any other reasonable definition of the term.
2
Desideratum 1. The scaling of ||∆Wℓ ||∗ can similarly be obtained by writing ∆hℓ (x) = ∆Wℓ hℓ−1 (x) + . . .
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 Desideratum 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 Condition 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” (µP) (Yang & Hu, 2021b).
1.2 Summary of contributions
Concretely, our contributions are as follows:
• We propose the spectral scaling condition (Condition 1) and show that it suffices to achieve feature
learning in neural networks even at large width.
• We show how Condition 1 may be implemented: either via direct spectral normalization, or by layer-
wise initialization scales {σℓ } and learning rates {ηℓ } that recover maximal update parametrization.2
• We show that other popular scaling rules, including so-called standard parameterization and neural
tangent parametrization, fail to satisfy Condition 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. Appendix B). 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 Condition 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 (n) = O(g(n)) means that f (n) “scales no faster than” g(n),
• f (n) = Θ(g(n)) means that f (n) “scales like” or “is order” g(n),
• f (n) = Ω(g(n)) means that f (n) “scales at least as fast as” g(n).
Formally, f (n) = Θ(g(n)) is equivalent to the statement that there exist constants c, C > 0 such that
c·g(n) ≤ f (n) ≤ C ·g(n) for all sufficiently large d. The weaker statements f (n) = O(g(n)) and f (n) = Ω(g(n))
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 ℓ2 -norm ||·||2 to assess the size of a vector. For
matrices, we will principally use the spectral norm ||·||∗ (a.k.a. operator norm) defined as follows:
3
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
2
spectral norm with the Frobenius norm ||·||F given by ||A||F = ij A2ij .
P
Properties of the spectral norm. Let A, B ∈ Rm×n be arbitrary matrices and v ∈ Rn be an arbitrary
vector. As with all norms, the spectral norm is subadditive, meaning that it obeys the triangle inequality
||A + B||∗ ≤ ||A||∗ + ||B||∗ . The spectral norm is also submultiplicative in the sense that ||Av||∗ ≤ ||A||∗ · ||v||2
and ||AB||∗ ≤ ||A||∗ · ||B||∗ . If we interpret a vector v ∈ Rn as a 1 × n matrix, then the ℓ2 , spectral and
Frobenius norms are equivalent: ||v||2 = ||v||∗ = ||v||F .
Special cases of the spectral norm. For a rank-one matrix A, which can be written as an outer-product
A = uv ⊤ , it holds that ||A||∗ = ||A||F = ||u||2 · ||v||2 . A matrix B ∈ Rm×n is semi-orthogonal if either
B ⊤ B = In or BB ⊤ = Im . A semi-orthogonal matrix has unit spectral norm: ||B||∗ = 1.
We let hL (x) ∈ RnL be the network output. We will keep the input dimension n0 and output dimension nL
fixed and consider scaling with respect to the hidden dimensions n1 , . . . , nL−1 .
We let the global loss be L = g(hL (x), y) where g and y are a loss function and target vector, respectively.
During training, we will take gradient steps at each layer as ∆Wℓ = −ηℓ · ∇Wℓ L, where ηℓ is a layerwise
learning rate. We will ultimately solve for the scale of ηℓ , but for now we will be content to discuss the
perturbation ∆Wℓ directly. By Equation (2), hidden vector updates at subsequent layers are related by:
hℓ (x) + ∆hℓ (x) = (Wℓ + ∆Wℓ )(hℓ−1 (x) + ∆hℓ−1 (x)). (3)
Hidden vector sizes. To reiterate Desideratum 1, we wish for features at the ℓth layer to have a norm
√
which scales as ||hℓ (x)||2 = Θ( nℓ ). Upon a gradient update, this feature vector should undergo an update of
√
size ||∆hℓ (x)||2 = Θ( nℓ ). We will show that Condition 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 ℓ − 1 and their updates
satisfy Desideratum 1, then so will those at layer ℓ. 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 Desideratum 1, and then show that they are in fact also not smaller than asked by Desideratum 1.
Hidden vector updates. By the subadditivity and submultiplicativity of the spectral norm, Equations (2)
and (3) imply that:
√
||hℓ (x)||2 ≤ ||Wℓ ||∗ · ||hℓ−1 (x)||2 = Θ( nℓ ); (4)
√
||∆hℓ (x)||2 ≤ ||∆Wℓ ||∗ · ||hℓ−1 (x)||2 + ||Wℓ ||∗ · ||∆hℓ−1 (x)||2 + ||∆Wℓ ||∗ · ||∆hℓ−1 (x)||2 = Θ( nℓ ), (5)
where on the right hand sides of the inequality we have inserted Desideratum 1 and Condition 1. The spectral
scaling condition thus gives features and feature updates obeying the correct upper bounds, and we need
merely show comparable lower bounds.
4
Tightness of bounds via matrix-vector alignment. The upper bound in the submultiplicativity property
||Av||2 ≤ ||A||∗ · ||v||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 hℓ−1 (x) ∈ Rnℓ−1 . Assume Wℓ in
Rnℓ ×nℓ−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 (nℓ ≥ nℓ−1 ), then with high probability:
||Wℓ hℓ−1 (x)||2 = Θ(||Wℓ ||∗ · ||hℓ−1 (x)||2 ).
Claim 2 (Alignment of updates). For an update ∆Wℓ given by gradient descent with batch size 1,
||∆Wℓ hℓ−1 (x)||2 = ||∆Wℓ ||∗ · ||hℓ−1 (x)||2 .
In words, Claim 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. Claim 2 states the same of weight updates, but the
proportionality constant is precisely one and requires no condition on dimensionality.3
We now justify these claims in turn. For Claim 1, first suppose that Wℓ is a random semi-orthogonal
matrix as is a popular initialization strategy. Then all singular values of Wℓ are one and, since fan-out
exceeds fan-in, the null-space of Wℓ is empty. Taken together, these observations imply the equality:
||Wℓ hℓ−1 (x)||2 = ||Wℓ ||∗ · ||hℓ−1 (x)||2 . Fortunately, if the elements of Wℓ are instead sampled i.i.d. from a
centered Gaussian distribution with standard deviation σℓ , then the situation is similar. It is easily shown
√
by the law of large numbers that ||Wℓ hℓ−1 (x)||2 ≈ σℓ nℓ ||hℓ−1 (x)||2 , and it is a standard result in random
√ √
matrix theory that ||Wℓ ||∗ ≈ σℓ ( nℓ−1 + nℓ ) (Rudelson & Vershynin, 2010; Vershynin, 2018). Claim 1 for
Gaussian intialization follows by combining these results.
Claim 2 is easily verified as follows. Observe that we can write the update at layer ℓ as the outer-product:
∆Wℓ = −ηℓ · ∇hℓ (x) L · hℓ−1 (x)⊤ . (6)
So the update ∆Wℓ is rank-one with right singular vector hℓ−1 (x). To verify the claim, observe that:
2
||∆Wℓ hℓ−1 (x)||2 = ηℓ · ∇hℓ (x) L 2
· ||hℓ−1 (x)||2 = ||∆Wℓ ||∗ · ||hℓ−1 (x)||2 . (7)
With the claims established, we can now get lower bounds on hidden vector size which serve to verify
Desideratum 1. The features at initialization scale correctly as:
√
||hℓ (x)||2 = Θ (||Wℓ ||∗ · ||hℓ−1 (x)||2 ) = Θ( nℓ ), (8)
where we have first used Claim 1 and then inserted Condition 1. To bound the size of ∆hℓ (x), let us first
observe from Equation (3) that
∆hℓ (x) = ∆Wℓ hℓ−1 (x) + Wℓ ∆hℓ−1 (x) + ∆Wℓ ∆hℓ−1 (x). (9)
So long as the first term ∆Wℓ hℓ−1 (x) does not perfectly cancel with the latter two, we have that:
√
||∆hℓ (x)||2 = Ω(||∆Wℓ ||∗ · ||hℓ ||2 ) = Ω( nℓ ), (10)
where in the last step we have inserted Condition 1. Combining Equation (10) with our matching upper
√
bound from Equation (5), we conclude that ||∆hℓ (x)||2 = Θ( nℓ ) as desired.
We have achieved both clauses of Desideratum 1 at layer ℓ, completing a recursive step from layer ℓ − 1. The
norm of the input 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 Note that x in Claim 2 is the same input that induced the gradient ∆W in the previous step. Claim 2 is generally not an
ℓ
equality if this not true.
5
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 (Yaida,
2022; Yang & Hu, 2021b) discusses layerwise initialization scales {σℓ } and learning rates {ηℓ } 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 Wℓ and ∆Wℓ . This enables us to
determine the sizes of hidden vectors and their updates quite easily, whereas the same calculation in terms of
σℓ and ηℓ is more involved. Section 4 shows how to recover σℓ and ηℓ from our spectral scaling condition.
The sort of perfect cancellation required to violate this assumption will be rare in practice (and adding a
small amount of randomness
p to the learning rate ηℓ will fix any occurrence with high probability). It follows
√
that ||Wℓ + ∆Wℓ ||∗ = Θ( nℓ /nℓ−1 ) and ||hℓ (x) + ∆hℓ (x)||2 = Θ( nℓ ). With these facts in place, the same
argument we used in Section 3.1 for the first step ensures that Desideratum 1 also holds at later steps.
3.2.2 Nonlinearities
We now add a nonlinearity ϕ to each layer of our MLP. The modified forward recursion relation is:
hℓ (x) = Wℓ h′ℓ−1 (x), h′ℓ (x) = ϕ(hℓ (x)); for ℓ = 2, . . . , L−1, (13)
with base case h1 (x) = W1 x and output hL (x) = WL h′L−1 (x). We assume that the hidden features before
and after the application of the nonlinearity are of the same scale:
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, ∆Wℓ will be rank-one and align to incoming signal as:
at each step. All our scaling arguments from the linear case therefore carry through: the term ∆Wℓ h′ℓ−1 (x) is
sufficient to induce a change ∆hℓ (x) satisfying Desideratum 1. (The other terms which depend on ∆h′ℓ−1 (x)
will be no larger.) Condition 1 therefore still achieves correctly-scaled feature evolution.
6
1.00
ReLU
Stable rank of ∆W`
Mean alignment of
102
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 n1 = n2 = 300, and ReLU and tanh activation functions. We then compute gradient
updates ∆Wℓ for layer ℓ = 2 on randomly-sampled size-B subsets of CIFAR-10. Left. As a measure of
effective rank, we report the stable rank srank(∆Wℓ ) := ||∆Wℓ ||F / ||∆Wℓ ||∗ of gradient updates. The stable
2 2
rank remains less than 10 even when B is large, which is much less than its maximal possible value of
min(n1 , n2 ) = 300. Right. We report the average alignment ||∆Wℓ hℓ−1 (x)||2 / ||∆Wℓ ||∗ h′ℓ−1 (x) 2 of the
weight update ∆Wℓ to the incoming vector h′ℓ−1 (x), averaged over 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.
We additionally make the assumption that the batch size is fixed and independent of width:
Combining these assumptions, we find that ||∆Wℓ hℓ (xi )||2 = Θ (||∆Wℓ ||∗ · ||hℓ (xi )||2 ): the batch update is
aligned with incoming signal in a scaling sense. Our previous scaling arguments in fact needed only alignment
in this “big-Θ” sense (as opposed to the perfect rank-one sense of Claim 2), and thus they still carry through:
our spectral scaling condition continues to suffice to achieve proper feature learning as per Desideratum 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.
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 ∆Wℓ calculated from these optimizers. In a gist, the main ingredients are 1)
7
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.
Matrix Shape Stable Rank Spectral Norm Frobenius Norm RMS Entry Size
√
1 1/n
p
Wℓ n×n n n
∆Wℓ n×n 1 1 1 1/n
q q q
min(nℓ ,nℓ−1 )
Wℓ nℓ × nℓ−1 min(nℓ , nℓ−1 ) nℓ
nℓ−1 min(nℓ , nℓ−1 ) · nℓ
nℓ−1 nℓ ×nℓ−1 · nℓ
nℓ−1
q q q
∆Wℓ nℓ × nℓ−1 1 nℓ
nℓ−1
nℓ
nℓ−1
1
nℓ ×nℓ−1 · nℓ
nℓ−1
Table 1: Conversion between matrix norms for initial weights Wℓ and updates ∆Wℓ . Results are
first given for square matrices and vectors, then finally generalized to matrices of all shapes. Table entries
denote the Θ-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 p square root of the number of entries. Notice that a n × n initial matrix Wℓ with Θ(1) spectral
norm has Θ( 1/n) entry √ size, whereas a n × n update ∆Wℓ with Θ(1) spectral norm has only Θ(1/n) entry
size, which is a factor n 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.
Parametrization 1 (Spectral parametrization). We claim that the spectral scaling condition (Condi-
tion 1) is satisfied and feature learning is achieved (as per Desideratum 1) if the initialization scale
and learning rate of each layer ℓ are chosen according to:
1
r
nℓ nℓ
σℓ = Θ √ min 1, ; ηℓ = Θ .
nℓ−1 nℓ−1 nℓ−1
Naïve method: direct spectral normalization. The most naïve way to impose Condition 1 is to
directly normalize the relevant quantities by their spectral
p norm. For instance, to initialize the weight matrix
Wℓ ∈ Rnℓ ×nℓ−1 at the ℓth layer to have spectral norm nℓ /nℓ−1 , one could sample a temporary matrix Wℓ′
8
using any standard initializer and then re-normalize according to:
Wℓ′
r
nℓ
Wℓ = σ × , (16)
nℓ−1 ||Wℓ′ ||∗
where σ = Θ(1) is a width-independent prefactor. Similarly, to ensure that the gradient step ∆Wℓ also has a
spectral norm of the proper size, one might spectrally normalize the gradient according to:
∇W ℓ L
r
nℓ
∆Wℓ = −η × , (17)
nℓ−1 ||∇Wℓ L||∗
where η = Θ(1) is a width-independent prefactor.
It is quick to check that normalizing as in Equations (16) and (17) will satisfy Condition 1, but computing
spectral norms is expensive and (as we will now show) can be avoided entirely by working out how ||Wℓ′ ||∗
and ||∇Wℓ L||∗ will scale and dividing by the appropriate factor.
Random initialization. Let us suppose that, as is common practice, Wℓ is initialized as Wℓ = σℓ · Wℓ′ ,
where all elements of Wℓ′ 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 ||Wℓ ||∗ ≈ p σℓ · ( nℓ + nℓ−1 ) (Rudelson &
Vershynin,
p 2010; Vershynin, 2018). To get the desired scaling ||Wℓ ||∗ = Θ( nℓ /nℓ−1 ), we need merely choose
√ √
σℓ = Θ( nℓ /nℓ−1 · ( nℓ + nℓ−1 )−1 ). Simplifying within the Θ(·), we arrive at σℓ scaled as in the spectral
parametrization (Parametrization 1). Initializing weights with a prefactor σℓ scaling in this manner achieves
the correct spectral norm of Wℓ . We note
√ that the constant factor suppressed by the Θ(·) here will usually
be small—for example, a prefactor of 2 agrees with typical practice for ReLU networks at most p layers. If
Wℓ′ is instead a random semi-orthogonal matrix, then we can simply use a prefactor σℓ = Θ( nℓ /nℓ−1 ).
Gradient updates. Here we give two methods for obtaining weight updates ∆Wℓ 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 ||∇Wℓ L||∗ ≈ ||∇Wℓ L||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
The main challenge lies in finding the scaling of the gradient ||∇Wℓ L||∗ . Note that we expect each gradient
√
update ∆Wℓ to induce a change ||∆hL (x)||2 = Θ( nL ) in the output which induces a change ∆L = Θ(1) for
common loss functions L. Taylor expanding the loss to first order, we also expect that
∆L = Θ(⟨∆Wℓ , ∇Wℓ L⟩) = Θ(||∆Wℓ ||F · ||∇Wℓ L||F ) = Θ ||∆Wℓ ||∗ · ||∇Wℓ L||∗ , (19)
where ⟨·, ·⟩ 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 ∆L = Θ(1) and the spectral scaling
condition ||∆Wℓ ||∗ = Θ( nℓ /nℓ−1 ), we can conclude that
p
This result may also be reached via direct layerwise-recursive analysis of the size of thepgradient. Returning
to Equation (18), we now see that we achieve a properly-scaled update ||∆Wℓ ||∗ = Θ( nℓ /nℓ−1 ) if we take
ηℓ = Θ (nℓ /nℓ−1 ) as prescribed by Parametrization 1.
In summary: training with layerwise initialization σℓ and learning rate ηℓ scaled as in our Parametrization 1
will implement the spectral scaling condition and give features and feature evolution of the correct size.
9
hyperparameters σℓ and ηℓ that suffice to put this condition into action in a standard deep learning paradigm
(Parametrization 1). Here we compare this “spectral parametrization” with popular parametrizations. We find
that our parametrization recovers the “maximal update parametrization” (µP) at all layers and is different to
other parametrizations.
5.1 Comparison with “maximal update parametrization”
Maximal update parametrization (µP) was recently proposed as a scaling rule that retains feature learning even
at infinite width. µP as given in Table 3 of Yang et al. (2021) may be recovered from our Parametrization 1 by
setting n0 = nL = 1 and n1 = n2 = ... = nL−1 . Our Parametrization 1 actually streamlines and generalizes
µ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 n0 and output dimension nL (as opposed to neglecting
them as agnostic Θ(1) quantities) and treats hidden widths of unequal dimension (nℓ−1 ̸= nℓ ).
5.2 Contrast to “standard parametrization”
At present, the vast majority of deep learning systems use either “Kaiming,” “Xavier,” or “LeCun” initialization
(Glorot & Bengio, 2010; He et al., 2015; 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:
√
σℓ = Θ(1/ nℓ−1 ) and ηℓ = Θ(1). (21)
Notice that SP initialization exceeds Parametrization 1 in any layer with fan-out smaller than fan-in. This
includes the final layer p
in sufficiently wide networks. So, while Parametrization 1 implies that weight matrices
have spectral norm Θ( fan-out/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 Wℓ / nℓ−1
where the entries of Wℓ are sampled iid standard normal, and applies gradient descent with a layer-independent
√
step-size. Since dividing a weight matrix by nℓ−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:
√
σℓ = Θ(1/ nℓ−1 ) and ηℓ = Θ(1/nℓ−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 σL is nL−1 larger than
Parametrization 1, so that by the linearity of backpropagation, the gradient to any middle layer Wℓ is also
√ √
nL−1 larger. Then the 1/nℓ−1 learning rate in NTP induces a change ∆Wℓ that is nL−1 /nℓ−1 (smaller)
compared to Parametrization 1, which prescribes a learning rate of ηℓ = 1. Because Parametrization 1
guarantees ||∆Wℓ ||∗ = Θ(1), NTP causes ∆Wℓ to vanish in spectral norm as hidden widths nℓ−1 , nL → ∞.
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) (Bordelon &
Pehlevan, 2022; Chizat et al., 2019). This operation transforms the NTP into µP.
5.4 Contrast to “Frobenius-normalized updates”
Condition 1 mandates that the spectral norm of the update at layer ℓ be proportional to the spectral norm
of the weight matrix to which it is applied: ||∆Wℓ ||∗ ∝ ||Wℓ ||∗ . This contrasts with a body of optimization
work (Bernstein et al., 2020a;b; Liu et al., 2021; Shazeer & Stern, 2018; You et al., 2017) that has suggested
instead that the Frobenius norm of the update at layer ℓ be proportional to the Frobenius norm of the weight
matrix to which it is applied: ||∆Wℓ ||F ∝ ||Wℓ ||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
10
100 A 100 B
||h 2 (x) − h 02 (x)||2
||W 2 − W 02 || ∗
10 1 10 1
||h 02 (x)||2
1/2
||W 02 || ∗
1/2
10 2 µP 10 2
NTP
10 3 10 3
C D
Final-layer alignment
100 100
||W 2 − W 02 ||F
10 1 1.0 10 1 1.0
||W 02 ||F
1/2 1/2
10 2 0.5 10 2 0.5
1
10 3 0.0 10 3 0.0
0.0 102 0.5 103 1.0 0.0 102 0.5 103 1.0
example
Width nd example
Width dn
Figure 2: Spectral quantities are Θ(1) under µ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.
1), when in reality weight matrices have high stable rank and updates have Θ(1) stable [Link] light of this
fact, the Frobenius-normalized proportionality rule should be modified to ||∆Wℓ ||F ∝ ||Wℓ ||F / min(nℓ , nℓ−1 )
in order to see proper feature evolution.
where h02 (x) and h2 (x) are the second (preactivation) hidden vector at initialization and after training
respectively, and the expectation is over samples x from the batch. As shown in Figure 2A, this feature
11
evolution ratio remains roughly
√ fixed as width d grows when using µP, in satisfaction of Desideratum 1. By
contrast, it decays as 1/ n 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:
W2 − W20 ∗ / W20 ∗ , where W20 and W2 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 µP, in accordance with
Condition 1. By contrast, it decays as 1/ n for the NTP.
Final-layer alignment. We measure the alignment of the final layer to incoming vectors as follows:
||W3 h′2 (x)||2
Final-layer alignment := Ex . (24)
||W3 ||∗ · ||h′2 (x)||2
√
This quantity√is Θ(1/ n) at initialization. As shown in Figure 2C, it grows to Θ(1) when using µP but
remains Θ(1/ n) 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
W2 − W20 F / W20 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
µ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 (Mei et al., 2019; Rotskoff & Vanden-Eijnden, 2022;
Sirignano & Spiliopoulos, 2022), and the related “feature learning” limit (Geiger et al., 2020; Yaida, 2022;
Yang & Hu, 2021b) (i.e. the µ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 (Condition 1), this method treats
only the weights Wℓ , not the updates ∆Wℓ . Furthermore, spectral normalization implementations typically
set the spectral norm of weight matrices
p either to one or to a tunable hyperparameter (Farnia et al., 2019),
and do not include the key factor of nℓ /nℓ−1 in our Condition 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 (Bernstein et al., 2020a;b; Carbonnelle
& Vleeschouwer, 2019; Liu et al., 2021; Shazeer & Stern, 2018; You et al., 2017). A particular focus has
12
been placed on setting the Frobenius norm of update steps to be small relative to the Frobenius norm of the
weight matrices (Bernstein et al., 2020a;b; Liu et al., 2021; You et al., 2017). 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 (Desideratum 1) and culminating in the demonstration that these
conditions may be achieved by simple scaling rules (Condition 1 and Parametrization 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 µ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 ℓ2 -norm of a vector v ∈ Rn incorporating a
normalization prefactor, all vector and matrix norms used in our spectral analysis become Θ(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 ||x||2 = Θ( n0 ).
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
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 1998. Cited on page 2.
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. Cited on
page 1.
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: the silent
alignment effect. In International Conference on Learning Representations, 2022. Cited on page 1.
Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv:1607.06450, 2016. Cited
on page 12.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD: Compressed
Optimisation for Non-Convex Problems, February 2018. URL [Link]
Cited on page 7.
13
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. Cited on pages 2, 6, 10, 12,
and 13.
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. Cited on pages 10, 12, and 13.
Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan, and Yisong Yue. Automatic Gradient Descent:
Deep Learning without Hyperparameters. arXiv:2304.05187, 2023. Cited on page 1.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural
networks. In Neural Information Processing Systems, 2022. Cited on page 10.
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. Cited on page 1.
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.
Cited on page 1.
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. Cited on page 12.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Neural
Information Processing Systems, 2019. Cited on pages 10 and 12.
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. Cited on page 1.
Inderjit S. Dhillon and Joel A. Tropp. Matrix nearness problems with Bregman divergences. SIAM Journal
on Matrix Analysis and Applications, 2008. Cited on page 2.
Farzan Farnia, Jesse Zhang, and David Tse. Generalizable adversarial training via spectral normalization. In
International Conference on Learning Representations, 2019. Cited on page 12.
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. Cited on pages 1
and 12.
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. Cited on page 12.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks.
In International Conference on Artificial Intelligence and Statistics, 2010. Cited on page 10.
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. Cited
on page 10.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing
internal covariate shift. In International Conference on Machine Learning, 2015. Cited on page 12.
Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Convergence and generalization
in neural networks. In Neural Information Processing Systems, 2018. Cited on pages 1, 10, 12, and 17.
14
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference
on Learning Representations, 2015. Cited on page 1.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of
Toronto, 2009. Cited on page 17.
Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient backprop. In Neural
Networks: Tricks of the Trade. Springer, 2002. Cited on pages 2 and 10.
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 Represen-
tations, 2018. Cited on page 1.
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. Cited on pages 1, 12, and 17.
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. Cited on page 1.
Yang Liu, Jeremy Bernstein, Markus Meister, and Yisong Yue. Learning by turning: Neural architecture
aware optimisation. In International Conference on Machine Learning, 2021. Cited on pages 2, 10, 12,
and 13.
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. Cited on page 12.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for
Generative Adversarial Networks. In International Conference on Learning Representations, 2018. Cited
on page 12.
Arkady S. Nemirovsky and David B. Yudin. Problem complexity and method efficiency in optimization. Wiley,
1983. Cited on page 2.
Chris Olah, Alexander Mordvintsev, and Ludwig Schubert. Feature visualization. Distill, 2017. Cited on
page 2.
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. Cited on page 21.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional
image generation with CLIP latents. arXiv:2204.06125, 2022. Cited on page 1.
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. Cited on
page 12.
Mark Rudelson and Roman Vershynin. Non-asymptotic theory of random matrices: Extreme singular values.
In International Congress of Mathematicians, 2010. Cited on pages 5 and 9.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In
International Conference on Machine Learning, 2018. Cited on pages 2, 10, and 12.
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. Cited on page 1.
15
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. Mathematics of
Operations Research, 2022. Cited on page 12.
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. Cited on page 1.
Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science.
Cambridge University Press, 2018. Cited on pages 5 and 9.
Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the NTK for understanding generalization
in deep learning. arXiv:2206.10012, 2022. Cited on pages 1 and 12.
Sho Yaida. Meta-principled family of hyperparameter scaling strategies. arXiv:2210.04909, 2022. Cited on
pages 6 and 12.
Greg Yang and Edward J. Hu. Feature learning in infinite-width neural networks. In ICML 2021 Work-
shop on Over-parameterization: Pitfalls and Opportunities, 2021a. [Link]
feature-learning-in-infinitewidth-neural-networks. Cited on page 12.
Greg Yang and Edward J. Hu. Tensor Programs IV: Feature learning in infinite-width neural networks. In
International Conference on Machine Learning, 2021b. Cited on pages 1, 3, 6, 12, and 18.
Greg Yang and Etai Littwin. Tensor programs IVb: Adaptive optimization in the ∞-width limit.
arXiv:2308.01814, 2023. Cited on page 18.
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. Cited on pages 1, 10, and 13.
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. Cited on pages 1, 2, 10, 12, and 13.
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. Cited on page 2.
Matthew D. Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In European
Conference on Computer Vision, 2014. Cited on page 2.
16
ReLU ReLU
Stable rank of H `−1
102 102
Stable rank of G `
tanh tanh
101 101
100 100
101 102 103 104 101 102 103 104
Batch size (B) Batch size (B)
Figure 3: Stable ranks for the set of forward hidden vectors Hℓ−1
′
and the set of backward hidden vectors Gℓ
for ℓ = 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
2 2
a matrix M is given by ||M ||F / ||M ||∗ .
A Experimental details
Experimental details for Figure 1. We examine randomly-initialized MLPs with depth L = 3, widths
n0 = 3072, n1 = n2 = d = 300, n3 = 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 ∆Wℓ at layer
ℓ = 2 with arbitrary learning rate. We then compute the stable rank of ∆Wℓ and the alignment metric
" #
||∆Wℓ hℓ−1 (xi )||2
Ex , (25)
||∆Wℓ ||∗ h′ℓ−1 (xi ) 2
where the expectation is taken over 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 ∆Wℓ . Let
′
Hℓ−1 = [h′1 (x1 ) . . . h′1 (xB )] ∈ Rd×B be the matrix formed from stacking the ℓ = 1 post-nonlinearity hidden
vectors from the full batch, and let Gℓ = [g2 (x1 ) . . . g2 (xB )] ∈ Rd×B with gℓ (x) = ∇hℓ (x) L be the matrix
⊤
formed from stacking the loss gradients at layer ℓ = 2 from the full batch. It is the case that ∆Wℓ ∝ Gℓ Hℓ−1
′
,
and if ∆Wℓ has low stable rank, it is a reasonable guess that either Hℓ−1 ′
and Gℓ also has low stable rank.
In fact, we find that both Hℓ−1 ′
and Gℓ have low stable rank, as shown in Figure 3.
Experimental details for Figure 2. We train MLPs with depth L = 3, widths n0 = 3072, n1 = n2 =
d, n3 = 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 ±1 targets.
We use two different hyperparameter schemes as follows. To implement µP, we take
√
2 2 2 1
r r
d
σ1 = ; σ2 = ; σ3 = ; η1 = η · ; η2 = η; η3 = η · , (µP)
n0 d d n0 d
with global learning rate η = 0.1. To implement NTP, we follow Jacot et al. (2018) and Lee et al. (2019):
s
2 1
σℓ = ; ηℓ = η · (NTP)
nℓ−1 nℓ−1
at all layers, again with η = 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 104 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
17
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.
Theorem 1. In µP, for almost every learning rate (in the measure-theoretic sense), Condition 1 is satisfied
at any time during training for sufficiently large width. µP is the unique stable abc-parametrization with this
property.
The main insight leading to Theorem 1 is that: ∆Wℓ / ||∆Wℓ ||F converges to a Hilbert-Schmidt operator (in
an appropriate space) as width goes to infinity. This in turn implies:
Lemma 1. Unless ∆Wℓ = 0, ||∆Wℓ ||F / ||∆Wℓ ||∗ = Θ(1) as width goes to infinity.4
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:
Random Initialization
The following is our main proposition for initialization:
18
Proof. This is a claim about random matrices and follows from classical random matrix theory. Here’s a
quick sketch: If σ is the standard
√ deviation of a matrix entry, then ||W ||F ≈ σ · n from the central limit
theorem, and ||W ||∗ ≈ 2σ · n 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 8x⟩ (called a ket) as width
goes to infinity in the sense that the scaled inner product ⟨x, y⟩/n converges almost surely to ⟨x8y⟩ = E 8x⟩8y⟩.
The kets form a Hilbert space Z. Then any weight parameter converges to a linear operator from Z to Z;
any vector parameter converges to a linear operator from Z to R or R to 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 W , at any step of
training, ∆W / ||∆W ||F converges to a Hilbert-Schmidt integral operator.
Proof. This is trivial for vector parameters. For matrix parameters, observe ∆W
g = ∆W / ||∆W || is always a
F
nonlinear outer product
∆W
g = Q(x; y; c)
for multi-vectors x and y and multi-scalars c and some nonlinearity Q. Then the limit 8∆W
g 8 acts on a ket
8z⟩ by
g 8z⟩ = E Q(8x⟩; 8y⟩ 1 ; c̊)8z⟩ 1
8∆W
1
Proposition 3. In Proposition 2, for large enough width, unless ∆W = 0, both spectral norm and Frobenius
norm of ∆W / ||∆W ||F are Θ(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 Z, one can construct a nonzero vector z in an extension of the program that defined
∆W , such that
8∆W
g 8z⟩ = θ 8∆W g 8 ||8z⟩|| ,
∗
for some θ ∈ (1/2, 1].5 This implies, by the Master Theorem, that
θ ∆W z
g
≤ ∆W
g
2 ||z|| ∗
19
Proof. In µP, it’s trivial to see that ||∆W ||F = Θ(1). So any nonzero ∆W (at any point of training) has both
spectral norm and Frobenius norm Θ(1) by the above.
√
Simple calculation then shows that W at any fixed time has Θ( n) since this is the case at initialization by
Proposition 1. This means the quadratic mean of the singular values of W is Θ(1), so its max singular value
must be Ω(1). But it furthermore must be Θ(1) because W at initialization and all of its updates have O(1)
spectral norm.
Theorem 2. In µP, for all but a measure-zero set of learning rates, Condition 1 is satisfied at any time
during training for sufficiently large width. µP is the unique stable and faithful abcd-parametrization with this
property.
Proof. In µP, by Proposition 4, all matrix parameters satisfy Condition 1 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, W has Θ(1) Frobenius norm at initialization and at all times during training.
By Proposition 3, W also has Θ(1) spectral norm, so Condition 1 is satisfied. A similar but easier argument
applies to all scalar parameters. This shows µP satisfies Condition 1.
Since any other stable and faithful parametrization essentially just rescales the initialization and the update,
we see that no other parametrization can satisfy Condition 1.
For SGD, all abc-parametrizations are equivalent to a faithful abcd-parametrization because Q is identity, so
Theorem 2 recovers Theorem 1.
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, Assumption 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 Assumption 1 also holds
when ∆Wℓ and ∆hℓ (x) are respectively replaced by Wℓ − Wℓ0 and hℓ (x) − h0ℓ (x), the total updates across
all of training.
20
1.00 1.00
||h ` (x)|| 2
||W` || ∗
`=1
0.25 0.5 `=2 0.25 0.5 `=1
||W`0 || ∗
`=3 `=2
0.00 0.01 0.00 0.0 1
100.0 100.5
2 101.0
3 100.0 100.5
2 101.0
3
example
Width dn example
Width n
d
Figure 4: Verification of Assumption 1. Subplots show the ratios between left- and right-hand sides
of conditions of Assumption 1 at various layers. Envelopes show variation over trials. The ratios are Θ(1)
(showing no obvious decay with network width), verifying Assumption 1.
Recall that h′ℓ (x) = ϕ(hℓ (x)), and that hℓ (x) contains elements of size Θ(1). Satisfaction of Assumption 2
merely requires that ϕ maps a nonvanishing fraction of preactivations to nonzero quantities. Its violation
would require preactivations to concentrate in regions of R which ϕ maps to zero. This might occur, for
example, in the unlikely scenario in which ϕ = ReLU and almost all preactivations are negative.6 For
activations that do not map nonzero inputs to zero (for example, ϕ = tanh), Assumption 2 may be dropped
altogether. Figure 5 verifies Assumption 2 in a deep ReLU MLP. See Poole et al. (2016) for a theoretical
framework sufficient to check Assumption 2 at initialization.
1.00
0.75
1.0
||h ` (x)|| 2
||h `0 (x)|| 2
0.50
0.25 0.5 ` = 1 (init) ` = 2 (init) 1
p
2
` = 1 (final) ` = 2 (final)
0.00 0.0
0.0 102 0.5 101.0
3
example
Width nd
Figure 5: Verification of Assumption 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/ 2, which is what one expects from an approximately-mean-zero random variable passed
through a ReLU nonlinearity.)
(i)
Assumption 3. ||∆Wℓ hℓ (xi )||2 = Θ(|| B1 ∆Wℓ hℓ (xi )||2 ).
6 It will also occur if the activation function is ϕ(z) = 0, but we can neglect this edge case.
21
` = 2 (init)
` = 2 (final)
0.2
(i)
0.1 0.5
0.0
0.0 102 0.5 101.0
3
example
Width nd
(i)
Figure 6: Verification of Assumption 3. We compute ∥∆Wℓ hℓ (xi )|2 / ||∆Wℓ hℓ (xi )||2 at layer ℓ = 2 for a
single step both at initialization and after training. As expected, this ratio remains Θ(1) as width
√ grows —
note the lack of decay with width. We note as a curiosity that this quantity hovers around 1/ B (dashed
line), which one expects from the central limit theorem if, for example, ∆Wℓ resembles a sum of B terms
(i)
similar to ∆Wℓ but with random signs.
7 Actually, a vector of d iid samples from a unbounded subgaussian or subexponential distribution will have its extreme
elements scale like polylog(d), so the really correct thing to say here is Θ̃(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.
22
E.2 Defining natural norms
Definition 2 (Natural ℓ2 -norm). The natural ℓ2 -norm for a dense vector v ∈ Rm is the RMS norm
1
||v||2̃ := √ ||v||2 . (28)
m
The natural ℓ2 -norm for a sparse vector v is simply the usual ℓ2 -norm:
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 ℓ
forms a “dense” hidden space.
Definition 3 (Natural spectral norm). Given a parameter matrix A, equip its input and output spaces with
their natural norms. Then the natural spectral norm ||A||∗˜ of 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 A ∈ Rm×n is
√
m
||A||∗˜ := √ ||A||∗ . (30)
n
Desideratum 2 (Feature learning, natural norms). Let hℓ (x) ∈ Rnℓ denote the features of input x
at layer ℓ of a deep neural network, and let ∆hℓ (x) ∈ Rnℓ denote their change after a gradient update.
We desire that:
||hℓ ||2̃ = Θ(1) and ||∆hℓ ||2̃ = Θ(1), at layers ℓ = 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 ∆Wℓ ∈ Rnℓ ×nℓ−1
to the ℓth weight matrix Wℓ ∈ Rnℓ ×nℓ−1 . The spectral norms of these matrices should satisfy:
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 Θ(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 Condition 1.
23