The Brute-Force Aesthetics of Matrix Function Approximation

By 苏剑林 | June 25, 2026

Previously, in "Efficient Calculation of Matrix Square Roots and Inverse Square Roots" and "Efficient Calculation of Matrix r-th Roots and Inverse r-th Roots", we explored efficient matrix power calculations. These generalized the Newton-Schulz iteration, originally used for $\newcommand{msign}{\mathop{\text{msign}}}\msign$, to matrix powers. More recently, the paper "Muonp: Muon with Fractional Spectral Powers" posted on arXiv provided another approach to constructing iterations for matrix powers, which gave me some inspiration.

However, whether it is my approach or the Muonp approach, generally speaking, they are "usable but not universal enough"—for example, they can only calculate rational powers of matrices, and the complexity increases as the numerator and denominator of the simplified fraction increase, which is clearly not scientific. To overcome these shortcomings, this article proposes a general approximation framework that can theoretically fit any matrix function.

Two Types of Functions

When speaking of matrix functions, they actually have two slightly different meanings. Let's briefly introduce them here.

The first type operates on eigenvalues (EIG-type). Suppose the eigenvalue decomposition of matrix $\boldsymbol{M}$ is $\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$ element-wise to the diagonal. This is actually what we usually call a matrix function; for instance, the matrix exponential and matrix logarithm belong to this category, and they are usually defined by power series.

The second type operates on singular values (SVD-type). Suppose the singular value decomposition of matrix $\boldsymbol{M}$ is $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, then we define $f\{\boldsymbol{M}\} \triangleq \boldsymbol{U}f(\boldsymbol{\Sigma})\boldsymbol{V}^{\top}$. The $\msign$ operation in Muon, and the $\newcommand{mclip}{\mathop{\text{mclip}}}\mclip$ operation discussed in "Calculating Singular Value Clipping mclip via msign (Part 1)" and "Calculating Singular Value Clipping mclip via msign (Part 2)", belong to this category.

EIG-type matrix functions are only valid for square matrices, and the existence of eigenvalue decomposition is only guaranteed in 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 can be defined for matrices of any shape; any real matrix has a singular value decomposition within the real domain, and singular values are always non-negative. These characteristics make the definition and analysis of SVD-type matrix functions relatively simpler.

These two types of matrix functions each have their own application scenarios. Their computational difficulty is not fundamentally different, but there are differences in the details of calculation. For powers, any positive integer power of eigenvalues can be calculated directly, but for singular values, we can only efficiently calculate odd powers:

That is, $[\boldsymbol{M}]^n$ is consistent with the $\boldsymbol{M}^n$ defined by matrix multiplication, while $\{\boldsymbol{M}\}^{2n+1}$ needs to be calculated indirectly via $(\boldsymbol{M}^{\top}\boldsymbol{M})^n$. Therefore, when we want to approximate a matrix function through polynomials, EIG-type functions can use polynomials of any degree, while SVD-type functions can only use odd-degree polynomials.

Existing Methods

In this section, we use the SVD-type fractional power $\{\boldsymbol{M}\}^{1/3}$ as an example to introduce two existing approximation schemes. For simplicity, assume the singular values of $\boldsymbol{M}$ are in the range $[0, 1]$.

The first idea is to use the identity $\{\boldsymbol{M}\}^{1/3} = \boldsymbol{M}(\boldsymbol{M}^{\top}\boldsymbol{M})^{-1/3}$, and then substitute the results from "Efficient Calculation of Matrix r-th Roots and Inverse r-th Roots", yielding:

Ultimately, $\lim_{t\to\infty} \boldsymbol{G}_t = \{\boldsymbol{M}\}^{1/3}$. Note that the iteration for $\boldsymbol{P}_t$ is independent. The principle of this scheme is to choose appropriate $a_t, b_t, c_t$ so that $\boldsymbol{P}_t \to \boldsymbol{I}$, and then the continuous multiplication of the isolated term $a_{t+1}\boldsymbol{I} + b_{t+1}\boldsymbol{P}_t + c_{t+1}\boldsymbol{P}_t^2$ tends toward $\boldsymbol{P}_0^{-1/3}$. Its main problem is the need to explicitly calculate $\boldsymbol{M}^{\top}\boldsymbol{M}$ first, which squares the condition number and leads to a decrease in computational precision.

The second idea comes from "Muonp: Muon with Fractional Spectral Powers". It treats finding $m^{1/3}$ as the root of the equation $x^3 - m = 0$. Then we only need to design an odd-degree polynomial iteration to find the root of that equation to write the corresponding matrix iteration. The paper considers a fixed-point iteration:

Simple analysis proves that the condition for convergence for any $m \in [0, 1]$ is $c \leq 2/3$; if taking a fixed value, $c=2/3$ can be considered. The biggest inspiration this idea gives me is to include the input $\boldsymbol{M}$ in each iteration step, rather than relying solely on the previous result $\boldsymbol{X}_t$. However, it still leaves many unresolved issues. For example, when the iteration format has multiple parameters, how to determine these parameters one by one, and whether different parameters can be chosen at each step to improve the convergence speed—the answers to these questions are unknown.

