The Aesthetics of Brute Force in Approximating Matrix Functions

矩阵函数近似中的暴力美学 苏剑林 · 2026-06-25 · Mathematics function iteration approximation matrix SVD

Previously, in Efficient Computation of Matrix Square Roots and Inverse Square Roots and Efficient Computation of Matrix r-th Roots and Inverse r-th Roots, we explored efficient ways to compute matrix powers, generalizing the Newton–Schulz iteration originally used to compute $\newcommand{msign}{\mathop{\text{msign}}}\msign$ to arbitrary matrix powers. Meanwhile, a recent arXiv paper, Muonp: Muon with Fractional Spectral Powers, offers yet another way of constructing iterations to compute matrix powers, which gave me some further inspiration.

But whether it's my own approach or Muonp's, both are, broadly speaking, "usable but not general enough" — for instance, they can only compute rational powers of a matrix, and their complexity grows with the numerator and denominator of the reduced fraction representing the power. That's clearly not very satisfying scientifically. To overcome these shortcomings, this post proposes a general approximation framework that can, in principle, fit any matrix function.

Two Kinds of Functions

When we talk about matrix functions, there are actually two somewhat different meanings involved, which I'll briefly introduce here.

The first kind acts on eigenvalues (the EIG type). Let the eigendecomposition of matrix $\boldsymbol{M}$ be $\boldsymbol{Q}\boldsymbol{\Lambda}\boldsymbol{Q}^{-1}$; then we define $f[\boldsymbol{M}]\triangleq\boldsymbol{Q}f(\boldsymbol{\Lambda})\boldsymbol{Q}^{-1}$, where $f(\boldsymbol{\Lambda})$ means applying the operation $f$ elementwise along the diagonal. This is essentially what we usually mean by "matrix function" — the matrix exponential and matrix logarithm both fall into this category, and they are typically defined via power series.

The second kind acts on singular values (the SVD type). Let the singular value decomposition of matrix $\boldsymbol{M}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$; then we define $f\{\boldsymbol{M}\} \triangleq \boldsymbol{U}f(\boldsymbol{\Sigma})\boldsymbol{V}^{\top}$. Muon's $\msign$ operation, as well as the $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$ operation discussed in Computing Singular Value Clipping mclip via msign (Part 1) and Part 2, both belong to this category.

EIG-type matrix functions only apply to square matrices, and the existence of the eigendecomposition can only be guaranteed over the complex field, so the definition of $f$ usually needs to be extended to complex numbers (unless the input matrix is restricted to real symmetric matrices). SVD-type matrix functions, on the other hand, can be defined for matrices of any shape — every real matrix has a real-valued singular value decomposition, and singular values are always non-negative. These properties make the definition and analysis of SVD-type matrix functions somewhat simpler.

These two types of matrix functions each have their own use cases, and there's no fundamental difference in how hard they are to compute, though the computational details differ. For powers, eigenvalues can be raised to any positive integer power directly, but for singular values we can only efficiently compute odd powers:

\begin{equation}\begin{aligned}[] [\boldsymbol{M}]^n =&\, \boldsymbol{Q}\boldsymbol{\Lambda}^n\boldsymbol{Q}^{-1} = (\boldsymbol{Q}\boldsymbol{\Lambda}\boldsymbol{Q}^{-1})^n = \boldsymbol{M}^n \\[4pt] \{\boldsymbol{M}\}^{2n+1} =&\, \boldsymbol{U}\boldsymbol{\Sigma}^{2n+1}\boldsymbol{V}^{\top} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}(\boldsymbol{V}\boldsymbol{\Sigma}^2\boldsymbol{V}^{\top})^n = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^n \end{aligned}\end{equation}

That is, $[\boldsymbol{M}]^n$ agrees with the $\boldsymbol{M}^n$ we'd get by defining it via ordinary matrix multiplication, whereas $\{\boldsymbol{M}\}^{2n+1}$ has to be computed indirectly via $(\boldsymbol{M}^{\top}\boldsymbol{M})^n$. So when we want to approximate a matrix function using a polynomial, EIG-type functions can use polynomials of any degree, while SVD-type functions are restricted to odd-degree polynomials.

