The Derivative of msign

msign的导数 苏剑林 · 2025-06-13 · Mathematics calculus matrix gradient muon

In this post we derive the differentiation formula for the $\newcommand{msign}{\mathop{\text{msign}}}\msign$ operator. If you'd like to combine TTT with Muon, the way Test-Time Training Done Right does, this article might be of help.

Two Definitions

This post again assumes readers already have some familiarity with $\msign$; if not, you might first check out Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices and The Newton-Schulz Iteration for the msign Operator (Part 1). Now suppose we have a matrix $\boldsymbol{M}\in\mathbb{R}^{n\times m}$, then

\begin{equation}\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}^{\top} = \text{SVD}(\boldsymbol{M}) \quad\Rightarrow\quad \msign(\boldsymbol{M}) = \boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}\end{equation}

where $\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m},\boldsymbol{V}\in\mathbb{R}^{m\times m}$, and $r$ is the rank of $\boldsymbol{M}$. In simple terms, $\msign$ is the matrix obtained by turning all nonzero singular values of the original matrix into 1. Building on the SVD, we can further show that

\begin{equation}\msign(\boldsymbol{M}) = (\boldsymbol{M}\boldsymbol{M}^{\top})^{-1/2}\boldsymbol{M}= \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/2}\end{equation}

Here the inverse of the $1/2$-th power of the matrix $^{-1/2}$ — since $\boldsymbol{M}\boldsymbol{M}^{\top}$ and $\boldsymbol{M}^{\top}\boldsymbol{M}$ are (semi-)positive-definite symmetric, the $1/2$-th power always exists, but the inverse may not, in which case we resort to the "pseudo-inverse." The name $\msign$ comes from the resemblance of the above expression to the real-valued sign function $\newcommand{sign}{\mathop{\text{sign}}}\sign(x) = x/\sqrt{x^2}$. However, as we've mentioned before, the sign function also has another matrix version, which we'll call $\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\newcommand{csgn}{\mathop{\text{csgn}}}\mcsgn$ here:

\begin{equation}\mcsgn(\boldsymbol{M}) = \boldsymbol{M}(\boldsymbol{M}^2)^{-1/2}\end{equation}

That is, $\boldsymbol{M}^{\top}\boldsymbol{M}$ in $\msign$ is replaced by $\boldsymbol{M}^2$. Since only square matrices can be squared, this definition applies only to square matrices. Introducing two similar-but-different definitions within a single article is prone to causing confusion, but unfortunately, both definitions are needed in the calculations that follow, so we have no choice but to let them coexist.

$\mcsgn$ has similarity invariance: if $\boldsymbol{M}=\boldsymbol{P}\boldsymbol{\Lambda}\boldsymbol{P}^{-1}$, then $\mcsgn(\boldsymbol{M})=\boldsymbol{P}\mcsgn(\boldsymbol{\Lambda})\boldsymbol{P}^{-1}$. Furthermore, if $\boldsymbol{\Lambda}$ is diagonalizable (which is almost always achievable over the complex field), then we have

\begin{equation}\mcsgn(\boldsymbol{M}) = \boldsymbol{P}\csgn(\boldsymbol{\Lambda})\boldsymbol{P}^{-1}\end{equation}

$\csgn(\boldsymbol{\Lambda})$ means every diagonal entry is replaced by $\csgn$, where $\csgn(z) = z/\sqrt{z^2}$ is the complex version of the sign function: if the real part of $z$ is nonzero, it equals $\sign(\mathop{\text{Re}}[z])$. Seen this way, the difference between $\msign$ and $\mcsgn$ is that the former applies the sign function to the singular values from the SVD, while the latter applies the sign function to the eigenvalues from the eigendecomposition. When $\boldsymbol{M}$ is a symmetric matrix, the two coincide.

The Same Computation

Currently, numerical computation of $\msign$ mainly relies on a "Newton-Schulz iteration" of the following form:

\begin{equation}\boldsymbol{X}_0 = \frac{\boldsymbol{M}}{\Vert\boldsymbol{M}\Vert_F},\qquad \boldsymbol{X}_{t+1} = a_{t+1}\boldsymbol{X}_t + b_{t+1}\boldsymbol{X}_t(\boldsymbol{X}_t^{\top}\boldsymbol{X}_t) + c_{t+1}\boldsymbol{X}_t(\boldsymbol{X}_t^{\top}\boldsymbol{X}_t)^2\end{equation}

As for the choice of coefficients, we've already discussed this in detail in The Newton-Schulz Iteration for the msign Operator (Part 1) and The Newton-Schulz Iteration for the msign Operator (Part 2). The more novel result from the latter is:

$$\begin{array}{c|ccc} \hline t & a\times 1.01 & b\times 1.01^3 & c\times 1.01^5 \\ \hline \quad 1\quad & 8.28721 & -23.5959 & 17.3004 \\ 2 & 4.10706 & -2.94785 & 0.544843 \\ 3 & 3.94869 & -2.9089 & 0.551819 \\ 4 & 3.31842 & -2.48849 & 0.510049 \\ 5 & 2.30065 & -1.6689 & 0.418807 \\ 6 & 1.8913 & 1.268 & 0.376804 \\ 7 & 1.875 & -1.25 & 0.375 \\ 8 & 1.875 & -1.25 & 0.375 \\ \hline \end{array}$$

The nice thing about this result is that it can be truncated or stacked arbitrarily: keeping only the first 5 rows gives the optimal 5-step iteration, keeping the first 6 rows gives the optimal 6-step iteration, and the approximation quality is guaranteed to be no worse than that of the 5-step iteration, and so on.

As for $\mcsgn$, it is simply $\msign$ with $\boldsymbol{M}^{\top}\boldsymbol{M}$ replaced by $\boldsymbol{M}^2$, so in principle it too could be computed via a Newton-Schulz iteration. But since eigenvalues can be complex, convergence in general is much harder to achieve. However, if we can confirm that the eigenvalues of the matrix $\boldsymbol{M}$ are all real (as is the case for the block-triangular matrices to which we'll apply $\mcsgn$ later in this post), then we can reuse the iteration and coefficients of $\msign$:

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\boldsymbol{X}_0 = \frac{\boldsymbol{M}}{\sqrt{\tr(\boldsymbol{M}^2)}},\qquad \boldsymbol{X}_{t+1} = a_{t+1}\boldsymbol{X}_t + b_{t+1}\boldsymbol{X}_t^3 + c_{t+1}\boldsymbol{X}_t^5\end{equation}

The Derivation

Now let's get to the main topic: finding the derivative of $\boldsymbol{O}=\msign(\boldsymbol{M})$. If you're only using Muon as an ordinary optimizer, then this post is probably not relevant to you. It's when we want to follow TTT and use the Muon optimizer to build an RNN model that we need the derivative of $\msign$: in that setting $\msign$ appears in the model's forward pass, and backpropagating through the whole model naturally requires the derivative of $\msign$.

Since $\msign$ is computed via a Newton-Schulz iteration, it can in fact be backpropagated through directly, so numerically differentiating $\msign$ is not itself a problem. But backpropagating through the iteration means storing a lot of intermediate states, which tends to blow up GPU memory — hence the desire for an analytical solution to simplify things. On the other hand, in The Derivative of SVD we already worked out the derivative of $\msign$, but that was based on the SVD expression, and SVD is not a GPU-efficient algorithm.

So our goal here is to find a result that doesn't rely on SVD and that can be computed efficiently. We start from the identity

\begin{equation}\boldsymbol{M} = \boldsymbol{O}\boldsymbol{M}^{\top}\boldsymbol{O}\end{equation}

(which follows directly from the definition of $\msign$), and differentiate both sides to get

\begin{equation}d\boldsymbol{M} = (d\boldsymbol{O})\boldsymbol{M}^{\top}\boldsymbol{O} + \boldsymbol{O}(d\boldsymbol{M}^{\top})\boldsymbol{O} + \boldsymbol{O}\boldsymbol{M}^{\top}(d\boldsymbol{O})\label{eq:dm-do}\end{equation}

The difficulty with this result is that we cannot simply isolate a term of the form $d\boldsymbol{M}=f(d\boldsymbol{O})$ or $d\boldsymbol{O}=f(d\boldsymbol{M})$, so it's not easy to see the relationship between $\nabla_{\boldsymbol{O}}\mathcal{L}$ and $\nabla_{\boldsymbol{W}}\mathcal{L}$ ($\mathcal{L}$ being the loss function). In this situation the best approach is to return to the fundamental idea behind matrix differentiation — the "trace trick":

Trace trick If we can find a matrix $\boldsymbol{G}$ of the same shape as $\boldsymbol{M}$ satisfying
\begin{equation}d\mathcal{L}=\langle \boldsymbol{G}, d\boldsymbol{M}\rangle_F = \tr(\boldsymbol{G}^{\top} (d\boldsymbol{M}))\end{equation}
then $\boldsymbol{G} = \nabla_{\boldsymbol{M}}\mathcal{L}$.

The essence of the trace trick is to reduce a matrix/vector expression to a scalar, then reduce the scalar to a trace, after which we can exploit trace identities:

\begin{equation}\tr(\boldsymbol{A}\boldsymbol{B}) = \tr(\boldsymbol{B}\boldsymbol{A}) = \tr(\boldsymbol{A}^{\top}\boldsymbol{B}^{\top}) = \tr(\boldsymbol{B}^{\top}\boldsymbol{A}^{\top})\end{equation}

Now let $\boldsymbol{X}$ be an arbitrary matrix of the same shape as $\boldsymbol{M}$. Multiplying both sides of equation $\eqref{eq:dm-do}$ by $\boldsymbol{X}^{\top}$ and taking the trace gives

\begin{equation}\begin{aligned} \tr(\boldsymbol{X}^{\top}(d\boldsymbol{M})) =&\, \tr(\boldsymbol{X}^{\top}(d\boldsymbol{O})\boldsymbol{M}^{\top}\boldsymbol{O}) + \tr(\boldsymbol{X}^{\top}\boldsymbol{O}(d\boldsymbol{M}^{\top})\boldsymbol{O}) + \tr(\boldsymbol{X}^{\top}\boldsymbol{O}\boldsymbol{M}^{\top}(d\boldsymbol{O})) \\[7pt] =&\, \tr(\boldsymbol{M}^{\top}\boldsymbol{O}\boldsymbol{X}^{\top}(d\boldsymbol{O})) + \tr(\boldsymbol{O}\boldsymbol{X}^{\top}\boldsymbol{O}(d\boldsymbol{M}^{\top})) + \tr(\boldsymbol{X}^{\top}\boldsymbol{O}\boldsymbol{M}^{\top}(d\boldsymbol{O})) \\[7pt] =&\, \tr(\boldsymbol{M}^{\top}\boldsymbol{O}\boldsymbol{X}^{\top}(d\boldsymbol{O})) + \tr(\boldsymbol{O}^{\top}\boldsymbol{X}\boldsymbol{O}^{\top}(d\boldsymbol{M})) + \tr(\boldsymbol{X}^{\top}\boldsymbol{O}\boldsymbol{M}^{\top}(d\boldsymbol{O})) \\[7pt] \end{aligned}\end{equation}

from which we obtain

\begin{equation}\tr((\boldsymbol{X}^{\top} - \boldsymbol{O}^{\top}\boldsymbol{X}\boldsymbol{O}^{\top})(d\boldsymbol{M})) = \tr((\boldsymbol{M}^{\top}\boldsymbol{O}\boldsymbol{X}^{\top} + \boldsymbol{X}^{\top}\boldsymbol{O}\boldsymbol{M}^{\top})(d\boldsymbol{O}))\end{equation}

If we let $\boldsymbol{M}^{\top}\boldsymbol{O}\boldsymbol{X}^{\top} + \boldsymbol{X}^{\top}\boldsymbol{O}\boldsymbol{M}^{\top}=(\nabla_{\boldsymbol{O}}\mathcal{L})^{\top}$, then the above equation takes on the meaning of $d\mathcal{L}$, so by the trace trick we have $\boldsymbol{X}^{\top} - \boldsymbol{O}^{\top}\boldsymbol{X}\boldsymbol{O}^{\top}=(\nabla_{\boldsymbol{M}}\mathcal{L})^{\top}$. This tells us the relationship between $\nabla_{\boldsymbol{M}}\mathcal{L}$ and $\nabla_{\boldsymbol{O}}\mathcal{L}$, described by the following system of equations:

\begin{gather}\boldsymbol{X} - \boldsymbol{O}\boldsymbol{X}^{\top}\boldsymbol{O} = \nabla_{\boldsymbol{M}}\mathcal{L} \label{eq:g-m}\\[7pt] \boldsymbol{X}\boldsymbol{O}^{\top}\boldsymbol{M} + \boldsymbol{M}\boldsymbol{O}^{\top}\boldsymbol{X} = \nabla_{\boldsymbol{O}}\mathcal{L}\label{eq:g-o}\end{gather}

The Theoretical Form

So the problem now becomes: solve for $\boldsymbol{X}$ from equation $\eqref{eq:g-o}$, and substitute it into equation $\eqref{eq:g-m}$ to get $\nabla_{\boldsymbol{M}}\mathcal{L}$ — that is, express $\nabla_{\boldsymbol{M}}\mathcal{L}$ as a function of $\nabla_{\boldsymbol{O}}\mathcal{L}$, avoiding a direct computation of $\nabla_{\boldsymbol{M}}\boldsymbol{O}$. Clearly, the only difficulty is solving equation $\eqref{eq:g-o}$.

In this section we'll first derive a not-so-practical theoretical solution based on SVD, which will help us understand the properties of equation $\eqref{eq:g-o}$ and check consistency with earlier results. Let $\boldsymbol{X}=\boldsymbol{U}\boldsymbol{Y}\boldsymbol{V}^{\top}$, and note also that $\boldsymbol{O}^{\top}\boldsymbol{M} = (\boldsymbol{M}^{\top}\boldsymbol{M})^{1/2} = \boldsymbol{V}(\boldsymbol{\Sigma}^{\top}\boldsymbol{\Sigma})^{1/2}\boldsymbol{V}^{\top}$ and $\boldsymbol{M}\boldsymbol{O}^{\top}=(\boldsymbol{M}\boldsymbol{M}^{\top})^{1/2} = \boldsymbol{U}(\boldsymbol{\Sigma}\boldsymbol{\Sigma}^{\top})^{1/2}\boldsymbol{U}^{\top}$. Substituting these identities into equation $\eqref{eq:g-o}$ gives

\begin{equation}\boldsymbol{U}\boldsymbol{Y}(\boldsymbol{\Sigma}^{\top}\boldsymbol{\Sigma})^{1/2}\boldsymbol{V}^{\top} + \boldsymbol{U}(\boldsymbol{\Sigma}\boldsymbol{\Sigma}^{\top})^{1/2}\boldsymbol{Y}\boldsymbol{V}^{\top} = \nabla_{\boldsymbol{O}}\mathcal{L}\end{equation}

that is

\begin{equation}\boldsymbol{Y}(\boldsymbol{\Sigma}^{\top}\boldsymbol{\Sigma})^{1/2} + (\boldsymbol{\Sigma}\boldsymbol{\Sigma}^{\top})^{1/2}\boldsymbol{Y} = \boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}\label{eq:g-o-2}\end{equation}

