Efficient Computation of Matrix r-th Roots and Inverse r-th Roots

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

In the previous post, Efficient Computation of Matrix Square Roots and Inverse Square Roots, I started from the $\newcommand{mcsgn}{\mathop{\text{mcsgn}}}\mcsgn$ operator and proposed a rather elegant method for computing matrix square roots and inverse square roots. What's curious is that, after simplification, the final formula no longer bears any resemblance to the original $\mcsgn$ form. This naturally raises a deeper question: what is the more fundamental mechanism underlying this scheme? And is it possible to generalize it to arbitrary $r$-th roots?

After analyzing the problem from this angle, I was pleasantly surprised to find that we can understand the previous iterative algorithm from a much simpler perspective — and from this new perspective, it generalizes easily to computing arbitrary $r$-th roots and inverse $r$-th roots. Let's walk through this process.

Recap

Let $\boldsymbol{G}\in\mathbb{R}^{m\times n}$ be an arbitrary matrix, and let $\boldsymbol{P}\in\mathbb{R}^{n\times n}$ be a matrix whose eigenvalues all lie within $[0,1]$. The previous post gave:

\begin{gather} \boldsymbol{G}_0 = \boldsymbol{G}, \quad \boldsymbol{P}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{G}_{t+1} = \boldsymbol{G}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) \label{eq:r2-rsqrt}\\[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 \label{eq:r3-rsqrt}\\[6pt] \lim_{t\to\infty} \boldsymbol{G}_t = \boldsymbol{G}\boldsymbol{P}^{-1/2}\notag \end{gather}more

Substituting $\boldsymbol{G}=\boldsymbol{P}$ gives us $\boldsymbol{P}^{1/2}$, and substituting $\boldsymbol{G}=\boldsymbol{I}$ gives us $\boldsymbol{P}^{-1/2}$. On closer inspection, we notice that the above iteration is actually a manifestation of the following limit:

\begin{equation} \prod_{t=0}^{\infty}(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) = \boldsymbol{P}^{-1/2}\label{eq:prod-rsqrt}\end{equation}

Interestingly, proving this limit directly is not complicated: just take the square root of both sides of equation $\eqref{eq:r3-rsqrt}$, then substitute into the equation above to get

\begin{equation} \prod_{t=0}^{\infty}(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2) = \prod_{t=0}^{\infty} \boldsymbol{P}_{t+1}^{1/2}\boldsymbol{P}_t^{-1/2} = \lim_{t\to\infty} \boldsymbol{P}_t^{1/2}\boldsymbol{P}_0^{-1/2} = \lim_{t\to\infty} \boldsymbol{P}_t^{1/2}\boldsymbol{P}^{-1/2}\end{equation}

From this we can see that as long as the sequence $\{\boldsymbol{P}_t\}$ remains invertible throughout and eventually converges to $\boldsymbol{I}$, the limit $\eqref{eq:prod-rsqrt}$ holds automatically. As for how the iteration $\eqref{eq:r3-rsqrt}$ ensures that $\{\boldsymbol{P}_t\}$ satisfies these two conditions, we'll discuss that shortly.

General Form

Let us consider the more general iteration

\begin{gather} \boldsymbol{G}_0 = \boldsymbol{G}, \quad \boldsymbol{P}_0 = \boldsymbol{P} \notag\\[6pt] \boldsymbol{G}_{t+1} = \boldsymbol{G}_t(a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^s\\[6pt] \boldsymbol{P}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^r\boldsymbol{P}_t \end{gather}

Similarly, if the sequence $\{\boldsymbol{P}_t\}$ remains invertible throughout and eventually converges to $\boldsymbol{I}$, then we can prove that

\begin{equation}\lim_{t\to\infty} \boldsymbol{G}_t = \boldsymbol{G}\boldsymbol{P}^{-s/r}\end{equation}

holds. This gives us a general iterative form for computing arbitrary $-s/r$-th powers of a matrix. Building on this result, we only need to choose $\boldsymbol{G}=\boldsymbol{P}, s=r-1$ to obtain $\boldsymbol{P}^{1/r}$, so all that remains is to focus on solving for the inverse $0\sim 1$-th power.

The problem then becomes: how do we choose an appropriate $\{a_t,b_t,c_t\}$ so that the sequence $\{\boldsymbol{P}_t\}$ converges to $\boldsymbol{I}$ as quickly as possible? The faster the convergence, the fewer iteration steps we need to reach a given precision.

Iteration Coefficients

By assumption, $\boldsymbol{P}_0 = \boldsymbol{P}$ is a matrix whose eigenvalues all lie within $[0,1]$, while the target matrix $\boldsymbol{I}$ has all eigenvalues equal to 1. So the sequence $\{\boldsymbol{P}_t\}$ is really the process of transforming eigenvalues from arbitrary values within $[0,1]$ to $1$ — which is exactly what $\mcsgn$ does!

Let us set $\boldsymbol{X}_t = \boldsymbol{P}_t^{1/r}$, so that $\boldsymbol{X}_0 = \boldsymbol{P}^{1/r}$ is likewise a matrix whose eigenvalues all lie within $[0,1]$, and the iterative equation becomes

\begin{equation}\boldsymbol{X}_{t+1} = a_{t+1}\boldsymbol{X}_t + b_{t+1}\boldsymbol{X}_t^{r+1} + c_{t+1}\boldsymbol{X}_t^{2r+1}\end{equation}

The problem now becomes how to make $\boldsymbol{X}_0$ approach $\boldsymbol{I}$ as quickly as possible — which is essentially the same problem we discussed in Newton–Schulz Iteration for the msign Operator (Part 1) and Newton–Schulz Iteration for the msign Operator (Part 2). There, the "second part" gave the theoretically optimal solution for $r=2$, but both the derivation process and the conclusion generalize to arbitrary $r$.

Specifically, we first convert the problem into a scalar iteration:

\begin{equation}x_{t+1} = f_t(x_t) = a_{t+1}x_t + b_{t+1}x_t^{r+1} + c_{t+1}x_t^{2r+1}\end{equation}

then show that the greedy solution is in fact optimal, and finding the greedy solution reduces to solving the equation

\begin{equation}\begin{gathered} f_t(l_t) = 1 - \mathcal{E}, \quad f_t(u_t) = 1 + \mathcal{E} \\ f_t(x_1) = 1 + \mathcal{E}, \quad f_t(x_2) = 1 - \mathcal{E} \\ f_t'(x_1) = 0, \quad f_t'(x_2) = 0 \end{gathered}\end{equation}

For simplicity, we parameterize $f_t$ as

\begin{equation}f_t'(x) = k(x^r-x_1^r)(x^r-x_2^r)\end{equation}

and then, just as in the "second part", we can solve it using Mathematica.

Initialization Analysis

Before proceeding to the formal solution, we need to analyze the initialization. In the previous post, Efficient Computation of Matrix Square Roots and Inverse Square Roots, we mentioned that under the assumption that all eigenvalues of $\boldsymbol{P}$ are non-negative, we can compress all the eigenvalues into $[0,1]$ by dividing by $\newcommand{tr}{\mathop{\text{tr}}}\tr(\boldsymbol{P})$. However, this compression ratio is often excessive; in this post, we instead use

\begin{equation}\boldsymbol{P}_0 = \frac{\boldsymbol{P}}{\sqrt{\tr(\boldsymbol{P}^2)}}\end{equation}

We know that $\tr(\boldsymbol{P}^2)$ equals the sum of squares of all eigenvalues, while $\tr(\boldsymbol{P})^2$ equals the square of the sum of all eigenvalues, and when the eigenvalues are non-negative, $\tr(\boldsymbol{P}^2)\leq\tr(\boldsymbol{P})^2$ always holds. So the above expression provides a tighter initial value. In particular, computing $\tr(\boldsymbol{P}^2)$ does not require explicitly computing $\boldsymbol{P}^2$, since we have the identity

\begin{equation}\tr(\boldsymbol{P}^2) = \langle \boldsymbol{P}, \boldsymbol{P}^{\top}\rangle_F\end{equation}

Next, we need to analyze how small the eigenvalues we need to handle can be — this is the same as the initial singular value analysis in Newton–Schulz Iteration for the msign Operator (Part 1). After dividing by $\sqrt{\tr(\boldsymbol{P}^2)}$, the eigenvalues of $\boldsymbol{P}_0$ form a unit vector; if all the eigenvalues were equal, each would be $1/\sqrt{n}$. By the pigeonhole principle, in general there must exist an eigenvalue smaller than $1/\sqrt{n}$, so to be conservative we accommodate down to $0.01/\sqrt{n}$.

For a sufficiently large LLM, $n$ is already on the order of $100^2$, so we need to accommodate down to $0.0001$. Note that this is just the eigenvalue of $\boldsymbol{P}_0$, and since $\boldsymbol{X}_0 = \boldsymbol{P}_0^{1/r}$, for $\boldsymbol{X}_0$ we only need to accommodate down to $0.0001^{1/r}$ — which is somewhat more favorable than the case of $\mcsgn$ and $\newcommand{msign}{\mathop{\text{msign}}}\msign$, because the input of $\mcsgn$ and $\msign$ is $\boldsymbol{X}_0$, requiring us to accommodate small eigenvalues down to $\boldsymbol{X}_0$, whereas here the input is $\boldsymbol{P}_0$, so we only need to start our consideration from $\boldsymbol{P}_0$.

Computation Results

Taking all of the above considerations into account, our final solution code is as follows:

r = 4;
df[x_] = k*(x^r - x1^r) (x^r - x2^r);
f[x_] = Integrate[df[x], {x, 0, x}];
sol[l_, u_] := 
 NSolve[{f[l] == 1 - e, f[x1] == 1 + e, f[x2] == 1 - e, f[u] == 1 + e,
    l < x1 < x2 < u, e > 0, k > 0}, {k, x1, x2, e}]
ff[x_, l_, u_] = f[x]*2/(f[l] + f[u]) // Expand;
lt = 0.0001^(1/r); ut = 1; lambda = 0.1;
While[1 - lt > 0.0001,
 fff[x_] = ff[x, lt, ut] /. sol[Max[lt, lambda*ut], ut][[1]];
 Print[fff[x]];
 lt = fff[lt]; ut = 2 - lt]
f[x] /. Solve[f[1] == 1, k][[1]] /. {x1 -> 1, x2 -> 1}

The computed results for $r=1\sim 5$ are as follows:

$$\begin{array}{c|ccc} \hline r & t & a & b & c \\ \hline & \quad 1\quad & 14.2975 & -31.2203 & 18.9214 \\ & 2 & 7.12258 & -7.78207 & 2.35989 \\ \quad 1\quad & 3 & 6.9396 & -7.61544 & 2.3195 \\ & 4 & 5.98456 & -6.77016 & 2.12571 \\ & 5 & 3.79109 & -4.18664 & 1.39555 \\ & \geq 6 & 3 & -3 & 1 \\ \hline & 1 & 7.42487 & -18.3958 & 12.8967 \\ & 2 & 3.48773 & -2.33004 & 0.440469 \\ 2 & 3 & 2.77661 & -2.07064 & 0.463023 \\ & 4 & 1.99131 & -1.37394 & 0.387593 \\ & \geq 5 & 15/8 & -5/4 & 3/8 \\ \hline & 1 & 5.05052 & -13.5427 & 10.2579 \\ & 2 & 2.31728 & -1.06581 & 0.144441 \\ 3 & 3 & 1.79293 & -0.913562 & 0.186699 \\ & 4 & 1.56683 & -0.786609 & 0.220008 \\ & \geq 5 & 14/9 & -7/9 & 2/9 \\ \hline & 1 & 3.85003 & -10.8539 & 8.61893 \\ 4 & 2 & 1.80992 & -0.587778 & 0.0647852 \\ & 3 & 1.50394 & -0.594516 & 0.121161 \\ & \geq 4 & 45/32 & -9/16 & 5/32 \\ \hline & 1 & 3.11194 & -8.28217 & 6.67716 \\ 5 & 2 & 1.5752 & -0.393327 & 0.0380364 \\ & 3 & 1.3736 & -0.44661 & 0.0911259 \\ & \geq 4 & 33/25 & -11/25 & 3/25 \\ \hline \end{array}$$

where the convergence value in the last step is derived from $x_1=x_2=1$ and $f(1)=1$.

Let's Test It

Here is a simple test script:

import numpy as np
import jax.numpy as jnp

coefs = [
    None,
    [
        (14.2975, -31.2203, 18.9214),
        (7.12258, -7.78207, 2.35989),
        (6.9396, -7.61544, 2.3195),
        (5.98456, -6.77016, 2.12571),
        (3.79109, -4.18664, 1.39555),
        (3, -3, 1),
    ],
    [
        (7.42487, -18.3958, 12.8967),
        (3.48773, -2.33004, 0.440469),
        (2.77661, -2.07064, 0.463023),
        (1.99131, -1.37394, 0.387593),
        (15 / 8, -5 / 4, 3 / 8),
    ],
    [
        (5.05052, -13.5427, 10.2579),
        (2.31728, -1.06581, 0.144441),
        (1.79293, -0.913562, 0.186699),
        (1.56683, -0.786609, 0.220008),
        (14 / 9, -7 / 9, 2 / 9),
    ],
    [
        (3.85003, -10.8539, 8.61893),
        (1.80992, -0.587778, 0.0647852),
        (1.50394, -0.594516, 0.121161),
        (45 / 32, -9 / 16, 5 / 32),
    ],
    [
        (3.11194, -8.28217, 6.67716),
        (1.5752, -0.393327, 0.0380364),
        (1.3736, -0.44661, 0.0911259),
        (33 / 25, -11 / 25, 3 / 25),
    ],
]

def abc(r=1, steps=None, scale=1):
    w, steps = coefs[r], steps or len(coefs[r])
    for a, b, c in w[:steps] + w[-1:] * max(steps - len(w), 0):
        yield a / scale, b / scale**(r + 1), c / scale**(2 * r + 1)

def matmul_invroot(G, P, r, s=1, steps=None, eps=1e-5):
    """return G @ P^(-s/r)
    """
    I = jnp.eye(P.shape[0], dtype=P.dtype)
    P = P / (t := (P * P.mT).sum()**0.5) + eps * I
    for a, b, c in abc(r, steps, 1.001):
        W = a * I + b * P + c * P @ P
        W1, W2 = jnp.linalg.matrix_power(W, s), jnp.linalg.matrix_power(W, r)
        G, P = G @ W1, P @ W2
    return G * t**(-s / r)

def matmul_invroot_by_eigh(G, P, r, s=1):
    """return G @ P^(-s/r)
    """
    S, Q = jnp.linalg.eigh(P)
    return G @ Q @ jnp.diag(S**(-s / r)) @ jnp.linalg.inv(Q)

d = 1000
s, r = 1, 4
G = np.random.randn(2 * d, d) / d**0.5
P = (x := np.random.randn(d, d) / d**0.5) @ x.T + 0.001 * np.eye(d)

X1 = matmul_invroot_by_eigh(G, P, r, s)
X2 = matmul_invroot(G, P, r, s, eps=0)
jnp.abs(X1 - X2).mean()  # ~= 1e-3

X2 = matmul_invroot(jnp.array(G, dtype='bfloat16'), jnp.array(P, dtype='bfloat16'), r, s, eps=0)
jnp.abs(X1 - X2).mean()  # ~= 2e-3

There are a few points worth noting here. First, the smallest eigenvalue of the input $\boldsymbol{P}$ must not be too small, or the iteration process can blow up very easily — even when we only want to compute a positive power such as $\boldsymbol{P}^{1/2}$. This is actually not hard to understand: $\sqrt{x}$ is somewhat ill-behaved near $x=0$ too, and if numerical error causes it to "accidentally" slip onto the negative real axis, it simply has no (real) solution there, at which point the iteration's behavior becomes unpredictable.

How small is "too small"? Roughly speaking, the smallest eigenvalue of $\boldsymbol{P}/\sqrt{\tr(\boldsymbol{P}^2)}$ clearly must not be smaller than the minimum eigenvalue we're accounting for, i.e., $0.0001$. If this can't be guaranteed, it is recommended to simply set

\begin{equation}\boldsymbol{P}_0 = \frac{\boldsymbol{P}}{\sqrt{\tr(\boldsymbol{P}^2)}} + \epsilon \cdot\boldsymbol{I} \end{equation}

where $\epsilon\sim 0.0001$. This sacrifices a bit of precision but noticeably improves numerical stability.

Additionally, in most cases the number of iteration steps does not need to exceed the recommended value len(coefs[r]), especially in low-precision computation scenarios, since more iteration steps make it easier for accumulated error to cause blow-up. In fact, as long as the eigenvalues stay within the range we've accounted for, the recommended number of steps is already sufficient to reach the desired precision — unless we're iterating in fp32 or even higher precision, in which case we might consider setting $\epsilon=0$, scale=1, and using more iteration steps.

Summary

In this post, we generalized the result from the previous post to the computation of arbitrary $r$-th roots and inverse $r$-th roots, arriving at a general iterative scheme for computing arbitrary $-1/r$-th powers of a matrix.