Computing the Singular Value Clip mclip via msign (Part 1)

通过msign来计算奇异值裁剪mclip（上） 苏剑林 · 2025-06-07 · Mathematics iteration approximation matrix SVD muon

In two earlier posts, Newton-Schulz Iteration for the msign Operator (Part 1) and Newton-Schulz Iteration for the msign Operator (Part 2), we discussed numerical methods for computing the matrix $\newcommand{msign}{\mathop{\text{msign}}}\newcommand{sign}{\mathop{\text{sign}}}\newcommand{clip}{\mathop{\text{clip}}}\newcommand{mclip}{\mathop{\text{mclip}}}\msign$ operator. In this post we turn our attention to the "Singular Value Clipping" operation, which recently sparked a lot of discussion on @_arohan_'s Twitter, and which we also mentioned earlier in Higher-Order MuP: A Simpler yet Smarter Spectral Conditioning Scaling. From here on we'll abbreviate it as $\mclip$.

Basic Concepts

For a scalar $x$, the $\clip$ operation is defined as

\begin{equation}\clip(x) = \max(\min(x, 1), -1) = \left\{\begin{aligned}1, &\quad x\geq 1 \\ x, &\quad x\in(-1, 1)\\ -1, &\quad x\leq -1 \end{aligned}\right.\end{equation}more

that is, anything greater than $1$ or less than $-1$ gets clipped, otherwise left unchanged. We define the $\mclip$ of a matrix $\boldsymbol{M}\in\mathbb{R}^{n\times m}$ as

\begin{equation}\mclip(\boldsymbol{M}) = \boldsymbol{U}\clip(\boldsymbol{\Sigma})\boldsymbol{V}^{\top} \end{equation}

where $\boldsymbol{M}=\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ is the SVD of matrix $\boldsymbol{M}$, $\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{V}\in\mathbb{R}^{m\times m}$ are orthogonal matrices, $\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m}$ is the diagonal matrix of singular values, and adding $\clip$ to a diagonal matrix means applying $\clip$ elementwise to its diagonal entries. Note that the matrix $\boldsymbol{\Sigma}$ is diagonal and always non-negative, so we also have

\begin{equation}\mclip(\boldsymbol{M}) = \boldsymbol{U}\min(\boldsymbol{\Sigma}, 1)\boldsymbol{V}^{\top} \end{equation}

SVD is of course the standard way to compute $\mclip$, but SVD isn't particularly efficient. Given our previous experience with $\msign$, it's natural to think that, just like with $\msign$, we could find a Newton-Schulz iteration for $\mclip$ as well. That approach is certainly viable, but building on $\msign$, there's actually an even smarter way to do it.

Standing on the Shoulders of Giants

This clever idea comes from @leloykun, who, in the blog post Numerically Stable Spectral Clipping Via Newton-Schulz Iteration, proposed standing on the shoulders of $\msign$ and expressing $\mclip$ in terms of $\msign$, so that there's no need to search for a separate Newton-Schulz iteration. In that post he also gives an ingenious solution, but personally I find it a bit unintuitive and not especially efficient either. Below I'll present my own approach.

My starting point is the scalar identity (which Kimi helped me find)

$$\min(x, 1) = \frac{1}{2} [x + 1 - (x-1)\sign(x-1)] $$

For simplicity, let's first assume $\boldsymbol{M}$ is a full-rank square matrix. Then