In addition, both of the above ideas share a common problem: the amount of calculation depends strongly on the simplest fraction form of the power sought. For example, if we want to find the $0.33$ power, the simplest fraction is $33/100$, so both ideas would require us to find the $100$-th power of some matrix, which is a considerable amount of calculation. But in reality, $0.33$ is almost the same as $1/3$; there shouldn't be such a big difference.

Brute-Force Aesthetics

Therefore, we need a more general framework capable of deriving a usable approximation for a specified matrix function, and providing similar results for similar matrix functions. Ideally, it should theoretically support any function besides power functions. This is the goal of this article.

Without loss of generality, consider the SVD-type matrix function $f\{\boldsymbol{M}\}$, which transforms the singular value $m$ into $f(m)$. We need to construct an odd-degree polynomial iteration to approximate $f(m)$. Assuming each step depends simultaneously on the current value $x_t$ and the input $m$ (generally, we could also consider adding $x_1, \dots, x_{t-1}$ to the iteration), and the degree of each iteration step does not exceed 3, we can construct a general iteration format:

where $\boldsymbol{c}_{t+1} = (c_{t+1,1}, \dots, c_{t+1,6})$ are the parameters to be found. There are 6 parameters per step, and we allow each step to take different values to improve the degree of approximation. The logic here is straightforward: include all terms we can think of without forcing interpretability; let the subsequent fitting results decide whether they are useful. While this lacks the elegance of a theoretical analytical solution, it possesses a brute-force aesthetic of universal computation.

How do we solve for $\boldsymbol{c}_{t+1}$ next? A naive approach would be like previous articles "Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices" and "Newton-Schulz Iteration for the msign Operator (Part 1)": first fix a number of iteration steps $T$, which allows the entire iteration to be viewed as a model, and then choose a regression target to train end-to-end using a gradient optimizer.

However, although this theoretically offers a chance to get closer to the global optimum, in practice it often faces many difficulties—the nonlinearity and non-convexity of the model after multiple iterations, as well as the randomness of initialization, which significantly affect the results.

Greedy Strategy

To stabilize the results, I propose a layer-by-layer greedy strategy—solving for $\boldsymbol{c}_{t+1}$ at each step assuming $\boldsymbol{c}_1, \boldsymbol{c}_2, \dots, \boldsymbol{c}_t$ are already known. The optimization target is:

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 reduction. We can also consider choosing different $B_i$ for each $c_{t+1,i}$. For the cost function $d(x_{t+1}, f)$, we have two choices:

where $[a,b]$ is the interval of singular values we care about. Obviously, $L_2$ focuses on the average error, while $L_{\infty}$ focuses on the maximum error. As we will see later, we can solve both choice cases exactly.

Although the greedy strategy cannot guarantee the optimum, it has unique advantages. First, Polar Express proved that if $f(m)=1$ (i.e., calculating $\msign$), the greedy solution is the optimal solution, indicating that in some cases greedy solutions can be very close to the optimum. Second, the greedy solution for each step is exact and progressive, ensuring every step gets closer to the target. We can pre-calculate enough $\boldsymbol{c}^*$ and truncate at any step to get a good approximation.

Finally, since each step of the greedy solution approximates the target, its fluctuations are controllable, and there won't be a situation where one step amplifies drastically only for the next step to shrink, which is particularly important for numerical calculations, especially low-precision ones. We can continue adding constraints to achieve desired properties, such as $\Vert \boldsymbol{c}_{t+1}\Vert_{\infty} \leq B$ being a very basic constraint. More can be added if necessary.

Solving One by One

The main reason for choosing these two cost functions is that their corresponding optimization problems can be solved exactly. Let's demonstrate the solving process.

First is $L_2$, which requires calculating an integral. After approximating it with discretization, it becomes a regression problem on finite samples. Noticing that $x_{t+1}$ is linear with respect to $\boldsymbol{c}_{t+1}$, this is simply a linear regression problem, exactly solvable! Even if we add some linear constraints on this basis, it's just a convex quadratic program, which remains exactly solvable with mature tools.

Then there is $L_{\infty}$, which requires taking the $\max$ error over the entire interval. We also discretize this. Next, we can transform it into a linear programming problem through a clever substitution: introducing a new variable $z$, we have:

If we discretize $[a,b]$ into $N$ points, these samples will translate into $2N$ linear inequality constraints involving 7 unknowns. We then minimize $z$ under these constraints. Solving linear programming is even simpler and more mature than quadratic programming, so there are no real difficulties.

