Beyond MuP: 3. Special Cases Require Special Handling

Responsible for the Loss

The answer is, the LM Head needs to be “responsible” for the Loss.

Keep in mind that steepest descent serves training. From the inference perspective, the model takes several tokens to predict the next token; but from the training perspective, the true “model” is: input several tokens together with the next token to compute the Loss. In other words, both the data and the label are essentially inputs, and the real output is the Loss. For earlier layers, we can ignore the label and the Loss, but the LM Head, being the last layer adjacent to the Loss, must account for the influence of the label and the Loss.

Thus, the input to the LM Head becomes \boldsymbol{x} and the next token ID t, and the output becomes the cross-entropy loss, i.e., \ell(\boldsymbol{x},t;\boldsymbol{W}) = \log\sum_{i=1}^{|V|} e^{\langle \boldsymbol{x},\boldsymbol{w}_i\rangle} - \langle \boldsymbol{x},\boldsymbol{w}_t\rangle = \log\sum_{i=1}^{|V|} e^{\langle \boldsymbol{x},\boldsymbol{w}_i - \boldsymbol{w}_t\rangle} where \boldsymbol{w}_i\triangleq \boldsymbol{W}_{:, i} is the i-th column of \boldsymbol{W}. Because \ell is a complex nonlinear function of \boldsymbol{x},t,\boldsymbol{W}, its three criteria cannot be computed exactly; our goal is to obtain a bound that is as tight as possible.

Forward Stability

First, the relatively simple forward stability; a simple bound gives \begin{aligned} \ell(\boldsymbol{x},t;\boldsymbol{W}) = \log\sum_{i=1}^{|V|} e^{\langle \boldsymbol{x},\boldsymbol{w}_i - \boldsymbol{w}_t\rangle} \leq&\, \log \left(|V| \max_i e^{\langle \boldsymbol{x},\boldsymbol{w}_i - \boldsymbol{w}_t\rangle}\right) \\ =&\, \log |V| + \max_i \langle \boldsymbol{x},\boldsymbol{w}_i - \boldsymbol{w}_t\rangle \\ \leq &\, \log |V| + \max_i \Vert\boldsymbol{x}\Vert_2 \Vert\boldsymbol{w}_i - \boldsymbol{w}_t\Vert_2 \end{aligned} Hence \begin{aligned} \text{Forward Stability:}\quad\max_{t, \Vert\boldsymbol{x}\Vert_{\mathrm{RMS}}=1} \ell(\boldsymbol{x},t;\boldsymbol{W}) \leq&\, \log |V| + d\max_{i,t} \Vert\boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{\mathrm{RMS}} \\ \leq&\, \log |V| + 2d\max_i \Vert\boldsymbol{w}_i\Vert_{\mathrm{RMS}} \end{aligned} If we drop the constant \log|V|, the inequality becomes a lower bound, so this bound is asymptotically fairly tight. To make it \Theta(1), the initialization variance of the LM Head should be chosen as \Theta(1/d^2).

An Important Inequality

For the remaining two criteria, because they involve differences, the computation is a bit more involved. We first prove an inequality we will need: \left|\log\sum_{i=1}^n e^{a_i} - \log\sum_{i=1}^n e^{b_i}\right| \leq \max_i |a_i - b_i|\label{leq:lse-ab} The proof is not difficult but requires a small trick: let the right-hand side be M. By monotonicity of \log,\sum,\exp, \log\sum_{i=1}^n e^{a_i} = \log\sum_{i=1}^n e^{(a_i - b_i)+b_i} \leq \log\sum_{i=1}^n e^{M + b_i} = M + \log\sum_{i=1}^n e^{b_i} This proves \log\sum_{i=1}^n e^{a_i} - \log\sum_{i=1}^n e^{b_i} \leq M By symmetry, swapping a_i and b_i also holds, thus proving the original inequality.

Dependence Stability

Using inequality [leq:lse-ab] and Cauchy-Schwarz, \begin{aligned} \frac{|\ell(\boldsymbol{x}_1,t;\boldsymbol{W}) - \ell(\boldsymbol{x}_2,t;\boldsymbol{W})|}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}}} \leq&\, \frac{\max_i |\langle \boldsymbol{x}_1,\boldsymbol{w}_i - \boldsymbol{w}_t\rangle - \langle \boldsymbol{x}_2,\boldsymbol{w}_i - \boldsymbol{w}_t\rangle|}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}}} \\ =&\, \frac{\max_i |\langle \boldsymbol{x}_1 - \boldsymbol{x}_2,\boldsymbol{w}_i - \boldsymbol{w}_t\rangle|}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}}} \\ \leq&\, \frac{\max_i d\, \Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}} \Vert \boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{\mathrm{RMS}}}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}}} \\ =&\, d\max_i \Vert \boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{\mathrm{RMS}} \end{aligned} Therefore \begin{aligned} \text{Dependence Stability:}\quad\max_{\begin{gathered}t \\ \Vert\boldsymbol{x}_1\Vert_{\mathrm{RMS}}=1 \\ \Vert\boldsymbol{x}_2\Vert_{\mathrm{RMS}}=1\end{gathered}} |\ell(\boldsymbol{x}_1,t;\boldsymbol{W}) - \ell(\boldsymbol{x}_2,t;\boldsymbol{W})| \leq&\, d\max_{i,t} \Vert\boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{\mathrm{RMS}}\\ \leq&\, 2d\max_i \Vert\boldsymbol{w}_i\Vert_{\mathrm{RMS}} \end{aligned} The result coincides with forward stability. The calculation here is similar to the Embedding layer: since the label t is a discrete ID, we do not consider its backpropagation, so the denominator only involves the difference of \boldsymbol{x}_1 and \boldsymbol{x}_2.

