Implementing Muon via Streaming Power Iteration: 1. First Encounter

基于流式幂迭代的Muon实现：1. 初识 苏剑林 · 2026-03-12 · Machine Learning iteration matrix SVD optimizer muon

The core operation of Muon is $\newcommand{msign}{\mathop{\text{msign}}}\msign$, and the current standard implementation is the Newton-Schulz iteration. It has to be said, this is indeed a very efficient and GPU-friendly algorithm, and Muon's popularity owes at least more than half its credit to this algorithm. However, this algorithm also gives an impression of being a "sole proprietor with no competitors", because it seems to be limited to computing $\msign$ — once we want to modify Muon in some way (for instance, swapping $\msign$ for the $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$ discussed here), the corresponding computation becomes cumbersome.

This post proposes a new implementation approach — approximating the SVD via Streaming Power Iteration. This isn't an entirely new idea; it has already appeared in some prior optimizer work, but here we extract it and present it as a standalone algorithm.

Recap

We won't go over the details of Muon again — readers can refer back to earlier posts such as Muon Optimizer Appreciation: The Essential Leap from Vectors to Matrices, Muon Sequel: Why Did We Choose to Try Muon?, and Muon Optimizer Guide: Getting Started Quickly and Key Details. Here we just give the formula directly:

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t [\msign(\boldsymbol{M}_t) + \lambda \boldsymbol{W}_{t-1}] \\ \end{aligned}\end{equation}more

where $\msign$ is

\begin{equation}\msign(\boldsymbol{M})=\boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}=\boldsymbol{U}_{[:, :r]}\boldsymbol{V}_{[:, :r]}^{\top}\end{equation}

Here $\boldsymbol{M}\in\mathbb{R}^{n\times m}$; without loss of generality we take $n\geq m$, and for simplicity, in most cases we assume $r=m$ (i.e., full rank), only discussing rank-deficient cases when it is really necessary.

Since SVD is fairly expensive, in most cases we use the Newton-Schulz iteration to compute $\msign$, which we have already discussed in detail in The Newton-Schulz Iteration for the msign Operator (Part 1) and The Newton-Schulz Iteration for the msign Operator (Part 2). Overall, the Newton-Schulz iteration is very clever and is the main contributor to Muon's success, but its extensibility is fairly weak.

To extend the range of applications of the Newton-Schulz iteration, I've previously done some work along these lines, such as Computing Singular Value Clipping mclip via msign (Part 1), Computing Singular Value Clipping mclip via msign (Part 2), Efficient Computation of Matrix Square Roots and Inverse Square Roots, and Efficient Computation of Matrix r-th Roots and Inverse r-th Roots, but overall there's still a fairly limited amount that can be done this way.

Clearly, the once-and-for-all approach is to compute the SVD directly, which is the direction we'll focus on next.

Power Iteration

In posts such as Lipschitz Constraints in Deep Learning: Generalization and Generative Models and From Spectral-Norm Gradients to Thoughts on a New Style of Weight Decay, we've already had a first encounter with Power Iteration, which we used to find the dominant eigenvector of $\boldsymbol{M}^{\top}\boldsymbol{M}$, or equivalently the dominant right singular vector of $\boldsymbol{M}$, with the iteration scheme:

\begin{equation}\boldsymbol{v}_1^{(t)} = \frac{\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_1^{(t-1)}}{\Vert\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_1^{(t-1)}\Vert_2}\end{equation}

Suppose we've already found the dominant eigenvector $\boldsymbol{v}_1$; we can add orthogonalization to power iteration to find the second eigenvector:

\begin{equation}\boldsymbol{v}_2^{(t)} = \frac{\tilde{\boldsymbol{v}}_2^{(t)} - \langle\tilde{\boldsymbol{v}}_2^{(t)},\boldsymbol{v}_1\rangle\boldsymbol{v}_1}{\Vert\tilde{\boldsymbol{v}}_2^{(t)} - \langle\tilde{\boldsymbol{v}}_2^{(t)},\boldsymbol{v}_1\rangle\boldsymbol{v}_1\Vert_2},\qquad \tilde{\boldsymbol{v}}_2^{(t)} = \boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_2^{(t-1)}\end{equation}

Since orthogonality with $\boldsymbol{v}_1$ is enforced, this will converge to the second eigenvector $\boldsymbol{v}_2$. Similarly, given $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_{k-1}$, we can combine this with Gram-Schmidt orthogonalization to find the $k+1$-th eigenvector:

\begin{equation}\boldsymbol{v}_k^{(t)} = \frac{\tilde{\boldsymbol{v}}_k^{(t)} - \langle\tilde{\boldsymbol{v}}_k^{(t)},\boldsymbol{v}_1\rangle\boldsymbol{v}_1 - \cdots - \langle\tilde{\boldsymbol{v}}_k^{(t)},\boldsymbol{v}_{k-1}\rangle\boldsymbol{v}_{k-1}}{\Vert\tilde{\boldsymbol{v}}_k^{(t)} - \langle\tilde{\boldsymbol{v}}_k^{(t)},\boldsymbol{v}_1\rangle\boldsymbol{v}_1 - \cdots - \langle\tilde{\boldsymbol{v}}_k^{(t)},\boldsymbol{v}_{k-1}\rangle\boldsymbol{v}_{k-1}\Vert_2},\qquad \tilde{\boldsymbol{v}}_k^{(t)} = \boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_k^{(t-1)}\label{eq:vk-pi}\end{equation}

In practice, we don't need to wait until all of $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_{k-1}$ have been computed before computing $\boldsymbol{v}_k$ — all of $\boldsymbol{V}=[\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_m]$ can be iterated in parallel. Specifically, starting from an existing approximation $\boldsymbol{V}_{t-1}$, we compute $\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{V}_{t-1}$ in bulk, then re-orthogonalize by columns (using QR decomposition), which yields a better approximation, which we denote as $\boldsymbol{V}_t$. Repeating this iteration, we eventually converge to our target $\boldsymbol{V}$:

\begin{equation}\newcommand{QR}{\mathop{\text{QR}}}\boldsymbol{V}_t = \QR(\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{V}_{t-1})\end{equation}

Here $\QR$ refers to the orthogonal matrix from the QR decomposition. Once we have $\boldsymbol{V}$, it's clear that $\newcommand{ColNorm}{\mathop{\text{ColNorm}}}\boldsymbol{U} = \ColNorm(\boldsymbol{M}\boldsymbol{V})$, where $\ColNorm$ denotes L2 normalizing each column (axis=0), together with $\newcommand{diag}{\mathop{\text{diag}}}\boldsymbol{\Sigma}=\diag(\boldsymbol{U}^{\top}\boldsymbol{M}\boldsymbol{V})$. This gives us an approximate computation scheme for the SVD based on power iteration and QR decomposition. Of course, when $n > m$, it only yields an incomplete decomposition, with $\boldsymbol{U}\in\mathbb{R}^{n\times m}$ and $\boldsymbol{\Sigma},\boldsymbol{V}\in\mathbb{R}^{m\times m}$, but this is already sufficient for our purposes.

Streaming Update

However, using power iteration to compute the SVD is extremely inefficient in practice, far slower than directly calling the SVD function built into a framework, so this isn't really practical. But considering that training itself is already a long-running iterative process, we can assume that $\boldsymbol{V}$ doesn't change much from one step to the next, so we can store the $\boldsymbol{V}$ from the previous step and use it as the initialization for the current step, then perform only a single power iteration per step:

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{V}_t =&\, \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}) \\[5pt] \boldsymbol{U}_t =&\, \ColNorm(\boldsymbol{M}_t\boldsymbol{V}_t) \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (\boldsymbol{U}_t\boldsymbol{V}_t^{\top} + \lambda \boldsymbol{W}_{t-1}) \\ \end{aligned}\label{eq:muon-qr}\end{equation}

where $\boldsymbol{V}_0=\boldsymbol{I}$. Empirically, Muon implemented via this kind of streaming power iteration does indeed produce an LM loss curve that nearly coincides with that of the Newton-Schulz version, which shows that it's indeed a viable approach. This is largely thanks to the momentum mechanism and small learning rates, which make the assumption that "$\boldsymbol{V}$ doesn't change much from step to step" approximately hold, thereby allowing the cost of power iteration to be "amortized" across steps.

Thanks to computing the SVD directly (even if only approximately), we can also apply various operations to the singular values and incorporate them into the optimizer, for example

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{V}_t =&\, \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}) \\[5pt] \boldsymbol{U}_t =&\, \ColNorm(\boldsymbol{M}_t\boldsymbol{V}_t) \\[5pt] \boldsymbol{\Sigma}_t =&\, \diag(\boldsymbol{U}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_t) \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (\boldsymbol{U}_t f(\boldsymbol{\Sigma}_t)\boldsymbol{V}_t^{\top} + \lambda \boldsymbol{W}_{t-1}) \\ \end{aligned}\end{equation}

