Computing Singular Value Clipping mclip via msign (Part 2)

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

Previously, in 《通过msign来计算奇异值裁剪mclip（上）》 we discussed the numerical computation of singular value clipping $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$, building on the ideas in @leloykun's post 《Numerically Stable Spectral Clipping Via Newton-Schulz Iteration》 (now substantially revised and renamed), which finds an expression based on $\newcommand{msign}{\mathop{\text{msign}}}\msign$ so as to avoid having to design yet another Newton-Schulz iteration. In that post I proposed a nested $\msign$ scheme with lower computational cost.

However, a couple of days ago @leloykun pointed out on Twitter that my scheme actually suffers from excessively large errors in practice. This post takes a closer look at that issue and presents a new scheme that is both more efficient and more accurate.more

Basic concepts

As usual, let's first lay out the basic concepts. First, the scalar $x$ operator $\newcommand{clip}{\mathop{\text{clip}}}\clip$, which this time we define in general as

\begin{equation}\clip\nolimits_{[\alpha,\beta]}(x) = \max(\min(x, \beta), \alpha) = \left\{\begin{aligned}\beta, &\quad \geq \beta \\ x, &\quad x\in(\alpha, \beta)\\ \alpha, &\quad x\leq \alpha \end{aligned}\right.\end{equation}

Unless otherwise noted, the interval defaults to $[-1,1]$, i.e., $\clip(x) = \clip_{[-1,1]}(x)$. Let the SVD of matrix $\boldsymbol{M}\in\mathbb{R}^{n\times m}$ be $\boldsymbol{M}=\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, where $\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{V}\in\mathbb{R}^{m\times m}$ are orthogonal matrices and $\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m}$ is the diagonal matrix of singular values. We then define

\begin{equation}\mclip\nolimits_{[\alpha,\beta]}(\boldsymbol{M}) = \boldsymbol{U}\clip\nolimits_{[\alpha,\beta]}(\boldsymbol{\Sigma})\boldsymbol{V}^{\top}\end{equation}

Adding $\clip$ to a diagonal matrix means applying $\clip$ to its diagonal entries individually. In plain terms, $\mclip_{[\alpha,\beta]}$ simply clips the singular values of $\boldsymbol{M}$ into $[\alpha,\beta]$.

Since singular values are non-negative, when $\alpha < 0$ we have $\mclip_{[\alpha,\beta]}(\boldsymbol{M})=\mclip_{[0,\beta]}(\boldsymbol{M})$. But as we'll see later, because of numerical errors in actual computation, allowing negative arguments in $\alpha$ can produce some interesting error-cancellation effects.

General theoretical solution

The goal of this section is to express $\mclip$ in terms of $\msign$, starting from the identity

\begin{equation}\newcommand{sign}{\mathop{\text{sign}}}\mclip\nolimits_{[\alpha,\beta]} (x) = \frac{\alpha + \beta + (\alpha - x)\sign(\alpha - x) - (\beta - x)\sign(\beta - x)}{2}\end{equation}

The key to finding this identity is to express $\clip$ as a linear combination of the absolute value and itself, and then use $|x|=x\sign(x)$ to transition to the $\sign$ operation; we won't go into the details here.

For simplicity, let's first assume $\boldsymbol{M}$ is a full-rank square matrix. Based on this identity, we have

\begin{equation}2\mclip\nolimits_{[\alpha,\beta]}(\boldsymbol{M}) = \boldsymbol{U}\Big((\alpha + \beta)\boldsymbol{I} + (\alpha \boldsymbol{I} - \boldsymbol{\Sigma})\sign(\alpha \boldsymbol{I} - \boldsymbol{\Sigma}) - (\beta \boldsymbol{I} - \boldsymbol{\Sigma})\sign(\beta \boldsymbol{I} - \boldsymbol{\Sigma})\Big)\boldsymbol{V}^{\top}\end{equation}

Expanding the right-hand side, we get several types of terms ($\gamma\in\{\alpha,\beta\}$):