Existing Methods

In this section, we take the SVD-type fractional power $\{\boldsymbol{M}\}^{1/3}$ as an example and introduce two existing approximation schemes. For simplicity, assume that the singular values of $\boldsymbol{M}$ all lie within $[0, 1]$.

The first approach makes use of the identity $\{\boldsymbol{M}\}^{1/3} = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/3}$, and then substitutes in the results from Efficient Computation of Matrix r-th Roots and Inverse r-th Roots, giving

\begin{gather} \boldsymbol{G}_0 = \boldsymbol{M}, \quad \boldsymbol{P}_0 = \boldsymbol{M}^{\top}\boldsymbol{M} \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)\\[6pt] \boldsymbol{P}_{t+1} = (a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2)^3\boldsymbol{P}_t \end{gather}

and ultimately $\lim_{t\to\infty} \boldsymbol{G}_t = \{\boldsymbol{M}\}^{1/3}$. Note that the iteration for $\boldsymbol{P}_t$ is independent — the idea behind this scheme is to choose $a_t,b_t,c_t$ appropriately so that $\boldsymbol{P}_t\to \boldsymbol{I}$, and then the separately isolated product of $a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2$ terms will converge to $\boldsymbol{P}_0^{-1/3}$. Its main drawback is that $\boldsymbol{M}^{\top}\boldsymbol{M}$ must first be computed explicitly, which squares the condition number and consequently degrades numerical precision.

The second approach comes from Muonp: Muon with Fractional Spectral Powers. It treats the computation of $m^{1/3}$ as finding the root of the equation $x^3 - m = 0$, so all we need is to design an odd-degree polynomial iteration that solves for this root, and we can then write down the corresponding matrix iteration. The paper considers the fixed-point iteration

\begin{equation}x_{t+1} = x_t + c (m - x_t^3) \qquad\Leftrightarrow\qquad \boldsymbol{X}_{t+1} = \boldsymbol{X}_t + c (\boldsymbol{M} - \boldsymbol{X}_t\boldsymbol{X}_t^{\top}\boldsymbol{X}_t)\end{equation}

A simple analysis shows that the condition guaranteeing convergence for any $m\in[0, 1]$ is $c\leq 2/3$; if we want a fixed value, we can take $c=2/3$. What I found most inspiring about this idea is that it feeds the input $\boldsymbol{M}$ into every step of the iteration, rather than depending solely on the previous result $\boldsymbol{X}_t$ — but it still leaves many open questions. For instance, when the iteration scheme has multiple parameters, how do we determine them one by one? Can we choose different parameters at each step to speed up convergence? The answers to these questions remain unclear.

Beyond that, both approaches share a common issue: the amount of computation depends heavily on the reduced-fraction form of the target power. For example, if we want the power $0.33$, whose reduced fraction is $33/100$, then both approaches require us to somehow compute the $100$-th power of some matrix — a substantial amount of computation, even though $0.33$ is barely different from $1/3$. There shouldn't be such a large discrepancy.

The Aesthetics of Brute Force

So what we need is a more general framework — one that can derive a usable approximation for any specified matrix function, automatically produces similar results for similar matrix functions, and ideally, in principle, supports arbitrary functions beyond just powers. That's the goal of this post.

Without loss of generality, consider an SVD-type matrix function $f\{\boldsymbol{M}\}$ that maps singular values $m$ to $f(m)$. We need to construct an odd-degree polynomial iteration that approximates $f(m)$. Suppose each step depends simultaneously on the current value $x_t$ and the input $m$ (more generally, we could also add $x_1,\cdots,x_{t-1}$ into the iteration), with each iteration step of degree at most 3. Then we can construct the general iteration scheme

\begin{equation}x_{t+1} = c_{t+1,1} x_t + c_{t+1,2} m + c_{t+1,3} x_t^3 + c_{t+1,4} m^3 + c_{t+1,5} x_t^2 m + c_{t+1,6} x_t m^2\end{equation}