\begin{equation}\begin{aligned} 2\mclip(\boldsymbol{M}) =&\, \boldsymbol{U} [2\min(\boldsymbol{\Sigma},1)] \boldsymbol{V}^{\top} \\[6pt] =&\, \boldsymbol{U} [\boldsymbol{\Sigma} + \boldsymbol{I} - (\boldsymbol{\Sigma} - \boldsymbol{I})\sign(\boldsymbol{\Sigma} - \boldsymbol{I})] \boldsymbol{V}^{\top} \\[6pt] =&\, \boldsymbol{U} [\boldsymbol{\Sigma} + \boldsymbol{I} - (\boldsymbol{\Sigma} - \boldsymbol{I})\msign(\boldsymbol{\Sigma} - \boldsymbol{I})] \boldsymbol{V}^{\top} \\[6pt] =&\, \boldsymbol{M} + \boldsymbol{U}\boldsymbol{V}^{\top} - \boldsymbol{U}(\boldsymbol{\Sigma} - \boldsymbol{I})\msign(\boldsymbol{\Sigma} - \boldsymbol{I}) \boldsymbol{V}^{\top} \end{aligned}\label{eq:2-mclip-M}\end{equation}

Note that

\begin{equation}\begin{aligned} &\,\boldsymbol{U}(\boldsymbol{\Sigma} - \boldsymbol{I})\msign(\boldsymbol{\Sigma} - \boldsymbol{I}) \boldsymbol{V}^{\top} \\[6pt] =&\, \boldsymbol{U}(\boldsymbol{\Sigma} - \boldsymbol{I}) \boldsymbol{U}^{\top} \boldsymbol{U}\msign(\boldsymbol{\Sigma} - \boldsymbol{I}) \boldsymbol{V}^{\top} \\[6pt] =&\, (\boldsymbol{U}\boldsymbol{\Sigma} \boldsymbol{U}^{\top} - \boldsymbol{I}) \msign(\boldsymbol{M} - \boldsymbol{U}\boldsymbol{V}^{\top}) \\[6pt] =&\, (\boldsymbol{U}\boldsymbol{\Sigma} \boldsymbol{V}^{\top} (\boldsymbol{U}\boldsymbol{V}^{\top})^{\top} - \boldsymbol{I}) \msign(\boldsymbol{M} - \boldsymbol{U}\boldsymbol{V}^{\top}) \\[6pt] =&\, (\boldsymbol{M} (\boldsymbol{U}\boldsymbol{V}^{\top})^{\top} - \boldsymbol{I}) \msign(\boldsymbol{M} - \boldsymbol{U}\boldsymbol{V}^{\top}) \\[6pt] \end{aligned}\end{equation}

where the second equality uses the fact that for any orthogonal matrix $\boldsymbol{P},\boldsymbol{Q}$, $\boldsymbol{P}\msign(\boldsymbol{R})\boldsymbol{Q} = \msign(\boldsymbol{P}\boldsymbol{R}\boldsymbol{Q})$ holds. Substituting this back into equation $\eqref{eq:2-mclip-M}$ gives

\begin{equation}2\mclip(\boldsymbol{M}) = \boldsymbol{M} + \boldsymbol{U}\boldsymbol{V}^{\top} + (\boldsymbol{I} - \boldsymbol{M}(\boldsymbol{U}\boldsymbol{V}^{\top})^{\top}) \msign(\boldsymbol{M} - \boldsymbol{U}\boldsymbol{V}^{\top})\label{eq:mclip-M-core}\end{equation}

If $\boldsymbol{M}$ is a general matrix of rank $r$, then $\boldsymbol{U}\boldsymbol{V}^{\top}$ should be replaced by $\boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}$, and we can directly substitute $\boldsymbol{M} = \boldsymbol{U}_{[:,:r]}\boldsymbol{\Sigma}_{[:r,:r]}\boldsymbol{V}_{[:,:r]}^{\top}$ into the above to verify that the equality still holds.

(Note: thanks to @YouJiacheng for discussion on this section.)

Reference Implementation

We know that $\boldsymbol{U}\boldsymbol{V}^{\top}=\msign(\boldsymbol{M})$, so computing $\mclip$ via equation $\eqref{eq:mclip-M-core}$ only requires evaluating $\msign$ twice:

\begin{equation}2\mclip(\boldsymbol{M}) = \boldsymbol{M} + \msign(\boldsymbol{M}) + (\boldsymbol{I} - \boldsymbol{M}\msign(\boldsymbol{M})^{\top}) \msign(\boldsymbol{M} - \msign(\boldsymbol{M}))\end{equation}