$$\begin{array}{c|c} \hline \text{original} & \text{simplify} \\ \hline \boldsymbol{U}\boldsymbol{V}^{\top} & \msign(\boldsymbol{M}) \\ \hline \boldsymbol{U}\sign(\gamma \boldsymbol{I} - \boldsymbol{\Sigma})\boldsymbol{V}^{\top} & \begin{aligned}&\, \msign(\gamma \boldsymbol{U}\boldsymbol{V}^{\top} - \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}) \\ =&\, \msign(\gamma \msign(\boldsymbol{M}) - \boldsymbol{M}) \end{aligned} \\ \hline \boldsymbol{U}\boldsymbol{\Sigma}\sign(\gamma \boldsymbol{I} - \boldsymbol{\Sigma})\boldsymbol{V}^{\top} & \begin{aligned}&\, \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}\boldsymbol{V}\boldsymbol{U}^{\top}\boldsymbol{U}\sign(\gamma \boldsymbol{I} - \boldsymbol{\Sigma})\boldsymbol{V}^{\top} \\ =&\, \boldsymbol{M}\msign(\boldsymbol{M})^{\top}\msign(\gamma \msign(\boldsymbol{M}) - \boldsymbol{M}) \end{aligned} \\ \hline \end{array}$$

Substituting and simplifying gives

\begin{equation}\mclip\nolimits_{[\alpha,\beta]}(\boldsymbol{M}) = \frac{1}{2}\left\{\begin{aligned}&\,(\alpha + \beta)\msign(\boldsymbol{M}) \\ + &\, (\alpha \boldsymbol{I} - \boldsymbol{M}\msign(\boldsymbol{M})^{\top})\msign(\alpha \msign(\boldsymbol{M}) - \boldsymbol{M})\\ - &\, (\beta \boldsymbol{I} - \boldsymbol{M}\msign(\boldsymbol{M})^{\top})\msign(\beta \msign(\boldsymbol{M}) - \boldsymbol{M}) \end{aligned}\right\}\label{eq:general}\end{equation}

For non-square or rank-deficient matrices, one can substitute $\msign(\boldsymbol{M})=\boldsymbol{U}_{[:,:r]}\boldsymbol{V}_{[:,:r]}^{\top}$ into the above and verify it still holds, so the above is the general theoretical solution for $\mclip$.

Initial form

Equation $\eqref{eq:general}$ appears to require computing $\msign$ at least three times, and the inputs to the latter two $\msign$ calls depend on the result of the first $\msign$, so it is formally a nested $\msign$. When we take $\alpha=0,\beta=1$, the number of $\msign$ calls can be reduced to two:

\begin{equation}\mclip(\boldsymbol{M}) = \frac{1}{2}\Big[\boldsymbol{M} + \msign(\boldsymbol{M}) + (\boldsymbol{I} - \boldsymbol{M}\msign(\boldsymbol{M})^{\top}) \msign(\boldsymbol{M} - \msign(\boldsymbol{M}))\Big]\label{eq:mclip-1}\end{equation}

This is exactly the result I gave in the previous post 《通过msign来计算奇异值裁剪mclip（上）》, requiring only two calls to $\msign$.

However, empirical testing shows that when the singular values of $\boldsymbol{M}$ are large and the computational precision of $\msign$ is low, this formula can produce fairly large errors — much larger than those in @leloykun's scheme. But @leloykun's scheme requires computing $\msign$ on a matrix $\begin{bmatrix}\boldsymbol{I} & \boldsymbol{M} \\ \boldsymbol{M}^{\top} & \boldsymbol{I}\end{bmatrix}$ about four times larger, which is not cheap, so it's still worth exploring what room for improvement exists here.

Removing the nesting

Intuitively, the source of the error is the accumulated error caused by the nested $\msign$ calls, so it's worth trying to eliminate the nesting. Fortunately, with a simple trick, we can indeed get rid of it!

First, we can prove

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

Then we have

\begin{equation}\msign(\boldsymbol{M})^{\top}\boldsymbol{M} - \boldsymbol{I} = \boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top} - \boldsymbol{I} = \boldsymbol{V}(\boldsymbol{\Sigma}-\boldsymbol{I})\boldsymbol{V}^{\top}\end{equation}

Based on the above, we claim

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

This makes use of a very simple property: $\forall x \geq 0, \sign(x-1) = \sign(x^2-1)$. Using this result, we obtain

\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]\label{eq:mclip-2}\end{equation}

Still two calls to $\msign$, but they no longer have a nested relationship, which means, in theory, there's no longer any accumulated error from nested $\msign$ calls. Empirical results show that equation $\eqref{eq:mclip-2}$'s error is indeed roughly half that of equation $\eqref{eq:mclip-1}$, but in extreme cases it still falls short of @leloykun's scheme, suggesting that nesting isn't actually the main source of error.

Mutual cancellation

So what other room for improvement is there? @leloykun's scheme requires the function to be odd, so it actually works with $\mclip_{[-1,1]}$ rather than $\mclip_{[0,1]}$. Could it be that this choice causes two error terms to cancel each other out, thereby achieving better numerical precision?

