Streaming Power Iteration for Muon: 5. Extensions
The theme of this series is "streaming power iteration," which, as the name suggests, combines "streaming" and "power iteration." Here, "power iteration" refers to a classic multi-step iterative scheme for computing matrix SVD, while "streaming" refers to amortizing an algorithm that would otherwise require many iterative steps across each training step, so that the computational cost becomes acceptable. The core idea is: rather than completing a complex computation all at once, it's better to keep approaching the target progressively over the course of training.
As an extension of this series, this post introduces a few more applications of the "streaming" idea, further demonstrating how a relatively expensive operation can be cleverly folded into the training process through streaming transformation.
Orthogonal Projection
In some scenarios, we may want to constrain certain parameter matrices to be orthogonal. Orthogonal matrices have good numerical stability, which can help avoid certain numerical blow-up or vanishing issues, and in some designs they also bring better theoretical guarantees. Of course, exactly where it's appropriate to constrain parameters to be orthogonal matrices needs to be analyzed case by case, which we won't go into here.
In Steepest Descent on Manifolds: 2. Muon + Orthogonal and Steepest Descent on Manifolds: 3. Muon + Stiefel, we explored the orthogonal (Stiefel) manifold to some extent, but the goal there was to combine steepest descent to derive new update rules, which ended up being fairly complex overall. Here we'll consider a much simpler approach: after every update step, re-project (retract) the parameters back onto the orthogonal manifold.
Without loss of generality, let the parameter be $\boldsymbol{W}\in\mathbb{R}^{n\times m}$, where $n \geq m$, so the operation we want to perform can be written as
\begin{equation}\newcommand{orth}{\mathop{\text{orth}}}\boldsymbol{W}_t = \boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t \qquad\to\qquad \boldsymbol{W}_t = \orth(\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t)\end{equation}
Here $\orth$ is defined as
\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}}\orth(\boldsymbol{W}) = \argmin_{\boldsymbol{O}^{\top}\boldsymbol{O}=\boldsymbol{I}} \Vert\boldsymbol{W} - \boldsymbol{O}\Vert_F\end{equation}
That is, finding the orthogonal matrix nearest to $\boldsymbol{W}$. This operation already appeared in the very first blog post in our Muon series, Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices — it's exactly the $\newcommand{msign}{\mathop{\text{msign}}}\msign$ inside Muon, so the operation we want to implement can also be written as
\begin{equation}\boldsymbol{W}_t = \msign(\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t)\label{eq:msign-u}\end{equation}
Streaming Iteration
In other words, after every update step, we need to run $\msign$ once more in order to project the parameters back onto the orthogonal manifold. However, the operation $\msign$, while not exactly expensive, isn't exactly cheap either. We know that the core operation in Muon is $\msign$, but Muon's $\msign$ runs in BF16, whereas the parameters are stored in FP32, so the parameters' $\msign$ needs to run in FP32, which is a fairly significant cost.
Efficient computation of $\msign$ is based on the Newton–Schulz iteration, which we discussed in detail in Newton–Schulz Iteration for the msign Operator (Part 1) and Newton–Schulz Iteration for the msign Operator (Part 2). Different Newton–Schulz iterations have different polynomial degrees and coefficients, corresponding to different convergence rates; a classic 3rd-order scheme is
\begin{equation}\boldsymbol{X}_t = \frac{3}{2}\boldsymbol{X}_{t-1} - \frac{1}{2}\boldsymbol{X}_{t-1}\boldsymbol{X}_{t-1}^{\top}\boldsymbol{X}_{t-1},\qquad \boldsymbol{X}_0 = \boldsymbol{W}\end{equation}
One can show that $\lim_{t\to\infty} \boldsymbol{X}_t = \msign(\boldsymbol{W})$. This iteration, while classic, converges slowly; in Muon we usually use a faster-converging 5th-order iteration. But regardless of which one is used, running the full iteration is a fairly significant cost.
But do we really need to run the full iteration at every single step on the parameter side? Suppose $\boldsymbol{W}_{t-1}$ is already orthogonal or nearly so; since the learning rate $\eta$ is small, $\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t$ won't deviate much from orthogonality either — and this is exactly where the streaming idea comes into play. Perhaps a single iteration per training step is enough, for example consider
\begin{equation}\boldsymbol{W}_t = \frac{3}{2}\tilde{\boldsymbol{W}}_t - \frac{1}{2}\tilde{\boldsymbol{W}}_t\tilde{\boldsymbol{W}}_t^{\top}\tilde{\boldsymbol{W}}_t,\qquad \tilde{\boldsymbol{W}}_t = \boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t\end{equation}
Asymptotically, this achieves the same effect as $\eqref{eq:msign-u}$, and each step only requires computing one extra term $\tilde{\boldsymbol{W}}_t\tilde{\boldsymbol{W}}_t^{\top}\tilde{\boldsymbol{W}}_t$, which is clearly cheaper than a full Newton–Schulz iteration.
Spectral Constraints
In general, orthogonality constraints can only be applied to certain special matrices, because they're too strict — they cut the matrix's degrees of freedom in half, which is effectively the same as cutting the parameter count in half. Often we may prefer a looser spectral constraint instead, such as the singular-value clipping mentioned in Higher-Order MuP: A Simpler yet Smarter Spectral-Condition Scaling.
If $\msign$ turns all of a matrix's singular values into 1, then singular-value clipping only turns singular values greater than 1 into 1, leaving the rest unchanged; in Computing Singular-Value Clipping mclip via msign (Part 1) and Computing Singular-Value Clipping mclip via msign (Part 2) we called this operator $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$. Suppose what we want to do is clip all singular values to no more than 1 after every parameter update step; this can be written as
\begin{equation}\boldsymbol{W}_t = \mclip(\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t)\label{eq:mclip-u}\end{equation}
However, computing $\mclip$ is considerably more troublesome than computing $\msign$. We previously discussed a scheme for implementing $\mclip$ based on $\msign$, such as the identity
\begin{equation}\mclip(\boldsymbol{M}) = \frac{1}{2}\Big[\boldsymbol{M} + \msign(\boldsymbol{M}) + (\msign(\boldsymbol{M}) - \boldsymbol{M}) \msign(\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I})\Big]\end{equation}
which requires two applications of $\msign$ to compute $\mclip$ — quite costly. Of course, we could alternatively run a separate streaming power iteration on $\boldsymbol{W}$ and then compute $\mclip$ via the SVD, but that requires introducing yet another cached variable, which also feels like a lot of extra trouble.
Clipping One at a Time
Let's take another angle on understanding $\mclip$: $\mclip$ turns every singular value greater than 1 into 1, so one necessary operation is to turn the leading (largest) singular value into 1 (if it's greater than 1). Once we've clipped the leading singular value, if there's still a singular value greater than 1, then the largest of the remaining ones becomes the new leading singular value. This means that by repeatedly "clipping the leading singular value to 1," we can eventually achieve $\mclip$.
The advantage of this "clip one at a time" strategy is that computing only the leading singular value and leading singular vector (let's call it $\mathop{\text{SVD1}}$) is far cheaper than performing a full SVD — it only requires vector-form power iteration
\begin{equation}\boldsymbol{v}\quad\leftarrow\quad \frac{\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}}{\Vert\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}\Vert}\end{equation}
to converge to the right leading singular vector $\boldsymbol{v}_1$ of $\boldsymbol{W}$, and then $\sigma_1 = \Vert\boldsymbol{W}\boldsymbol{v}_1\Vert$ and $\boldsymbol{u}_1 = \boldsymbol{W}\boldsymbol{v}_1/\sigma_1$ follow; in practice we just fix the number of iteration steps and take an approximate value. Then, sticking with the "streaming" idea, at every update step we perform just one leading-singular-value clipping, i.e.,
\begin{equation}\boldsymbol{W}_t = \tilde{\boldsymbol{W}}_t - \max(\sigma_1 - 1, 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\end{equation}
Under long-term training, this keeps the singular values of $\boldsymbol{W}$ controlled near 1 (due to the approximate nature of power iteration and the streaming approach, in practice they'll end up slightly larger than 1). This "streaming singular-value clipping" can be viewed as a variant of the "Spectral Weight Decay" introduced in From Spectral-Norm Gradients to New Forms of Weight Decay; in the paper Training Transformers with Enforced Lipschitz Constants, it's called the "Spectral Hammer."
Other Examples
When is the streaming idea applicable? The typical scenario is: we can anticipate that, under long-term training, some variable changes slowly, so we can try running only a small number of iterations per training step, hoping that's enough to correct for the change brought about by the parameter update. Besides the two new examples above, we've already used the "streaming" idea in some earlier posts, which we'll briefly revisit here.
In Steepest Descent on Manifolds: 3. Muon + Stiefel and Steepest Descent on Manifolds: 4. Muon + Spectral Sphere, when solving for steepest descent on the corresponding manifolds, we needed to solve a nonlinear equation; at the time we used fixed-point iteration to solve it. Later, in Steepest Descent on Manifolds: 5. Dual Gradient Descent, we discussed the dual gradient descent solution and proposed amortizing the solving process across each training step via the streaming idea.
Coincidentally, in A Tour of MoE: 6. Optimal Transport for Balanced Routing, we approached the MoE load balancing problem from the perspective of optimal assignment. When solving the dual problem for the optimal assignment, we originally needed to alternately optimize $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ to convergence at every step, but considering that $\boldsymbol{\beta}$ shouldn't change much from one step to the next, we again relied on the streaming idea: performing just one update step each for $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ per training step turned out to achieve good load-balancing results.
In short, once we realize that a variable changes slowly, we can consider running only a small number of iteration steps at each step, in order to amortize the computational cost. To achieve this, sometimes we need to introduce an extra cached variable (such as $\boldsymbol{V}$ in streaming power iteration), but sometimes we don't. Designing good streaming iterations sometimes calls for a bit of cleverness, and it's a topic well worth digging into and savoring in detail.
Summary
This post has mainly introduced two more examples of applying the "streaming" idea: projecting parameters onto the orthogonal (Stiefel) manifold, and clipping the spectral norm of parameters to no more than 1, both at very little extra cost. These examples once again illustrate the elegance of the streaming idea: many computations that seem to require being done all at once in a complex fashion can in fact be decomposed into gradual, incremental adjustments and woven into every single step of training.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.