% Source: https://kexue.fm/archives/11654
\documentclass[11pt]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{lmodern}
\usepackage{geometry}
\usepackage{hyperref}
\usepackage{graphicx}
\usepackage{amsmath,amssymb}
\usepackage{booktabs,longtable,array}
\usepackage{xcolor}
\usepackage{listings}

\title{A Streaming Power Iteration-Based Implementation Approach for Muon}
\author{Su Jianlin}
\date{2026-03-12}

\begin{document}

\maketitle

\newcommand{\msign}{\mathop{\text{msign}}}
\newcommand{\mclip}{\mathop{\text{mclip}}}
\newcommand{\QR}{\mathop{\text{QR}}}
\newcommand{\ColNorm}{\mathop{\text{ColNorm}}}
\newcommand{\diag}{\mathop{\text{diag}}}

\section*{Introduction}

The core operation of \href{https://kexue.fm/archives/10592}{Muon} is $\msign$. The current standard implementation uses the \href{https://kexue.fm/archives/10996}{Newton-Schulz iteration}. Undoubtedly, this is a highly efficient and GPU-friendly algorithm, and its contribution accounts for at least half of Muon's popularity. However, this algorithm gives the impression of being a one-of-a-kind solution—it appears confined solely to computing $\msign$. Once we attempt to modify Muon (e.g., replacing $\msign$ with $\mclip$ as in \href{https://kexue.fm/archives/11059}{this article}), the corresponding computations become significantly more complicated.

This article proposes a new implementation idea—approximating SVD via Streaming Power Iteration. This is not an entirely novel concept, as it has appeared in some previous optimizer-related works; however, here we distill it into an independent algorithm for standalone use.

\section{Content Review \label{sec:content-review}}

We will not elaborate on the details of Muon here. Readers are encouraged to refer to earlier articles such as \href{https://kexue.fm/archives/10592}{\textit{Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices}}, \href{https://kexue.fm/archives/10739}{\textit{Sequel to Muon: Why Did We Choose to Try Muon?}}, and \href{https://kexue.fm/archives/11416}{\textit{Guide to the Muon Optimizer: Quick Start and Key Details}}. Here, we present Muon's 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}
where $\msign$ is defined as
\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}$, and without loss of generality, we assume $n\geq m$. For simplicity, in most cases we assume $r=m$ (i.e., full rank), and only discuss reduced-rank situations when absolutely necessary.

Since SVD is computationally expensive, the Newton-Schulz iteration is commonly used to compute $\msign$, which we have detailed in \href{https://kexue.fm/archives/10922}{\textit{Newton-Schulz Iteration for the msign Operator (Part 1)}} and \href{https://kexue.fm/archives/10996}{\textit{Newton-Schulz Iteration for the msign Operator (Part 2)}}. In summary, the Newton-Schulz iteration is remarkably clever and has been instrumental in Muon's success, but its extensibility is relatively weak.

To extend the applicability of Newton-Schulz iteration, the author has previously conducted several studies, such as \href{https://kexue.fm/archives/11006}{\textit{Computing Singular Value Clipping (mclip) via msign (Part 1)}}, \href{https://kexue.fm/archives/11059}{\textit{Computing Singular Value Clipping (mclip) via msign (Part 2)}}, \href{https://kexue.fm/archives/11158}{\textit{Efficient Computation of Matrix Square Root and Inverse Square Root}}, and \href{https://kexue.fm/archives/11175}{\textit{Efficient Computation of Matrix r-th Root and Inverse r-th Root}}. However, the range of such extensions remains limited.

Clearly, the most straightforward solution would be to compute SVD directly—this is the approach we focus on next.

\section{Power Iteration \label{sec:power-iteration}}

In articles such as \href{https://kexue.fm/archives/6051}{\textit{Lipschitz Constraints in Deep Learning: Generalization and Generative Models}} and \href{https://kexue.fm/archives/10648}{\textit{From Spectral Norm Gradient to a New Form of Weight Decay}}, we have already encountered \href{https://en.wikipedia.org/wiki/Power_iteration}{Power Iteration}, which we use to find the principal eigenvector of $\boldsymbol{M}^{\top}\boldsymbol{M}$, or equivalently, the right principal singular vector of $\boldsymbol{M}$. The iteration format is:
\begin{equation}
\boldsymbol{v}_1^{(t)} = \frac{\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_1^{(t-1)}}{\|\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_1^{(t-1)}\|_2}
\end{equation}
Assuming we have obtained the principal eigenvector $\boldsymbol{v}_1$, we can add orthogonalization to the power iteration to compute 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}{\|\tilde{\boldsymbol{v}}_2^{(t)} - \langle\tilde{\boldsymbol{v}}_2^{(t)},\boldsymbol{v}_1\rangle\boldsymbol{v}_1\|_2},\qquad \tilde{\boldsymbol{v}}_2^{(t)} = \boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_2^{(t-1)}
\end{equation}
By ensuring orthogonality to $\boldsymbol{v}_1$, this converges to the second eigenvector $\boldsymbol{v}_2$. Similarly, assuming $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_{k-1}$ are known, we can use Gram-Schmidt orthogonalization to compute 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}}{\|\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}\|_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 do not need to wait until $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_{k-1}$ are fully computed before starting $\boldsymbol{v}_k$. The entire matrix $\boldsymbol{V}=[\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_m]$ can be updated in parallel. Specifically, starting from an existing approximation $\boldsymbol{V}_{t-1}$, we compute $\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{V}_{t-1}$ in batch and then re-orthonormalize column-wise (using QR decomposition), resulting in an improved approximation denoted $\boldsymbol{V}_t$. Repeating this process eventually converges to the desired $\boldsymbol{V}$:
\begin{equation}
\boldsymbol{V}_t = \QR(\boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{V}_{t-1})
\end{equation}
With $\boldsymbol{V}$ obtained, we clearly have $\boldsymbol{U} = \ColNorm(\boldsymbol{M}\boldsymbol{V})$, where $\ColNorm$ denotes L2 normalization along columns (axis=0), and $\boldsymbol{\Sigma}=\diag(\boldsymbol{U}^{\top}\boldsymbol{M}\boldsymbol{V})$. This gives us an approximate SVD computation scheme based on power iteration and QR decomposition. Of course, when $n > m$, this 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 sufficient for practical purposes.

\section{Streaming Update \label{sec:streaming-update}}

However, computing SVD via power iteration is extremely inefficient in practice—far slower than calling built-in SVD functions from deep learning frameworks—making it impractical. However, considering that training itself is a long iterative process, we can assume that $\boldsymbol{V}$ changes only slightly at each step. Thus, we can store the previous step's $\boldsymbol{V}$ and use it as initialization for the current step, performing only a single power iteration per update:
\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}$. Empirical results show that this streaming power iteration implementation of Muon produces loss convergence curves nearly identical to those of the Newton-Schulz version on language modeling tasks, confirming its feasibility. This is largely due to the momentum mechanism and small learning rates, which make the assumption of minor changes in $\boldsymbol{V}$ per step approximately valid, thereby allowing the power iteration cost to be “amortized” over each step.