Written in component form, the left-hand side above is $\boldsymbol{Y}_{i,j}\sigma_j + \sigma_i \boldsymbol{Y}_{i,j} = (\sigma_i + \sigma_j)\boldsymbol{Y}_{i,j}$, where $\sigma_1,\sigma_2,\cdots,\sigma_r$ are the nonzero singular values of $\boldsymbol{M}$, and $0=\sigma_{r+1}=\sigma_{r+2}=\cdots$. Clearly, when $\boldsymbol{M}$ is a full-rank square matrix, we can solve to get

\begin{equation}\boldsymbol{Y} = (\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}) \oslash \boldsymbol{S}\end{equation}

where $\boldsymbol{S}_{i,j} = \sigma_i+\sigma_j$ and $\oslash$ denotes Hadamard division (element-wise division). Substituting $\boldsymbol{X}=\boldsymbol{U}\boldsymbol{Y}\boldsymbol{V}^{\top}$ into equation $\eqref{eq:g-m}$ then yields the result consistent with the one in The Derivative of SVD. This convergence of independent routes to the same answer boosts our confidence that our derivation so far is correct.

What if $\boldsymbol{M}$ is not full-rank or not square? In that case, if the right-hand side's $\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}$ "doesn't cooperate," equation $\eqref{eq:g-o-2}$ has no solution. But equation $\eqref{eq:g-o-2}$ arises from a genuine problem, so it must have a solution — meaning the right-hand side "has to cooperate"! What does cooperation look like here? If $\boldsymbol{M}$ has rank $r$, then in the matrix $\boldsymbol{S}$ only $\boldsymbol{S}_{[:r,:r]}$ is nonzero, and for equation $\eqref{eq:g-o-2}$ to have a solution, the part outside $(\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V})_{[:r,:r]}$ must be zero. Under this condition we can write

