Muon Implementation Based on Streaming Power Iteration: 3. Refinement
Looking back at the previous two posts, Muon Implementation Based on Streaming Power Iteration: 1. First Encounter and Muon Implementation Based on Streaming Power Iteration: 2. Speedup, we introduced a Streaming Power Iteration implementation scheme for Muon, gave an initial validation of its feasibility, and then further discussed accelerating the core operation—QR decomposition—bringing it close to the efficiency of the Newton-Schulz iteration implementation.
In this post, we no longer restrict ourselves to optimizing a single step of the QR decomposition. Instead, we take a more holistic view of streaming power iteration and, combined with the specific computational context, further "refine" the implementation details, minimizing computational bottlenecks as much as possible so that its efficiency approaches the theoretical limit.
Existing Results
Streaming power iteration is essentially "computing an SVD while training." The idea is to use power iteration to compute the SVD, and by caching the result from the previous step, spread the computation evenly across every training step, making it feasible to embed an SVD inside the optimizer. As for Muon, it's simply one basic application of this—because the most fundamental way to implement Muon's core operation $\newcommand{msign}{\mathop{\text{msign}}}\msign$ is precisely via SVD. Specifically, Muon's update formula is
\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}
Here all matrices have shape $n\times m$, and we adopt the convention $n\geq m$. Let the SVD of $\boldsymbol{M}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ (where $\boldsymbol{U}\in\mathbb{R}^{n\times m}$ and $\boldsymbol{\Sigma},\boldsymbol{V}\in\mathbb{R}^{m\times m}$), so that $\msign(\boldsymbol{M})=\boldsymbol{U}\boldsymbol{V}^{\top}$; hence implementing SVD amounts to implementing $\msign$. Of course, direct SVD is usually quite expensive, but streaming power iteration makes it feasible.
In the previous post, we also discussed four ideas for speeding up streaming power iteration. The first—enabling full-precision FP32 multiplication—is generic, while the other three are, to some extent, mutually exclusive, and we can only pick one. My recommendation is to go with the second, since it has a higher theoretical ceiling, and the refinements in this post are also based on that second approach. Substituting the second idea into the streaming power iteration version and applying it to Muon, the iteration formula becomes
\begin{equation}\newcommand{QR}{\mathop{\text{QR}}}\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}\QR(\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}
Clearly, the most expensive operation now is $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1}))$, and that is exactly the target of our next optimization.
Speeding Up the Decomposition
To ensure efficiency, we don't call the framework's built-in QR decomposition function for $\QR$ here; instead we use "SCQR" (Shifted Cholesky QR), which splits the QR decomposition of matrix $\boldsymbol{A}$ into two steps: 1. Perform a Cholesky decomposition on $\boldsymbol{A}^{\top}\boldsymbol{A} + \lambda \boldsymbol{I}$ 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 theoretically very efficient, but they don't always succeed, so we need an extra check step—falling back to the built-in standard QR function on failure, which almost always succeeds. However, once this fallback is triggered, end-to-end efficiency takes a big hit. The main reason SCQR fails is that Cholesky decomposition is extremely sensitive to the condition number of the matrix; the regularization term $\lambda\boldsymbol{I}$ exists precisely to reduce the condition number of $\boldsymbol{A}^{\top}\boldsymbol{A}$.
However, this presents a dilemma: the larger $\lambda$ is, the more likely SCQR succeeds, but the final result deviates further from orthogonality (i.e., larger error), degrading performance; the smaller $\lambda$ is, the higher the precision, of course, but the probability of falling back to standard QR also rises, hurting efficiency. In practice, we found that setting $\lambda=\epsilon \Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$ with $\epsilon=10^{-9}$ strikes a good balance between effectiveness and efficiency.
The speedup ideas in the previous post all revolved around reducing the condition number. The first version of streaming power iteration was $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1})$, where the matrix we need to Cholesky-decompose is $\boldsymbol{V}_{t-1}^{\top}(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t)^2\boldsymbol{V}_{t-1}$—that is, the condition number of $\boldsymbol{M}_t$ gets raised to the fourth power, which clearly explodes. After switching to $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1}))$, although $\QR$ now happens twice, the condition number of the matrix in each Cholesky decomposition is only the square of that of $\boldsymbol{M}_t$—a marked reduction—so SCQR's success rate improves considerably, and speed actually increases as a result.
Reordering the Computation
Everything above is essentially a recap of the previous two posts (sorry for the long lead-in, but sharpening the axe never wastes time cutting wood). In this section we finally get to the new optimization idea. Looking carefully, one notices that we currently introduce $\QR$ twice, but these two occurrences of $\QR$ are considered independently. @YouJiacheng and @Kimi discovered that if we consider them jointly, we can find some acceleration tricks.
Under the default ordering, the computation flow for $\boldsymbol{V}_t = \QR(\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1}))$ is
\begin{equation}\begin{aligned} \boldsymbol{A}_{(1), t} =&\, \boldsymbol{M}_t\boldsymbol{V}_{t-1} \\ \boldsymbol{R}_{(1), t}^{\top}\boldsymbol{R}_{(1), t} =&\, \boldsymbol{A}_{(1), t}^{\top}\boldsymbol{A}_{(1), t} + \lambda \boldsymbol{I}\qquad(\text{Cholesky decomposition}) \\ \boldsymbol{Q}_{(1), t} =&\, \boldsymbol{A}_{(1), t} \boldsymbol{R}_{(1), t}^{-1} \qquad(\text{Triangular Solve}) \\ \boldsymbol{A}_{(2), t} =&\, \boldsymbol{M}_t^{\top}\boldsymbol{Q}_{(1), t} \\ \boldsymbol{R}_{(2), t}^{\top}\boldsymbol{R}_{(2), t} =&\, \boldsymbol{A}_{(2), t}^{\top}\boldsymbol{A}_{(2), t} + \lambda \boldsymbol{I}\quad(\text{Cholesky decomposition}) \\ \boldsymbol{Q}_{(2), t} =&\, \boldsymbol{A}_{(2), t} \boldsymbol{R}_{(2), t}^{-1} \qquad(\text{Triangular Solve}) \\ \end{aligned}\label{eq:qr2}\end{equation}
Among these, the four steps $\boldsymbol{M}_t\boldsymbol{V}_{t-1}$, $\boldsymbol{A}_{(1), t}^{\top}\boldsymbol{A}_{(1), t}$, $\boldsymbol{A}_{(1), t} \boldsymbol{R}_{(1), t}^{-1}$, $\boldsymbol{M}_t^{\top}\boldsymbol{Q}_{(1), t}$ all have complexity $\mathcal{O}(nm^2)$, while the rest have complexity $\mathcal{O}(m^3)$; when $n \gg m$, $\mathcal{O}(nm^2)$ can become the bottleneck. Interestingly, through an algebraic identity we can make $\mathcal{O}(nm^2)$ appear only once!
\begin{equation}\begin{aligned} \boldsymbol{A}_{(1), t} =&\, (\boldsymbol{M}_t^{\top}\boldsymbol{M}_t)\boldsymbol{V}_{t-1} \\ \boldsymbol{R}_{(1), t}^{\top}\boldsymbol{R}_{(1), t} =&\, \boldsymbol{V}_{t-1}^{\top}\boldsymbol{A}_{(1), t} + \lambda \boldsymbol{I}\qquad(\text{Cholesky decomposition}) \\ \boldsymbol{A}_{(2), t} =&\, \boldsymbol{A}_{(1), t} \boldsymbol{R}_{(1), t}^{-1} \qquad(\text{Triangular Solve}) \\ \boldsymbol{R}_{(2), t}^{\top}\boldsymbol{R}_{(2), t} =&\, \boldsymbol{A}_{(2), t}^{\top}\boldsymbol{A}_{(2), t} + \lambda \boldsymbol{I}\qquad(\text{Cholesky decomposition}) \\ \boldsymbol{Q}_{(2), t} =&\, \boldsymbol{A}_{(2), t} \boldsymbol{R}_{(2), t}^{-1} \qquad(\text{Triangular Solve}) \\ \end{aligned}\label{eq:qr2-sim}\end{equation}
This equivalent version really deserves careful appreciation! First, one can prove that it is theoretically completely equivalent to the original, and this equivalence does not depend on the exact orthogonality of $\boldsymbol{V}_{t-1}$ and $\boldsymbol{Q}_{(1), t}$. After the transformation, only the step $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t$ has complexity $\mathcal{O}(nm^2)$, and everything else is $\mathcal{O}(m^3)$—and the total number of steps is even reduced by one (merging what were originally $\boldsymbol{Q}_{(1), t} = \boldsymbol{A}_{(1), t} \boldsymbol{R}_{(1), t}^{-1}$ and $\boldsymbol{A}_{(2), t} = \boldsymbol{M}_t^{\top}\boldsymbol{Q}_{(1), t}$ into a single step)!
Note 1: According to @YouJiacheng, this elegant transformation was automatically discovered by Kimi after he explained equation $\eqref{eq:qr2}$ to it.
Note 2: Equations $\eqref{eq:qr2}$ and $\eqref{eq:qr2-sim}$ actually differ subtly—in equation $\eqref{eq:qr2}$ the first step is $(\boldsymbol{M}_t\boldsymbol{V}_{t-1})^{\top}(\boldsymbol{M}_t\boldsymbol{V}_{t-1})$, whereas in equation $\eqref{eq:qr2-sim}$ it is $\boldsymbol{V}_{t-1}^{\top}(\boldsymbol{M}_t^{\top}\boldsymbol{M}_t)\boldsymbol{V}_{t-1}$. These two algorithms are mathematically completely equivalent, but they behave differently under finite-precision floating-point arithmetic: the matrix produced by the latter's multiplication has a larger condition number, so the Cholesky decomposition may need a slightly larger regularization term.
Simplifying the Regularization
For the matrix $\boldsymbol{A}^{\top}\boldsymbol{A}$, the regularization term we add during Cholesky decomposition is $\lambda\boldsymbol{I}$, where $\lambda=\epsilon \Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$ and $\epsilon=10^{-9}$. We never gave a detailed explanation of why this particular form was chosen, so let's expand on that here, and derive a more concise regularization term based on the specifics of the problem.
Since $\boldsymbol{A}^{\top}\boldsymbol{A}$ is symmetric positive definite, its SVD coincides with its eigendecomposition. Suppose its SVD is $\boldsymbol{A}^{\top}\boldsymbol{A} = \boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$; then $\boldsymbol{A}^{\top}\boldsymbol{A} + \lambda\boldsymbol{I} = \boldsymbol{V}(\boldsymbol{\Sigma} + \lambda\boldsymbol{I})\boldsymbol{V}^{\top}$. Let the largest and smallest singular values of $\boldsymbol{A}^{\top}\boldsymbol{A}$ be $\sigma_{\max},\sigma_{\min}$, respectively; then the largest and smallest singular values of $\boldsymbol{A}^{\top}\boldsymbol{A} + \lambda\boldsymbol{I}$ are $\sigma_{\max} + \lambda,\sigma_{\min} + \lambda$, and the condition number—the ratio of the largest to smallest singular value—drops from $\sigma_{\max}/\sigma_{\min}$ down to
\begin{equation}\frac{\sigma_{\max} + \lambda}{\sigma_{\min} + \lambda} < \frac{\sigma_{\max} + \lambda}{\lambda} = \frac{\sigma_{\max}}{\lambda} + 1\end{equation}
If we want to keep the condition number bounded by $1/\epsilon + 1$, then $\lambda \geq \epsilon \sigma_{\max}$, which shows that ideally we should use the largest singular value of $\boldsymbol{A}^{\top}\boldsymbol{A}$—i.e., its spectral norm—as the basis for tuning $\lambda$. But computing the spectral norm is relatively costly, so instead we used the simpler Frobenius norm $\Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$, which is where $\lambda=\epsilon \Vert\boldsymbol{A}^{\top}\boldsymbol{A}\Vert_F$ comes from; as for $\epsilon=10^{-9}$, that is purely an empirical finding.
However, "the spectral norm is costly to compute, so use the Frobenius norm instead" is only a general conclusion applicable to arbitrary matrices. Here, the streaming power iteration we're performing is itself being used to compute an SVD, and as training proceeds, $\boldsymbol{V}_t$ increasingly converges to the right singular matrix of $\boldsymbol{M}_t$. Since $\boldsymbol{M}_t$ changes slowly, $\boldsymbol{V}_{t-1}$ also changes only slightly, so theoretically $\boldsymbol{V}_{t-1}^{\top}\boldsymbol{M}_t^{\top}\boldsymbol{M}_t\boldsymbol{V}_{t-1}$ becomes progressively closer to a diagonal matrix, and its top-left element gets progressively closer to its spectral norm!
Likewise, $\tilde{\boldsymbol{U}}_t = \QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})$ becomes progressively closer to the left singular matrix of $\boldsymbol{M}_t$, so $\tilde{\boldsymbol{U}}_t^{\top}\boldsymbol{M}_t\boldsymbol{M}_t^{\top}\tilde{\boldsymbol{U}}_t$ also becomes progressively closer to a diagonal matrix, whose top-left element gets progressively closer to its spectral norm. Therefore, in our specific setting, the simplest and most accurate basis is just to directly use $(\boldsymbol{A}^{\top}\boldsymbol{A})_{[0,0]}$ as an approximation to the spectral norm, i.e., $\lambda=\epsilon \cdot (\boldsymbol{A}^{\top}\boldsymbol{A})_{[0,0]}$. In practice, $\epsilon=10^{-7}$ strikes a good balance between effectiveness and efficiency.
Reference Implementation
Combining the modifications from the two sections above, the reference implementation of the iteration from $\boldsymbol{V}_{t-1}$ to $\boldsymbol{V}_t$ is:
import jax.numpy as jnp
from jax.scipy.linalg import solve_triangular
from jax import lax
def shift_old(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_old(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])
def v_step_old(M, V, eps=1e-9):
return scqr(M.mT @ scqr(M @ V, eps), eps)
def shift(A, eps=1e-7):
return A + eps * A[..., :1, :1] * jnp.eye(A.shape[-1])
def v_step(M, V, eps=1e-7):
A = (M.mT @ M) @ V
R = jnp.linalg.cholesky(shift(V.mT @ A, eps), upper=True)
B = solve_triangular(R.mT, A.mT, lower=True).mT
R = jnp.linalg.cholesky(shift(B.mT @ B, eps), upper=True)
Q = solve_triangular(R.mT, B.mT, lower=True).mT
return lax.cond(jnp.isfinite(Q).all(), lambda: Q, lambda: jnp.linalg.qr(A)[0])
Concurrent Work
Between the release of Muon Implementation Based on Streaming Power Iteration: 2. Speedup and the publication of this post, some interesting optimization efforts also appeared externally, sharing similar optimization ideas with the two modifications proposed in this post; it's worth studying them together.
First, after the previous post was published, @Ji_Ha_Kim also proposed some improvement ideas. For instance, he mentioned that in a conversation with GPT he found (see the link) that we might be able to skip one occurrence of $\text{Triangular Solve}$! Specifically, we have
\begin{equation}\begin{aligned} \boldsymbol{V}_t =&\, \QR(\boldsymbol{M}_t^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})) \\ =&\, \QR(\boldsymbol{V}_{t-1}(\boldsymbol{M}_t \boldsymbol{V}_{t-1})^{\top}\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})) \\ =&\, \QR(\boldsymbol{V}_{t-1}(\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})^{\top}\boldsymbol{M}_t \boldsymbol{V}_{t-1})^{\top}) \\ =&\, \boldsymbol{V}_{t-1}\QR((\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})^{\top}\boldsymbol{M}_t \boldsymbol{V}_{t-1})^{\top}) \\ \end{aligned}\end{equation}
It's easy to see that $\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})^{\top}\boldsymbol{M}_t \boldsymbol{V}_{t-1}$ is exactly the R obtained from QR-decomposing $\boldsymbol{M}_t\boldsymbol{V}_{t-1}$, which can be obtained directly via Cholesky decomposition.
In other words, in theory, once R is obtained from the first Cholesky decomposition, the second $\QR$ can proceed directly, skipping one occurrence of $\text{Triangular Solve}$. However, this only holds theoretically, because the result depends on the strict orthogonality of $\boldsymbol{V}_{t-1}$ and $\QR(\boldsymbol{M}_t\boldsymbol{V}_{t-1})$, which only holds under exact QR (i.e., $\lambda=0$). In practice, to ensure efficiency we must use SCQR, whose result is not strictly orthogonal, so exploiting orthogonality prematurely in the algebraic transformation would instead cause error accumulation in the iterative algorithm.
Around the same time, the Tri Dao team published Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon, proposing an acceleration idea for the $\msign$ operator. By definition, $\msign(\boldsymbol{M}) = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$; the team wanted to use Newton-Schulz iteration to compute $(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}$ of $m\times m$ rather than $\msign$, which would noticeably reduce the computation when $n\gg m$. In fact, many researchers had tried this idea before, but all had failed—yet Tri Dao's team cleverly solved the problem via Restart.
Clearly, this optimization direction for $\msign$ is consistent with streaming power iteration's shift from $\eqref{eq:qr2}$ to $\eqref{eq:qr2-sim}$. Coincidentally, @Ji_Ha_Kim suggested changing the Newton-Schulz iteration for $\msign$ from a polynomial form to a rational form, which achieves the same good results with fewer iterations. The problem with rational iteration is that it requires computing a matrix inverse; however, combined with the specific context of $\msign$, it only needs to invert a symmetric positive-definite matrix of size $m\times m$, which can be computed via Cholesky decomposition plus two applications of $\text{Triangular Solve}$—still acceptable.
That said, this makes the computational flow of the rational iteration highly overlap with that of streaming power iteration, and each step additionally requires one more $\text{Triangular Solve}$, so it doesn't appear that its speed can surpass streaming power iteration.
Summary
This post further "refined" the implementation details of streaming power iteration, with the main improvements being: 1. Reordering the computation to reduce the number of $\mathcal{O}(nm^2)$-complexity operations from four down to one; 2. Simplifying the regularization term by exploiting the specific context of streaming power iteration. These optimizations further reduce the computational bottlenecks of streaming power iteration, pushing its efficiency close to the theoretical limit.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.