Muon Implementation Based on Streaming Power Iteration: 2. Acceleration

基于流式幂迭代的Muon实现：2. 加速 苏剑林 · 2026-03-26 · Machine Learning iteration matrix SVD optimizer muon

In the first post, A Streaming-Power-Iteration Implementation of Muon: 1. Getting Acquainted], I isolated Streaming Power Iteration as a standalone concept and used it to build a new implementation of Muon. Since this new scheme directly approximates the SVD, it offers much richer room for extension compared to the standard implementation based on Newton–Schulz iteration, and it is worth digging into further.

Computationally, the main change in the new scheme is replacing Newton–Schulz iteration with a $\newcommand{QR}{\mathop{\text{QR}}}\QR$ decomposition, which introduces some slowdown. The previous post already discussed some basic speedup tricks, but they still didn't match the standard implementation in speed. In this post, we continue exploring ways to accelerate $\QR$, trying to close the gap as much as possible.

Streaming Iteration

We'll continue to use all the concepts and notation from the first post — readers who have doubts should go back and review it first. To recap, the update formula for Muon is

\begin{equation}\newcommand{msign}{\mathop{\text{msign}}}\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 the standard implementation of $\msign$ is Newton–Schulz iteration], which is also the most expensive computation in the Muon optimizer. By contrast, the update formula for the streaming power iteration scheme is

\begin{equation}\newcommand{ColNorm}{\mathop{\text{ColNorm}}}\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}\end{equation}

If we repeatedly execute $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1})$, this is exactly standard power iteration, and the result will converge to the right singular matrix of $\boldsymbol{M}_t$, thereby achieving the SVD of $\boldsymbol{M}_t$, from which we can compute $\msign$. However, running a full power iteration at every step is too expensive, so instead we cache the result from the previous step, $\boldsymbol{V}_{t-1}$, and perform just one iteration of $\QR$ per step as an approximation — this is what "streaming" means here.

Now the most expensive operation becomes the $\QR$ decomposition. The most naive implementation is naturally to call the framework's built-in QR function, which is based on Householder transformations — very stable, but relatively slow.

First Speedup

To accelerate things, in the previous post we introduced Cholesky QR, which splits the QR decomposition of a matrix $\boldsymbol{A}$ into two steps: 1. Perform a Cholesky decomposition of $\boldsymbol{A}^{\top}\boldsymbol{A}$ to obtain the upper-triangular matrix $\boldsymbol{R}$; 2. Solve the equation $\boldsymbol{Q}\boldsymbol{R}=\boldsymbol{A}$ to obtain the orthogonal matrix $\boldsymbol{Q}$. Both steps are, in theory, extremely efficient, but in practice the computation can fail if the condition number is too large. To address this, we also introduced the Shift trick], which adds a regularization term $\lambda \boldsymbol{I}$ ($\lambda=\epsilon \Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$) to $\boldsymbol{A}^{\top}\boldsymbol{A}$ to reduce the condition number.

Combining the two, we call this "SCQR" (Shifted Cholesky QR). A 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])

Note that the smaller $\lambda$ is, the more likely SCQR is to fail, while the larger $\lambda$ is, the further the result deviates from orthogonality, degrading performance — so $\lambda$ needs to be "just right." This means that this scheme still has a fairly high chance of falling back to standard QR. Additionally, the previous post also mentioned that adding $\ColNorm$ to the power iteration (i.e., changing it to $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\ColNorm(\boldsymbol{M}_t\boldsymbol{V}_{t-1}))$) can stabilize training, and this effect is even more pronounced under SCQR.

Full Precision

The above essentially summarizes the entirety of the first post — it got the whole pipeline running and validated feasibility, and compared with directly calling the framework's built-in QR decomposition, SCQR did provide some speedup. But it was still noticeably slower than the $\msign$ implementation based on Newton–Schulz iteration, so we still need to find ways to speed things up further.

This section introduces the first acceleration trick: turning on "full" FP32-precision matrix multiplication. First, it's worth noting that the newly added steps in streaming power iteration are all computed at FP32 precision. However, starting with the A100's introduction of the TF32 format, some frameworks (such as Jax, which I use for my small-scale experiments, or certain versions of Torch) will, by default, convert FP32 arrays to TF32 format for matrix multiplication in order to speed things up — you need to manually enable it to get true FP32-precision multiplication.

Some readers might wonder: shouldn't increasing multiplication precision slow things down — why does it speed things up instead? This is indeed counterintuitive, but not hard to understand once you think about it: lowering the precision of a matrix tends to increase its condition number, which raises the probability of SCQR failing and falling back to standard QR, thereby increasing runtime. Conversely, raising the precision increases the success rate of SCQR, and QR is precisely the most time-consuming part — so the total runtime actually decreases.

According to available information, Jax has always defaulted to using TF32 for FP32 multiplication, so it needs jax.config.update('jax_default_matmul_precision', 'highest') to manually enable full precision. Torch is a bit more complicated: versions 1.7 through 1.11 default to TF32 multiplication, but starting from 1.12, FP32 multiplication is the default. Given that Torch is now at version 2.11, presumably most users no longer need to enable this manually.