\begin{equation}\boldsymbol{Y} = \lim_{\epsilon\to 0}\,\, (\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}) \oslash (\boldsymbol{S} + \epsilon) \end{equation}

This amounts to saying that we can perturb the singular values slightly, turning the problem into one where all singular values are nonzero; after completing the computation we let the perturbation go to zero, thereby recovering the correct result.

An Efficient Solution

The SVD-based solution from the previous section is mostly of theoretical value. To compute things efficiently on a GPU, we need a solution of a different form. Introduce the notation $\boldsymbol{M}\boldsymbol{O}^{\top}=\boldsymbol{A},\boldsymbol{O}^{\top}\boldsymbol{M}=\boldsymbol{B},\nabla_{\boldsymbol{O}}\mathcal{L}=\boldsymbol{C}$; then equation $\eqref{eq:g-o}$ is actually a Sylvester equation:

\begin{equation}\boldsymbol{A}\boldsymbol{X}+\boldsymbol{X}\boldsymbol{B} = \boldsymbol{C}\end{equation}

There are many methods for solving Sylvester equations, and perhaps the most elegant and GPU-efficient one is based on $\mcsgn$ (not $\msign$) — an approach adapted here from Fast Differentiable Matrix Square Root. First, starting from the equation above, we can verify that the following holds:

\begin{equation}\begin{bmatrix} \boldsymbol{A} & -\boldsymbol{C} \\ \boldsymbol{0} & -\boldsymbol{B}\end{bmatrix} = \begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}\begin{bmatrix} \boldsymbol{A} & \boldsymbol{0} \\ \boldsymbol{0} & -\boldsymbol{B}\end{bmatrix}\begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}^{-1} \end{equation}

Taking $\mcsgn$ of both sides, and using the properties of $\mcsgn$, we get

\begin{equation}\mcsgn\left(\begin{bmatrix} \boldsymbol{A} & -\boldsymbol{C} \\ \boldsymbol{0} & -\boldsymbol{B}\end{bmatrix}\right) = \begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}\begin{bmatrix} \mcsgn(\boldsymbol{A}) & \boldsymbol{0} \\ \boldsymbol{0} & -\mcsgn(\boldsymbol{B})\end{bmatrix}\begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}^{-1} \end{equation}

Note that $\boldsymbol{A}=\boldsymbol{M}\boldsymbol{O}^{\top}=(\boldsymbol{M}\boldsymbol{M}^{\top})^{1/2}, \boldsymbol{B}=\boldsymbol{O}^{\top}\boldsymbol{M}=(\boldsymbol{M}^{\top}\boldsymbol{M})^{1/2}$: assuming $\boldsymbol{M}$ is a full-rank square matrix, then $\boldsymbol{A},\boldsymbol{B}$ are both positive-definite symmetric, and the $\mcsgn$ of a positive-definite symmetric matrix is always a square matrix, so

\begin{equation}\mcsgn\left(\begin{bmatrix} \boldsymbol{A} & -\boldsymbol{C} \\ \boldsymbol{0} & -\boldsymbol{B}\end{bmatrix}\right) = \begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}\begin{bmatrix} \boldsymbol{I} & \boldsymbol{0} \\ \boldsymbol{0} & -\boldsymbol{I}\end{bmatrix}\begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}^{-1} = \begin{bmatrix} \boldsymbol{I} & -2\boldsymbol{X} \\ \boldsymbol{0} & -\boldsymbol{I}\end{bmatrix} \end{equation}

The last simplification uses the identity $\begin{bmatrix} \boldsymbol{I} & \boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}^{-1}=\begin{bmatrix} \boldsymbol{I} & -\boldsymbol{X} \\ \boldsymbol{0} & \boldsymbol{I}\end{bmatrix}$. From this result we see that we only need to compute $\mcsgn$ for the block matrix $\begin{bmatrix} \boldsymbol{A} & -\boldsymbol{C} \\ \boldsymbol{0} & -\boldsymbol{B}\end{bmatrix}$, and then we can read $\boldsymbol{X}$ off the upper-right block of the result. $\mcsgn$ can be computed efficiently via a Newton-Schulz iteration, so this scheme is GPU-friendly.

When $\boldsymbol{M}$ is not full-rank or not square, $\boldsymbol{A},\boldsymbol{B}$ is only positive-semidefinite, in which case $\mcsgn$ is no longer $\boldsymbol{I}$. However, the experience from the previous section tells us that since $\nabla_{\boldsymbol{O}}\mathcal{L}$ "has to cooperate," we only need to perturb $\boldsymbol{\Sigma}$ slightly to turn it into the positive-definite case, which is solvable. Perturbing $\boldsymbol{\Sigma}$ in this way is equivalent to adding $\epsilon \boldsymbol{I}$ to $\boldsymbol{A},\boldsymbol{B}$, so

\begin{equation}\boldsymbol{X} = -\frac{1}{2} \left(\lim_{\epsilon\to 0}\,\, \mcsgn\left(\begin{bmatrix} \boldsymbol{A} + \epsilon \boldsymbol{I} & -\boldsymbol{C} \\ \boldsymbol{0} & -\boldsymbol{B} - \epsilon \boldsymbol{I}\end{bmatrix}\right)\right)_{[:n,n:]} \end{equation}

In actual computation, we simply have to pick a reasonably small $\epsilon > 0$ for the approximation; we could consider $\epsilon=10^{-3}$, which falls within the lower bound range we previously found for the Newton-Schulz iteration.

Summary

This post discussed how to compute the derivative of the $\msign$ operator. If you're interested in combining "TTT + Muon," this article may be of help.