where $\boldsymbol{c}_{t+1} = (c_{t+1,1},\cdots,c_{t+1,6})$ are the parameters to be determined — 6 parameters per step, allowed to take different values at each step to improve the approximation. The idea here is quite direct: throw in every term we can think of, without insisting on interpretability, and let the subsequent fitting results decide which terms are actually useful. Although this lacks the elegance of an analytically derived closed-form solution, it has a certain brute-force beauty of general-purpose computation.

So how do we solve for $\boldsymbol{c}_{t+1}$? The naive approach, as in earlier posts like Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices and Newton–Schulz Iteration for the msign Operator (Part 1), is to fix a number of iteration steps $T$, treat the whole iteration as a model, pick a regression target, and then train it end-to-end with a gradient-based optimizer.

But even though this approach theoretically has a chance of getting closer to the global optimum, in practice it often runs into all sorts of difficulties — the non-linearity and non-convexity of the model after many iteration steps, combined with the randomness of initialization, can significantly affect the outcome.

A Greedy Strategy

To make things more stable, I propose a layer-by-layer greedy strategy: at each step, assuming $\boldsymbol{c}_1,\boldsymbol{c}_2,\cdots,\boldsymbol{c}_t$ is already known, we solve for $\boldsymbol{c}_{t+1}$, with the optimization objective

\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}}\boldsymbol{c}_{t+1}^* = \argmin_{\boldsymbol{c}_{t+1}} d(x_{t+1}, f)\qquad\text{s.t.}\qquad \Vert \boldsymbol{c}_{t+1}\Vert_{\infty} \leq B\end{equation}

Here $\Vert\boldsymbol{c}_{t+1}\Vert_{\infty} \leq B$ constrains the maximum absolute value of $\boldsymbol{c}_{t+1}$ to not exceed $B$, avoiding precision loss caused by extreme amplification followed by shrinkage; we could also allow a different $B_i$ for each $c_{t+1,i}$. $d(x_{t+1}, f)$ is the cost function, and we have two choices:

\begin{align}L_2:&\,\qquad \int_a^b (x_{t+1} - f(m))^2 dm \\ L_{\infty}:&\,\qquad \max_{m\in [a, b]} |x_{t+1} - f(m)| \\ \end{align}

where $[a,b]$ is the interval of singular values we care about. Clearly, $L_2$ cares about the average error, while $L_{\infty}$ cares about the maximum error. As we'll see below, both choices admit exact solutions.

Although the greedy strategy can't guarantee global optimality, it has its own advantages. First, Polar Express proves that when $f(m)=1$ (i.e., computing $\msign$), the greedy solution is the optimal solution, which suggests that in some cases the greedy solution can be quite close to optimal. Second, each step of the greedy solution is exact and progressive, ensuring that every step brings us closer to the target — we can precompute a sufficient number of $\boldsymbol{c}^*$ steps and then, as needed, truncate to any number of leading steps and still get a good approximation.

Finally, since each step of the greedy solution is approximating the target, its fluctuations remain controlled — there's no phenomenon of one step drastically amplifying something only for the next step to shrink it back down, which is especially important for numerical computation, particularly at low precision. We can also add further constraints to enforce desired properties; for instance, $\Vert \boldsymbol{c}_{t+1}\Vert_{\infty} \leq B$ is a very basic constraint, and more can be added if necessary.

Solving Step by Step

The main reason we chose these two particular cost functions is that the optimization problems they lead to can both be solved exactly. Let's walk through the solution process for each.

First, $L_2$: it requires computing an integral, which we approximate via discretization, turning it into a finite-sample regression problem. Since $x_{t+1}$ is linear in $\boldsymbol{c}_{t+1}$, this is just a linear regression problem — exactly solvable! Even adding some linear constraints on top of this only turns it into a convex quadratic program, which remains exactly solvable, and for which mature solvers are readily available.

