Efficient Computation of Matrix Square Roots and Inverse Square Roots

矩阵平方根和逆平方根的高效计算 苏剑林 · 2025-07-19 · Mathematics algebra iteration matrix linear

Let $\boldsymbol{P}\in\mathbb{R}^{n\times n}$ be an $n$-th order square matrix all of whose eigenvalues are non-negative real numbers. This post discusses how to compute its square root $\boldsymbol{P}^{1/2}$ and inverse square root $\boldsymbol{P}^{-1/2}$.

Basic Concepts

The square root of a matrix $\boldsymbol{P}$ refers to a matrix $\boldsymbol{X}$ satisfying $\boldsymbol{X}^2=\boldsymbol{P}$. We know that positive numbers have two square roots, so it's not hard to imagine that matrix square roots are generally not unique either. However, the "arithmetic square root" is unique: for a positive number, the arithmetic square root is the positive one of its two square roots; similarly, we call the square root of $\boldsymbol{P}$ whose eigenvalues are all non-negative the arithmetic square root. Throughout this post, "matrix square root" defaults to meaning the arithmetic square root. more

The computation here relies on the matrix sign function we discussed in What Can the Matrix Sign Function mcsgn Compute?:

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

Simply put, it maps the eigenvalues of an arbitrary matrix $\boldsymbol{M}\in\mathbb{R}^{n\times n}$ to the corresponding values of the sign function, giving a new matrix. Assuming the eigenvalues of $\boldsymbol{M}$ are all real, $\mcsgn$ can be computed efficiently via a Newton-Schulz iteration:

\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}

where $\frac{\boldsymbol{M}}{\sqrt{\tr(\boldsymbol{M}^2)}}$ is chosen so as to scale the eigenvalues of $\boldsymbol{X}_0$ into $[-1,1]$, and $a_t,b_t,c_t$ are the coefficients derived in Newton-Schulz Iterations for the msign Operator (Part 2):

$$\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}$$

In fact, when the eigenvalues of $\boldsymbol{M}$ are all real, the computation principle of $\mcsgn$ is essentially the same as that of another matrix sign function $\newcommand{msign}{\mathop{\text{msign}}}\msign$.

The Underlying Principle

The starting point for what follows is the identity

\begin{equation}\mcsgn\left(\begin{bmatrix}\boldsymbol{0} & \boldsymbol{A} \\ \boldsymbol{B} & \boldsymbol{0}\end{bmatrix}\right)=\begin{bmatrix}\boldsymbol{0} & \boldsymbol{A}(\boldsymbol{B}\boldsymbol{A})^{-1/2} \\ \boldsymbol{B}(\boldsymbol{A}\boldsymbol{B})^{-1/2} & \boldsymbol{0}\end{bmatrix}\label{eq:core}\end{equation}

This can be verified directly by substituting the definition of $\mcsgn$ (note: $\boldsymbol{A},\boldsymbol{B}$ need not be a square matrix). Next we need to determine under what conditions the eigenvalues of the matrix on the left-hand side that we're taking the $\mcsgn$ of are all real. Let $\lambda$ be one of its nonzero eigenvalues; then

\begin{equation}0=\det\left(\lambda\boldsymbol{I} - \begin{bmatrix}\boldsymbol{0} & \boldsymbol{A} \\ \boldsymbol{B} & \boldsymbol{0} \end{bmatrix}\right) = \det\left(\begin{bmatrix}\lambda\boldsymbol{I} & -\boldsymbol{A} \\ -\boldsymbol{B} & \lambda\boldsymbol{I} \end{bmatrix}\right) = \det(\lambda^2 \boldsymbol{I} - \boldsymbol{A}\boldsymbol{B})\end{equation}

that is, $\lambda^2$ is an eigenvalue of the matrix $\boldsymbol{A}\boldsymbol{B}$. This means that all the eigenvalues of the above block matrix are real if and only if all the eigenvalues of $\boldsymbol{A}\boldsymbol{B}$ are non-negative.

It's certainly possible to iterate directly on the original matrix, but that would be computationally wasteful. We can exploit its anti-diagonal structure to reduce the amount of computation. Since