Double Orthogonalization

The second acceleration trick I came up with is adding an extra orthogonalization step for the left singular matrix, i.e., changing the power iteration step to

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

This step is likewise counterintuitive — adding an extra $\QR$ operation ends up making things faster overall, for reasons similar to the previous section: both reduce the condition number of the matrix being decomposed, thereby increasing the success rate of SCQR. Since SCQR itself is very fast, running it twice doesn't add much time, while significantly reducing the number of fallbacks to standard QR — resulting in a noticeable overall speedup.

Understanding this speedup requires two steps: first, showing that adding this extra step $\QR$ does not, in theory, change the power iteration; second, showing that adding this step $\QR$ does indeed lower the condition number. The first point is easy to understand: if $\boldsymbol{A}=\boldsymbol{Q}\boldsymbol{R}$, then $\boldsymbol{Q}=\boldsymbol{A}\boldsymbol{R}^{-1}$, where $\boldsymbol{R}^{-1}$ is also an upper-triangular matrix — that is, the QR decomposition can be written as a right-multiplication by an upper-triangular matrix. Then

\begin{equation}\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1}) = \boldsymbol{M}_t^{\top}(\boldsymbol{M}_t\boldsymbol{V}_{t-1}\times \text{some upper triangular matrix}) = \boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}\times \text{some upper triangular matrix} \end{equation}

By the uniqueness of QR decomposition, right-multiplying by an upper-triangular matrix doesn't change the result of $\QR$, so in theory this is equivalent to $\QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1})$.

As for the condition number, it equals the ratio of the largest to the smallest singular value. If we perform a single $\QR$, then the matrix to be Cholesky-decomposed is $\boldsymbol{V}_{t-1}^{\top}(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t)^2\boldsymbol{V}_{t-1}$ — note that orthogonal transformations don't change singular values, and hence don't change the condition number, so at this point the condition number of the matrix to be decomposed reaches the fourth power of the condition number of $\boldsymbol{M}_t$! If we instead perform two rounds of $\QR$, the matrix to be Cholesky-decomposed becomes $\boldsymbol{Q}_t^{\top}(\boldsymbol{M}_t\boldsymbol{M}_t^{\top})\boldsymbol{Q}_t$, where $\boldsymbol{Q}_t$ is the orthogonal matrix from the first $\QR$ — in which case the condition number is only the square of $\boldsymbol{M}_t$, a significant reduction.

Shift Invariance

The third acceleration trick came out of a discussion with @YouJiacheng], and it exploits the shift invariance of the eigenmatrix. As we know, the power iteration $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1})$ can also be understood as computing the eigenmatrix of the positive-definite matrix $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t$, and positive-definite matrices have the property that adding any multiple of the identity matrix leaves the eigenmatrix unchanged.

In other words, $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t$ and $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t + \lambda \boldsymbol{I}$ share the same eigenmatrix, so we can rewrite the power iteration as

\begin{equation}\boldsymbol{V}_t = \QR((\boldsymbol{M}_t^{\top}\boldsymbol{M}_t + \lambda \boldsymbol{I})\boldsymbol{V}_{t-1}) = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1} + \lambda \boldsymbol{V}_{t-1})\end{equation}

without changing the convergence result of the power iteration. So what's the benefit of adding $\lambda \boldsymbol{I}$ to $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t$? Again, the answer is to reduce the condition number, i.e., $(\sigma_{\max} + \lambda)/(\sigma_{\min} + \lambda) < \sigma_{\max}/\sigma_{\min}$, which also improves the success rate of Cholesky QR. Note that here we're talking about Cholesky QR rather than SCQR, because by setting an appropriate $\lambda$ externally, we can already guarantee the condition number without needing a Shift — so the resulting output is guaranteed to be orthogonal, which is a nice property in itself.

But don't celebrate too soon. The larger $\lambda$ is, the easier it naturally becomes for Cholesky QR to succeed, but this also slows down the convergence of the power iteration! This is because the convergence speed of power iteration depends on the ratio of adjacent singular values — the smaller $\sigma_{i+1}/\sigma_i$ is, the faster the convergence (with singular values sorted from largest to smallest), and $(\sigma_{i+1} + \lambda)/(\sigma_i + \lambda) > \sigma_{i+1}/\sigma_i$, so the larger $\lambda$ is, the slower the power iteration converges, and the worse the final result becomes.

So we must carefully tune the value of $\lambda$ to balance the success rate of Cholesky QR against the convergence speed of power iteration. Through testing, I found that taking $\lambda = \epsilon\Vert\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\Vert_F$ with $\epsilon=10^{-4}$ gives fairly good results. Another approach is to use a larger $\lambda$ to guarantee the success rate of Cholesky QR, and then perform two iterations to speed up convergence of the power iteration, i.e.,

