Above MuP: 2. Linear Layers and Steepest Descent

MuP之上：2. 线性层与最速下降 苏剑林 · 2026-02-15 · Machine Learning matrix linear optimizer muon MuP

In the previous post, Above MuP: 1. Three Properties of a Good Model, we proposed three core metrics—forward stability, dependence stability, and update stability—and gave their corresponding mathematical definitions. We also proposed characterizing how good a model is by whether it satisfies $\Theta(1)$, which will serve as the theoretical cornerstone for our subsequent analysis and computation. In this post, we will combine these metrics with the idea of steepest descent to design a "fastest-while-stable" update rule tailored to each parameter.

\begin{align} &\text{forward stability:}\quad\max_{\boldsymbol{x}} \Vert \boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})\Vert_{RMS} = \Theta(1) \label{eq:c1} \\[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:c2} \\[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:c3} \end{align}

We take the linear layer as our first example. The result should be familiar to some readers—it's the Muon optimizer, which gradually rose to prominence last year. Of course, our goal is not to rediscover Muon, but to demonstrate the process of designing models and optimizers from first principles, which will provide a unified methodology for handling other parameters later on. more

Linear Transformations

For a linear layer, the input is a vector $\boldsymbol{x}\in\mathbb{R}^{d_{in}}$, the parameter is a matrix $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$, and the model is $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W})=\boldsymbol{x}\boldsymbol{W}$. Note that in none of the three metric definitions did we require $\boldsymbol{x}$ to be bounded, so for a plain linear layer, none of the three metrics is guaranteed to exist—for instance, $\max\limits_{\boldsymbol{x}}\Vert\boldsymbol{x}\boldsymbol{W}\Vert_{RMS}$ is generally infinite. To fix this, we simply need to supplement the model with some operation that makes the result bounded, for example:

\begin{align} \newcommand{Norm}{\mathop{\text{Norm}}} &\text{In Norm:}\quad \Norm(\boldsymbol{x})\boldsymbol{W} \\[5pt] &\text{Out Norm:}\quad \Norm(\boldsymbol{x}\boldsymbol{W}) \end{align}

where $\Norm(\boldsymbol{x}) = \boldsymbol{x} / \Vert\boldsymbol{x}\Vert_{RMS}$, and here we omit the gamma parameter of RMS Norm, assuming its effect is secondary. As we know, residual connections commonly come in two flavors, Pre Norm and Post Norm. Pre Norm obviously corresponds to In Norm, but it's worth pointing out here that Post Norm is actually also In Norm:

\begin{align} \newcommand{Norm}{\mathop{\text{Norm}}} &\text{Pre Norm:}\quad \boldsymbol{x}_{t+1} = \boldsymbol{x}_t + \boldsymbol{F}_t(\Norm(\boldsymbol{x}_t)) \\[5pt] &\text{Post Norm:} \quad \boldsymbol{x}_{t+1} = \Norm(\underbrace{\boldsymbol{x}_t + \boldsymbol{F}_t(\boldsymbol{x}_t)}_{\text{denote}\boldsymbol{y}_{t+1}}) \quad \Rightarrow\quad \boldsymbol{y}_{t+1} = \Norm(\boldsymbol{y}_t) + \boldsymbol{F}_t(\Norm(\boldsymbol{y}_t)) \end{align}

So Post Norm, compared to Pre Norm, is simply replacing $\boldsymbol{x}_t + \boldsymbol{F}_t(\Norm(\boldsymbol{x}_t))$ with $\Norm(\boldsymbol{x}_t) + \boldsymbol{F}_t(\Norm(\boldsymbol{x}_t))$; for $\boldsymbol{F}_t$, both are In Norm, and this post will use In Norm as our running example.

Compared to Out Norm, In Norm has another advantage: more room for speed-up, since $(\boldsymbol{x} / \Vert\boldsymbol{x}\Vert_{RMS})\boldsymbol{W}=\boldsymbol{x}\boldsymbol{W} / \Vert\boldsymbol{x}\Vert_{RMS}$—in principle $\boldsymbol{x}\boldsymbol{W}$ and $\Vert\boldsymbol{x}\Vert_{RMS}$ can be computed in parallel and only divided at the end, reducing latency. This idea shows up in works such as FlashNorm: fast normalization for LLMs, Block-level AI Operator Fusion, and Superoptimizing RMSNorm and Linear.

Initial Variance

Following the discussion in the previous section, let's agree to only consider linear layers equipped with In Norm. Then, from the definition of the spectral norm, we can compute the three metrics as:

\begin{align} &\text{forward stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{RMS}=1} \Vert \boldsymbol{x}\boldsymbol{W}\Vert_{RMS} = \sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{W}\Vert_2 \\[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}} = \sqrt{\frac{d_{in}}{d_{out}}}\Vert\boldsymbol{W}\Vert_2 \\[5pt] &\text{update stability:}\quad\max_{\Vert\boldsymbol{x}\Vert_{RMS}=1} \Vert \boldsymbol{x}(\boldsymbol{W} + \Delta\boldsymbol{W}) - \boldsymbol{x}\boldsymbol{W}\Vert_{RMS} = \sqrt{\frac{d_{in}}{d_{out}}}\Vert\Delta\boldsymbol{W}\Vert_2 \end{align}

where taking $\Vert\cdot\Vert_2$ of a matrix denotes its spectral norm. As we can see, all three metrics are variants of the spectral norm—or more precisely, the three metrics the author proposed are themselves generalizations built starting from the spectral norm.

The first two metrics are functions of $\boldsymbol{W}$, and for a linear layer they happen to coincide. If we want them to be $\Theta(1)$, then we need $\Vert\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$, which at the very least imposes a requirement on the initialization of $\boldsymbol{W}$. According to Fast Estimation of the Spectral Norm of a Random Matrix, a standard normal matrix of size $d_{in}\times d_{out}$ has a spectral norm of roughly $\sqrt{d_{in}} + \sqrt{d_{out}}$, so in order for the initialization to satisfy $\Vert\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$, the initial variance $\sigma^2$ should satisfy

\begin{equation}\sigma = \Theta\left(\sqrt{\frac{d_{out}}{d_{in}}}\frac{1}{\sqrt{d_{in}} + \sqrt{d_{out}}}\right)\end{equation}

Additionally, we could also consider maintaining the constraint $\Vert\boldsymbol{W}\Vert_2$ throughout training, which has inspired some works such as Steepest Descent on Manifolds: 4. Muon + Spectral Sphere and Controlled LLM Training on Spectral Sphere. We'll return to this topic in a later post.

Steepest Descent

Next, let's focus on the "update stability" metric $\sqrt{d_{in}/d_{out}}\Vert\Delta\boldsymbol{W}\Vert_2$, which is the spectral-norm variant of the parameter increment $\Delta\boldsymbol{W}$. As we know, the update quantity is determined by the optimizer, so this part offers guidance for designing the optimizer. Following the "fastest-while-stable" principle, now that we have "stable," when is it "fastest"?

This is precisely the question steepest descent aims to answer. We've discussed this before in Muon Sequel: Why Did We Choose to Try Muon?, Steepest Descent on Manifolds: 1. SGD + Hypersphere, and Steepest Descent on Manifolds: 2. Muon + Orthogonality, but for the sake of completeness of this series, let's go through it once more without complaint. Steepest descent refers to the update that decreases the loss the fastest under some constraint, formally defined as

\begin{equation}\min_{\Delta \boldsymbol{W}} \mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W}) \qquad \text{s.t.}\qquad \rho(\Delta\boldsymbol{W})\leq \eta\end{equation}

where $\mathcal{L}$ is the loss function and $\rho(\Delta\boldsymbol{W})$ is the stability metric of the increment $\Delta\boldsymbol{W}$, which we already have—namely $\sqrt{d_{in}/d_{out}}\Vert\Delta\boldsymbol{W}\Vert_2$. However, solving this problem directly is still too complex; we need to replace $\mathcal{L}(\boldsymbol{W} +\Delta\boldsymbol{W})$ with the first-order approximation $\mathcal{L}(\boldsymbol{W}

) + \langle \boldsymbol{G}, \Delta\boldsymbol{W}\rangle_F$ in order to make the problem tractable. The problem to be solved is then equivalent to

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\min_{\Delta \boldsymbol{W}} \tr(\boldsymbol{G}^{\top}\Delta\boldsymbol{W}) \qquad \text{s.t.}\qquad \Vert\Delta\boldsymbol{W}\Vert_2\leq\eta\sqrt{\frac{d_{out}}{d_{in}}}\end{equation}

where $\boldsymbol{G}=\nabla_{\boldsymbol{W}}\mathcal{L}(\boldsymbol{W})$ is the gradient of the loss function, and we have used the identity $\langle \boldsymbol{G}, \Delta\boldsymbol{W}\rangle_F=\tr(\boldsymbol{G}^{\top}\Delta\boldsymbol{W})$.

Solution Process

Continuing, let $\Delta\boldsymbol{W}=-\kappa \boldsymbol{\Phi}$, and rewrite the optimization objective as

\begin{equation}\max_{\kappa,\boldsymbol{\Phi}}\kappa\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad 0\leq \kappa \leq \eta\sqrt{\frac{d_{out}}{d_{in}}}, \quad\Vert\boldsymbol{\Phi}\Vert_2=1\end{equation}

Clearly, the optimization over $\kappa$ can be carried out on its own, with the maximum attained at $\kappa = \eta\sqrt{d_{out}/d_{in}}$, so we only need to solve

\begin{equation}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2=1\end{equation}

Next, let $\boldsymbol{G}$ have SVD $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top} = \sum\limits_{i=1}^r \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top}$, where $r$ is the rank of $\boldsymbol{G}$. We have