\begin{equation} \begin{bmatrix}\boldsymbol{0} & \boldsymbol{Y} \\ \boldsymbol{Z} & \boldsymbol{0}\end{bmatrix}^3 = \begin{bmatrix}\boldsymbol{0} & (\boldsymbol{Y}\boldsymbol{Z})\boldsymbol{Y} \\ \boldsymbol{Z}(\boldsymbol{Y}\boldsymbol{Z}) & \boldsymbol{0}\end{bmatrix},\quad \begin{bmatrix}\boldsymbol{0} & \boldsymbol{Y} \\ \boldsymbol{Z} & \boldsymbol{0}\end{bmatrix}^5 = \begin{bmatrix}\boldsymbol{0} & (\boldsymbol{Y}\boldsymbol{Z})^2\boldsymbol{Y} \\ \boldsymbol{Z}(\boldsymbol{Y}\boldsymbol{Z})^2 & \boldsymbol{0}\end{bmatrix} \\ \end{equation}

we can obtain the iteration

\begin{gather} \boldsymbol{Y}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)\boldsymbol{Y}_t \label{eq:r1} \\[6pt] \boldsymbol{Z}_{t+1} = \boldsymbol{Z}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2) \label{eq:r2} \end{gather}

so that $\boldsymbol{Y}_t\to \boldsymbol{A}(\boldsymbol{B}\boldsymbol{A})^{-1/2},\boldsymbol{Z}_t\to \boldsymbol{B}(\boldsymbol{A}\boldsymbol{B})^{-1/2}$. In particular, multiplying the two equations above together gives the recursion for $\boldsymbol{Y}_t\boldsymbol{Z}_t$:

\begin{equation}\boldsymbol{Y}_{t+1}\boldsymbol{Z}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)^2\boldsymbol{Y}_t\boldsymbol{Z}_t\label{eq:r3}\end{equation}

Computing the Square Root

Now we get to the actual computation of the square root. Since we've assumed that the eigenvalues of $\boldsymbol{P}$ are non-negative, we can always divide by $\tr(\boldsymbol{P})$ to further compress its eigenvalues into $0\sim 1$. So without loss of generality, assume the eigenvalues of $\boldsymbol{P}$ all lie within $[0,1]$, so that we can directly compute $\mcsgn$ using the Newton-Schulz iteration.

Substituting $\boldsymbol{A}=\boldsymbol{P},\boldsymbol{B}=\boldsymbol{I}$ into equation $\eqref{eq:core}$ gives

\begin{equation}\mcsgn\left(\begin{bmatrix}\boldsymbol{0} & \boldsymbol{P} \\ \boldsymbol{I} & \boldsymbol{0}\end{bmatrix}\right)=\begin{bmatrix}\boldsymbol{0} & \boldsymbol{P}^{1/2} \\ \boldsymbol{P}^{-1/2} & \boldsymbol{0}\end{bmatrix}\end{equation}

Remarkably, in theory we only need to run $\mcsgn$ once to obtain both the square root and the inverse square root — that is, by iterating according to equations $\eqref{eq:r1}$ and $\eqref{eq:r2}$, we can accomplish both tasks simultaneously!

In practice, however, things aren't so ideal. If $\boldsymbol{P}$ has singular values extremely close to zero, then $\boldsymbol{P}^{-1/2}$ will blow up numerically (this is equivalent to encountering $1/\sqrt{0}$), whereas $\boldsymbol{P}^{1/2}$ will not. So if we only care about the value of $\boldsymbol{P}^{1/2}$, simultaneously computing $\boldsymbol{P}^{1/2},\boldsymbol{P}^{-1/2}$ actually increases numerical instability. In that case, a better approach is to iterate via equations $\eqref{eq:r1}$ and $\eqref{eq:r3}$, computing only $\boldsymbol{P}^{1/2}$:

\begin{gather} \boldsymbol{Y}_0 = \boldsymbol{P}, \quad \boldsymbol{Y}_0\boldsymbol{Z}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{Y}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)\boldsymbol{Y}_t \\[6pt] \boldsymbol{Y}_{t+1}\boldsymbol{Z}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)^2\boldsymbol{Y}_t\boldsymbol{Z}_t \\[6pt] \lim_{t\to\infty} \boldsymbol{Y}_t = \boldsymbol{P}^{1/2}\notag \end{gather}

Since the limit of $\boldsymbol{Z}_t$ is $\boldsymbol{P}^{-1/2}$, the limit of $\boldsymbol{Y}_t\boldsymbol{Z}_t$ is $\boldsymbol{I}$. Therefore iterating $\boldsymbol{Y}_t\boldsymbol{Z}_t$ is less prone to numerical risk. Reference code is as follows:

import numpy as np