\begin{equation}\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\tilde{\boldsymbol{V}}_t + \lambda \tilde{\boldsymbol{V}}_t),\qquad \tilde{\boldsymbol{V}}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1} + \lambda \boldsymbol{V}_{t-1})\end{equation}

This lets us have the best of both worlds for Cholesky QR and power iteration convergence, at the cost of needing two $\QR$ operations per step.

Multi-Step Correction

The fourth acceleration trick is called "SCQR2," a general multi-step correction technique for SCQR. Let's revisit the two steps of SCQR (given the matrix $\boldsymbol{A}$ to be decomposed):

\begin{align}1)\quad&\, \boldsymbol{R}^{\top}\boldsymbol{R}= \boldsymbol{A}^{\top}\boldsymbol{A} + \lambda \boldsymbol{I} &\,(\text{correct}\boldsymbol{A}^{\top}\boldsymbol{A}+\lambda\boldsymbol{I}\text{Cholesky decomposition}) \\[5pt] 2)\quad&\, \boldsymbol{Q} = \boldsymbol{A}\boldsymbol{R}^{-1}&\,(\text{solve triangular linear equation}\boldsymbol{Q}\boldsymbol{R}=\boldsymbol{A})\end{align}

The problem with SCQR is that the larger $\lambda$ is, the easier the Cholesky decomposition succeeds, but the less orthogonal $\boldsymbol{Q} = \boldsymbol{A}\boldsymbol{R}^{-1}$ becomes. The idea behind SCQR2 is: first perform one round of SCQR with a larger $\lambda$. At this point, although the result isn't orthogonal, it is closer to orthogonal than the original $\boldsymbol{A}$ — indicating that the condition number has already been reduced. We can then apply SCQR again with a smaller $\lambda$ to correct the orthogonality. A rough implementation is as follows:

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
    """
    R = jnp.linalg.cholesky(shift(A.mT @ A, eps), upper=True)
    return solve_triangular(R.mT, A.mT, lower=True).mT

def scqr2(A, eps1=1e-4, eps2=1e-8):
    """SCQR两次，失败则回退到默认QR
    """
    Q = scqr(scqr(A, eps1), eps2)
    return lax.cond(jnp.isfinite(Q).all(), lambda: Q, lambda: jnp.linalg.qr(A)[0])

To understand why this second correction step works, suppose the first round of SCQR yields $\boldsymbol{Q}_1 = \boldsymbol{A}\boldsymbol{R}_1^{-1}$. Although it deviates from orthogonality, it has the form "$\boldsymbol{A}\times \text{upper triangular matrix}$" — and as noted earlier, right-multiplying by an upper-triangular matrix doesn't change the QR result, so this allows us to apply SCQR again on top of the first round's result. In principle, we could of course perform even more correction steps.

Summary of Methods

We've now discussed four acceleration tricks; here's a brief summary of their properties.

The first trick — raising the precision of FP32 matrix multiplication — is general-purpose: Jax needs it enabled manually, while newer versions of Torch already do so by default. The second, third, and fourth tricks are each standalone and cannot be combined with one another. Intuitively, trick two has the higher ceiling, because tricks three and four both take $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}$ as input — meaning the condition number has already been amplified, and they then try to remedy it after the fact — whereas trick two modifies the input to be $\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})$, reducing the condition number at the source.

Interestingly, tricks two, three, and four all seem to converge on requiring two rounds of $\QR$. Except for trick three, which — with careful tuning of $\lambda$ — can get by with just one $\QR$, all the others need at least two $\QR$. This really does seem to be the safest choice. In terms of speed, if trick three can be tuned to use only one $\QR$, it's the fastest; otherwise it's about as fast as trick two. Trick four is somewhat unstable: applying SCQR2 to $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}$ is fast but doesn't guarantee good results, while applying it to $\boldsymbol{M}_t^{\top}\ColNorm(\boldsymbol{M}_t\boldsymbol{V}_{t-1})$ guarantees good results but at reduced speed.

I recommend combining tricks one and two, which offers good guarantees on both effectiveness and efficiency. In isolated testing, its speed comes out to roughly half that of the $\msign$ Newton–Schulz iteration. Readers might think, "all that effort for only half the speed?" But this is actually pretty good, considering we're computing everything in FP32 and need two rounds of $\QR$. On the other hand, the end-to-end time spent on the $\msign$ step is only around 1% of the total, so doubling it only adds about 1% more overall time — an acceptable tradeoff.

Furthermore, the efficiency of Newton–Schulz iteration depends on the number of iteration steps. If we use the coefficients from Polar Express] to further increase the number of steps to improve precision, the speed gap with our approach here will narrow further. In short, streaming power iteration is indeed slower, but it also yields richer, more accurate results (the SVD), which enables many more possibilities.

Summary

This post introduced further acceleration techniques for streaming power iteration, whose essence lies in finding ways to lower the condition number of the matrix, thereby improving the success rate of Cholesky QR.