The computational cost is roughly twice that of $\msign$; by comparison, Numerically Stable Spectral Clipping Via Newton-Schulz Iteration requires evaluating $\msign$ once on a matrix about 4 times larger, which comes to roughly 8 times the cost of $\msign$.

Building on $\msign$, implementing equation $\eqref{eq:mclip-M-core}$ takes as few as two lines of code, as shown below:

import numpy as np

def msign(m):
    u, s, vh = np.linalg.svd(m, full_matrices=False)
    return u @ vh

def mclip(m):
    ms2 = msign(m - (ms := msign(m)))
    return (m + ms + ms2 - m @ ms.mT @ ms2) / 2

m = np.random.randn(10, 20)
u, s, vh = np.linalg.svd(m, full_matrices=False)

result1 = u @ np.diag(s.clip(0, 1)) @ vh
result2 = mclip(m)
np.abs(result1 - result2).mean()

Here we directly use SVD to compute $\msign$, just to quickly verify the correctness of equation $\eqref{eq:mclip-M-core}$. In actual computation, readers are free to replace the $\msign$ function with the corresponding Newton-Schulz iteration.

Other Functions

We can use the same approach to compute matrix versions of other functions as well, such as the step function. Let's define the scalar step function $\newcommand{mstep}{\mathop{\text{mstep}}}\newcommand{step}{\mathop{\text{step}}}$ as

\begin{equation}\step(x) = \frac{1}{2}[\sign(x - 1) + 1]\end{equation}

meaning it becomes 1 above 1 and 0 below 1. We can then define

\begin{equation}\mstep(\boldsymbol{M}) = \boldsymbol{U}\step(\boldsymbol{\Sigma})\boldsymbol{V}^{\top}\end{equation}

that is, singular values greater than 1 are kept and clipped to 1, while those less than 1 are set to zero directly. Following the same steps, we get

\begin{equation}\mstep(\boldsymbol{M}) = \frac{1}{2}[\msign(\boldsymbol{M}) + \msign(\boldsymbol{M} - \msign(\boldsymbol{M}))]\end{equation}

We can even express even functions this way — for instance, defining

\begin{equation}\mathop{\text{msquare}}(\boldsymbol{M}) = \boldsymbol{U} \boldsymbol{\Sigma}^2\boldsymbol{V}^{\top} = \boldsymbol{U}\boldsymbol{V}^{\top}(\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{U}^{\top})(\boldsymbol{U} \boldsymbol{\Sigma}\boldsymbol{V}^{\top}) = \msign(\boldsymbol{M})\boldsymbol{M}^{\top}\boldsymbol{M}\end{equation}

This differs from the matrix square defined directly by $\boldsymbol{M}^2$, since the latter only applies to square matrices and squares the eigenvalues under eigendecomposition, whereas the expression above squares the singular values under SVD. More generally, we have

\begin{equation}\boldsymbol{U} \boldsymbol{\Sigma}^{2n}\boldsymbol{V}^{\top} = \msign(\boldsymbol{M})(\boldsymbol{M}^{\top}\boldsymbol{M})^n,\quad \boldsymbol{U} \boldsymbol{\Sigma}^{2n+1}\boldsymbol{V}^{\top} = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^n\end{equation}

This shows that for any polynomial $f(x)$ (not just odd polynomials), $\boldsymbol{U}f(\boldsymbol{\Sigma})\boldsymbol{V}^{\top}$ can be obtained from $\boldsymbol{M}$ and $\msign(\boldsymbol{M})$ through a finite number of matrix additions and multiplications.

Summary

This post introduced an approach for performing general operations on the singular values of a matrix — using the matrix itself together with its $\msign$ — covering singular value clipping, the step function, and arbitrary-degree polynomials (not just odd ones).