def abc(steps):
    coefs = [
        (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)
    ]
    for a, b, c in coefs[:steps] + max(steps - 7, 0) * coefs[-1:]:
        yield a / 1.01, b / 1.01**3, c / 1.01**5

def msqrt(P, steps=6):
    Y = YZ = P / (t := np.trace(P))
    I = np.eye(P.shape[0])
    for a, b, c in abc(steps):
        W = a * I + b * YZ + c * YZ @ YZ
        Y, YZ = W @ Y, W @ W @ YZ
    return Y * t**0.5

d = 100
P = (x := np.random.randn(d, d) / d**0.5) @ x.T
np.abs(msqrt(P) @ msqrt(P) - P).mean()  # ~= 2e-4

The Inverse Square Root

If we absolutely need to explicitly obtain the inverse square root $\boldsymbol{P}^{-1/2}$, then there isn't really a good way around it — whatever is going to blow up will blow up regardless. In this case, whether we use the combination in equation $\eqref{eq:r2},\eqref{eq:r1}$ or equation $\eqref{eq:r2},\eqref{eq:r3}$, the effect should be about the same, though the latter should be relatively more stable:

\begin{gather} \boldsymbol{Z}_0 = \boldsymbol{I}, \quad \boldsymbol{Y}_0\boldsymbol{Z}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{Z}_{t+1} = \boldsymbol{Z}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)\label{eq:r2-rsqrt} \\[6pt] \boldsymbol{Y}_{t+1}\boldsymbol{Z}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)^2\boldsymbol{Y}_t\boldsymbol{Z}_t\label{eq:r3-rsqrt} \\[6pt] \lim_{t\to\infty} \boldsymbol{Z}_t = \boldsymbol{P}^{-1/2}\notag \end{gather}

Reference code is as follows:

def mrsqrt(P, steps=6):
    YZ = P / (t := np.trace(P))
    Z = I = np.eye(P.shape[0])
    for a, b, c in abc(steps):
        W = a * I + b * YZ + c * YZ @ YZ
        Z, YZ = Z @ W, W @ W @ YZ
    return Z / t**0.5

d = 100
P = (x := np.random.randn(d, d) / d**0.5) @ x.T
np.abs(mrsqrt(P) @ mrsqrt(P) @ P - np.eye(d)).mean()  # ~= 5e-4

Matrix Multiplication

In most cases, however, computing $\boldsymbol{P}^{-1/2}$ is just an intermediate step; afterward we typically need to multiply by another matrix. Let $\boldsymbol{G}\in\mathbb{R}^{m\times n}$ be a matrix, and suppose we need to compute $\boldsymbol{G}\boldsymbol{P}^{-1/2}$. If we can treat $\boldsymbol{G}\boldsymbol{P}^{-1/2}$ as a single object to be iterated on as a whole, this often gives better numerical stability than computing $\boldsymbol{P}^{-1/2}$ separately and then performing the matrix multiplication.

Let's look carefully at equations $\eqref{eq:r2-rsqrt}$ and $\eqref{eq:r3-rsqrt}$. It's not hard to see that when we treat $\boldsymbol{Y}_t\boldsymbol{Z}_t$ as a whole, its iteration $\eqref{eq:r3-rsqrt}$ is actually independent of $\boldsymbol{Z}_t$, so equation $\eqref{eq:r2-rsqrt}$ for $\boldsymbol{Z}_t$ is essentially just a linear recursion! Multiplying it by a matrix on the left does not change the form of the iteration — we only need to modify the initial value — giving us

\begin{gather} \boldsymbol{Z}_0 = \boldsymbol{G}, \quad \boldsymbol{Y}_0\boldsymbol{Z}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{Z}_{t+1} = \boldsymbol{Z}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2) \label{eq:r2-final} \\[6pt] \boldsymbol{Y}_{t+1}\boldsymbol{Z}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Y}_t\boldsymbol{Z}_t + c_{t+1}(\boldsymbol{Y}_t\boldsymbol{Z}_t)^2)^2\boldsymbol{Y}_t\boldsymbol{Z}_t \label{eq:r3-final}\\[6pt] \lim_{t\to\infty} \boldsymbol{Z}_t = \boldsymbol{G}\boldsymbol{P}^{-1/2}\notag \end{gather}

Reference code:

import scipy as sp

