% Source: https://kexue.fm/archives/11719
\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}

\geometry{a4paper, margin=1in}
\hypersetup{colorlinks=true, linkcolor=blue, citecolor=blue, urlcolor=blue}

\title{Muon Implementation Based on Streaming Power Iteration: 5. Extensions}
\author{Su Jianlin}
\date{2026-04-17}

\begin{document}

\maketitle

\newcommand{\orth}{\mathop{\mathrm{orth}}}
\newcommand{\argmin}{\mathop{\mathrm{argmin}}}
\newcommand{\msign}{\mathop{\mathrm{msign}}}
\newcommand{\mclip}{\mathop{\mathrm{mclip}}}
\newcommand{\SVDone}{\mathop{\mathrm{SVD1}}}

The central theme of this series is ``streaming power iteration''. As the name suggests, it combines two components: ``streaming'' and ``power iteration''. Here, ``power iteration'' refers to a classical multi-step iterative method for computing the singular value decomposition (SVD) of a matrix, while ``streaming'' means distributing the originally multi-step algorithm across each training step, making the computational cost acceptable. The core idea is: instead of performing a complex calculation in one go, continuously approximate the target during training.

As an extension of this series, this article introduces several other applications of the ``streaming'' idea, further illustrating how computationally expensive operations can be cleverly integrated into the training process through streaming transformations.

\section{Orthogonal Projection}\label{sec:orthogonal-projection}

In certain scenarios, we may want to enforce orthogonality constraints on specific parameter matrices. Orthogonal matrices offer good numerical stability, avoiding issues such as numerical explosion or vanishing, and can provide better theoretical guarantees in some designs. Of course, determining where it is appropriate to impose orthogonality constraints requires case-by-case analysis, which we will not elaborate on here.