The successful transformation of the $L_{\infty}$ case into linear programming is one of the advantages of the greedy strategy. If we tried to use gradient descent-like optimizers to jointly optimize all $\boldsymbol{c}$, the $\max$ in $L_{\infty}$ would be a core obstacle to optimization because gradients cannot propagate effectively through $\max$. Combined with the extreme nonlinearity of the function after multiple iterations, efforts to find a global solution would likely fail. But through the "greedy solution + linear programming" transformation, we can stably obtain an effective solution.

Reference Implementation

If restricted to Numpy and Scipy, for linear regression with boundary constraints, we can use scipy.optimize.lsq_linear; for general linear programming, we can use scipy.optimize.linprog. However, considering the generality and conciseness of the code, as well as the possibility of adding custom constraints at any time, we suggest solving based on CVXPY and other specialized convex optimization tools.

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

import cvxpy as cp
import numpy as np

def solve_step_linf(m_samples, x_t_samples, target_f_samples, B=10.0):
    """
    Solve for c_{t+1} using L_infinity norm via Linear Programming.
    """
    n_samples = len(m_samples)
    # The 6 terms in the iteration
    features = np.stack([
        x_t_samples,                           # x_t
        m_samples,                             # m
        x_t_samples**3,                        # x_t^3
        m_samples**3,                          # m^3
        (x_t_samples**2) * m_samples,          # x_t^2 * m
        x_t_samples * (m_samples**2)           # x_t * m^2
    ], axis=1) # (N, 6)

    c = cp.Variable(6)
    z = cp.Variable()

    # x_{t+1} = features @ c
    x_next = features @ c

    constraints = [
        x_next - target_f_samples <= z,
        target_f_samples - x_next <= z,
        c <= B,
        c >= -B
    ]

    prob = cp.Problem(cp.Minimize(z), constraints)
    prob.solve(solver=cp.ECOS)

    return c.value, z.value

# Example setup for m^(1/3)
m = np.linspace(0, 1.01, 1000)
f_m = m**(1/3)
x_t = m.copy() # x_0 = m
T = 10
coeffs = []

for t in range(T):
    c_val, err = solve_step_linf(m, x_t, f_m)
    coeffs.append(c_val)
    # Update x_t for next step
    feat = np.stack([x_t, m, x_t**3, m**3, x_t**2 * m, x_t * m**2], axis=1)
    x_t = feat @ c_val
    print(f"Step {t+1}, Max Error: {err:.6e}")

Comparison of Results

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

$L_2$ greedy solution, 10 steps: Max error approx $8.5\times 10^{-2}$, Mean Squared Error approx $5.4\times 10^{-5}$;

$L_{\infty}$ greedy solution, 10 steps: Max error approx $4.6\times 10^{-2}$, Mean Squared Error approx $8.2\times 10^{-4}$;

Muonp iteration, 10 steps: Max error approx $1.4\times 10^{-1}$, Mean Squared Error approx $6.3\times 10^{-4}$;

Muonp iteration, 20 steps: Max error approx $1.0\times 10^{-1}$, Mean Squared Error approx $1.3\times 10^{-4}$.

Comparison chart follows:

Comparison of Cube Root Approximation Methods

Numerically, the $L_{\infty}$ greedy solution has a clear advantage in maximum error, followed by $L_2$. The fixed-step Muonp scheme at $T=20$ still has a larger maximum error than the $L_2$ greedy solution at $T=10$. As seen in the chart, the Muonp scheme's disadvantage is mainly in the region near 0. This is because $c=2/3$, while balancing the entire $[0,1]$ interval, is too small for the vicinity of 0—this is the primary flaw of static step sizes.

By changing $1/3$ to $1/5$ in the code, we can obtain an approximation for the 5th root without changing the iteration order. However, using the Muonp scheme or my previous $r$-th root algorithm would require increasing the iteration order to at least 5. This demonstrates the versatility of the framework presented here.

Some Results

Below are the iteration coefficients for the $L_{\infty}$ greedy solutions 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 boundary $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 1), the interval is $[0.001, 1.01]$.

The iteration format for each step is:

where $\boldsymbol{X}_0 = \boldsymbol{M}$, assuming the singular values of $\boldsymbol{M}$ are normalized to $[0,1]$. Note that $\boldsymbol{M}^{\top} \boldsymbol{M}$ and $\boldsymbol{M} \boldsymbol{M}^{\top} \boldsymbol{M}$ must be calculated in the first step and can be stored. $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})$. Thus, while this iteration has six terms, it actually only adds one matrix multiplication step $\boldsymbol{X}_t \boldsymbol{M}^{\top}\boldsymbol{M}$ compared to Muonp's iteration.

The coefficients are as follows:

Article Summary

This article proposes a general matrix function approximation framework based on a greedy strategy. Instead of pursuing strict interpretability, it directly constructs simple polynomial iterations and uses a greedy strategy at each step to regress to the target. By choosing appropriate cost functions, the problem is transformed into quadratic or linear programming, allowing for the exact and stable solution of iteration parameters, ultimately obtaining an effective computational scheme.