sMuon: Muon with low-rank adapters

We recently introduced an algorithm for low-rank Muon training in our paper - we give a more concise introduction in this post. For code see here.

tldr

Our method, sMuon, performs competitively with other Muon techniques while remaining fast for larger adapters (see figure and table).

high-level idea

During "standard" full-parameter training, Muon tells us to orthogonalize weight updates (i.e. flatten the singular values). This looks like, $$ \begin{align*} \delta W = -\eta\,\operatorname{msign}(G) = -\eta\,UV^\top, \end{align*} $$ where $U \Sigma V^\top = \mathrm{svd}(G)$, and $\eta$ is the learning rate.

During LoRA training we instead have an adapter $\Delta W = BA$ with gradient $G_{\Delta W}$. Again, Muon tells us to make the following adapter descent step, $$ \begin{align*} \delta \Delta W = -\eta\,\operatorname{msign}(G_{\Delta W}), \end{align*} $$ but this is not possible in practice because we update $B$ and $A$, rather than the adapter as a whole.

In our paper, we use a simple workaround: select updates $\delta B$ and $\delta A$ whose induced adapter change is as close as possible to the ideal update. That is, solve, $$ \begin{align*} \mathrm{argmin}_{\delta B,\,\delta A} \|-\eta\,\operatorname{msign}(G_{\Delta W}) + [(B + \delta B)(A + \delta A) - BA]\|_F^2. \end{align*} $$ This is a least-squares fit to the ideal step, with the ugly-looking $(B + \delta B)(A + \delta A) - BA$ term being the change to the adapter after updating $B$ and $A$.

We make two simplifying steps. First, we linearize; this is easily justified since we expect $\delta B$ and $\delta A$ to be small. Second, we replace $G_{\Delta W}$ with its low-rank projection $H$, $$ H=G_{\Delta W} - (I - P_B)G_{\Delta W}(I - P_A), $$ where $P_B$ and $P_A$ are projection matrices onto the column space of $B$ and row space of $A$ respectively. Remember that the gradient $G_{\Delta W}$ will generally have a much larger rank than the adapter $\Delta W$ itself. We can therefore understand $H$ as retaining components of the gradient that align with either the input or output directions represented by the LoRA adapter.

After making these two approximations, one solution to the least-squares problem is, $$ \begin{align} \delta A &= -\eta\, B^\dagger \, \operatorname{msign}(H)\label{eq:dA},\\ \delta B &= -\eta\, (I - BB^\dagger)\,\operatorname{msign}(H) A^\dagger\label{eq:dB}, \end{align} $$ where $BB^\dagger=P_B$, and $B^\dagger=(B^\top B)^{-1}B^\top$ and $A^\dagger=A^\top(AA^\top)^{-1}$ are Moore-Penrose inverses (assuming $A$ and $B$ have full rank). See the paper for the derivation of this solution (it is straightforward - we find $\delta B$ and $\delta A$ which make the gradient equal to zero).

To summarize, our target adapter descent step is $-\eta \operatorname{msign}(H)$. Using Eqs. \eqref{eq:dA} and \eqref{eq:dB} gives us an adapter step that is very close to the target.

Note: Using $H$ instead of $G_{\Delta W}$ at first glance seems arbitrary and unjustified. It is actually another linearization: Muon can be viewed as minimizing a spectral-regularized gradient descent objective, and swapping the adapter gradient for the projected gradient $H$ gives the same objective to first order. However, the main reason to use $H$ is that it enables an efficient implementation, which we discuss next.

efficient computation

We seek a cheap way to calculate Eqs. $\eqref{eq:dA}$ and $\eqref{eq:dB}$.

Observation 1 ($H$ is low rank). By the chain rule, $G_{B} = G_{\Delta W} A^\top$ and $G_{A} = B^\top G_{\Delta W}$. This allows us to rewrite $H$, $$ \begin{align} H &= G_{\Delta W} - (I - P_B)G_{\Delta W}(I - P_A)\notag\\ &= P_B G_{\Delta W} + (I - P_B) G_{\Delta W} P_A\notag\\ &= B(B^\top B)^{-1}B^\top G_{\Delta W} + (I - P_B) G_{\Delta W} A^\top (AA^\top)^{-1}A\notag\\ &= (B^\dagger)^\top G_{A} + (I - P_B) G_B (A^\dagger)^{\top}\label{eq:H_low_rank}. \end{align} $$ From Eq. \eqref{eq:H_low_rank}, $H$ is the sum of two rank $r$ matrices, thus $H$ has rank at most $2r.$