This makes it much easier to implement $\mclip$, or Muon variants based on general Schatten norms. In short, once we have an explicit result for $\boldsymbol{U}_t,\boldsymbol{\Sigma}_t,\boldsymbol{V}_t$ (even an approximate one), it becomes easy to try small modifications, significantly boosting both extensibility and "playability".

Speeding Up the Decomposition

Now the pressure shifts to the QR decomposition. The most time-consuming step in $\eqref{eq:muon-qr}$ is the QR decomposition, and the standard implementation is Householder QR. Although this is already considerably faster than SVD, it's still slower than computing $\msign$ via Newton-Schulz iteration (a polynomial iteration that allows BF16 multiplication — practically a cheat code). So, to make this new approach more competitive, we still need to speed up the QR decomposition.

For a given matrix $\boldsymbol{A}\in\mathbb{R}^{n\times m}$ ($n\geq m$), QR decomposition seeks an orthogonal matrix $\boldsymbol{Q}\in\mathbb{R}^{n\times m}$ and upper-triangular matrix $\mathbb{R}\in\mathbb{R}^{m\times m}$ such that $\boldsymbol{A}=\boldsymbol{Q}\boldsymbol{R}$ (here the orthogonal matrix only needs to satisfy $\boldsymbol{Q}^{\top}\boldsymbol{Q}=\boldsymbol{I}$; more precisely it's what's called a Stiefel matrix). Notice that $\boldsymbol{A}^{\top}\boldsymbol{A}=\boldsymbol{R}^{\top}\boldsymbol{R}$ — that is, we only need to decompose $\boldsymbol{A}^{\top}\boldsymbol{A}$ into the product of a lower-triangular matrix and its transpose in order to obtain $\boldsymbol{R}$, and this is exactly what Cholesky decomposition does!

Cholesky decomposition is very efficient, so the first step can use it to obtain $\boldsymbol{R}$, then we can solve the equation $\boldsymbol{Q}\boldsymbol{R}=\boldsymbol{A}$ to get $\boldsymbol{Q}$. This equation can be rewritten as $\boldsymbol{R}^{\top}\boldsymbol{Q}^{\top}=\boldsymbol{A}^{\top}$, which can be solved using solve_triangular, also very efficient. Together, these two steps make up a QR decomposition algorithm known as "Cholesky QR". If we ignore numerical stability, it may well be the fastest QR decomposition method there is.

Unfortunately, compared to standard QR decomposition, Cholesky QR is very unstable — it is extremely sensitive to the condition number of $\boldsymbol{A}^{\top}\boldsymbol{A}$. To address this, Shifted CholeskyQR for computing the QR factorization of ill-conditioned matrices (abbreviated "SCQR") proposes adding a regularization term $\lambda \boldsymbol{I}$ ($\lambda=\epsilon \Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$) to $\boldsymbol{A}^{\top}\boldsymbol{A}$ to alleviate this problem. But this is a double-edged sword: the larger $\epsilon$ is, the more stable SCQR becomes, but the resulting $\boldsymbol{Q}$ deviates further from orthogonality, and the final performance degrades accordingly.

Furthermore, even after introducing $\epsilon$, there's no guarantee that SCQR will succeed, so we still need an extra check, falling back to standard QR if it fails.

Reference Implementation

A simple reference implementation based on Jax is as follows:

import jax.numpy as jnp
from jax.scipy.linalg import solve_triangular
from jax import lax

def shift(A, eps=1e-9):
    return A + eps * jnp.linalg.matrix_norm(A, keepdims=True) * jnp.eye(A.shape[-1])

def scqr(A, eps=1e-9):
    """先按Shifted Cholesky QR算，失败则回退到默认QR
    """
    R = jnp.linalg.cholesky(shift(A.mT @ A, eps), upper=True)
    Q = solve_triangular(R.mT, A.mT, lower=True).mT
    return lax.cond(jnp.isfinite(Q).all(), lambda: Q, lambda: jnp.linalg.qr(A)[0])

Simple tests show that when it executes successfully, SCQR is about as efficient as the Newton-Schulz version $\msign$. However, to maintain a good approximation, $\epsilon$ can't be set too small, otherwise the performance drops noticeably. In practice, it usually needs to be set to around $\epsilon=10^{-9}$ for the results to be reasonably reliable, at which point SCQR still has a fairly high probability of falling back to standard QR, so the final speed still lags somewhat behind Newton-Schulz iteration.

Besides directly improving the QR decomposition algorithm, there are other acceleration tricks — for example, keeping only the top $k$ eigenvectors, so $\boldsymbol{V}$ only needs to be initialized at size $m\times k$ rather than $m\times m$, which also reduces the computation somewhat. Further acceleration of the QR decomposition is left for readers to explore further; we won't dwell on it here.