Thanks to the direct approximation of SVD, we can now manipulate singular values and incorporate these 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 \href{https://en.wikipedia.org/wiki/Schatten_norm}{Schatten norms}. In short, having explicit (even if approximate) results for $\boldsymbol{U}_t, \boldsymbol{\Sigma}_t, \boldsymbol{V}_t$ significantly increases flexibility and enables more creative experimentation.

\section{Accelerated Decomposition \label{sec:accel-decomposition}}

Now the bottleneck shifts to QR decomposition. In equation~\eqref{eq:muon-qr}, the most time-consuming step is QR decomposition. Standard implementations use Householder QR, which, although faster than full SVD, is still slower than computing $\msign$ via Newton-Schulz iteration (which benefits from polynomial iterations and BF16 multiplication, giving it an almost unfair advantage). To make this new approach competitive, we need to accelerate QR decomposition.

Given a 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 an upper-triangular matrix $\boldsymbol{R}\in\mathbb{R}^{m\times m}$ such that $\boldsymbol{A}=\boldsymbol{Q}\boldsymbol{R}$ (here, “orthogonal” means $\boldsymbol{Q}^{\top}\boldsymbol{Q}=\boldsymbol{I}$; more precisely, $\boldsymbol{Q}$ is a Stiefel matrix). Noting that $\boldsymbol{A}^{\top}\boldsymbol{A}=\boldsymbol{R}^{\top}\boldsymbol{R}$, we only need to decompose $\boldsymbol{A}^{\top}\boldsymbol{A}$ into the product of a lower triangular matrix and its transpose to obtain $\boldsymbol{R}$. This is precisely what the \href{https://en.wikipedia.org/wiki/Cholesky_decomposition}{Cholesky decomposition} achieves.

Cholesky decomposition is highly efficient, so we can first use it to compute $\boldsymbol{R}$, then solve $\boldsymbol{Q}\boldsymbol{R}=\boldsymbol{A}$ for $\boldsymbol{Q}$. Rewriting this as $\boldsymbol{R}^{\top}\boldsymbol{Q}^{\top}=\boldsymbol{A}^{\top}$, we can solve it using \href{https://docs.jax.dev/en/latest/_autosummary/jax.scipy.linalg.solve_triangular.html}{solve\_triangular}, which is also very efficient. Combining these two steps yields an algorithm known as “Cholesky QR,” which, neglecting numerical stability, might be the fastest QR decomposition method available.

Unfortunately, Cholesky QR is much less stable than standard QR decomposition, being extremely sensitive to the \href{https://en.wikipedia.org/wiki/Condition_number}{condition number} of $\boldsymbol{A}^{\top}\boldsymbol{A}$. To mitigate this, \href{https://papers.cool/arxiv/1809.11085}{\textit{Shifted CholeskyQR for computing the QR factorization of ill-conditioned matrices}} suggests adding $\lambda \boldsymbol{I}$ to $\boldsymbol{A}^{\top}\boldsymbol{A}$, where $\lambda=\epsilon \|\boldsymbol{A}^{\top}\boldsymbol{A}\|_F$. However, this is a double-edged sword: larger $\epsilon$ makes Cholesky QR more stable, but results in a less orthogonal $\boldsymbol{Q}$, degrading final performance.

Furthermore, even with $\epsilon$, we cannot guarantee success, necessitating a fallback to standard QR in case of failure.

\section{Reference Implementation \label{sec:implementation}}

A simple reference implementation in JAX is provided below:

\begin{lstlisting}[language=Python, basicstyle=\ttfamily\footnotesize, numbers=left, numberstyle=\tiny, frame=single, breaklines=true]
import jax.numpy as jnp
from jax.scipy.linalg import solve_triangular
from jax import lax

def cholesky_qr(A, eps=1e-9):
    """Compute QR via Cholesky; fall back to default QR if failed
    """
    B = A.mT @ A
    B_norm = (B**2).sum(axis=(-1, -2), keepdims=True)**0.5
    B = B + eps * B_norm * jnp.eye(A.shape[-1])
    R = jnp.linalg.cholesky(B, 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])
\end{lstlisting}

Simple benchmarks show that when Cholesky QR succeeds, its efficiency is comparable to that of the Newton-Schulz $\msign$ implementation. However, to maintain accuracy, $\epsilon$ must not be too small; empirically, $\epsilon=10^{-9}$ is typically required. Even then, Cholesky QR frequently falls back to standard QR, making it slightly slower overall than Newton-Schulz iteration.

Besides improving the QR algorithm, other acceleration techniques exist—e.g., retaining only the first $k$ singular vectors reduces $\boldsymbol{V}$ from $m\times m$ to $m\times k$, lowering computational cost. Further exploration of QR acceleration is left as future work.

\section{Additional Details \label{sec:details}}

Several subtleties deserve special attention, as they closely impact training stability and final performance.

First, per the convention $\boldsymbol{M}_t\in\mathbb{R}^{n\times m}$, we must ensure $n\geq m$; otherwise, transpose the matrix. If $n < m$, $\boldsymbol{M}_t^{\top}\boldsymbol{M}_t \boldsymbol{V}_{t-1}$ is necessarily rank-deficient, and QR decomposition on such matrices is ill-posed. This is especially problematic for Cholesky QR, which is prone to pathological behavior, ultimately degrading performance. Ensuring $n\geq m$ thus improves numerical stability, accelerates training, and enhances results—a win-win-win.

Second, empirical testing reveals that adding $\ColNorm$ before QR significantly improves results:
\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 changing $\tilde{\boldsymbol{v}}_k^{(t)} = \boldsymbol{M}^{\top}\boldsymbol{M}\boldsymbol{v}_k^{(t-1)}$ in equation~\eqref{eq:vk-pi} to $\tilde{\boldsymbol{v}}_k^{(t)} = \boldsymbol{M}^{\top}(\boldsymbol{M}\boldsymbol{v}_k^{(t-1)} / \|\boldsymbol{M}\boldsymbol{v}_k^{(t-1)}\|_2)$, without affecting the convergence of power iteration. Experimentally, this additional $\ColNorm$ step clearly improves training performance, particularly under Cholesky QR, significantly narrowing the performance gap with standard QR.

The underlying mechanism for this improvement remains unclear. Empirical observations suggest that this step slightly increases the likelihood of Cholesky QR falling back to standard QR, and since it also benefits standard QR, it may not be specific to Cholesky QR alone. It is speculated that this step accelerates the convergence of power iteration itself.

\section{Related Work \label{sec:related-work}}

As noted at the beginning, the idea of streaming power iteration has already appeared in several optimizer-related works, such as \href{https://papers.cool/arxiv/2405.18144}{\textit{4-bit Shampoo for Memory-Efficient Network Training}}, \href{https://papers.cool/arxiv/2409.11321}{\textit{SOAP: Improving and Stabilizing Shampoo using Adam}}, \href{https://papers.cool/arxiv/2502.17410}{\textit{COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs}}, and \href{https://papers.cool/arxiv/2504.05295}{\textit{Dion: Distributed Orthonormalized Updates}}.

Indeed, adapting a multi-step iterative algorithm into a single-step streaming version—leveraging the fact that model training involves long-term updates to amortize cost—is not an overly complex idea. We have previously explored similar ideas in \href{https://kexue.fm/archives/11388}{\textit{Steepest Descent on Manifold: 5. Dual Gradient Descent}}. Thus, the existence of many related works is not surprising.

The primary reference for this article is the recent paper \href{https://papers.cool/arxiv/2602.09006}{\textit{ARO: A New Lens On Matrix Optimization For Large Models}}, which essentially contains most of the content presented here and even extends it further. The extension is particularly insightful: notice that Muon's key update term can be written as
\begin{equation}
\ColNorm(\boldsymbol{M}_t\boldsymbol{V}_t)\boldsymbol{V}_t^{\top}
\end{equation}
Here, $\ColNorm(\boldsymbol{M}_t)$ can be viewed as a base optimizer—applying simple column normalization to momentum, which is not highly competitive. By finding a new orthogonal basis $\boldsymbol{V}_t$, applying the base optimizer in this new basis, and transforming back, we get Muon, which is much stronger than direct $\ColNorm$. This naturally leads to the next question: can $\ColNorm$ be replaced with other base optimizers?

In response, ARO proposes a general optimizer framework—Rotated 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 any matrix function, and $\boldsymbol{V}$ is now denoted $\boldsymbol{R}$ (for Rotation). We will discuss Rotated Steepest Descent in detail in a future article.

\section{Summary \label{sec:summary}}

This article introduced an approach to implement Muon by computing SVD via Streaming Power Iteration. Each step requires only one QR decomposition, offering greater flexibility and extensibility compared to the standard Newton-Schulz iteration-based implementation.

\vspace{30pt}
\noindent\color{orange}
\textit{\textbf{When reproducing this article, please include the following link:} \url{https://kexue.fm/archives/11654}} \\
\textit{\textbf{For further details on reproduction, see:} \url{https://kexue.fm/archives/6508}}

\end{document}