Observation 2 (change orthogonalization bases). We can transform the row and column bases, perform orthogonalization, then transform back. Mathematically, for (square) orthogonal $U$ and $V$, \begin{align}\label{eq:msign_transform_bases} \operatorname{msign}(H) = U \,\operatorname{msign}(U^\top H V ) V^\top. \end{align} A stronger version also holds when $U$ and $V$ are semi-orthogonal (i.e. not necessarily square): \begin{align} \mathrm{col}(U) \supseteq \mathrm{col}(H),\,\mathrm{col}(V)\supseteq\mathrm{col}(H^\top) \implies \operatorname{msign}(H) = U \,\operatorname{msign}(U^\top H V ) V^\top\label{eq:msign_lr}. \end{align} Assume the adapter is $\Delta W\in\mathbb{R}^{d_1\times d_2}$, then we can use $U\in\mathbb{R}^{d_1\times 2r}$ and $V\in\mathbb{R}^{d_2\times 2r}$, since we know that $H$ has rank at most $2r$. This allows us to trade an $\operatorname{msign}$ on $H$, a $d_1\times d_2$ matrix for an $\operatorname{msign}$ on $U^\top HV$, a $2r\times 2r$ matrix.

Finding appropriate $U$ and $V$. We still need bases for the column and row spaces of $H$. We can do this by looking at $H$, $$ \begin{align*} H= B \underbrace{(B^\top B)^{-1} G_{A}}_{:=X} + \underbrace{(I - P_B) G_B (AA^\top)^{-1}}_{:= Y} A^{\top}. \end{align*} $$ It is enough if the columns of $U$ span $\operatorname{col}(B) + \operatorname{col}(Y)$, and the columns of $V$ span $\operatorname{col}(A^\top) + \operatorname{col}(X^\top)$! All that remains is to ensure that $U$ and $V$ are semi-orthogonal, which is achieved by, $$ \begin{aligned} U &= [ U_1 \quad U_2]\\ V &= [ V_1 \quad V_2] \end{aligned} $$ where, $$ \begin{aligned} U_1&=\operatorname{msign}(B),&U_2&=\operatorname{msign}(Y),\\ V_1&=\operatorname{msign}(A^{\top}),& V_2&=\operatorname{msign}((I-P_A)X^\top). \end{aligned} $$ Here, the $\operatorname{msign}$ operations orthogonalize the columns of each submatrix ($U/V_{1/2}$), and the $(I-P_A)$ term ensures that the columns of $V_2$ are orthogonal to the columns of $V_1$.

Note 1: in practice we avoid instantiating large matrices such as $P_A$ and $P_B$ directly. For example, if we must compute $P_B C$, then we can compute this as two matmuls, $P_B C = B(B^\top B)^{-1}B^\top C = (B(B^\top B)^{-1})(B^\top C)$.

Note 2: in the above derivation we need to assume $B$ and $A$ both have full rank (so that $B^\top B$ and $AA^\top$ are non-singular). In the implementation, we compute inverse powers of $B^\top B + \epsilon I$ and $AA^\top + \epsilon I$ to avoid issues and numerical instability.

Summary. Finally, substituting the expression for $\operatorname{msign}(H)$ into the formulae for $\delta B$ and $\delta A$, after a bit of algebra, gives an efficient low-rank optimizer (see the paper for the full algorithm). The crucial point is that the matrix sign of the large matrix $H$ can be replaced by four smaller matrix-sign computations, allowing for a more memory efficient computation than naively calculating $\operatorname{msign}(H)$.

You may have noticed that we still need to compute terms like $(B^\top B)^{-1/2}$ and $(AA^\top)^{-1/2}$. We can compute these inverse square roots and the $\operatorname{msign}$ operations using Newton–Schulz-style iterations. Since these iterations use matmuls, the resulting optimizer can be implemented using matmuls alone. Thus we avoid complex linear algebra routines. Avoiding QR and SVD decompositions is what enables sMuon to be much faster than Riemannion at larger $r$.

Citation

If you want to cite this work, please cite the original paper,

@article{anson2026approximate,
  title={Approximate Muon with low-rank adapters},
  author={Anson, Ben and Houghton, Conor and Milsom, Edward},
  journal={arXiv preprint arXiv:2608.14492},
  year={2026}
}