\begin{equation}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi})=\tr\left(\sum_{i=1}^r \sigma_i \boldsymbol{v}_i \boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\right) = \sum_{i=1}^r \sigma_i \boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i\end{equation}

By definition, when $\Vert\boldsymbol{\Phi}\Vert_2=1$ we have $\Vert\boldsymbol{\Phi}\boldsymbol{v}_i\Vert_2\leq \Vert\boldsymbol{v}_i\Vert_2=1$, hence $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i\leq 1$, so

\begin{equation}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi})\leq \sum_{i=1}^r \sigma_i = \Vert \boldsymbol{G}\Vert_*\end{equation}

where $\Vert\cdot\Vert_*$ is called the matrix's nuclear norm, and equality holds when all $\boldsymbol{u}_i^{\top}\boldsymbol{\Phi}\boldsymbol{v}_i$ equal 1, in which case

\begin{equation}\newcommand{msign}{\mathop{\text{msign}}}\boldsymbol{\Phi} = \sum_{i=1}^r \boldsymbol{u}_i \boldsymbol{v}_i^{\top} = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top} = \msign(\boldsymbol{G})\end{equation}

Summary of Results

To summarize briefly: starting from the three stability metrics, we've obtained at least two conclusions so far. First, the initialization variance $\sigma^2$ of the parameter $\boldsymbol{W}$ should satisfy

\begin{equation}\sigma = \Theta\left(\sqrt{\frac{d_{out}}{d_{in}}}\frac{1}{\sqrt{d_{in}} + \sqrt{d_{out}}}\right)\end{equation}

Second, its increment $\Delta\boldsymbol{W}$ should take the following form:

\begin{equation}\Delta\boldsymbol{W} = -\eta\sqrt{\frac{d_{out}}{d_{in}}}\msign(\boldsymbol{G})\end{equation}

This is exactly the MuP version of Muon (for the differences among various versions, see A Guide to the Muon Optimizer: Getting Started Quickly and Key Details; standard Muon replaces $\boldsymbol{G}$ with its momentum, which can be viewed as a smoother estimate of the gradient). There's also still work to be done regarding the constraint on $\boldsymbol{W}$, which we'll leave for a later post.

Since we've already given full introductions to MuP and Muon across several previous blog posts, neither of these two results is new here. So this post serves only as a first case study, demonstrating the plausibility of the metrics $\eqref{eq:c1},\eqref{eq:c2},\eqref{eq:c3}$, which will provide a unified formula for the stability metrics of the parameters and increments of any layer, thereby generalizing the conclusions of Muon.

An Open Question

Before moving on to generalization, there's one more question we need to answer: the derivation above is based on the In Norm design—does that mean we need to add In Norm to every single linear layer? Can Muon still be used without In Norm? To answer this, let's first borrow a passage from the previous post:

Here, $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ can be a single layer, a block composed of several layers, or even the entire model. In principle, the coarser the granularity, the looser—or one might say the more accurate—the resulting constraint, but solving for $\max$ also becomes harder. So it depends on our ability to compute $\max$.

Put simply, the more accurately we compute the stability metrics, the better, but approximation is allowed. So without In Norm, the extent to which Muon remains usable depends on the extent to which "$\Vert\boldsymbol{x}\Vert_{RMS}=\text{some const}$" holds. Take the FFN layer $\boldsymbol{y}=\phi(\boldsymbol{x}\boldsymbol{W}_{up})\boldsymbol{W}_{down}$ as an example: if we assume the activation function $\phi$ has Lipschitz constant 1, then it still holds that

\begin{equation}\Vert\boldsymbol{y}\Vert_{RMS} \leq \Vert\boldsymbol{x}\Vert_{RMS} \times\sqrt{\frac{d_{in}}{d_{mid}}}\Vert\boldsymbol{W}_{up}\Vert_2\times \sqrt{\frac{d_{mid}}{d_{out}}}\Vert\boldsymbol{W}_{down}\Vert_2\end{equation}

where $\boldsymbol{W}_{up}\in\mathbb{R}^{d_{in}\times d_{mid}},\boldsymbol{W}_{down}\in\mathbb{R}^{d_{mid}\times d_{out}}$. This means that even if we only apply RMS Norm to $\boldsymbol{x}$, the same stability metric still approximately holds for the second parameter $\boldsymbol{W}_{down}$, and hence Muon remains usable there too.

Similarly, even without any RMS Norm at all, as long as we still believe that "$\Vert\boldsymbol{x}\Vert_{RMS}=\text{some const}$" holds to some extent, we can still try using the Muon optimizer for the subsequent linear layers.

Conclusion

Starting from the three stability metrics introduced in the previous post, this article demonstrated the process of "reproducing" the results related to MuP and Muon for linear layers. Next, we will apply this same methodology to "tailor" initializations, optimizers, and more for parameters beyond the linear layer.