An Initial Look at MuP: Scaling Laws for Hyperparameter Transfer Across Model Sizes
As we all know, training a large LLM to completion even once is expensive, which makes it infeasible to repeatedly test hyperparameters directly on large LLMs. A natural idea is to carefully search for hyperparameters on a small model with the same architecture, and then, once the optimal combination has been found, transfer it directly to the large model. Simple as this idea sounds, realizing it is not trivial — it requires us to understand the scaling laws that relate common hyperparameters to model scale, and MuP is exactly one practical realization of this idea.
MuP, sometimes also written $\mu P$, stands for Maximal Update Parametrization, and comes from the paper Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. As LLM training has become more widespread, it has gradually become one of the de facto standards of scientific model training.
Gist of the Method
Before getting into the main topic, I have to complain a bit: the original MuP paper is written in a way that's far too obscure, and its conclusions aren't stated clearly either, which unnecessarily adds to the difficulty of understanding it. So in what follows I'll try, as best I can, to reproduce MuP's conclusions in a (hopefully) concise and clear manner. more
Let's start with the conclusion: MuP mainly studies how hyperparameters should transfer across model scale. There are a few key terms here:
1. Hyperparameters, which for now mainly means the learning rate;
2. Model scale, which for now mainly means model width;
3. The core idea here is "transfer".
Note that MuP does not study what the optimal hyperparameters are; it only studies how the optimal hyperparameters change as model scale changes. So we need to search for the optimal hyperparameter combination on some small model, and then transfer it to the large model — that's the use case and usage method of MuP.
The principle behind deriving MuP is to ensure that the model's forward pass, backward pass, loss increment, and feature changes do not vary noticeably as model scale changes:
1. Concretely, this is done by analyzing the order of magnitude at initialization, and then assuming that the resulting conclusion also holds for the subsequent optimization dynamics;
2. In plain terms: assume you get the initialization right, and everything afterward will automatically follow the correct trajectory (well begun is half done?);
3. Of course you could also tell a story about the law of large numbers or the central limit theorem to justify this assumption, but personally I don't think that's necessary.
Forward Pass
Let's start the discussion with the forward pass, since this is the relatively simple and mature part. First, consider the linear layer $\boldsymbol{Y}=\boldsymbol{X}\boldsymbol{W}$, where $\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}},\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$. We use RMS (Root Mean Square) as an indicator of the scale of a matrix, e.g.,
\begin{equation}\text{RMS}(\boldsymbol{W}) = \sqrt{\frac{1}{d_{in} d_{out}}\sum_{i=1}^{d_{in}} \sum_{j=1}^{d_{out}} W_{i,j}^2}\end{equation}
We know that, to make the RMS of $\boldsymbol{X}$ at initialization roughly equal to the RMS of $\boldsymbol{Y}$ (call this "stable"), $\boldsymbol{W}$ should use:
LeCun initialization: random initialization with "mean 0, variance $1/d_{in}$".
This is already one of the basic results of deep learning, so I won't re-derive it here; readers who aren't yet familiar with it can refer to earlier posts such as Understanding Parameter Initialization Strategies from a Geometric Perspective and A Brief Discussion of Initialization, Parametrization, and Normalization in Transformers.
Next, consider the nonlinear layer $\boldsymbol{Y}=\phi(\boldsymbol{X}\boldsymbol{W})$, where $\phi$ is an element-wise activation function. If we still want to maintain that the RMS of $\boldsymbol{X}$ is approximately equal to the RMS of $\boldsymbol{Y}$, the result changes slightly. For instance, for the $\text{relu}$ activation we obtain
Kaiming initialization: random initialization with "mean 0, variance $2/d_{in}$".
It's easy to see that Kaiming initialization differs from LeCun initialization only by a (scale-independent) constant factor of 2 in the variance, and one can show similar results hold for other activation functions. So we can draw the conclusion:
fan_in initialization: to ensure stability of the forward pass, one should use random initialization with "mean 0, variance proportional to $1/d_{in}$".
This conclusion can also be understood as "the effect of the activation function is independent of model scale," so if we only want to analyze scale effects, we can ignore the presence of (element-wise) activation functions, and directly obtain the scaling law $\propto 1/d_{in}$ from LeCun initialization.
Backward Pass
Now let's move on to analyzing the backward pass (gradients). Note that here we adopt the convention that a variable and its gradient share the same shape, so we can compute
\begin{align} \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}} =&\, \boldsymbol{X}^{\top}\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\otimes \phi'(\boldsymbol{X}\boldsymbol{W})\right) \\[5pt] \frac{\partial\mathcal{L}}{\partial \boldsymbol{X}} =&\, \left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\otimes \phi'(\boldsymbol{X}\boldsymbol{W})\right)\boldsymbol{W}^{\top} \end{align}
The first formula gives the gradient of the parameter within the current layer, and the second gives the gradient that propagates back to the previous layer; $\otimes$ denotes the Hadamard product, and $\phi'$ is the derivative of $\phi$.
Note the fact that the derivatives of the activation functions we commonly use can all be bounded by a (scale-independent) constant, so at least in terms of order of magnitude we can write
\begin{align} \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}} =&\, \boldsymbol{X}^{\top}\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\otimes \phi'(\boldsymbol{X}\boldsymbol{W})\right) \sim \boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}} \label{eq:grad-w}\\[5pt] \frac{\partial\mathcal{L}}{\partial \boldsymbol{X}} =&\, \left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\otimes \phi'(\boldsymbol{X}\boldsymbol{W})\right)\boldsymbol{W}^{\top}\sim \frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\boldsymbol{W}^{\top}\label{eq:grad-x} \end{align}
Let's first look at the second formula. Compared with $\boldsymbol{Y}=\boldsymbol{X}\boldsymbol{W}$, the matrix multiplied on the right becomes $\boldsymbol{W}^{\top}$, so according to the conclusion of the previous section, if we want to keep the RMS of the backward pass stable, then $\boldsymbol{W}$ should be initialized as:
fan_out initialization: random initialization with "mean 0, variance $1/d_{out}$".
When $d_{in}\neq d_{out}$, the requirements of the forward pass and the backward pass come into conflict, and at this point someone proposed a compromise:
Xavier initialization: random initialization with "mean 0, variance $2/(d_{in} + d_{out})$".
This is also called "fan_avg initialization," since it's just a simple average of $d_{in}$ and $d_{out}$; other averaging schemes are also worth considering, see Reflections on Dimension-Averaging Strategies for Non-Square Matrices in Initialization Methods. Xavier initialization appears to accommodate both forward and backward passes at once, but you could also say it accommodates neither perfectly; a better approach is to design the model so that most parameters are square matrices, as in the model family $\eqref{eq:model}$ discussed later.
Loss Increment
With the groundwork of forward and backward pass laid, we can now try to analyze the increment of the loss function. Consider the change in the loss function under $\boldsymbol{W}\to \boldsymbol{W} + \Delta\boldsymbol{W}$:
\begin{equation}\Delta \mathcal{L} = \mathcal{L}(\boldsymbol{W} + \Delta\boldsymbol{W}) - \mathcal{L}(\boldsymbol{W})\approx \left\langle\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}, \Delta\boldsymbol{W}\right\rangle_F\end{equation}
Here $\langle\cdot,\cdot\rangle_F$ denotes the Frobenius inner product, i.e., flattening the matrices into vectors and taking their inner product. Consider gradient descent $\Delta\boldsymbol{W} = -\eta \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}$, where $\eta$ is naturally the learning rate; combining this with equation $\eqref{eq:grad-w}$, we get
$$\begin{equation}\Delta \mathcal{L}\approx -\eta\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right\Vert_F^2\sim -\eta \left\Vert\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\right\Vert_F^2\end{equation}$$
In fact, this equation already tells us why the same learning rate $\eta$ can't be used across different model scales:
1. $\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ is a $d_{in}\times d_{out}$ matrix;
2. $\left\Vert\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\right\Vert_F^2$ is the sum of squares of $d_{in}\times d_{out}$ numbers;
3. $\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ is exactly the product of the forward and backward quantities;
4. if both the forward and backward passes are stable, then every element of $\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ is $\Theta(1)$ ($\Theta$ being "Big Theta Notation");
5. so $\left\Vert\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}\right\Vert_F^2$ is $\Theta(d_{in} d_{out})$.
Point 4 deserves a bit more elaboration. $\boldsymbol{X}^{\top}$ is a $d_{in}\times b$ matrix and $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ is a $b\times d_{out}$ matrix; multiplying the two amounts to taking the inner product of $d_{in} d_{out}$ pairs of $b$-dimensional vectors, and an inner product is a sum over $b$ terms. Meanwhile the loss $\mathcal{L}$ is usually an average over samples (i.e., includes dividing by $b$). So if $\boldsymbol{X}^{\top}$ and $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ are both scale-independent, then their product is also basically scale-independent [i.e., their RMS values are all $\Theta(1)$].
This final conclusion tells us that if we directly use a small model's learning rate for a large model, then for a sufficiently large model, the loss increment at each step will explode as the parameter scale (i.e., $d_{in} d_{out}$) grows. This means we can't replicate the convergence behavior of the small model, and might even fail to converge at all because the step size is too large.
At this point what many people might think of is scaling $\Delta\mathcal{L}$ with $\eta\propto 1/(d_{in} d_{out})$, and indeed this idea is already in line with MuP's approach. But in real scenarios, because of the aforementioned incompatibility between forward and backward requirements, point 4 — "if both the forward and backward passes are stable, then every element of $\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}}$ is $\Theta(1)$" — doesn't always hold, so the actual situation is somewhat more complicated.
Model Assumptions
Now let's consider a scenario closer to practice. Our task is to train a model of $\mathbb{R}^{d_{in}}\mapsto \mathbb{R}^{d_{out}}$, where $d_{in},d_{out}$ is determined by the data and cannot be changed. As stated at the outset, MuP aims to study how the optimal hyperparameters scale with model size, so any fixed, unchanging quantity is effectively a constant, or $\Theta(1)$; for example, an initialization variance of $1/d_{in}$ is equivalent to saying the initialization variance is $\Theta(1)$.
What we can change is the model's architecture, parameter count, etc., but MuP mainly considers the scaling law with respect to width, so let's fix the model architecture. The model family we mainly consider here is:
\begin{equation}\begin{gathered} \boldsymbol{Y}_{in} = \boldsymbol{X} \boldsymbol{W}_{in} \\[5pt] \boldsymbol{Y}_{out} = \text{NN}(\boldsymbol{Y}_{in},\boldsymbol{\Omega}) \\[5pt] \boldsymbol{Z} = \boldsymbol{Y}_{out} \boldsymbol{W}_{out} \end{gathered}\label{eq:model}\end{equation}
where:
1. $\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$ (including the batch size);
2. $\boldsymbol{W}_{in} \in \mathbb{R}^{d_{in}\times d}, \boldsymbol{W}_{out} \in \mathbb{R}^{d\times d_{out}}$;
3. $\text{NN}$ is an arbitrary $\mathbb{R}^d\mapsto \mathbb{R}^d$ neural network;
4. here $d$ is exactly what we usually call the hidden size;
5. we can arbitrarily increase $d$ to boost the model's parameter count and capacity;
6. MuP is precisely about studying how the hyperparameters vary with $d$.
To be more concrete, here the $\text{NN}$ we consider is an $K$-layer MLP:
\begin{equation}\begin{aligned} \boldsymbol{Y}_0 =&\, Y_{in} \\[5pt] \boldsymbol{Y}_{k+1} =&\, \phi(\boldsymbol{Y}_k \boldsymbol{W}_{k+1}) \\[5pt] \boldsymbol{Y}_{out} =&\, \boldsymbol{Y}_K \end{aligned}\end{equation}
Here $\boldsymbol{\Omega}=\{\boldsymbol{W}_1,\boldsymbol{W}_2,\cdots,\boldsymbol{W}_K\}$, $\boldsymbol{W}_k\in\mathbb{R}^{d\times d}$, i.e., both are $d\times d$ square matrices, and both use fan_in initialization (equivalently, also fan_out initialization).
To add a bit more: the convention here that all the parameter matrices are $d\times d$ square matrices is purely to simplify the analysis, not a strict requirement. Because the real purpose here is to assume that among the parameters of $\text{NN}$, there is no scale-independent shape — for instance, a shape like $d\times 64$ is not allowed, because $64$ is a constant — but a shape like $d\times 4d$ is allowed, because regardless of whether you use fan_in, fan_out, or fan_avg initialization, the variance is always proportional to $1/d$.
Putting It All Together
Once the specific model has been fixed, we can now assemble all the previous conclusions. The parameters to be updated are divided into three parts, $\boldsymbol{W}_{in},\boldsymbol{\Omega},\boldsymbol{W}_{out}$, and we compute their gradients respectively:
\begin{align} \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}} =&\, \boldsymbol{Y}_{out}^{\top}\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}} \\[6pt] \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k} =&\, \frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k} \cdot\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}_{out}} = \frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k} \cdot\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}\boldsymbol{W}_{out}^{\top}\right) \\[6pt] \frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}} =&\, \boldsymbol{X}^{\top} \frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}_{in}} = \boldsymbol{X}^{\top} \left(\frac{\partial\boldsymbol{Y}_{out}}{\partial \boldsymbol{Y}_{in}}\cdot\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}_{out}}\right) = \boldsymbol{X}^{\top} \left(\frac{\partial\boldsymbol{Y}_{out}}{\partial \boldsymbol{Y}_{in}}\cdot\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}\boldsymbol{W}_{out}^{\top}\right)\right) \\[6pt] \end{align}
The operation $\cdot$ here needs a bit of explanation: since $\boldsymbol{Y}_{in},\boldsymbol{Y}_{out}$ are each matrices, $\frac{\partial\boldsymbol{Y}_{out}}{\partial \boldsymbol{Y}_{in}}$ is in principle a fourth-order tensor, and the chain rule $\frac{\partial\boldsymbol{Y}_{out}}{\partial \boldsymbol{Y}_{in}}\cdot\frac{\partial\mathcal{L}}{\partial \boldsymbol{Y}_{out}}$ is really a multiplication of higher-order tensors. We won't expand on this in detail here, so we simply denote it by $\cdot$; the reader just needs to know that it's a general extension of matrix multiplication.
Now let's observe the pattern:
1. all three formulas contain $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}$;
2. the latter two both contain $\boldsymbol{W}_{out}^{\top}$;
3. $\boldsymbol{W}_k$ are all square matrices, and $\frac{\partial\boldsymbol{Y}_{out}}{\partial \boldsymbol{Y}_{in}}$ and $\frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k}$ are both stable [RMS is $\Theta(1)$];
4. if $\boldsymbol{W}_{in}$ also uses fan_in initialization, then $\boldsymbol{Y}_{out}$ is also stable;
5. for $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}\boldsymbol{W}_{out}^{\top}$ to be stable, the initialization variance must be $1/d_{out}$, but $d_{out}$ is scale-independent, essentially a constant.
Given this:
1. the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}$ is $\Theta(1)$, and $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}\right\Vert_F^2$ is the sum of squares of $d\times d_{out}$ numbers, so its magnitude is $\Theta(d\times d_{out})$ — and remembering that $d_{out}$ is a constant, this is effectively $\Theta(d)$. So, in order to get $\Delta\mathcal{L}$ for $\Theta(1)$, its learning rate must satisfy $\eta_{out}\propto 1/d$;
2. $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}\right\Vert_F^2$ is a sum over $d^2$ numbers, and both $\frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k}$ and $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}$ have RMS $\Theta(1)$; if we directly set the initialization variance of $\boldsymbol{W}_{out}$ to $\propto 1/d^2$, then the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}$ will be $\Theta(1/d)$, and summing the squares gives exactly $\Theta(1)$, so the learning rate doesn't need to change;
3. in this case the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}$ is also $\Theta(1/d)$, but $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}\right\Vert_F^2$ is only a sum of squares over $d_{in}\times d$ numbers, so the result is $\Theta(1/d)$. To get $\Delta\mathcal{L}$ for $\Theta(1)$, the learning rate instead needs to be scaled up by a factor of $d$ to cancel out this effect, i.e., $\eta_{in}\propto d$.
Feature Changes
The result above is not wrong, but on closer inspection we notice a problem with the derivation: points 2 and 3 above are both built on the setting "we directly set the initialization variance of $\boldsymbol{W}_{out}$ to $\propto 1/d^2$" — yet so far this setting has had no direct justification. If we don't further explain this, the derivation remains incomplete.
In fact, if we only look at the requirement $\Delta \mathcal{L}=\Theta(1)$, we indeed cannot rule out other possible choices. For instance, we could instead set the initialization variance of $\boldsymbol{W}_{out}$ to $\propto 1/d$; then the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}$ would be $\Theta(1/\sqrt{d})$, and summing squares gives $\Theta(d)$, so as long as the learning rate is $\eta\propto 1/d$, we could likewise achieve $\Delta \mathcal{L}=\Theta(1)$. So, to justify the necessity of "setting the initialization variance of $\boldsymbol{W}_{out}$ to $\propto 1/d^2$," we need to introduce a new condition.
The loss function $\mathcal{L}$ is a macroscopic, or external, indicator of the model — looking at its change alone is no longer sufficient to explain all the results, so we need to dig deeper into the model's internals. Concretely, we want the change in the output of each layer of the model (commonly referred to as the feature, or sometimes the activation) to also be scale-invariant. For example, for the linear layer $\boldsymbol{Y}_k = \boldsymbol{Y}_{k-1} \boldsymbol{W}_k$, the change in output caused by the parameter update $\boldsymbol{W}_k\to \boldsymbol{W}_k + \Delta \boldsymbol{W}_k$ is
\begin{equation}\Delta\boldsymbol{Y}_k = \boldsymbol{Y}_{k-1} (\boldsymbol{W}_k + \Delta \boldsymbol{W}_k) - \boldsymbol{Y}_{k-1} \boldsymbol{W}_k = \boldsymbol{Y}_{k-1} \Delta\boldsymbol{W}_k\end{equation}
Note that $\boldsymbol{Y}_{k-1}\in\mathbb{R}^{b\times d},\Delta\boldsymbol{W}_k\in\mathbb{R}^{d\times d}$, so $\boldsymbol{Y}_{k-1} \Delta\boldsymbol{W}_k$ is the inner product of $b\times d$ pairs of $d$-dimensional vectors. Note that here $\Delta\boldsymbol{W}_k$ is a carefully designed update quantity, which is unlikely to be independent of $\boldsymbol{Y}_{k-1}$ the way it is at initialization, so the "inner product of $d$-dimensional vector pairs" is more likely to be $\Theta(d)$ (a $d$-dimensional inner product has a sum of $d$ terms). Therefore if the RMS of $\Delta\boldsymbol{Y}_{k-1}$ is $\Theta(1)$, we can consider the RMS of $\Delta\boldsymbol{Y}_k$ to be $\Theta(d\times \text{RMS}(\Delta \boldsymbol{W}_k))$.
So, in order for the RMS of $\Delta\boldsymbol{Y}_k$ to be $\Theta(1)$, we obtain an additional requirement on $\Delta \boldsymbol{W}_k$:
\begin{equation}\text{RMS}(\Delta \boldsymbol{W}_k) = \Theta(1 / d)\label{eq:dw-rms}\end{equation}
Combining $\Delta \boldsymbol{W}_k = -\eta\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}$ and $\Delta\mathcal{L}=\Theta(1)$, we then arrive at the result "the initialization variance of $\boldsymbol{W}_{out}$ is set to $\propto 1/d^2$."
(Note: this section relies on a suggestion from @Chenyu Zheng — many thanks!)
The Adam Version
The above is the MuP analysis for SGD. For Adam, we typically use SignSGD as an approximation for the order-of-magnitude analysis:
1. $\Delta \boldsymbol{W} = -\eta \mathop{\text{sign}}\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right)$;
2. $\Delta \mathcal{L} \approx -\eta \left|\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right|_1$;
3. here $|\cdot|_1$ means taking the absolute value of each element and summing.
For more on the SignSGD approximation itself, readers can refer to articles such as How Should the Learning Rate Change as Batch Size Increases? and How Does Adam's Epsilon Affect the Scaling Law of the Learning Rate?; we won't go into detail here. In short, SignSGD is a commonly used approximation for analyzing Adam-related scaling laws.
Now we can mimic the SGD analysis:
1. the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}$ is $\Theta(1)$, and $\left|\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}\right|_1$ is a sum over $d\times d_{out}$ numbers, of magnitude $\Theta(d\times d_{out}) = \Theta(d)$, so its learning rate must satisfy $\eta_{out}\propto 1/d$ to cancel out the scale effect;
2. $\left|\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}\right|_1$ is a sum over $d^2$ numbers, and both $\frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k}$ and $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}$ have RMS $\Theta(1)$; setting the initial variance of $\boldsymbol{W}_{out}$ to $\propto 1/d^2$, the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}$ becomes $\Theta(1/d)$, and summing over $d^2$ numbers gives $\Theta(d)$, so the learning rate must transform as $\eta_k\propto 1/d$ to cancel out the scale effect;
3. in this case the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}$ is also $\Theta(1/d)$, but $\left|\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}\right|_1$ is just a sum over $d_{in}\times d$ numbers, so it's already $\Theta(1)$, and thus the learning rate doesn't need to change with scale.
(Note: readers can verify for themselves that equation $\eqref{eq:dw-rms}$ is satisfied.)
The Muon Version
Naturally, we can't skip an analysis of Muon. For Muon itself, we already gave a detailed introduction in Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices and Muon Sequel: Why Did We Choose to Try Muon?, so we won't repeat that here. Similar to how we use SignSGD for Adam, we use MSignSGD as an approximation for Muon:
1. $\Delta \boldsymbol{W} = -\eta \mathop{\text{msign}}\left(\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right)$;
2. $\Delta \mathcal{L} \approx -\eta \left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right\Vert_*$ (see the proof in Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices);
3. here $\Vert\cdot\Vert_*$ denotes the Nuclear norm, i.e., the sum of all the singular values of the matrix;
4. the Nuclear norm is not easy to compute, but the $F$ norm is, being equal to the square root of the sum of squares of all singular values;
5. we use the $F$ norm as an approximation to the Nuclear norm, hence $\Delta \mathcal{L} \approx -\eta \left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right\Vert_*\approx -\eta \left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}}\right\Vert_F$;
6. the $F$ norm is also equal to the square root of the sum of squares of all the matrix's elements.
Now we can begin the analysis:
1. the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}$ is $\Theta(1)$, so the magnitude of $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{out}}\right\Vert_*$ is $\Theta(\sqrt{d\times d_{out}}) = \Theta(\sqrt{d})$; to cancel the scale effect, its learning rate must satisfy $\eta_{out}\propto 1/\sqrt{d}$;
2. $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}\right\Vert_F$ is the square root of the sum of squares of $d^2$ numbers, and both $\frac{\partial \boldsymbol{Y}_{out}}{\partial \boldsymbol{W}_k}$ and $\frac{\partial\mathcal{L}}{\partial \boldsymbol{Z}}$ have RMS $\Theta(1)$; setting the initial variance of $\boldsymbol{W}_{out}$ to $\propto 1/d^2$, the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_k}$ becomes $\Theta(1/d)$, and after summing squares and taking the square root the result is $\Theta(1)$, so the learning rate doesn't need to change;
3. in this case the RMS of $\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}$ is also $\Theta(1/d)$, but $\left\Vert\frac{\partial\mathcal{L}}{\partial \boldsymbol{W}_{in}}\right\Vert_F$ is just the square root of a sum of squares over $d_{in}\times d$ numbers, so it's of order $\Theta(1/\sqrt{d})$, and the learning rate instead needs to be scaled up by a factor of $\sqrt{d}$ to cancel this effect, i.e., $\eta_{in}\propto \sqrt{d}$.
(Note: the conclusion for Muon here is correct, but it does not satisfy condition $\eqref{eq:dw-rms}$, because equation $\eqref{eq:dw-rms}$, strictly speaking, also depends on the assumption that the update quantity is element-wise, which Muon does not satisfy; so in fact it isn't directly applicable. We haven't gone into detailed discussion of this here, and instead directly reused the conclusion "the initialization variance of $\boldsymbol{W}_{out}$ is set to $\propto 1/d^2$," sidestepping equation $\eqref{eq:dw-rms}$.)
Summary of Conclusions
Putting all the conclusions above together:
$$\begin{array}{c|c|c|c|c|c|c} \hline & \boldsymbol{W}_{in}\text{variance} & \boldsymbol{W}_{in}\text{learning rate} & \boldsymbol{W}_k\text{variance} & \boldsymbol{W}_k\text{learning rate} & \boldsymbol{W}_{out}\text{variance} & \boldsymbol{W}_{out}\text{learning rate} \\ \hline \text{SGD} & 1/d_{in} & d & 1 / d & 1 & 1/d^2 & 1 / d\\ \text{Adam} & 1/d_{in} & 1 & 1 / d & 1 / d & 1/d^2 & 1 / d\\ \text{Muon} & 1/d_{in} & \sqrt{d} & 1 / d & 1 & 1/d^2 & 1 / \sqrt{d} \\ \hline \end{array}$$
Here $\boldsymbol{W}_k$ refers to all parameters except $\boldsymbol{W}_{in},\boldsymbol{W}_{out}$, and it should also be emphasized that these relationships are all "proportional to" rather than "equal to". Additionally, in practice things can be adjusted somewhat according to specific needs. For example, in practice when we use Muon, the optimization of $\boldsymbol{W}_{in}$ and $\boldsymbol{W}_{out}$ usually doesn't use Muon but rather Adam, which leads to two changes:
1. $\eta_{out}\propto 1/d$;
2. $\eta_{in}$ stays unchanged.
If we combine this with the Adjust LR approach proposed in Muon is Scalable for LLM Training, then the learning rate needs an extra factor of $\sqrt{\max(n, m)}$, where $n\times m$ is the shape of the parameter matrix. We've already assumed that the parameters in the $\text{NN}$ part are all scaled proportionally, so $\sqrt{\max(n, m)}\propto \sqrt{d}$. Therefore, to cancel out the scale effect introduced by Adjust LR, we need:
3. $\eta_k\propto 1/\sqrt{d}$.
Summary
This post has tried to introduce MuP (Maximal Update Parametrization) as concisely and clearly as possible — a body of work aimed at studying how hyperparameters should transfer across model scale. Based on MuP, we can search for the optimal hyperparameters (mainly learning rate and initialization here) at relatively low cost on a small model, and then transfer them to a large model, lowering the cost of training large models.
To be honest, the discussion and analysis here is still fairly preliminary — for instance, we haven't considered bias terms, haven't evaluated whether the conclusions generalize beyond MLP architectures, and haven't carefully considered the role of normalization and residual connections. Not considering bias terms is purely laziness on my part, and I'll leave it as an exercise for the reader; as for MuP under different architectures, the analysis is generally more involved, but thanks to the similarity between neural network architectures, the conclusions are largely the same, and we can use them without proof. Personally, I think the more critical points for improvement are the effects of normalization and residual connections — especially normalization, since it allows the forward pass to be stabilized without relying on special initializations, which brings much greater freedom and possibility.
Of course, all of that is left for future analysis.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.