In the articles \href{https://kexue.fm/archives/11215}{``Steepest Descent on Manifolds: 2. Muon + Orthogonal''} and \href{https://kexue.fm/archives/11221}{``Steepest Descent on Manifolds: 3. Muon + Stiefel''} we have explored the orthogonal (Stiefel) manifold, although that approach aimed to derive new update rules by combining steepest descent, making it more complex. Here we consider a simpler approach: reprojecting (retracting) parameters back onto the orthogonal manifold after each update step.

Without loss of generality, suppose a parameter matrix $\boldsymbol{W}\in\mathbb{R}^{n\times m}$ with $n \geq m$; the operation we need can be written as
\begin{equation}
\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}
where $\orth$ is defined as
\begin{equation}
\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 closest to $\boldsymbol{W}$. This operation has already appeared in our first blog post on Muon \href{https://kexue.fm/archives/10592}{``Appreciating the Muon Optimizer: From Vector to Matrix Essentials''}, and it is essentially the $\msign$ operation in Muon. Thus, the update can alternatively be written as
\begin{equation}
\boldsymbol{W}_t = \msign(\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t)
\label{eq:msign-u}
\end{equation}

\section{Streaming Iteration}\label{sec:streaming-iteration}

In other words, after each update step, we perform a $\msign$ operation to project the parameters back onto the orthogonal manifold. However, computing $\msign$ is neither cheap nor negligible. While we know that $\msign$ is central to Muon, Muon runs its $\msign$ in BF16, whereas model parameters are stored in FP32, so computing $\msign$ on parameters in FP32 incurs a significant cost.

Efficient computation of $\msign$ is based on Newton-Schulz iteration, which we have discussed in detail in \href{https://kexue.fm/archives/10922}{``Newton-Schulz Iteration for msign Operators (Part 1)''} and \href{https://kexue.fm/archives/10996}{``Newton-Schulz Iteration for msign Operators (Part 2)''}. Different Newton-Schulz iterations use different polynomials with varying degrees and coefficients, resulting in different convergence speeds. A classical third-order format 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}
which provably satisfies $\lim_{t\to\infty} \boldsymbol{X}_t = \msign(\boldsymbol{W})$. Although classic, this iteration converges slowly; in Muon we typically use faster fifth-order versions. Regardless, full iteration incurs significant cost.

However, in the context of parameter updates, is it truly necessary to perform the complete iteration at every step? Assuming $\boldsymbol{W}_{t-1}$ is already orthogonal or nearly so, and given a small learning rate $\eta$, the updated $\boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t$ will not deviate significantly from orthogonality. This is where the streaming idea becomes effective---perhaps one iteration suffices per training step. Consider instead
\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 requires only one additional computation $\tilde{\boldsymbol{W}}_t\tilde{\boldsymbol{W}}_t^{\top}\tilde{\boldsymbol{W}}_t$, which is significantly cheaper than a full Newton-Schulz iteration.

\section{Spectral Constraints}\label{sec:spectral-constraints}

Generally, orthogonal constraints can only be applied to certain special matrices, as they are too restrictive, reducing the matrix's degrees of freedom by half---effectively halving the number of free parameters. Often, one might prefer a looser spectral constraint, such as the singular value clipping mentioned in \href{https://kexue.fm/archives/10795}{``Higher-Order MuP: A More Concise and More Refined Spectral Conditioning Scaling''}.

If $\msign$ sets all singular values of a matrix to 1, then singular value clipping only reduces singular values greater than 1 to 1, leaving others unchanged. In \href{https://kexue.fm/archives/11006}{``Computing Singular Value Clipping mclip via msign (Part 1)''} and \href{https://kexue.fm/archives/11059}{``Computing Singular Value Clipping mclip via msign (Part 2)''}, this is referred to as $\mclip$. If what we want is to clip singular values to lie within 1 after each parameter update, we write
\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 expensive than $\msign$. Previous work explored using $\msign$ to compute $\mclip$, for example via 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 $\msign$ computations and is thus quite costly. Alternatively, one could run an additional streaming power iteration on $\boldsymbol{W}$ to compute $\mclip$ via SVD, but that introduces another cache variable and becomes cumbersome.

\section{Sequential Clipping}\label{sec:sequential-clipping}

Let us reinterpret $\mclip$ from a different angle: $\mclip$ converts all singular values greater than 1 to 1, so a necessary step is to clip the largest singular value to 1 (if greater than 1). After clipping, if other singular values exceed 1, the largest remaining becomes the new principal singular value. This suggests that repeatedly applying ``clipping the principal singular value to 1'' suffices to achieve $\mclip$.

The advantage of this ``sequential clipping'' strategy is that computing only the principal singular value and corresponding vectors (denote this operation as $\SVDone$) is far cheaper than a full SVD. It can be efficiently approximated via vectorized power iteration
\begin{equation}
\boldsymbol{v} \leftarrow \frac{\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}}{\Vert\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}\Vert}
\end{equation}
which converges to the right principal singular vector $\boldsymbol{v}_1$ of $\boldsymbol{W}$, from which we obtain $\sigma_1 = \Vert\boldsymbol{W}\boldsymbol{v}_1\Vert$ and $\boldsymbol{u}_1 = \boldsymbol{W}\boldsymbol{v}_1/\sigma_1$. In practice, we use a fixed number of iterations to get an approximation. Then, by maintaining the ``streaming'' philosophy, we perform only a single principal singular value clipping per update:
\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 = \SVDone(\tilde{\boldsymbol{W}}_t),\quad\tilde{\boldsymbol{W}}_t = \boldsymbol{W}_{t-1} - \eta \boldsymbol{\Phi}_t
\end{equation}
Over long-term training, this keeps the singular values of $\boldsymbol{W}$ near 1 (slightly above due to approximation and streaming effects). This ``streaming singular value clipping'' can be seen as a variant of the ``Spectral Weight Decay'' introduced in \href{https://kexue.fm/archives/10648}{``From Spectral Norm Gradients to Novel Weight Decay Ideas''}, and is dubbed ``Spectral Hammer'' in the paper \href{https://papers.cool/arxiv/2507.13338}{``Training Transformers with Enforced Lipschitz Constants''}.

\section{Other Examples}\label{sec:other-examples}

When is the streaming idea applicable? A typical scenario is when we expect a variable to change slowly over long training periods, suggesting that minimal iterations per step may suffice to correct updates. Beyond the two new examples above, we have already employed the ``streaming'' idea in earlier articles---a few are recalled briefly.

In \href{https://kexue.fm/archives/11221}{``Steepest Descent on Manifolds: 3. Muon + Stiefel''} and \href{https://kexue.fm/archives/11241}{``Steepest Descent on Manifolds: 4. Muon + Spectral Spheres''}, solving for steepest descent on corresponding manifolds required solving nonlinear equations. There, we mentioned fixed-point iteration; later in \href{https://kexue.fm/archives/11388}{``Steepest Descent on Manifolds: 5. Dual Gradient Descent''}, we discussed a dual gradient descent solution and proposed distributing the solution process across each training step via streaming.

Similarly, in \href{https://kexue.fm/archives/11619}{``MoE Tours: 6. Optimal Allocation Promotes Balance''}, we viewed MoE's load balancing from an optimal assignment perspective. Solving the dual problem originally required alternating optimization of $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ to convergence at each step. However, since $\boldsymbol{\beta}$ should change only slightly between steps, the streaming idea suggests that updating $\boldsymbol{\alpha}$ and $\boldsymbol{\beta}$ just once per step suffices to achieve good load balancing.

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

This article introduced two additional applications of the ``streaming'' idea: achieving parameter projection onto orthogonal (Stiefel) manifolds with minimal overhead, and clipping parameter spectral norms to stay at or below 1. These examples again illustrate the power of the streaming concept: many operations seemingly requiring one-off, complex computations can be decomposed into incremental adjustments integrated into each training step.

\end{document}