Other Details

There are also a few other details worth paying special attention to, which are closely related to both training stability and final performance.

First, following the convention of $\boldsymbol{M}_t\in\mathbb{R}^{n\times m}$, we need to ensure $n\geq m$, otherwise we should transpose. If $n < m$, then the matrix $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t \boldsymbol{V}_{t-1}$ is necessarily rank-deficient, and performing QR decomposition on a rank-deficient matrix is ill-posed — SCQR in particular is much more prone to various pathological behaviors, ultimately hurting performance. So ensuring $n\geq m$ both guarantees numerical stability and improves performance, while also speeding things up — killing several birds with one stone.

Second, in practice we found that adding a $\ColNorm$ to the $\QR$ step is also very helpful for improving performance:

\begin{equation}\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}) \qquad\to\qquad \boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\ColNorm(\boldsymbol{M}_t\boldsymbol{V}_{t-1}))\end{equation}

This is equivalent to simply changing $\tilde{\boldsymbol{v}}_k^{(t)} = \boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_k^{(t-1)}$ to $\tilde{\boldsymbol{v}}_k^{(t)} = \boldsymbol{M}^{\top}(\boldsymbol{M}\boldsymbol{v}_k^{(t-1)} / \Vert\boldsymbol{M}\boldsymbol{v}_k^{(t-1)}\Vert_2)$ in formula $\eqref{eq:vk-pi}$, without altering the convergence of the power iteration itself. But experiments show that this extra step $\ColNorm$ noticeably helps training performance, and the effect is even more pronounced with SCQR, where it clearly narrows the performance gap relative to standard QR.

It can be shown that in theory this operation doesn't change the power iteration itself; it only affects the numerical computation of the QR decomposition. Based on experimental observation, it actually makes SCQR fall back to standard QR with somewhat higher probability (though not by much, so it doesn't slow things down noticeably), and it also gives a slight boost to standard QR — so it seems that it improves the conditioning of the matrix being QR-decomposed, resulting in better QR quality overall.

Related Work

As mentioned at the start of this post, the idea of streaming power iteration has actually appeared multiple times already in various optimizer works, such as 4-bit Shampoo for Memory-Efficient Network Training, SOAP: Improving and Stabilizing Shampoo using Adam, COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs, and Dion: Distributed Orthonormalized Updates.

In fact, taking an algorithm that requires multiple iterations to converge and, exploiting the fact that model training itself already involves long-running updates, turning it into a streaming version that performs only one iteration per training step to amortize the cost, isn't too hard an idea to arrive at. We ourselves already tried something along these lines earlier, in Steepest Descent on Manifolds: 5. Dual Gradient Descent. So it's not really surprising that so much related work already exists.

The main reference for this post is ARO: A New Lens On Matrix Optimization For Large Models, published last month, which actually already covers most of the content of this post, and generalizes it further. The generalization is also well worth thinking about. Notice that Muon's update can currently be written as

\begin{equation}\ColNorm(\boldsymbol{M}_t\boldsymbol{V}_t)\boldsymbol{V}_t^{\top}\end{equation}

where $\ColNorm(\boldsymbol{M}_t)$ can be regarded as a base optimizer that simply performs column-wise normalization on the momentum — not particularly competitive on its own. So we find it a new set of orthogonal bases $\boldsymbol{V}_t$, apply the base optimizer in this new basis, and then transform back. We now recognize that this is exactly Muon, and it is indeed considerably stronger than directly applying $\ColNorm$. The natural next question, then, is: can $\ColNorm$ be replaced by some other base optimizer? To address this, ARO proposes a general optimizer framework (rotational steepest descent):

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta\boldsymbol{M}_{t-1} + \boldsymbol{G}_t \\[5pt] \boldsymbol{R}_t =&\, \QR(\boldsymbol{M}_t^{\top}f(\boldsymbol{M}_t\boldsymbol{R}_{t-1})) \\[5pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta_t (f(\boldsymbol{M}_t\boldsymbol{R}_t)\boldsymbol{R}_t^{\top} + \lambda \boldsymbol{W}_{t-1}) \\ \end{aligned}\end{equation}

where $f$ represents an arbitrary matrix function, and the original notation $\boldsymbol{V}$ has been replaced by $\boldsymbol{R}$ (Rotation). We'll leave further discussion of rotational steepest descent for another opportunity.

Summary

This post mainly introduced the idea of computing the SVD via Streaming Power Iteration, and using it to implement Muon, requiring only a single QR decomposition per step. Compared to the standard Newton-Schulz iteration implementation, this approach offers greater flexibility for extension.