Beyond MuP: 4. Holding the Line on Parameter Stability
Through the derivations and calculations of the previous posts, we've seen that the three stability metrics proposed in the first post, Beyond MuP: 1. Three Characteristics of a Good Model, generally split into two parts: "parameter stability" and "increment stability." In Beyond MuP: 2. Linear Layers and Steepest Descent and Beyond MuP: 3. Special Cases, Special Treatment, we demonstrated how combining increment stability with steepest descent yields new update rules (optimizers).
However, for parameter stability, we've so far only dealt with initialization. The task of this post is to explore how to maintain parameter stability throughout the entire training process, completing the practical side of the theory.
Background
Take Beyond MuP: 2. Linear Layers and Steepest Descent as an example: the three stability metrics are
\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 $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ is the parameter of the linear layer. We want all three metrics to be $\Theta(1)$, which means we want the parameter and its increment to satisfy $\Vert\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$ and $\Vert\Delta\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$ respectively. In Beyond MuP: 3. Special Cases, Special Treatment we did the same computation for layers like the Embedding and LM Head, arriving at similar conclusions, just with different corresponding norms.
We use the increment condition as a stability metric and, following the "stable yet fast" principle of steepest descent, derive the theoretically optimal update rule — for instance the Muon optimizer for linear layers:
\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}}\newcommand{tr}{\mathop{\text{tr}}}\newcommand{msign}{\mathop{\text{msign}}}\argmin_{\Vert\Delta\boldsymbol{W}\Vert_2\leq\eta\sqrt{\frac{d_{out}}{d_{in}}}} \tr(\boldsymbol{G}^{\top}\Delta\boldsymbol{W}) \qquad \Rightarrow \qquad \Delta\boldsymbol{W} = -\eta\sqrt{\frac{d_{out}}{d_{in}}}\msign(\boldsymbol{G})\end{equation}
As for the parameter stability part, we previously only required the initialization to satisfy $\Vert\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$, without knowing how to guarantee the same parameter stability throughout the entire training process.
A General Framework
How can we ensure that $\boldsymbol{W}$ keeps satisfying $\Vert\boldsymbol{W}\Vert_2 = \Theta(\sqrt{d_{out}/d_{in}})$? More generally, given a parameter $\boldsymbol{\omega}$ — which could be a vector, a matrix, or even a higher-order tensor — plus a norm $\Vert\cdot\Vert$ (typically induced by forward stability or dependency stability), and finally a target scale $\tau$, the question is: how do we keep $\boldsymbol{\omega}$ satisfying $\Vert\boldsymbol{\omega}\Vert=\Theta(\tau)$ throughout training?
Initial Thoughts
A naive idea is to directly enforce $\Vert\boldsymbol{\omega}\Vert=\tau$ ($\tau$ could also be replaced by a constant multiple of itself, though this doesn't affect the discussion). The simplest implementation is to rescale the norm back to $\tau$ via normalization after every optimization step (as in Hyperball and Nemotron-Flash). Another approach is to reparametrize the model directly using normalization, i.e., replacing $\boldsymbol{f}(\boldsymbol{x};\boldsymbol{\omega})$ with $\boldsymbol{f}(\boldsymbol{x};\tau\boldsymbol{\omega}/\Vert\boldsymbol{\omega}\Vert)$, which in theory achieves a similar effect.
A more refined approach is to adjust the update rule by combining it with the idea of steepest descent, as discussed in Steepest Descent on Manifolds: 1. SGD + Hypersphere, Steepest Descent on Manifolds: 4. Muon + Spectral Sphere, and the paper Controlled LLM Training on Spectral Sphere. Methodologically this is more elegant, but in practice it's more complicated, usually requiring the solution of a nonlinear equation to obtain the exact update.
But should we really be pinning some norm of a parameter to a fixed value so strictly? Intuitively, the norm of a parameter should be determined by the training process itself — at most we should set a prior range for it. Although some work has shown that, if configured properly, fixing the parameter norm to a preset value doesn't hurt performance, doing so still disrupts the original training dynamics, and may require considerable extra effort to understand and adapt to.
So the position this post takes is that we only need to guarantee $\Vert\boldsymbol{\omega}\Vert = \mathcal{O}(\tau)$ — specifically, we try to ensure $\Vert\boldsymbol{\omega}\Vert \leq \tau$ holds at every step. Whatever specific value it ends up taking, and whether it actually reaches $\Theta(\tau)$, is left entirely to the training algorithm itself, with no further intervention.
Post-hoc Clipping
The natural next question is: how do we achieve $\Vert\boldsymbol{\omega}\Vert \leq \tau$? More concretely, suppose the original update rule for $\boldsymbol{\omega}$ is
\begin{equation}\boldsymbol{\omega}_t = \boldsymbol{\omega}_{t-1} - \eta \boldsymbol{\phi}_t\label{eq:base-update}\end{equation}
How should we modify it so that $\boldsymbol{\omega}_t$ always satisfies $\Vert\boldsymbol{\omega}_t\Vert\leq\tau$? There are of course many possible methods — the normalization mentioned in the previous section is one such scheme. Given that, we'd like to pick out the scheme that has the least impact on the optimization process: given a parameter $\boldsymbol{\omega}$ and norm $\Vert\cdot\Vert$, we want, with the smallest possible modification, to bring its norm down to no more than $\tau$, formally defined as
\begin{equation}\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq\tau}} = \argmin_{\Vert\tilde{\boldsymbol{\omega}}\Vert\leq\tau} \Vert \boldsymbol{\omega} - \tilde{\boldsymbol{\omega}}\Vert_{RMS}\label{eq:nclip}\end{equation}
Readers familiar with convex optimization will readily recognize this as the projection of $\boldsymbol{\omega}$ onto a hypersphere of radius no greater than $\tau$ under some norm. The key point here is that we want to achieve the goal of a norm no larger than $\tau$, while minimizing the impact on the original parameter $\boldsymbol{\omega}$ — so we minimize the discrepancy measure $\Vert \boldsymbol{\omega} - \tilde{\boldsymbol{\omega}}\Vert_{RMS}$, which induces a specific projection, or clipping, operation.
As for how to compute $\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq\tau}}$, that depends on the specific norm — we'll get to that shortly. With this operation in hand, one scheme we can consider is truncating the parameter norm after every update step, i.e., changing equation $\eqref{eq:base-update}$ to
\begin{equation}\boldsymbol{\omega}_t = \color{skyblue}{\lfloor}\boldsymbol{\omega}_{t-1} - \eta \boldsymbol{\phi}_t\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq\tau}}\end{equation}
Let's call this scheme "Post Clip" — simple and intuitive, but it may feel a bit "non-smooth." This is easy to understand: suppose we initialize with a radius smaller than $\tau$, and as training begins, the parameter radius slowly increases; once it reaches $\tau$, the clipping "suddenly" kicks in. This process, while continuous, is not smooth — similar to a $\max(x,0)$ function.
Pre Decay
If this lack of smoothness bothers us, we can mimic weight decay and spread the penalty across every update step. Starting again from the update rule $\eqref{eq:base-update}$, suppose $\boldsymbol{\phi}_t$ satisfies $\Vert\boldsymbol{\phi}_t\Vert\leq\tau$. Then by the triangle inequality, $\Vert\boldsymbol{\omega}_t\Vert = \Vert\boldsymbol{\omega}_{t-1} - \eta \boldsymbol{\phi}_t\Vert\leq \Vert\boldsymbol{\omega}_{t-1}\Vert + \eta \tau$ — meaning that, in the worst case, the norm increases by $\eta\tau$ every step, which will "run out of control" if accumulated over the long run.
To prevent this, we can, right before $- \eta \boldsymbol{\phi}_t$, apply a bit of preprocessing to $\boldsymbol{\omega}_{t-1}$ to shrink its norm just enough to offset the growth introduced by the update. Following the spirit of weight decay, we can consider
\begin{equation}\boldsymbol{\omega}_t = \color{skyblue}{\lfloor}\boldsymbol{\omega}_{t-1}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq (1-\eta)\Vert\boldsymbol{\omega}_{t-1}\Vert}} - \eta \boldsymbol{\phi}_t\label{eq:pre-decay}\end{equation}
That is, we first shrink the norm of $\boldsymbol{\omega}_{t-1}$ to $1-\eta$ times its original value, then perform the update. This gives us
\begin{equation}\Vert\boldsymbol{\omega}_t\Vert \leq (1-\eta)\Vert\boldsymbol{\omega}_{t-1}\Vert + \eta \tau \leq \max(\Vert\boldsymbol{\omega}_{t-1}\Vert,\tau)\end{equation}
Propagating this forward, we get $\Vert\boldsymbol{\omega}_t\Vert \leq \max(\Vert\boldsymbol{\omega}_{t-1}\Vert,\tau) \leq \cdots \leq \max(\Vert\boldsymbol{\omega}_0\Vert,\tau)$ — meaning that as long as the initialization satisfies $\Vert\boldsymbol{\omega}_0\Vert\leq\tau$, the entire update chain automatically satisfies $\Vert\boldsymbol{\omega}_t\Vert\leq \tau$. This conclusion holds regardless of the specific norm chosen; it only relies on the triangle inequality of the norm. And the minimal-change operation that reduces the norm is precisely the clipping operator defined by equation $\eqref{eq:nclip}$, so using it to shrink the norm is the natural choice.
We'll call this scheme "Pre Decay." The difference from "Post Clip" is that the latter's threshold is static ($\tau$), so clipping isn't guaranteed to trigger; but the former's threshold is dynamic ($(1-\eta)\Vert\boldsymbol{\omega}_{t-1}\Vert$), and clipping is guaranteed to trigger at every step. This process is smoother, which is why we call it a "decay" rather than a "clip" — it is a generalization of weight decay.
Basic Results
So far we've established a general framework for constraining parameter norms, with two schemes — "Post Clip" and "Pre Decay" — whose core operation is the clipping operator $\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq\tau}}$ defined by $\eqref{eq:nclip}$. So far this has only been a formal definition; in practice we need to compute it for specific norms. Below are some basic results.
A Simple Example
Let's start with a simple example, choosing the norm $\Vert\cdot\Vert_{RMS}$, which for vectors is equivalent to the L2 norm, and for matrices is equivalent to the Frobenius norm. It's not hard to show that
\begin{equation}\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_{RMS}\leq\tau}} = \argmin_{\Vert\tilde{\boldsymbol{\omega}}\Vert_{RMS}\leq\tau} \Vert \boldsymbol{\omega} - \tilde{\boldsymbol{\omega}}\Vert_{RMS} = \min\left(1,\,\frac{\tau}{\Vert\boldsymbol{\omega}\Vert_{RMS}}\right)\boldsymbol{\omega}\end{equation}
We leave the proof to the reader (if you really can't figure it out, ask Kimi). In particular, substituting $\tau = (1 - \eta) \Vert\omega\Vert_{RMS}$ gives
\begin{equation}\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_{RMS}\leq (1 - \eta) \Vert\omega\Vert_{RMS}}} = \min\left(1,\,\frac{(1 - \eta) \Vert\omega\Vert_{RMS}}{\Vert\boldsymbol{\omega}\Vert_{RMS}}\right)\boldsymbol{\omega} = (1-\eta)\boldsymbol{\omega}\end{equation}
Then substituting this into equation $\eqref{eq:pre-decay}$ gives
\begin{equation}\boldsymbol{\omega}_t = (1-\eta)\boldsymbol{\omega}_{t-1} - \eta \boldsymbol{\phi}_t\end{equation}
It's easy to see that this is just ordinary weight decay. In other words, Pre Decay under the RMS norm is exactly the weight decay we commonly use — it is the Pre Decay scheme that produces the minimal change to the original parameter while keeping the RMS norm (equivalently, the L2 norm of a vector, or the Frobenius norm of a matrix) constrained.
Singular Value Clipping
Now we move on to the "main event" of this post — matrix parameters and Muon. Here we reuse the notation $\boldsymbol{W}$, and write Muon's original update rule as
\begin{equation}\boldsymbol{W}_t = \boldsymbol{W}_{t-1} - \eta\lambda\boldsymbol{\Phi}_t,\quad \boldsymbol{\Phi}_t=\frac{1}{\lambda}\sqrt{\frac{d_{out}}{d_{in}}}\msign(\boldsymbol{G}_t)\end{equation}
Let $\tau = \frac{1}{\lambda}\sqrt{\frac{d_{out}}{d_{in}}}$, so that $\Vert\boldsymbol{\Phi}_t\Vert_2=\tau$. The two schemes for making $\boldsymbol{W}_t$ satisfy $\Vert\boldsymbol{W}_t\Vert_2\leq\tau$ are:
\begin{align} \text{Post Clip:}\quad\boldsymbol{W}_t =&\, \color{skyblue}{\lfloor}\boldsymbol{W}_{t-1} - \eta\lambda\boldsymbol{\Phi}_t\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_2\leq\tau}} \\[5pt] \text{Pre Decay:}\quad\boldsymbol{W}_t =&\, \color{skyblue}{\lfloor}\boldsymbol{W}_{t-1}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_2\leq(1-\eta\lambda)\Vert\boldsymbol{W}_{t-1}\Vert_2}} - \eta\lambda\boldsymbol{\Phi}_t \\ \end{align}
The next task is to compute $\color{skyblue}{\lfloor}\boldsymbol{W}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_2\leq\tau}}$, which, by the equivalence between the RMS and Frobenius norms, equals
\begin{equation}\color{skyblue}{\lfloor}\boldsymbol{W}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_2\leq\tau}} = \argmin_{\Vert\tilde{\boldsymbol{W}}\Vert_2\leq\tau} \Vert\boldsymbol{W} - \tilde{\boldsymbol{W}}\Vert_F\label{eq:mclip-loss}\end{equation}
The optimal solution to this problem should be familiar to some readers — it's the "Singular Value Clipping (SVC)" we discussed in Higher-Order MuP: A Simpler yet Cleverer Way to Scale the Spectral Condition, referred to as $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$ in Computing Singular Value Clipping (mclip) via msign (Part 1) and Computing Singular Value Clipping (mclip) via msign (Part 2):
\begin{equation}\color{skyblue}{\lfloor}\boldsymbol{W}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert_2\leq\tau}} = \mclip(\boldsymbol{W};\tau) = \boldsymbol{U}\min(\boldsymbol{\Sigma},\tau)\boldsymbol{V}^{\top}\label{eq:2-to-mclip}\end{equation}
where $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ is the SVD of $\boldsymbol{W}$, and $\min(\boldsymbol{\Sigma},\tau)$ truncates the singular values to be no larger than $\tau$. We'll walk through the proof in the next section. With this notation, the two schemes can be written as
\begin{align} \text{Post Clip:}\quad\boldsymbol{W}_t =&\, \mclip(\boldsymbol{W}_{t-1} - \eta\lambda\boldsymbol{\Phi}_t;\tau) \\[5pt] \text{Pre Decay:}\quad\boldsymbol{W}_t =&\, \mclip(\boldsymbol{W}_{t-1};(1-\eta\lambda)\Vert\boldsymbol{W}_{t-1}\Vert_2) - \eta\lambda\boldsymbol{\Phi}_t \\ \end{align}
Derivation
In this section we prove the result $\eqref{eq:2-to-mclip}$. Let the SVD of $\boldsymbol{W}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, where $\boldsymbol{U}\in\mathbb{R}^{d_{in}\times d_{in}}$, $\boldsymbol{\Sigma}\in\mathbb{R}^{d_{in}\times d_{out}}$, $\boldsymbol{V}\in\mathbb{R}^{d_{out}\times d_{out}}$. Then
\begin{equation}\Vert\boldsymbol{W} - \tilde{\boldsymbol{W}}\Vert_F = \Vert\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top} - \tilde{\boldsymbol{W}}\Vert_F = \Vert\boldsymbol{U}(\boldsymbol{\Sigma} - \boldsymbol{U}^{\top}\tilde{\boldsymbol{W}}\boldsymbol{V})\boldsymbol{V}^{\top}\Vert_F = \Vert\boldsymbol{\Sigma} - \boldsymbol{U}^{\top}\tilde{\boldsymbol{W}}\boldsymbol{V}\Vert_F\end{equation}
The last equality holds because orthogonal matrices leave the Frobenius norm unchanged. Orthogonal matrices also leave the spectral norm unchanged, so letting $\tilde{\boldsymbol{\Sigma}}=\boldsymbol{U}^{\top}\tilde{\boldsymbol{W}}\boldsymbol{V}$, the objective $\eqref{eq:mclip-loss}$ can equivalently be simplified to
\begin{equation}\argmin_{\Vert\tilde{\boldsymbol{\Sigma}}\Vert_2\leq\tau} \Vert\boldsymbol{\Sigma} - \tilde{\boldsymbol{\Sigma}}\Vert_F\end{equation}
Note that here $\boldsymbol{\Sigma}$ is diagonal, with diagonal entries denoted $\sigma_1,\sigma_2,\cdots \geq 0$, but $\tilde{\boldsymbol{\Sigma}}$ is not yet fixed to be diagonal — in the proof we must allow it to be a general matrix. Written out in components, this gives
\begin{equation}\Vert\boldsymbol{\Sigma} - \tilde{\boldsymbol{\Sigma}}\Vert_F^2 = \sum_i \sigma_i^2 + \sum_{i,j} \tilde{\Sigma}_{i,j}^2 - 2\sum_i \sigma_i \tilde{\Sigma}_{i,i} \geq \sum_i \sigma_i^2 + \sum_i (\tilde{\Sigma}_{i,i}^2 - 2 \sigma_i \tilde{\Sigma}_{i,i})\end{equation}
Looking at each term individually, $\tilde{\Sigma}_{i,i}^2 - 2 \sigma_i \tilde{\Sigma}_{i,i}$ is simply a quadratic function of $\tilde{\Sigma}_{i,i}$, whose minimum is attained at $\tilde{\Sigma}_{i,i}=\sigma_i$. But we also have the constraint $\Vert\tilde{\boldsymbol{\Sigma}}\Vert_2\leq\tau$; since the spectral norm is at least as large as the absolute value of any matrix entry, we have at least the constraint $\tilde{\Sigma}_{i,i}\leq\tau$, under which the minimum of $\tilde{\Sigma}_{i,i}^2 - 2 \sigma_i \tilde{\Sigma}_{i,i}$ is attained at $\tilde{\Sigma}_{i,i}^* = \min(\sigma_i,\tau)$.
Considering that we want all the equalities to hold simultaneously, we obtain $\tilde{\Sigma}_{i,j}^*=0(i\neq j)$, which shows that $\tilde{\boldsymbol{\Sigma}}^*$ is also diagonal, and can be written compactly as $\tilde{\boldsymbol{\Sigma}}^*=\min(\boldsymbol{\Sigma},\tau)$, which in turn corresponds to $\tilde{\boldsymbol{W}}^*=\boldsymbol{U}\min(\boldsymbol{\Sigma},\tau)\boldsymbol{V}^{\top}$. This proves the result $\eqref{eq:2-to-mclip}$.
Clipping the Dominant Term
So how do we efficiently compute $\mclip$? Performing a full SVD at every training step is clearly too expensive. In Computing Singular Value Clipping (mclip) via msign (Part 1) and Computing Singular Value Clipping (mclip) via msign (Part 2), we actually explored this problem systematically; back then, the approach was to rely on $\msign$, but this requires 2–3 calls to $\msign$, which is fairly costly. For example, an identity discovered in the second of those posts is
\begin{equation}\mclip(\boldsymbol{W};\tau) =\frac{1}{2}\Bigl\{\boldsymbol{W}+\tau\msign(\boldsymbol{W})-(\tau\boldsymbol{I}-\boldsymbol{W}\msign(\boldsymbol{W})^{\top})\msign(\tau\msign(\boldsymbol{W})-\boldsymbol{W})\Bigr\}\end{equation}
which requires two calls to $\msign$. Since parameter computations are usually done in FP32, running $\msign$ twice is still fairly expensive, so this isn't especially practical.
Here we mainly consider the term-by-term clipping approach discussed in Streaming Power Iteration-Based Muon Implementation: 5. Extensions. Specifically, $\mclip$ turns every singular value greater than $\tau$ into $\tau$, so the necessary operation is to turn the dominant singular value into $\tau$ (if it exceeds $\tau$). After clipping the dominant singular value, if there are still singular values greater than $\tau$, then the largest of these becomes the new dominant singular value. So, by repeatedly "clipping the dominant singular value down to $\tau$," we can implement $\mclip$.
Since the dominant singular value and dominant singular vector can be efficiently computed via power iteration (denoted $\mathop{\text{SVD1}}$), clipping the dominant singular value can be considered efficient. Furthermore, if we assume that training is smooth enough, we can perform just one round of dominant-singular-value clipping per step and still approximately achieve the same effect. Based on this strategy, the two singular-value-constraint schemes can be written as
\begin{align} \text{Post Clip:}\quad\boldsymbol{W}_t =&\, \tilde{\boldsymbol{W}}_t - \max(\sigma_1 - \tau, 0) \boldsymbol{u}_1 \boldsymbol{v}_1^{\top},\quad\sigma_1, \boldsymbol{u}_1, \boldsymbol{v}_1 = \mathop{\text{SVD1}}(\tilde{\boldsymbol{W}}_t),\quad\tilde{\boldsymbol{W}}_t = \boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t \\[5pt] \text{Pre Decay:}\quad\boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \lambda\eta\sigma_1 \boldsymbol{u}_1 \boldsymbol{v}_1^{\top} - \eta \boldsymbol{\Phi}_t,\quad\sigma_1, \boldsymbol{u}_1, \boldsymbol{v}_1 = \mathop{\text{SVD1}}(\boldsymbol{W}_{t-1}) \end{align}
The "Pre Decay" version here is exactly the spectral weight decay introduced in From Spectral Norm Gradients to a New Kind of Weight Decay — more than a year later, we've arrived at the same result via a different route. As for the "Post Clip" version, @_arohan_ once mentioned it on X, calling it "Wion" at the time. In practice, since only one singular value is clipped per step, there may be some "overachieving" matrices whose spectral norm noticeably exceeds the set threshold — this is normal, and it will gradually come down during the LR decay phase.
Other Details
For more precise clipping, we can also use power iteration to compute the top-$k$ singular values and vectors simultaneously, clipping up to $k$ singular values per step. The cost is that the L2 normalization step in power iteration must be replaced with a QR decomposition, which itself has various acceleration tricks — see the streaming power iteration series for the relevant principles, e.g. Streaming Power Iteration-Based Muon Implementation: 1. First Encounter.
Besides the spectral norm of linear layer matrices, in Beyond MuP: 3. Special Cases, Special Treatment we encountered various other norms for other layers — for example, the Embedding and LM Head layers correspond respectively to the maximum row and column RMS, while the gamma parameter of RMS Norm layers corresponds to the maximum absolute value, also known as the infinity norm of a vector.
Fortunately, the clipping operators $\color{skyblue}{\lfloor}\boldsymbol{\omega}\color{skyblue}{\rfloor}_{\color{skyblue}{\Vert\cdot\Vert\leq\tau}}$ under these norms are all fairly easy to compute. For instance, the norm of the Embedding layer is the maximum row RMS, so the clipping operator simply clips the RMS of each row vector to no more than $\tau$; the LM Head is analogous, with rows swapped for columns; and the gamma parameter is even simpler — it's just dominant-term clipping $\mathop{\text{clip}}(\boldsymbol{\gamma};-\tau,\tau) = \max(\min(\boldsymbol{\gamma},\tau),-\tau)$.
These results are all intuitive, and their proofs are fairly simple, so we won't go through them here — we'll leave them as an exercise for the reader.
Why This Is Necessary
Some readers might wonder: does it really have to be this complicated? Can't we just use ordinary weight decay like in Training Deep Learning Models with Norm-Constrained LMOs? For example,
\begin{equation}\boldsymbol{W}_t = (1-\eta\lambda)\boldsymbol{W}_{t-1} - \eta\sqrt{\frac{d_{out}}{d_{in}}}\msign(\boldsymbol{G}_t)\label{eq:muon-wd}\end{equation}
This can also keep the spectral norm bounded within $\tau = \frac{1}{\lambda}\sqrt{\frac{d_{out}}{d_{in}}}$ — so why not use this simpler form?
The answer is: to avoid over-intervention. From the definition $\eqref{eq:nclip}$, we can see that our clipping operator is the operation that, subject to achieving the same effect, produces the smallest possible change to the original parameter. For the spectral norm, simply multiplying by $1-\eta\lambda$ can also bring the spectral norm of $\boldsymbol{W}_{t-1}$ down to no more than $(1-\eta\lambda)\Vert\boldsymbol{W}_{t-1}\Vert_2$, but since this differs from the minimal-change operator $\mclip$, it necessarily involves some degree of "over-intervention."
Over-intervention has two possible consequences: either we choose a small $\lambda$ to preserve performance, in which case $\tau$ is too large — meaning we can't actually guarantee the spectral norm stays within our expected range — or we choose a large $\lambda$ to guarantee control of the spectral norm, but this noticeably hurts performance. For example, if we want the spectral norm to stay within $5$ when $d_{in}=d_{out}$, then according to formula $\lambda=0.2$, for the Muon update in equation $\eqref{eq:muon-wd}$, a weight decay coefficient of 0.2 would be extremely large (typical values are around 0.01).
Note our repeated emphasis on "guarantee" — this is the crucial point. Suppose we use weight decay with a coefficient of 0.01; in theory the spectral norm could reach as high as 100, but in experiments on small models we might find it never even reaches 5. This is quite common. However, small models being safe doesn't mean large models are safe. As we've said before, large models are powerful enough to amplify even the tiniest bug — if the theoretical upper bound is 100, a small model may never get anywhere near it, but a large model genuinely might.
So it's really necessary to keep a reasonable theoretical bound on the key norms of the parameters — this is exactly the "stability" part of the "stable yet fast" principle. And the clipping operator defined by equation $\eqref{eq:nclip}$ is the "lightest" operation that guarantees such a bound — in other words, it's likely the operation with the least performance cost among all those that guarantee the same bound.
Summary
Building on the idea of minimal change, this post proposed a general framework for maintaining parameter stability throughout training, comprising two schemes: Post Clip and Pre Decay. Under the spectral norm, these further specialize into singular value clipping and spectral weight decay. These operations aim to keep the key norms of the parameters bounded while minimizing interference with training dynamics.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.