To test this, we substitute $\alpha=-1,\beta=1$ into equation $\eqref{eq:general}$, obtaining

\begin{equation}\mclip(\boldsymbol{M}) = \frac{1}{2}\left\{\begin{aligned} &\,(\boldsymbol{I} + \boldsymbol{M}\msign(\boldsymbol{M})^{\top})\msign(\msign(\boldsymbol{M}) + \boldsymbol{M}) \\ - &\,(\boldsymbol{I} - \boldsymbol{M}\msign(\boldsymbol{M})^{\top})\msign(\msign(\boldsymbol{M}) - \boldsymbol{M}) \end{aligned}\right\}\end{equation}

Applying the same de-nesting trick as in the previous section, we get

\begin{equation}\mclip(\boldsymbol{M}) = \frac{1}{2}\left\{\begin{aligned} &\,(\msign(\boldsymbol{M}) + \boldsymbol{M})\msign(\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I}) \\ + &\,(\msign(\boldsymbol{M}) - \boldsymbol{M})\msign(\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I}) \end{aligned}\right\}\label{eq:mclip-3}\end{equation}

Note that $\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I}$ is always a positive-definite symmetric matrix, so in theory $\msign(\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I})=\boldsymbol{I}$, which would bring us back to equation $\eqref{eq:mclip-2}$. But in actual computation, the error between $\msign(\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I})$ and $\boldsymbol{I}$ may cancel out the error introduced by $\msign(\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I})$, so we decide whether to keep it based on experiments.

As expected, the numerical error of equation $\eqref{eq:mclip-3}$ turns out to be even smaller than that of @leloykun's scheme! This confirms our conjecture: setting $\alpha=-1$ and $\beta=1$ so that $\mclip$ becomes an odd function helps cancel out errors.

A brief reflection on the cause

Why does this cancellation work out so neatly? We can do a simple quantitative analysis. Large errors arise under two conditions: first, $\boldsymbol{M}$ has very large singular values, and second, the number of iteration steps for $\msign$ is small, so $\msign$ itself isn't very accurate.

Looking at equation $\eqref{eq:mclip-3}$, it can be split into a sum of four terms. In fact, the two terms $\msign(\boldsymbol{M})\msign(\boldsymbol{M}^{\top}\boldsymbol{M} \pm \boldsymbol{I})$ are bounded, so even if $\msign$ isn't very precise, they essentially can't diverge. So the main source of error is

\begin{equation}\boldsymbol{M}\msign(\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I}) - \boldsymbol{M}\msign(\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I})\label{eq:error-1}\end{equation}

which is proportional to $\boldsymbol{M}$ and is the term most likely to amplify errors. Correspondingly, the main error term in equation $\eqref{eq:mclip-2}$ is

\begin{equation}\boldsymbol{M} - \boldsymbol{M}\msign(\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I})\label{eq:error-2}\end{equation}

Consider a singular value much greater than 1. If $\msign$ were exact, then the result of $\msign$ would be 1, and in both of the expressions above, the parts corresponding to large singular values would come out to the 0 we expect.

However, if $\msign$ has only a small number of iteration steps, it might instead come out as something like $0.6$ or $1.4$, in which case the corresponding part of equation $\eqref{eq:error-2}$ would produce a huge error like $\sim\pm 0.4 \boldsymbol{M}$. But in equation $\eqref{eq:error-1}$, when the singular value is large, the relative difference between $\boldsymbol{M}^{\top}\boldsymbol{M} - \boldsymbol{I}$ and $\boldsymbol{M}^{\top}\boldsymbol{M} + \boldsymbol{I}$ is not large, so the difference in $\msign(\boldsymbol{M}^{\top}\boldsymbol{M} \pm \boldsymbol{I})$ is small — meaning equation $\eqref{eq:error-1}$ can still cancel out most of the error.

But keep in mind, this all rests on the premise that $\boldsymbol{M}$ has singular values noticeably greater than 1, and that the number of iteration steps is small. If these two conditions aren't met, then equation $\eqref{eq:mclip-2}$'s original error isn't large to begin with, and equation $\eqref{eq:mclip-3}$ would instead end up increasing the error by computing $\msign$ an extra time. So which formula performs best in practice really depends on the specific situation.

Comparison code

Let's construct a set of singular values that includes some greater than 1 and some less than 1, with the largest singular value close to 1000, and then test each algorithm under bfloat16 precision. Reference code is given below (approximate results have been noted in the comments):

import numpy as np
import jax.numpy as jnp
import jax.lax as lax