Update Stability

Finally, the update stability, again using inequality [leq:lse-ab] and Cauchy-Schwarz: \begin{aligned} |\ell(\boldsymbol{x},t;\boldsymbol{W} + \Delta\boldsymbol{W}) - \ell(\boldsymbol{x},t;\boldsymbol{W})| \leq&\, \max_i |\langle \boldsymbol{x},\boldsymbol{w}_i + \Delta\boldsymbol{w}_i - \boldsymbol{w}_t - \Delta\boldsymbol{w}_t\rangle - \langle \boldsymbol{x},\boldsymbol{w}_i - \boldsymbol{w}_t\rangle| \\ =&\, \max_i |\langle \boldsymbol{x},\Delta\boldsymbol{w}_i - \Delta\boldsymbol{w}_t\rangle| \\ \leq &\, d \max_i \Vert\boldsymbol{x}\Vert_{\mathrm{RMS}} \Vert\Delta\boldsymbol{w}_i - \Delta\boldsymbol{w}_t\Vert_{\mathrm{RMS}} \end{aligned} Thus \begin{aligned} \text{Update Stability:}\quad\max_{t,\Vert\boldsymbol{x}\Vert_{\mathrm{RMS}}=1} |\ell(\boldsymbol{x},t;\boldsymbol{W} + \Delta\boldsymbol{W}) - \ell(\boldsymbol{x},t;\boldsymbol{W})| \leq&\, d \max_{i, t} \Vert\Delta\boldsymbol{w}_i - \Delta\boldsymbol{w}_t\Vert_{\mathrm{RMS}} \\ \leq&\, 2d\max_i \Vert\Delta\boldsymbol{w}_i\Vert_{\mathrm{RMS}} \end{aligned}

It is easy to see that the three stability criteria of the LM Head are essentially the same as those of the Embedding layer: they are all the maximum row/column norm of the parameter matrix or its update, which means the steepest descent for the LM Head is also Normalized SGD—the difference being that it is column-wise RMS Norm. Moreover, all three criteria for the LM Head carry a factor of d, so its initialization standard deviation and learning rate are scaled by \Theta(1/d), while for Embedding it is \Theta(1). This implies slightly different behavior when scaling across width.

Hadamard Product

As we know, RMS Norm is often followed by a multiplicative vector \boldsymbol{\gamma} (Hadamard product), i.e., (\boldsymbol{x} / \Vert\boldsymbol{x}\Vert_{\mathrm{RMS}})\odot\boldsymbol{\gamma}, to control the output scale. This parameter is not a matrix, so the preceding Muon or Normalized SGD does not directly apply.

One can follow the original definition to compute the three stability criteria for \boldsymbol{\gamma} step by step and then analyze its initialization and steepest descent. But there is a more clever way: notice that \boldsymbol{x}\odot\boldsymbol{\gamma}=\boldsymbol{x}\mathop{\mathrm{diag}}(\boldsymbol{\gamma}), i.e., the Hadamard product of \boldsymbol{x} and \boldsymbol{\gamma} equals the matrix product of \boldsymbol{x} and the diagonal matrix \mathop{\mathrm{diag}}(\boldsymbol{\gamma}). This turns it into a special linear layer with \boldsymbol{W}=\mathop{\mathrm{diag}}(\boldsymbol{\gamma}), allowing us to reuse the conclusions for linear layers.

According to the previous article, the initial spectral norm of \boldsymbol{W} should be \Theta(\sqrt{d_{\mathrm{out}}/d_{\mathrm{in}}}). Here \boldsymbol{W} is a square matrix, so this gives exactly \Theta(1). Since \boldsymbol{W} is a diagonal matrix, we can simply initialize it as the identity matrix to satisfy this requirement, which corresponds to initializing \boldsymbol{\gamma} to all ones.

As for the optimizer, let the gradient of \boldsymbol{\gamma} be \boldsymbol{g}; then the gradient of \boldsymbol{W} is \boldsymbol{G}=\mathop{\mathrm{diag}}(\boldsymbol{g}). We already know that the steepest descent for a linear layer is Muon, i.e., \Delta\boldsymbol{W}=-\eta\mathop{\mathrm{msign}}(\boldsymbol{G}). For a diagonal matrix, \mathop{\mathrm{msign}}(\boldsymbol{G})=\mathop{\mathrm{sign}}(\boldsymbol{G})=\mathop{\mathrm{diag}}(\mathop{\mathrm{sign}}(\boldsymbol{g})), so the steepest descent for the \boldsymbol{\gamma} parameter is SignSGD.