Next, $L_{\infty}$: it requires taking the $\max$ of the error over the entire interval, which we again approximate via discretization. We can then apply a clever transformation to turn this into a linear program: introducing a new variable $z$, we have

\begin{equation}\min_{\boldsymbol{c}_{t+1}} \max_{m\in [a, b]} |x_{t+1} - f(m)| = \min_{\boldsymbol{c}_{t+1}, z} \big\{z \,\big| -z \leq x_{t+1} - f(m) \leq z, \forall m\in[a, b]\big\}\end{equation}

If we discretize $[a,b]$ into $N$ points, these samples turn into $2N$ linear inequality constraints over $7$ unknowns, under which we minimize $z$. Linear programming is even simpler and more mature to solve than quadratic programming, so this poses no real difficulty either.

Being able to convert the $L_{\infty}$ case into a linear program is also one of the advantages of the greedy strategy. If we tried to jointly optimize all the $\boldsymbol{c}$ using a gradient-descent-style optimizer, the $\max$ of $L_{\infty}$ would be a major optimization obstacle, since gradients don't propagate well through $\max$, and combined with the extreme non-linearity of the function after many iteration steps, efforts to find a global solution would likely be futile. But through the "greedy solution + linear programming" transformation, we can stably obtain an effective solution.

Reference Implementation

If we restrict ourselves to NumPy and SciPy, we can solve the bounded linear regression using scipy.optimize.lsq_linear, and for the general linear program, scipy.optimize.linprog. But given the desire for generality, simplicity of code, and the ability to add custom constraints at will, I'd recommend using a dedicated convex optimization tool such as CVXPY.

Below, taking $f(m) = m^{1/3}$ as an example, here's a reference implementation based on CVXPY:

import numpy as np
import cvxpy as cp

N = 10000  # 离散化点数
B = 10  # 参数边界
m = np.linspace(0, 1.01, N)  # 我们关心的奇异值范围是0～1，但参数估计时考虑大一点，以保证稳定性

x, f = m, m**(1 / 3)
coefs = []
for t in range(10):
    A = np.array([x, m, x**3, m**3, x**2 * m, x * m**2]).T
    c = cp.Variable(6)
    objective = cp.Minimize(cp.sum_squares(A @ c - f))  # L2
    # objective = cp.Minimize(cp.max(cp.abs(A @ c - f)))  # L∞
    constraints = [cp.abs(c) <= B]
    problem = cp.Problem(objective, constraints)
    result = problem.solve()
    coefs.append(c.value.round(3))
    x = A @ c.value.round(3)
    print(f'iter {t + 1}, max error:', np.abs(x - f).max())
    print(f'iter {t + 1}, mse error:', np.square(x - f).mean())

coefs = np.array(coefs)

Comparison of Results

Below we compare the greedy solution for $L_2$, the greedy solution for $L_{\infty}$, and Muonp's fixed-step iteration $x_{t+1} = x_t + \frac{2}{3}(m - x_t^3)$, all applied to the cube root $f(m)=m^{1/3}$. All methods start from $x_0=m$, with discretization interval $[0, 1.01]$ and parameter bound $B=10$; $T$ is the total number of iteration steps.

$L_2$ greedy solution, 10 iterations: max error about $8.5\times 10^{-2}$, mean squared error about $5.4\times 10^{-5}$;
$L_{\infty}$ greedy solution, 10 iterations: max error about $4.6\times 10^{-2}$, mean squared error about $8.2\times 10^{-4}$;
Muonp, 10 iterations: max error about $1.4\times 10^{-1}$, mean squared error about $6.3\times 10^{-4}$;
Muonp, 20 iterations: max error about $1.0\times 10^{-1}$, mean squared error about $1.3\times 10^{-4}$.

Here's a comparison plot:

Comparison of cube-root approximation methods

Numerically, the $L_{\infty}$ greedy solution has a clear advantage in terms of maximum error, with $L_2$ coming in second; the fixed-step Muonp scheme, even at $T=20$, has a worse maximum error than the $L_2$ greedy solution at $T=10$. From the plot, we can see that Muonp's weak region is mainly near zero — this is because although $c=2/3$ takes the whole interval $[0,1]$ into account, it's too small near zero, which is the main drawback of a static step size.