def matmul_mrsqrt(G, P, steps=6):
    YZ = P / (t := np.trace(P))
    Z, I = G, np.eye(P.shape[0])
    for a, b, c in abc(steps):
        W = a * I + b * YZ + c * YZ @ YZ
        Z, YZ = Z @ W, W @ W @ YZ
    return Z / t**0.5

d = 100
P = (x := np.random.randn(d, d) / d**0.5) @ x.T
G = np.random.randn(2 * d, d) / d**0.5
X = matmul_mrsqrt(G, P)
np.abs(X @ sp.linalg.sqrtm(P) - G).mean()  # ~= 1e-4

Now, looking back at the algorithm for computing the square root, it's not hard to see that it is in fact just another equivalent way of writing this section's iteration when $\boldsymbol{G}=\boldsymbol{P}$, i.e. $\boldsymbol{P}^{1/2}=\boldsymbol{P}\boldsymbol{P}^{-1/2}$. So although we appear to have split things into three sections discussing three separate iterations, they are all fundamentally special cases of this last iteration!

The Ultimate Generalization

Finally, we can generalize this further to the computation of $\boldsymbol{Q}^{-1/2}\boldsymbol{G}\boldsymbol{P}^{-1/2}$, where $\boldsymbol{Q}\in\mathbb{R}^{m\times m}$ is another matrix with non-negative eigenvalues. The result is as follows:

\begin{gather} \boldsymbol{G}_0 = \boldsymbol{G}, \quad \boldsymbol{Q}_0 = \boldsymbol{Q},\quad \boldsymbol{P}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{G}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Q}_t + c_{t+1}\boldsymbol{Q}_t^2)\boldsymbol{G}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) \\[6pt] \boldsymbol{Q}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{Q}_t + c_{t+1}\boldsymbol{Q}_t^2)^2\boldsymbol{Q}_t \\[6pt] \boldsymbol{P}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^2\boldsymbol{P}_t \\[6pt] \lim_{t\to\infty} \boldsymbol{G}_t = \boldsymbol{Q}^{-1/2}\boldsymbol{G}\boldsymbol{P}^{-1/2}\notag \end{gather}

Reference code:

def mrsqrt_matmul_mrsqrt(Q, G, P, steps=6):
    Q = Q / (t1 := np.trace(Q))
    P = P / (t2 := np.trace(P))
    I1, I2 = np.eye(Q.shape[0]), np.eye(P.shape[0])
    for a, b, c in abc(steps):
        W1 = a * I1 + b * Q + c * Q @ Q
        W2 = a * I2 + b * P + c * P @ P
        G, Q, P = W1 @ G @ W2, W1 @ W1 @ Q, W2 @ W2 @ P
    return G / (t1 * t2) **0.5

d = 100
Q = (x := np.random.randn(2 * d, 2 * d) / (2 * d)**0.5) @ x.T
P = (x := np.random.randn(d, d) / d**0.5) @ x.T
G = np.random.randn(2 * d, d) / d**0.5
X = mrsqrt_matmul_mrsqrt(Q, G, P)
np.abs(sp.linalg.sqrtm(Q) @ X @ sp.linalg.sqrtm(P) - G).mean()  # ~= 2e-3

I leave it to readers to complete the proof themselves based on the results of the preceding sections.

For the Shampoo optimizer, we need to compute $\boldsymbol{Q}^{-1/4}\boldsymbol{G}\boldsymbol{P}^{-1/4}$. At present, the most feasible approach seems to be to first compute $\boldsymbol{Q}^{1/2}$ and $\boldsymbol{P}^{1/2}$ separately, and then substitute them into the above iteration to obtain $(\boldsymbol{Q}^{1/2})^{-1/2}\boldsymbol{G}(\boldsymbol{P}^{1/2})^{-1/2}$. This looks computationally expensive, but in practice, during the update stage of the optimizer, compute is often not the bottleneck — as long as the algorithm can be sufficiently parallelized, the wall-clock time need not increase significantly. Conveniently, the computations of $\boldsymbol{Q}^{1/2}$ and $\boldsymbol{P}^{1/2}$ can be parallelized, and during the iteration the two matrices W1 and W2 can also be computed in parallel, so this should still be acceptable.

Of course, it will certainly be slower than Muon — after all, Shampoo's complexity has increased substantially, so it can't come entirely for free (see the follow-up post Efficient Computation of Matrix r-th Roots and Inverse r-th Roots).

Summary

This post has shown how to recast the computation of matrix square roots and inverse square roots into the form $\mcsgn$, enabling efficient computation via its Newton-Schulz iteration.