Streaming Power Iteration-based Muon Implementation: 4. Principles
Having gone through the three previous articles, Streaming Power Iteration-based Muon Implementation: 1. First Encounter, Streaming Power Iteration-based Muon Implementation: 2. Acceleration, and Streaming Power Iteration-based Muon Implementation: 3. Refinement, you should now have a fairly good grasp of the idea, implementation, and acceleration details of Streaming Power Iteration. All told, this is a genuinely competitive way of implementing Muon, and since it directly approximates the SVD computation, it also enjoys better extensibility.
Due to space constraints, we described the underlying mathematics of the relevant operations rather sketchily in those posts. So in this article we fill in some of the mathematical derivations behind power iteration and QR decomposition, to build a more complete theoretical picture. That said, the derivations here still emphasize intuition over rigor — the main goal is to help readers (including the author) get their thinking straight, so professional readers, please bear with us.
Coaxial Equivalence
Before diving into the derivation, we need to introduce the notion of "coaxial equivalence." For a matrix $\boldsymbol{A},\boldsymbol{B}\in\mathbb{R}^{n\times m}$, if there exists a signature matrix $\boldsymbol{S}$ satisfying $\boldsymbol{A} = \boldsymbol{B}\boldsymbol{S}$, then we say $\boldsymbol{A}$ and $\boldsymbol{B}$ are "coaxial equivalent," and they are each other's "coaxial matrix." Here, a "signature matrix" refers to a diagonal matrix whose diagonal entries are $\pm 1$, i.e., $\newcommand{diag}{\mathop{\text{diag}}}\diag(\pm 1, \pm 1, \cdots, \pm 1)$.
It's worth noting that "coaxial" is a term the author coined himself to describe this equivalence relation. From the perspective of coordinate systems, matrices satisfying this condition $\boldsymbol{A},\boldsymbol{B}$ actually describe the same coordinate system, just with different choices of positive direction for some of the axes. Some literature seems to call this "sign equivalence," but the author feels "coaxial" is more intuitive.
The reason for introducing the concept of coaxiality is that many matrix decompositions are only unique up to coaxial equivalence (which makes the author wonder why such a commonly used equivalence relation doesn't have a standard name). Take QR decomposition, for example: let matrix $\boldsymbol{A}\in\mathbb{R}^{n\times m}(n\geq m)$ be full rank, and let $\boldsymbol{Q}_1\boldsymbol{R}_1$ and $\boldsymbol{Q}_2\boldsymbol{R}_2$ be two of its QR decompositions. Then $\boldsymbol{Q}_1$ and $\boldsymbol{Q}_2$ are coaxial, and $\boldsymbol{R}_1$ and $\boldsymbol{R}_2$ are coaxial.
This uniqueness is not hard to understand: during Gram-Schmidt orthogonalization, we can freely flip the sign of any orthogonalized vector without breaking orthogonality. Most textbooks establish the uniqueness of QR decomposition by requiring the diagonal entries of $\boldsymbol{R}$ to be positive — that's one valid approach, but it can sometimes introduce unnecessary complications. Likewise, the uniqueness of SVD also only holds up to coaxial equivalence.
The Iteration of Powers
In Streaming Power Iteration-based Muon Implementation: 1. First Encounter, we introduced power iteration by solving for eigenvectors one at a time, then switched directly to the parallel version, claiming it converges to the same result. Here, we'll start directly from the parallel version and prove its convergence.
Again, let the matrix be $\boldsymbol{A}\in\mathbb{R}^{n\times m}(n\geq m)$, and assume it is full rank (rank $m$). Consider the power iteration
\begin{equation}\newcommand{QR}{\mathop{\text{QR}}}\boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\boldsymbol{A} \boldsymbol{V}_{t-1}),\qquad \boldsymbol{V}_0 = \boldsymbol{I}\end{equation}
Let the QR decomposition of $\boldsymbol{A}^{\top}\boldsymbol{A} \boldsymbol{V}_{t-1}$ be $\boldsymbol{Q}_t \boldsymbol{R}_t$, so that $\QR(\boldsymbol{A}^{\top}\boldsymbol{A} \boldsymbol{V}_{t-1}) = \boldsymbol{A}^{\top}\boldsymbol{A} \boldsymbol{V}_{t-1}\boldsymbol{R}_t^{-1}$. Iterating step by step, we get
\begin{equation}\boldsymbol{V}_t = (\boldsymbol{A}^{\top}\boldsymbol{A})^t \boldsymbol{R}_1^{-1}\boldsymbol{R}_2^{-1} \cdots \boldsymbol{R}_t^{-1}\end{equation}
Next, let the SVD of $\boldsymbol{A}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, so that $\boldsymbol{A}^{\top}\boldsymbol{A}=\boldsymbol{V}\boldsymbol{\Sigma}^2\boldsymbol{V}^{\top}$. Substituting this into the above equation gives
\begin{equation}\boldsymbol{V}_t = \boldsymbol{V}\boldsymbol{\Sigma}^{2t}\boldsymbol{V}^{\top} \boldsymbol{R}_1^{-1}\boldsymbol{R}_2^{-1} \cdots \boldsymbol{R}_t^{-1} = \boldsymbol{V}\underbrace{(\boldsymbol{\Sigma}^{2t}\boldsymbol{V}^{\top}\boldsymbol{\Sigma}^{-2t})}_{\boldsymbol{B}}\,\underbrace{(\boldsymbol{\Sigma}^{2t} \boldsymbol{R}_1^{-1}\boldsymbol{R}_2^{-1} \cdots \boldsymbol{R}_t^{-1})}_{\boldsymbol{C}}\end{equation}
The result splits into three parts, $\boldsymbol{V},\boldsymbol{B},\boldsymbol{C}$, and what we want to prove is that, in the coaxial sense, $\boldsymbol{V}_t\to \boldsymbol{V}$. We already know that $\boldsymbol{V}_t,\boldsymbol{V}$ are orthogonal matrices, and by the closure of upper triangular matrices under multiplication, $\boldsymbol{C}$ is itself an upper triangular matrix. If we can show that $\boldsymbol{B}$ tends toward an upper triangular matrix, then by the uniqueness of QR decomposition we can conclude $\boldsymbol{V}_t\to \boldsymbol{V}$. Writing $\boldsymbol{B}=\boldsymbol{\Sigma}^{2t}\boldsymbol{V}^{\top}\boldsymbol{\Sigma}^{-2t}$ out in component form gives
\begin{equation}B_{i,j} = V_{j,i} (\sigma_i / \sigma_j)^{2t}\end{equation}
Assuming $\sigma_1 > \sigma_2 > \cdots > \sigma_m$, then as $i > j$ we have $(\sigma_i / \sigma_j)^{2t}\to 0$, meaning the lower-triangular part of $\boldsymbol{B}$ tends to 0 — in other words, $\boldsymbol{B}$ tends to an upper triangular matrix, which proves the claim. From this we can also see that the convergence rate of power iteration depends on the ratio of adjacent singular values: the larger $\sigma_i / \sigma_{i+1}$ is, the faster the convergence. Once two singular values become equal, power iteration breaks down — but in practice we can assume the probability of two singular values being exactly equal is zero, sidestepping the issue.
Getting to the Essence
Let's think carefully through the whole proof and understand what it's really doing at its core.
First, the block $\boldsymbol{C} = \boldsymbol{\Sigma}^{2t} \boldsymbol{R}_1^{-1}\boldsymbol{R}_2^{-1} \cdots \boldsymbol{R}_t^{-1}$ looks complicated, but its only role is to be "some upper triangular matrix" — what it actually equals doesn't matter at all, as long as it retains the upper-triangular form. What really does the heavy lifting is the fact that $\boldsymbol{B}=\boldsymbol{\Sigma}^{2t}\boldsymbol{V}^{\top}\boldsymbol{\Sigma}^{-2t}$ tends toward an upper triangular matrix, which is what lets us extract $\boldsymbol{V}$ via QR decomposition.
What actually achieves this effect is, in fact, the single step $(\boldsymbol{A}^{\top}\boldsymbol{A})^t$! In other words, in theory we have $\boldsymbol{V}_t = \QR((\boldsymbol{A}^{\top}\boldsymbol{A})^t)$ — meaning that, theoretically, all those repeated applications of $\QR$ earlier on are redundant, and it would suffice to perform $\QR$, i.e., orthogonalization, just once at the very end. Of course, this equivalence is only of theoretical value: computing $(\boldsymbol{A}^{\top}\boldsymbol{A})^t$ directly would blow up or collapse numerically, rendering the result useless. So performing $\QR$ periodically (rather than only once at the end) is essentially about maintaining numerical stability.
We can also understand the operation of repeated QR steps from the perspective of a "preconditioner." A preconditioner refers to some preprocessing applied to the input that, in theory, doesn't change the output, but is generally beneficial for numerical computation. For $\QR$, right-multiplying by any full-rank upper triangular matrix $\boldsymbol{R}$ doesn't change the result, i.e., $\QR(\boldsymbol{A}) = \QR(\boldsymbol{A}\boldsymbol{R})$. So any full-rank upper triangular matrix that we right-multiply by can be called a preconditioner for $\QR$.
We know that $\QR$ itself can also be written as a right-multiplication by an upper triangular matrix, so $\QR$ is itself a preconditioner for $\QR$. This means we can insert some copies of $\QR$ into $(\boldsymbol{A}^{\top}\boldsymbol{A})^t$, turning it into
\begin{equation}(\boldsymbol{A}^{\top}\boldsymbol{A})^t\qquad\to\qquad \boldsymbol{A}^{\top}\boldsymbol{A}\QR(\boldsymbol{A}^{\top}\boldsymbol{A}\QR(\boldsymbol{A}^{\top}\boldsymbol{A}\cdots\QR(\boldsymbol{A}^{\top}\boldsymbol{A})))\end{equation}
Applying $\QR$ to both sides at the end doesn't change the result, but since each $\QR$ step produces an orthogonal matrix, and right-multiplying by an orthogonal matrix carries almost no risk of numerical blow-up, $\QR$ turns out to be its own excellent preconditioner.
Related Variants
From the preconditioner point of view, we can also make sense of a number of power-iteration variants. For instance, the $\newcommand{ColNorm}{\mathop{\text{ColNorm}}}\ColNorm$ version we already used back in the opening post, Streaming Power Iteration-based Muon Implementation: 1. First Encounter:
\begin{equation}\boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\boldsymbol{A}\boldsymbol{V}_{t-1}) \qquad\to\qquad \boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\ColNorm(\boldsymbol{A}\boldsymbol{V}_{t-1}))\end{equation}
The two instances of $\boldsymbol{V}_t$ on the left and right are equal, because applying $\ColNorm$ to a matrix is equivalent to right-multiplying it by a diagonal matrix:
\begin{equation}\underbrace{\left[\frac{\boldsymbol{a}_1}{\Vert\boldsymbol{a}_1\Vert},\frac{\boldsymbol{a}_2}{\Vert\boldsymbol{a}_2\Vert},\cdots,\frac{\boldsymbol{a}_m}{\Vert\boldsymbol{a}_m\Vert}\right]}_{\ColNorm([\boldsymbol{a}_1,\boldsymbol{a}_2,\cdots,\boldsymbol{a}_m])} = [\boldsymbol{a}_1,\boldsymbol{a}_2,\cdots,\boldsymbol{a}_m]\begin{bmatrix}\Vert\boldsymbol{a}_1\Vert^{-1} & 0 & \cdots & 0 \\ 0 & \Vert\boldsymbol{a}_2\Vert^{-1} & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & \Vert\boldsymbol{a}_m\Vert^{-1}\end{bmatrix}\end{equation}
A diagonal matrix is a special case of an upper triangular matrix, so $\ColNorm$ is also a preconditioner for $\QR$, meaning it doesn't change the result. By the same reasoning, we can also shift $\ColNorm$ one step further outward, which likewise doesn't change the result of $\QR$:
\begin{equation}\boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\boldsymbol{A}\boldsymbol{V}_{t-1}) \qquad\to\qquad \boldsymbol{V}_t = \QR(\ColNorm(\boldsymbol{A}^{\top}\boldsymbol{A}\boldsymbol{V}_{t-1}))\end{equation}
Additionally, the star of Streaming Power Iteration-based Muon Implementation: 2. Acceleration and Streaming Power Iteration-based Muon Implementation: 3. Refinement was the doubly $\QR$ power iteration:
\begin{equation}\boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\boldsymbol{A}\boldsymbol{V}_{t-1}) \qquad\to\qquad \boldsymbol{V}_t = \QR(\boldsymbol{A}^{\top}\QR(\boldsymbol{A}\boldsymbol{V}_{t-1}))\end{equation}
Based on the discussion in those two sections, it's easy to show that adding one extra step of $\QR$ to $\boldsymbol{A}\boldsymbol{V}_{t-1}$ likewise doesn't, in theory, change the result of $\boldsymbol{V}_t$.
Finite Error
We've repeatedly used phrases like "in theory," because many of these results only hold under infinite-precision, exact $\QR$. In actual computation, $\QR$ is finite-precision, and for efficiency's sake we often use the less accurate "SCQR (Shifted Cholesky QR)" — so it's important to analyze how errors accumulate.
By now, after reading through this series, SCQR should be familiar: it splits the QR decomposition of a matrix $\boldsymbol{A}$ into two steps: (1) perform a Cholesky decomposition on $\boldsymbol{A}^{\top}\boldsymbol{A}+\lambda \boldsymbol{I}$ to obtain $\boldsymbol{R}$, and (2) obtain the orthogonal matrix $\boldsymbol{Q}$ via $\boldsymbol{Q}=\boldsymbol{A}\boldsymbol{R}^{-1}$. If $\lambda=0$, then SCQR is equivalent to exact QR; however, in practice, $\lambda=0$ almost always fails due to an excessively large condition number, so we need to set an appropriate $\lambda > 0$, which introduces error.
Fortunately, this error does not accumulate! Whether it's $\boldsymbol{A}^{\top}\boldsymbol{A}$ or $\boldsymbol{A}^{\top}\boldsymbol{A}+\lambda \boldsymbol{I}$, the result of their Cholesky decomposition is always an upper triangular matrix, and $\boldsymbol{A}$ ends up being right-multiplied by the inverse of this triangular matrix. And we already know that right-multiplying by an upper triangular matrix doesn't change the result of $\QR$. Therefore, ignoring the error introduced by matrix multiplication itself, we can say that SCQR is "lossless" with respect to $\QR$ itself.
In other words, as long as the very last step uses an exact $\QR$, then no matter how many inexact SCQR steps came before it, the final result will still be exact. Put differently, the total error is equivalent to the error of a single SCQR step (the last one), rather than accumulating over iterations. The "SCQR2" acceleration trick mentioned in Streaming Power Iteration-based Muon Implementation: 2. Acceleration also relies on this same principle.
Conversely, if some transformation we apply cannot be written in the form of "right-multiplication by an upper triangular matrix," then it will lossily corrupt the information in the original matrix, and errors will accumulate over the course of long iterations until the method completely fails. The simplification proposed by @Ji_Ha_Kim, which we discussed in the previous post, Streaming Power Iteration-based Muon Implementation: 3. Refinement, is a classic counterexample of exactly this kind.
Cholesky Decomposition
Finally, let's briefly review "Cholesky decomposition." Its goal is to express a given positive-definite symmetric matrix $\boldsymbol{B}$ in the form $\boldsymbol{L}\boldsymbol{L}^{\top}$, where $\boldsymbol{L}$ is a lower triangular matrix. If we then let $\boldsymbol{R}=\boldsymbol{L}^{\top}$, we obtain the upper-triangular form $\boldsymbol{R}^{\top}\boldsymbol{R}$.
Cholesky decomposition is fast because it admits an analytical solution that can be computed directly via recursion. Specifically, starting from the equality
\begin{equation}\begin{bmatrix} a_{1,1} & a_{1,2} & \cdots & a_{1,m} \\ a_{2,1} & a_{2,2} & \cdots & a_{2,m} \\ \vdots & \vdots & \ddots & \vdots \\ a_{m,1} & a_{m,2} & \cdots & a_{m,m} \end{bmatrix}=\begin{bmatrix} l_{1,1} & & & \\ l_{2,1} & l_{2,2} & & \\ \vdots & \vdots & \ddots & \\ l_{m,1} & l_{m,2} & \cdots & l_{m,m} \end{bmatrix}\begin{bmatrix} l_{1,1} & l_{2,1} & \cdots & l_{m,1} \\ & l_{2,2} & \cdots & l_{m,2} \\ & & \ddots & \vdots \\ & & & l_{m,m} \end{bmatrix}\end{equation}
and matching terms one by one, we can derive the recursive formula
\begin{equation}l_{j,j}=\pm\sqrt {a_{j,j}-\sum_{k=1}^{j-1}l_{j,k}^2},\qquad l_{i,j}=\frac{1}{l_{j,j}}\left(a_{i,j}-\sum_{k=1}^{j-1}l_{i,k}l_{j,k}\right)\quad \text{for}i > j\end{equation}
Cholesky QR is one of the classic applications of Cholesky decomposition. It's similar to Gram-Schmidt orthogonalization, but lets us skip the orthogonalization process and go directly to the triangular matrix $\boldsymbol{R}$. They also share the same numerical pitfalls: Gram-Schmidt orthogonalizes vectors one at a time, and once it encounters linearly dependent or zero vectors, the process breaks down — which, in terms of singular values, shows up as a low effective rank or a large condition number. Cholesky decomposition is just as prone to failure under those same conditions.
Another classic application of Cholesky decomposition is inverting positive-definite symmetric matrices. For instance, to solve the equation $\boldsymbol{B}\boldsymbol{X}=\boldsymbol{C}$, where $\boldsymbol{B}$ is positive-definite symmetric, we can first decompose it as $\boldsymbol{L}\boldsymbol{L}^{\top}$, and then rewrite the problem as $\boldsymbol{L}(\boldsymbol{L}^{\top}\boldsymbol{X})=\boldsymbol{C}$, which only requires solving two triangular systems — each step being fairly efficient. In particular, the fractional iteration for computing $\newcommand{msign}{\mathop{\text{msign}}}\msign$ proposed by @Ji_Ha_Kim finds the inverse matrix based on exactly this idea.
Summary
This article mainly filled in some of the mathematical derivations behind streaming power iteration, including the convergence of power iteration, the preconditioner view of QR decomposition, and a brief introduction to Cholesky decomposition. Hopefully this helps you understand the method more thoroughly at the level of first principles.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.