Above MuP: 3. Special Cases Deserve Special Treatment
After so many blog posts on the topic, most readers should by now be quite familiar with the Muon optimizer—even without knowing the theoretical details, you probably have the impression that it's "an optimizer tailor-made for matrix parameters." However, this description isn't entirely accurate—for instance, the Embedding layer at the input end and the LM Head at the output end both have matrix parameters, yet Muon is not well suited to them (see Muon Optimizer Guide: Quick Start and Key Details).
Why should they be "treated differently"? This post will continue using the three stability metrics introduced in the first installment to examine the initialization rules of different layer types and their corresponding steepest-descent directions, in order to answer this question.
Recap
In the first post, Above MuP: 1. Three Characteristics of a Good Model, we proposed three stability metrics:
\begin{align} &\text{forward stability:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:cc1} \\[5pt] &\text{dependency stability:}\quad\max_{\boldsymbol{x}_1,\boldsymbol{x}_2} \frac{\Vert \boldsymbol{f}(\boldsymbol{x}_1;\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x}_2;\boldsymbol{\omega})\Vert_{RMS}}{\Vert\boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{RMS}} = \Theta(1) \label{eq:cc2} \\[5pt] &\text{update stability:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega} + \Delta\boldsymbol{\omega}) - \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:cc3} \end{align}
All three metrics share a unified format: take the RMS over the output, then take the $\max$ over the input. Here $\boldsymbol{x}$ denotes the input, $\boldsymbol{\omega}$ denotes the parameters, and $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ can represent a single layer, a block, or even the whole model, depending on our ability to compute $\max$.
Since we haven't restricted the range of $\boldsymbol{x}$, the maximum doesn't always exist; sometimes we need to add extra operations to the model to make it exist, which in turn guides the model's design. For instance, in the previous post, Above MuP: 2. Linear Layers and Steepest Descent, in order to compute the stability metric for a linear layer, we appended an In Norm to it. Building on the idea of steepest descent, we also reproduced the derivation of the Muon optimizer.
Steepest descent is not a new concept—it answers the question "given a stability metric, what optimizer should be used?" The core contribution of the "Above MuP" series is instead to answer "what stability metric should be used?" It provides a formula for computing the stability metric that applies to any layer.
The Embedding Layer
Now let's consider the Embedding layer, arguably the simplest layer of all. Its input is an index $i$, and its output is the corresponding vector, i.e. $\boldsymbol{f}(i;\boldsymbol{E}) = \boldsymbol{E}_i$, where $\boldsymbol{E}$ is a $|V|\times d$ matrix and $\boldsymbol{E}_i \triangleq \boldsymbol{E}_{i,:}$ denotes the $i$-th row of $\boldsymbol{E}$. It's straightforward to compute
\begin{align} &\text{forward stability:}\quad\max_i \Vert\boldsymbol{E}_i\Vert_{RMS} = \Theta(1)\\[5pt] &\text{update stability:}\quad\max_i \Vert \Delta \boldsymbol{E}_i\Vert_{RMS} = \Theta(1) \label{eq:ec3} \end{align}
Note that there is no "dependency stability" here, because the input to the Embedding layer consists of discrete token IDs, and you can't subtract two IDs. Of course, you could force the computation by converting to one-hot vectors, but this doesn't yield any new signal, and from the perspective of backpropagation, we don't actually need to continue backpropagating into the token ID anyway—so there's no need to consider its stability either.
The remaining two stability results are the maximum row norm of $\boldsymbol{E}$ or $\Delta\boldsymbol{E}$ (times a factor of $1/\sqrt{d}$). Forward stability and dependency stability are used only to guide initialization here, telling us to initialize $\boldsymbol{E}$ with zero mean and variance $\Theta(1)$. As for update stability, equation $\eqref{eq:ec3}$ tells us that although both are matrices, the metric of "stability" appropriate for the Embedding layer should not be the spectral norm, but rather the maximum row norm—which means its steepest descent is not Muon.
To find the steepest descent for the Embedding layer, we need to solve the optimization problem
\begin{equation}\min_{\Delta \boldsymbol{E}} \langle\boldsymbol{G},\Delta\boldsymbol{E}\rangle \qquad \text{s.t.}\qquad \max_i \underbrace{\Vert\Delta\boldsymbol{E}_i\Vert_{RMS}}_{\Vert\Delta\boldsymbol{E}_i\Vert_2/\sqrt{d}}\leq\eta\end{equation}
This problem isn't hard to solve; we just need the Cauchy–Schwarz inequality:
\begin{equation}\langle\boldsymbol{G},\Delta\boldsymbol{E}\rangle = \sum_{i=1}^{|V|}\langle\boldsymbol{G}_i,\Delta\boldsymbol{E}_i\rangle \geq -\sum_{i=1}^{|V|}\Vert\boldsymbol{G}_i\Vert_2 \times \Vert\Delta\boldsymbol{E}_i\Vert_2 \geq -\eta\sqrt{d}\sum_{i=1}^{|V|}\Vert\boldsymbol{G}_i\Vert_2\end{equation}
with equality holding at $\Delta\boldsymbol{E}_i = - \eta\boldsymbol{G}_i / \Vert\boldsymbol{G}_i\Vert_{RMS}$. That is to say, the steepest descent appropriate for the Embedding layer is row-wise RMS Norm applied to the gradient (Normalized SGD).
The Output Head
Next let's look at the LM Head. On the surface, this also looks like a linear layer: the input is $\boldsymbol{x}\in\mathbb{R}^d$, the weight is $\boldsymbol{W}\in\mathbb{R}^{d\times |V|}$, the output is $\boldsymbol{x}\boldsymbol{W}\in\mathbb{R}^{|V|}$, and $\boldsymbol{x}$ typically also comes with an RMS Norm—in every respect it looks like a linear layer. So why doesn't Muon apply to it?
Answerable to the Loss
The answer is that the LM Head has to be "answerable" to the loss.
Note that steepest descent serves training. From the inference perspective, the model takes in some tokens to predict the next token; but from the training perspective, the "model" in the true sense is: taking in some tokens together with the next token to compute the loss. In other words, both the data and the label are actually inputs, and the real output is the loss. For earlier layers, we don't need to consider the label or the loss, but the LM Head, being the final layer that "borders" the loss, has no choice but to take the label and the loss into account.
So, the input to the LM Head becomes $\boldsymbol{x}$ together with the index of the next token, $t$, and the output becomes the cross-entropy loss, i.e.
\begin{equation}\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}\end{equation}
where $\boldsymbol{w}_i\triangleq \boldsymbol{W}_{:, i}$ is the $i$-th column of $\boldsymbol{W}$. Since $\ell$ is a complicated nonlinear function of $\boldsymbol{x},t,\boldsymbol{W}$, its three metrics cannot be computed exactly; our goal is to find a bound that is as tight as possible.
Forward Stability
First, the relatively simple case of forward stability. A simple bound gives
\begin{equation}\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}\end{equation}
hence
\begin{equation}\begin{aligned} \text{forward stability:}\quad\max_{t, \Vert\boldsymbol{x}\Vert_{RMS}=1} \ell(\boldsymbol{x},t;\boldsymbol{W}) \leq&\, \log |V| + d\max_{i,t} \Vert\boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{RMS} \\ \leq&\, \log |V| + 2d\max_i \Vert\boldsymbol{w}_i\Vert_{RMS} \end{aligned}\end{equation}
If we drop the constant $\log|V|$, this becomes a lower bound too, so this bound is fairly tight in the asymptotic sense. To make it equal to $\Theta(1)$, the initialization variance of the LM Head should be chosen as $\Theta(1/d^2)$.
A Key Inequality
For the remaining two metrics, since they involve taking differences, the computation is somewhat more involved. Let's first prove an inequality we'll need:
\begin{equation}\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}\end{equation}
The proof is not hard, but requires a small trick: denote the right-hand side by $M$; then by the monotonicity of $\log,\sum,\exp$ we readily get
\begin{equation}\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}\end{equation}
which proves that
\begin{equation}\log\sum_{i=1}^n e^{a_i} - \log\sum_{i=1}^n e^{b_i} \leq M\end{equation}
By symmetry, swapping $a_i$ and $b_i$ also holds, which proves the original inequality.
Dependency Stability
Using inequality $\eqref{leq:lse-ab}$ together with the Cauchy–Schwarz inequality, we get
\begin{equation}\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_{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_{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_{RMS}} \\ \leq&\, \frac{\max_i d\, \Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{RMS} \Vert \boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{RMS}}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{RMS}} \\ =&\, d\max_i \Vert \boldsymbol{w}_i - \boldsymbol{w}_t\Vert_{RMS} \end{aligned}\end{equation}
hence
\begin{equation}\begin{aligned} \text{dependency stability:}\quad\max_{\begin{gathered}t \\ \Vert\boldsymbol{x}_1\Vert_{RMS}=1 \\ \Vert\boldsymbol{x}_2\Vert_{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_{RMS}\\ \leq&\, 2d\max_i \Vert\boldsymbol{w}_i\Vert_{RMS} \end{aligned}\end{equation}
The result parallels forward stability exactly. The computation here is similar to that for the Embedding layer, because the label $t$ is a discrete ID, so we don't need to consider its backpropagation; the denominator only needs to account for the difference in $\boldsymbol{x}_1,\boldsymbol{x}_2$.
Update Stability
Finally, update stability—again using inequality $\eqref{leq:lse-ab}$ and the Cauchy–Schwarz inequality:
\begin{equation}\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_{RMS} \Vert\Delta\boldsymbol{w}_i - \Delta\boldsymbol{w}_t\Vert_{RMS} \end{aligned}\end{equation}
hence
\begin{equation}\begin{aligned} \text{update stability:}\quad\max_{t,\Vert\boldsymbol{x}\Vert_{RMS}=1} |\ell(\boldsymbol{x},t;\boldsymbol{W} + \Delta\boldsymbol{W}) - \ell(\boldsymbol{x}_2,t_2;\boldsymbol{W})| \leq&\, d \max_{i, t} \Vert\Delta\boldsymbol{w}_i - \Delta\boldsymbol{w}_t\Vert_{RMS} \\ \leq&\, 2d\max_i \Vert\Delta\boldsymbol{w}_i\Vert_{RMS} \end{aligned}\end{equation}
It's not hard to see that the three stability metrics for the LM Head are, in essence, the same as those for the Embedding layer: they're all about the maximum row/column norm of the parameter matrix or its increment. This means the steepest descent for the LM Head is also Normalized SGD, the only difference being that for the LM Head it's applied column-wise instead of row-wise. Moreover, all three metrics of the LM Head take the form of multiplying by $d$, so both its initialization standard deviation and its learning rate scale as $\Theta(1/d)$, whereas for the Embedding layer it's $\Theta(1)$—meaning they differ slightly when transferred across widths.
Other Modules
Besides linear layers, Embedding, and the LM Head, common Transformer models usually have a few other parameters or layers that deserve separate analysis. Let's go through them one by one.
Hadamard Products
As we know, after RMS Norm there is usually a multiplication by a $\boldsymbol{\gamma}$ vector (a Hadamard product), i.e. $(\boldsymbol{x} / \Vert\boldsymbol{x}\Vert_{RMS})\odot\boldsymbol{\gamma}$, in order to rescale the output. This parameter is not a matrix, so the earlier Muon or Normalized SGD results don't directly apply to it.
We could go through the original definitions step by step to compute the three stability metrics for $\boldsymbol{\gamma}$, and then analyze its initialization and steepest descent, but there's a slicker way here: note that $\newcommand{diag}{\mathop{\text{diag}}}\boldsymbol{x}\odot\boldsymbol{\gamma}=\boldsymbol{x}\diag(\boldsymbol{\gamma})$, i.e. multiplying $\boldsymbol{x}$ and $\boldsymbol{\gamma}$ via a Hadamard product is equivalent to matrix-multiplying $\boldsymbol{x}$ by the diagonal matrix $\diag(\boldsymbol{\gamma})$. This turns it into a special kind of linear layer, namely $\boldsymbol{W}=\diag(\boldsymbol{\gamma})$, and we can reuse the conclusions for linear layers.
According to the previous post, the initial spectral norm of $\boldsymbol{W}$ should be $\Theta(\sqrt{d_{out}/d_{in}})$; here $\boldsymbol{W}$ is a square matrix, so this is exactly $\Theta(1)$. Since $\boldsymbol{W}$ is itself a diagonal matrix, we can simply initialize $\boldsymbol{W}$ as the identity matrix to satisfy this requirement, which corresponds to initializing $\boldsymbol{\gamma}$ to all ones.
As for the optimizer: suppose the gradient of $\boldsymbol{\gamma}$ is $\boldsymbol{g}$, then the gradient of $\boldsymbol{W}$ is $\boldsymbol{G}=\diag(\boldsymbol{g})$. We know that the steepest descent for a linear layer is Muon, i.e. $\newcommand{msign}{\mathop{\text{msign}}}\Delta\boldsymbol{W}=-\eta\msign(\boldsymbol{G})$, and since for a diagonal matrix we have $\newcommand{sign}{\mathop{\text{sign}}}\msign(\boldsymbol{G})=\sign(\boldsymbol{G})=\diag(\sign(\boldsymbol{g}))$, the steepest descent for the $\boldsymbol{\gamma}$ parameters is SignSGD.
Linear Bias Terms
A conventional linear layer usually also has a bias vector $\boldsymbol{b}$, so the full linear operation is $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W},\boldsymbol{b}) = \boldsymbol{x}\boldsymbol{W}+\boldsymbol{b}$. However, in recent years, most open-source models have dropped the bias term, so it has become rather insignificant. Still, for completeness, let's discuss it here too.
After adding the bias vector, the three stability metrics become
\begin{align} &\text{forward stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{RMS}=1} \Vert \boldsymbol{x}\boldsymbol{W} + \boldsymbol{b}\Vert_{RMS} \\[5pt] &\text{dependency stability:}\quad\max_{\Vert\boldsymbol{x}_1\Vert_{RMS}=\Vert\boldsymbol{x}_2\Vert_{RMS}=1} \frac{\Vert \boldsymbol{x}_1\boldsymbol{W} - \boldsymbol{x}_2\boldsymbol{W}\Vert_{RMS}}{\Vert \boldsymbol{x}_1 - \boldsymbol{x}_2\Vert_{RMS}}\\[5pt] &\text{update stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{RMS}=1} \Vert \boldsymbol{x} \Delta\boldsymbol{W} + \Delta\boldsymbol{b}\Vert_{RMS} \end{align}
where dependency stability is the same as in the bias-free case, so we only need to look at forward stability and update stability. For simplicity, let's use inequality $\Vert \boldsymbol{x}\boldsymbol{W} + \boldsymbol{b}\Vert_{RMS}\leq \Vert \boldsymbol{x}\boldsymbol{W}\Vert_{RMS} + \Vert\boldsymbol{b}\Vert_{RMS}$. Suppose $\boldsymbol{W}$ retains its original initialization; then the $\Vert \boldsymbol{x}\boldsymbol{W}\Vert_{RMS}$ part already achieves $\Theta(1)$, so we only need $\Vert\boldsymbol{b}\Vert_{RMS}=\mathcal{O}(1)$—for simplicity, in practice $\boldsymbol{b}$ is usually initialized to all zeros.
Similarly, $\Vert \boldsymbol{x}\Delta\boldsymbol{W} + \Delta\boldsymbol{b}\Vert_{RMS}\leq \Vert \boldsymbol{x}\Delta\boldsymbol{W}\Vert_{RMS} + \Vert\Delta\boldsymbol{b}\Vert_{RMS}$; if we require $\Vert\Delta\boldsymbol{b}\Vert_{RMS}=\mathcal{O}(1)$, then the $\boldsymbol{b}$ parameters will undergo steepest descent with respect to the stability metric $\Vert\Delta\boldsymbol{b}\Vert_{RMS}$, and again the result is Normalized SGD.
Attention Scaling
Using the forward-stability metric, we can also re-derive the scaling factor in 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$ is treated as a linear layer, we can assume $\Vert\boldsymbol{q}\Vert_{RMS}=\Theta(1)$ and $\Vert\boldsymbol{k}\Vert_{RMS}=\Theta(1)$ have already been achieved. Then, by the Cauchy–Schwarz inequality,
\begin{equation}|\langle\boldsymbol{q},\boldsymbol{k}\rangle| \leq \Vert\boldsymbol{q}\Vert_2 \Vert\boldsymbol{k}\Vert_2 = d\Vert\boldsymbol{q}\Vert_{RMS} \Vert\boldsymbol{k}\Vert_{RMS} \end{equation}
where $d$ is the dimension of $\boldsymbol{q},\boldsymbol{k}$, i.e. the head dimension. Clearly, the above expression 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)$, which differs from the previous $1/\sqrt{d}$ (see A Brief Discussion of Initialization, Parameterization, and Normalization in Transformers).
So which one is correct? Actually, both are. $1/\sqrt{d}$ is the average result under random initialization, while $\Theta(1/d)$ is a limiting value that holds over the entire course of training. This doesn't mean we should directly change the scaling factor to $1/d$; rather, it means that a scaling factor inversely proportional to $d$ may provide better transferability, and the two are actually compatible. For example, suppose that at $d=128$, using $1/\sqrt{128}$ as the scaling factor works well; then when transferring to $d=256$, one might consider changing the scaling factor to $1/2\sqrt{128}$ rather than $1/\sqrt{256}$.
In practice, though, because of Flash Attention, the choice of head dimension is quite constrained—usually 128, and at most around 256—so in practice there's essentially never a need to transfer parameters across different head dimensions. This result is therefore more of theoretical interest.
Summary
Finally, let's summarize the main results of these two posts:
$$\begin{array}{|c|c|} \hline & \text{input} & \text{parameter} & \text{output} & \text{initial std} & \text{steepest descent} \\ \hline \text{Linear} & \boldsymbol{x} & \begin{aligned}\boldsymbol{W}\in&\,\mathbb{R}^{d_{in}\times d_{out}} \\ \boldsymbol{b}\in&\,\mathbb{R}^{d_{out}} \end{aligned} & \boldsymbol{x}\boldsymbol{W} + \boldsymbol{b} & \begin{aligned} \boldsymbol{W}:&\, \small{\sqrt{\frac{d_{out}}{d_{in}}}\frac{1}{\sqrt{d_{in}} + \sqrt{d_{out}}}} \\ \boldsymbol{b}:&\, 0\end{aligned}& \begin{aligned} \Delta\boldsymbol{W} =&\, \small{-\eta\sqrt{\frac{d_{out}}{d_{in}}}\msign(\boldsymbol{G})} \\ \Delta\boldsymbol{b} =&\, \small{-\eta \frac{\boldsymbol{g}}{\Vert\boldsymbol{g}\Vert_{RMS}}}\end{aligned} \\ \hline \text{Embedding} & i & \boldsymbol{E} \in\mathbb{R}^{|V|\times d} & \boldsymbol{E}_{i,:} & 1 & \Delta\boldsymbol{E}_{i,:} = -\eta \frac{\boldsymbol{G}_{i,:}}{\Vert\boldsymbol{G}_{i,:} \Vert_{RMS}} \\ \hline \text{LM Head} & \boldsymbol{x}, t & \boldsymbol{W} \in\mathbb{R}^{d \times |V|} & \log\sum\limits_{i=1}^{|V|} e^{\langle \boldsymbol{x},\boldsymbol{W}_{:,i} - \boldsymbol{W}_{:,t}\rangle} & \frac{1}{d} & \Delta\boldsymbol{W}_{:,i} = -\frac{\eta}{d} \frac{\boldsymbol{G}_{:,i}}{\Vert\boldsymbol{G}_{:,i}\Vert_{RMS}} \\ \hline \text{RMS Norm} & \boldsymbol{x} & \boldsymbol{\gamma} \in\mathbb{R}^d & \frac{\boldsymbol{x}}{\Vert\boldsymbol{x}\Vert_{RMS}}\odot\boldsymbol{\gamma} & 1 & \Delta\boldsymbol{\gamma} = -\eta \sign(\boldsymbol{g})\\ \hline \end{array}$$
Here, the steepest-descent directions for the Embedding and LM Head layers are row-wise/column-wise Normalized SGD respectively, consistent with works such as Scion; as for the scaling laws for variance and learning rate, they agree with the conclusions of MuP. In both of these posts, these results were derived from our proposed "three stability metrics," which shows that we have indeed found a unified form for the stability measure applicable to any layer.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.