def msign(x, steps=4, eps=1e-20):
    """The coefficients come from https://kexue.fm/archives/10996
    """
    abc = [
        (8.287212018145622, -23.59588651909882, 17.300387312530923),
        (4.107059111542197, -2.9478499167379084, 0.54484310829266),
        (3.9486908534822938, -2.908902115962947, 0.5518191394370131),
        (3.3184196573706055, -2.488488024314878, 0.5100489401237208),
        (2.3006520199548186, -1.6689039845747518, 0.4188073119525678),
        (1.8913014077874002, -1.2679958271945908, 0.37680408948524996),
        (1.875, -1.25, 0.375)
    ]
    y = x.mT if x.shape[-2] > x.shape[-1] else x
    y = y * lax.rsqrt((y**2).sum(axis=[-2, -1], keepdims=True) + eps)
    for a, b, c in abc[:steps] + max(steps - 7, 0) * abc[-1:]:
        a, b, c = a / 1.01, b / 1.01**3, c / 1.01**5
        y = a * y + (b * (u := y @ y.mT) + c * u @ u) @ y
    return y.mT if x.shape[-2] > x.shape[-1] else y

def mclip1(m):
    """1st version (2 nested msign)
    """
    ms2 = msign(m - (ms1 := msign(m)))
    return (m + ms1 + ms2 - m @ ms1.mT @ ms2) / 2

def mclip2(m):
    """2nd version (2 non-nested msign)
    """
    ms1 = msign(m)
    ms2 = msign(m.mT @ m - jnp.eye(m.shape[-1]))
    return (m + ms1 + (ms1 - m) @ ms2) / 2

def mclip3(m):
    """3rd version (3 non-nested msign)
    """
    ms1 = msign(m)
    ms2 = msign(m.mT @ m + jnp.eye(m.shape[-1]))
    ms3 = msign(m.mT @ m - jnp.eye(m.shape[-1]))
    return ((ms1 + m) @ ms2  + (ms1 - m) @ ms3) / 2

def spectral_clip(W):
    """@leloykun verision: https://leloykun.github.io/ponder/spectral-clipping/
    """
    m, n = W.shape
    H = jnp.block([[jnp.eye(m), W], [W.T, jnp.eye(n)]])
    OH = msign(H)
    P, Q = OH[:m, :m], OH[:m, m:]
    return Q + P @ W

m = np.random.randn(4096, 1024)
u, s, vh = jnp.linalg.svd(m, full_matrices=False)
s = np.concatenate([np.linspace(1, 1000, 128), np.linspace(0, 1, 896)])
s = np.sort(s)[::-1]
m = u @ jnp.diag(s) @ vh  # matrix with large singular values

result0 = u @ np.diag(s.clip(0, 1)) @ vh  # exact result via SVD
result1 = mclip1(m.astype('bfloat16'))
result2 = mclip2(m.astype('bfloat16'))
result3 = mclip3(m.astype('bfloat16'))
result4 = spectral_clip(m.astype('bfloat16'))

# spectral norm of the resulting matrix, closer to 1 is better.
jnp.linalg.svd(result0.astype('float32'))[1][0]  # = 1
jnp.linalg.svd(result1.astype('float32'))[1][0]  # ≈ 700
jnp.linalg.svd(result2.astype('float32'))[1][0]  # ≈ 250
jnp.linalg.svd(result3.astype('float32'))[1][0]  # ≈ 1.5
jnp.linalg.svd(result4.astype('float32'))[1][0]  # ≈ 13

# mean absolute error of singular values, closer to 0 is better.
jnp.abs(jnp.linalg.svd(result1.astype('float32'))[1] - s.clip(0, 1)).mean()  # ≈ 20
jnp.abs(jnp.linalg.svd(result2.astype('float32'))[1] - s.clip(0, 1)).mean()  # ≈ 10
jnp.abs(jnp.linalg.svd(result3.astype('float32'))[1] - s.clip(0, 1)).mean()  # ≈ 0.5
jnp.abs(jnp.linalg.svd(result4.astype('float32'))[1] - s.clip(0, 1)).mean()  # ≈ 0.7

# mean absolute error of total matrix, closer to 0 is better.
jnp.abs(result0 - result1).mean()  # ≈ 1
jnp.abs(result0 - result2).mean()  # ≈ 0.5
jnp.abs(result0 - result3).mean()  # ≈ 0.01
jnp.abs(result0 - result4).mean()  # ≈ 0.02

Summary

This post has further refined the previous post's scheme for computing $\mclip$ using $\msign$, successfully reducing the computational error by removing the nesting of $\msign$ and introducing an extra correction term.