By simply swapping $1/3$ for $1/5$ in the code, we can get an approximation for the fifth root, keeping the same iteration order — whereas with either the Muonp scheme or my earlier $r$-th root algorithm, the iteration order would need to be raised to at least 5. This illustrates just how general this framework is.

Some Results

Below are the iteration coefficients for the greedy solution $L_{\infty}$, computed for $\{\boldsymbol{M}\}^0$ ($\msign$), $\{\boldsymbol{M}\}^{1/2}$, $\{\boldsymbol{M}\}^{1/3}$, and $\{\boldsymbol{M}\}^{1/4}$. All coefficients are rounded to three decimal places, with parameter bound $B=10$. The fitting interval for $m^{1/2}, m^{1/3}, m^{1/4}$ is $[0, 1.01]$; for $m^0$ (i.e., the constant function 1), the fitting interval is $[0.001, 1.01]$.

Each iteration step takes the form

\begin{equation}\begin{aligned} \boldsymbol{X}_{t+1} =&\, c_{t+1, 1} \boldsymbol{X}_t + c_{t+1, 2} \boldsymbol{M} + c_{t+1, 3} \boldsymbol{X}_t \boldsymbol{X}_t^{\top} \boldsymbol{X}_t \\[4pt] &\, + c_{t+1, 4} \boldsymbol{M} \boldsymbol{M}^{\top} \boldsymbol{M} + c_{t+1, 5} \boldsymbol{X}_t \boldsymbol{X}_t^{\top}\boldsymbol{M} + c_{t+1, 6} \boldsymbol{X}_t \boldsymbol{M}^{\top}\boldsymbol{M} \end{aligned}\end{equation}

where $\boldsymbol{X}_0 = \boldsymbol{M}$, assuming that the singular values of $\boldsymbol{M}$ have already been normalized to lie within $[0,1]$. Note that $\boldsymbol{M}^{\top} \boldsymbol{M}$ and $\boldsymbol{M} \boldsymbol{M}^{\top} \boldsymbol{M}$ must be computed in the first step and can be cached; $c_{t+1, 3} \boldsymbol{X}_t \boldsymbol{X}_t^{\top} \boldsymbol{X}_t$ and $c_{t+1, 5} \boldsymbol{X}_t \boldsymbol{X}_t^{\top}\boldsymbol{M}$ can be merged into $ \boldsymbol{X}_t \boldsymbol{X}_t^{\top} (c_{t+1, 3}\boldsymbol{X}_t + c_{t+1, 5}\boldsymbol{M})$，所以这个迭代虽然有六项，但实际相比Muonp的迭代只增加一步矩阵乘法$\boldsymbol{X}_t \boldsymbol{M}^{\top}\boldsymbol{M}$.

The coefficients are as follows:

$$\begin{array}{c|c|cccccc} \hline & t & c_{t,1} & c_{t,2} & c_{t,3} & c_{t,4} & c_{t,5} & c_{t,6} \\ \hline & 1 & 2.564 & 2.564 & -1.256 & -1.256 & -1.256 & -1.256 \\ & 2 & 2.668 & 1.243 & -1.987 & 0.715 & 7.158 & -10.000 \\ & 3 & 2.670 & -0.793 & -0.791 & 3.138 & 1.727 & -4.170 \\ & 4 & 2.487 & 0.018 & -0.635 & 0.033 & 0.019 & -0.060 \\ \{\boldsymbol{M}\}^0 & 5 & 2.350 & 0.017 & -0.616 & -0.013 & -0.003 & 0.002 \\ & 6 & 2.095 & 0.001 & -0.582 & 0.001 & 0.000 & -0.002 \\ & 7 & 1.761 & -0.000 & -0.537 & 0.001 & 0.001 & -0.002 \\ & 8 & 1.547 & -0.001 & -0.508 & 0.001 & 0.001 & -0.001 \\ & 9 & 1.503 & -0.000 & -0.501 & -0.000 & 0.000 & 0.000 \\ & 10 & 1.499 & -0.000 & -0.499 & 0.000 & 0.000 & -0.000 \\ \hline & 1 & 0.931 & 0.931 & -0.245 & -0.245 & -0.245 & -0.245 \\ & 2 & 1.366 & 0.368 & -5.209 & 0.216 & 10.000 & -5.723 \\ & 3 & 1.495 & -0.070 & -5.241 & 0.379 & 10.000 & -5.557 \\ & 4 & 1.261 & 0.402 & -5.486 & -0.066 & 10.000 & -5.143 \\ \{\boldsymbol{M}\}^{1/2} & 5 & 1.145 & 0.370 & -4.675 & 1.648 & 10.000 & -7.530 \\ & 6 & 1.112 & 0.429 & -4.937 & 1.000 & 10.000 & -6.635 \\ & 7 & 1.071 & 0.535 & -5.013 & 1.129 & 10.000 & -6.762 \\ & 8 & 1.060 & 0.427 & -4.697 & 1.472 & 10.000 & -7.288 \\ & 9 & 1.035 & 0.651 & -5.296 & 0.695 & 10.000 & -6.123 \\ & 10 & 1.027 & 0.506 & -4.613 & 1.888 & 10.000 & -7.839 \\ \hline & 1 & 1.199 & 1.199 & -0.405 & -0.405 & -0.405 & -0.405 \\ & 2 & 1.168 & 1.717 & -4.360 & 0.767 & 10.000 & -8.393 \\ & 3 & 1.897 & -0.485 & -4.022 & 2.809 & 10.000 & -9.167 \\ & 4 & 1.702 & -0.140 & -4.187 & 2.457 & 10.000 & -8.854 \\ \{\boldsymbol{M}\}^{1/3} & 5 & 1.525 & 0.087 & -4.142 & 2.415 & 10.000 & -8.921 \\ & 6 & 1.392 & 0.406 & -4.142 & 2.750 & 10.000 & -9.500 \\ & 7 & 1.292 & 0.384 & -3.870 & 3.122 & 10.000 & -10.000 \\ & 8 & 1.245 & 0.809 & -4.187 & 3.013 & 10.000 & -10.000 \\ & 9 & 1.173 & 0.517 & -3.433 & 3.409 & 9.255 & -10.000 \\ & 10 & 1.180 & 1.091 & -4.355 & 2.894 & 10.000 & -9.933 \\ \hline & 1 & 1.385 & 1.385 & -0.517 & -0.517 & -0.517 & -0.517 \\ & 2 & 1.398 & 1.910 & -3.967 & 1.031 & 10.000 & -9.554 \\ & 3 & 2.155 & -0.986 & -3.272 & 3.804 & 9.311 & -10.000 \\ & 4 & 1.957 & -0.663 & -3.329 & 3.625 & 9.393 & -10.000 \\ \{\boldsymbol{M}\}^{1/4} & 5 & 1.789 & -0.099 & -3.540 & 3.340 & 9.381 & -10.000 \\ & 6 & 1.640 & 0.020 & -3.559 & 3.267 & 9.537 & -10.000 \\ & 7 & 1.517 & 0.324 & -3.590 & 3.190 & 9.429 & -10.000 \\ & 8 & 1.405 & 0.343 & -3.395 & 3.280 & 9.260 & -10.000 \\ & 9 & 1.345 & 0.651 & -3.514 & 3.224 & 9.145 & -10.000 \\ & 10 & 1.258 & 0.748 & -3.030 & 3.640 & 8.201 & -10.000 \\ \hline \end{array}$$

Summary

This post has proposed a general framework for approximating matrix functions based on a greedy strategy — one that doesn't strive for strict interpretability, but instead directly constructs simple polynomial iterations. Using the greedy strategy, each step regresses toward the target; by choosing a suitable cost function, the problem is turned into a quadratic or linear program, which can be solved exactly and stably for the iteration parameters, ultimately yielding an effective computational scheme.