Linear Bias Term

Traditional linear layers often include a bias vector \boldsymbol{b}, i.e., the full linear operation is \boldsymbol{f}(\boldsymbol{x};\boldsymbol{W},\boldsymbol{b}) = \boldsymbol{x}\boldsymbol{W}+\boldsymbol{b}. However, open-source models in recent years have mostly removed the bias term, so it has little presence nowadays. For completeness, we still supplement its discussion.

After adding the bias vector, the three stability criteria become \begin{aligned} &\text{Forward Stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{\mathrm{RMS}}=1} \Vert \boldsymbol{x}\boldsymbol{W} + \boldsymbol{b}\Vert_{\mathrm{RMS}} \\[5pt] &\text{Dependence Stability:}\quad\max_{\Vert\boldsymbol{x}_1\Vert_{\mathrm{RMS}}=\Vert\boldsymbol{x}_2\Vert_{\mathrm{RMS}}=1} \frac{\Vert \boldsymbol{x}_1\boldsymbol{W} - \boldsymbol{x}_2\boldsymbol{W}\Vert_{\mathrm{RMS}}}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{\mathrm{RMS}}}\\[5pt] &\text{Update Stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{\mathrm{RMS}}=1} \Vert \boldsymbol{x} \Delta\boldsymbol{W} + \Delta\boldsymbol{b}\Vert_{\mathrm{RMS}} \end{aligned} The dependence stability is the same as without bias, so we only need to examine forward and update stability. For simplicity, we use the inequality \Vert \boldsymbol{x}\boldsymbol{W} + \boldsymbol{b}\Vert_{\mathrm{RMS}}\leq \Vert \boldsymbol{x}\boldsymbol{W}\Vert_{\mathrm{RMS}} + \Vert\boldsymbol{b}\Vert_{\mathrm{RMS}}. Assuming \boldsymbol{W} continues to use its original initialization, the \Vert \boldsymbol{x}\boldsymbol{W}\Vert_{\mathrm{RMS}} part already achieves \Theta(1); thus we only need \Vert\boldsymbol{b}\Vert_{\mathrm{RMS}}=\mathcal{O}(1). In practice, for simplicity, \boldsymbol{b} is usually initialized to zeros.

Similarly, \Vert \boldsymbol{x}\Delta\boldsymbol{W} + \Delta\boldsymbol{b}\Vert_{\mathrm{RMS}}\leq \Vert \boldsymbol{x}\Delta\boldsymbol{W}\Vert_{\mathrm{RMS}} + \Vert\Delta\boldsymbol{b}\Vert_{\mathrm{RMS}}. Let \Vert\Delta\boldsymbol{b}\Vert_{\mathrm{RMS}}=\mathcal{O}(1); then the \boldsymbol{b} parameter will perform steepest descent with \Vert\Delta\boldsymbol{b}\Vert_{\mathrm{RMS}} as the stability metric, which also results in Normalized SGD.

Attention Scaling

Using the forward stability criterion, we can also re-derive the scaling factor of the Attention mechanism. Let \boldsymbol{q}=\boldsymbol{x}\boldsymbol{W}_q,\boldsymbol{k}=\boldsymbol{x}\boldsymbol{W}_k. If \boldsymbol{W}_q,\boldsymbol{W}_k are treated as linear layers, we can assume they already achieve \Vert\boldsymbol{q}\Vert_{\mathrm{RMS}}=\Theta(1) and \Vert\boldsymbol{k}\Vert_{\mathrm{RMS}}=\Theta(1). Then by Cauchy-Schwarz, |\langle\boldsymbol{q},\boldsymbol{k}\rangle| \leq \Vert\boldsymbol{q}\Vert_2 \Vert\boldsymbol{k}\Vert_2 = d\Vert\boldsymbol{q}\Vert_{\mathrm{RMS}} \Vert\boldsymbol{k}\Vert_{\mathrm{RMS}} where d is the dimensionality of \boldsymbol{q},\boldsymbol{k}, i.e., the head dimension. Clearly the right-hand side is \Theta(d); to make it \Theta(1), we need to multiply \boldsymbol{q}\cdot\boldsymbol{k} by a scaling factor on the order of \Theta(1/d). This differs from the previously common 1/\sqrt{d} (refer to A Brief Discussion on Transformer Initialization, Parameterization and Normalization).

Which one is correct? Actually both are. 1/\sqrt{d} is the average result under random initialization, while \Theta(1/d) is the extreme value applicable throughout training. It does not mean we should simply change the scaling factor to 1/d, but rather that scaling inversely with d may yield better transferability; the two are compatible. For example, suppose at d=128 a scaling factor of 1/\sqrt{128} works well; when transferring to d=256, one may consider changing the factor to 1/2\sqrt{128} instead of 1/\sqrt{256}.

In fact, constrained by Flash Attention, the choice of head dimension is not large, typically 128 and at most 256, so in practice there is almost no parameter transfer across head dimensions. Therefore this result is more of theoretical interest.