Randomized trace estimation and low-rank kernel approximation

Estimate the trace of a matrix you can only multiply by, and build a low-rank factor of a kernel matrix from a few of its columns. Hutch++ (vp.randla.trace_hutchpp, Meyer, Musco, Musco and Woodruff, 2021) spends a third of its matrix products on a sketch of the top of the spectrum and the rest on Hutchinson's estimator for what is left. Its error is guaranteed to fall at least like 1/m in the number of products m, where Hutchinson's falls like 1/sqrt(m). Randomly pivoted Cholesky (vp.randla.rp_cholesky, Chen, Epperly, Tropp and Webber, 2023) picks columns with probability proportional to the diagonal of what is still unexplained.

A matrix with a decaying spectrum

A 400 x 400 symmetric matrix with eigenvalues j^-2 in a random orthonormal basis: most of its trace sits in a few large eigenvalues, the case where Hutch++ helps most. Hutchinson's estimator is written here in NumPy with Rademacher probes, tr(A) ~ mean of g^T A g. Hutchinson's error follows 1/sqrt(m). Hutch++ beats its 1/m guarantee here: with eigenvalues falling like j^-2, the part its sketch misses shrinks as the sketch grows, and the fitted slope is about -2.

import numpy as np
import vipyr as vp
from vipyr.figure import figure

rng = np.random.default_rng(0)
n = 400
Q, _ = np.linalg.qr(rng.standard_normal((n, n)))
lam = np.arange(1, n + 1) ** -2.0
A = (Q * lam) @ Q.T
exact = lam.sum()
print("exact trace %.10f" % exact)


def hutchinson(A, m, seed):
    g = np.random.default_rng(seed).choice([-1.0, 1.0], size=(A.shape[0], m))
    return np.sum(g * (A @ g)) / m


ms = [12, 24, 48, 96, 192, 384]
seeds = range(20)
err_pp, err_h = [], []
for m in ms:
    e_pp = [
        abs(vp.randla.trace_hutchpp(A=A, m=m, seed=s).trace - exact) / exact
        for s in seeds
    ]
    e_h = [abs(hutchinson(A, m, s) - exact) / exact for s in seeds]
    err_pp.append(np.mean(e_pp))
    err_h.append(np.mean(e_h))
    print(
        "m = %3d   Hutch++ mean relative error %.2e   Hutchinson %.2e"
        % (m, err_pp[-1], err_h[-1])
    )
for name, e in (("Hutch++", err_pp), ("Hutchinson", err_h)):
    print(
        "%-10s fitted slope of log error against log m: %.2f"
        % (name, np.polyfit(np.log(ms), np.log(e), 1)[0])
    )

fig = figure(
    "Trace estimation error, mean of 20 seeds",
    "matrix-vector products m",
    "relative error",
)
fig.plot(ms, err_pp, "Hutch++ (vipyr)")
fig.plot(ms, err_h, "Hutchinson (NumPy)")
fig.plot(ms, [err_pp[0] * ms[0] / m for m in ms], "1/m")
fig.plot(ms, [err_h[0] * np.sqrt(ms[0] / m) for m in ms], "$1/\\sqrt{m}$")
fig.style("1/m", "s1 dash thin nomarker")
fig.style("$1/\\sqrt{m}$", "s2 dash thin nomarker")
fig.xscale("log")
fig.yscale("log")

exact trace 1.6424371892
m =  12   Hutch++ mean relative error 3.58e-02   Hutchinson 1.95e-01
m =  24   Hutch++ mean relative error 9.46e-03   Hutchinson 1.78e-01
m =  48   Hutch++ mean relative error 2.17e-03   Hutchinson 1.05e-01
m =  96   Hutch++ mean relative error 5.32e-04   Hutchinson 9.12e-02
m = 192   Hutch++ mean relative error 1.39e-04   Hutchinson 5.97e-02
m = 384   Hutch++ mean relative error 2.87e-05   Hutchinson 3.47e-02
Hutch++    fitted slope of log error against log m: -2.05
Hutchinson fitted slope of log error against log m: -0.50

A low-rank factor of a Gaussian kernel

The kernel matrix K_ij = exp(-|x_i - x_j|^2 / (2 h^2)) of 800 points in the unit square, h = 0.15. rp_cholesky returns F with K ~ F F^T after k pivots. The best possible rank-k error in the trace norm is the sum of the eigenvalues past the k-th, which needs the full eigendecomposition. RPCholesky reads only k columns of K and its diagonal. Its error stays within a factor of 2 to 7 of the best while both fall by four orders of magnitude between k = 10 and k = 120.

pts = np.random.default_rng(1).random((800, 2))
d2 = ((pts[:, None, :] - pts[None, :, :]) ** 2).sum(-1)
K = np.exp(-d2 / (2 * 0.15**2))
ev = np.sort(np.linalg.eigvalsh(K))[::-1]
ks = [10, 20, 40, 60, 80, 120]
rp, best = [], []
for k in ks:
    F = vp.randla.rp_cholesky(A=K, k=k, seed=7).F
    rp.append((np.trace(K) - np.sum(F * F)) / np.trace(K))
    best.append(ev[k:].sum() / np.trace(K))
    print(
        "k = %3d   RPCholesky trace error %.2e   best rank-k %.2e   ratio %.2f"
        % (k, rp[-1], best[-1], rp[-1] / best[-1])
    )

fig = figure(
    "Gaussian kernel, 800 points: trace error of the rank-k factor",
    "rank k",
    "$\\mathrm{tr}(K - FF^\\top)/\\mathrm{tr}(K)$",
)
fig.plot(ks, rp, "RPCholesky (vipyr)")
fig.plot(ks, best, "best rank k (eigendecomposition)")
fig.style("best rank k (eigendecomposition)", "s2 dash")
fig.yscale("log")

k =  10   RPCholesky trace error 5.31e-01   best rank-k 3.09e-01   ratio 1.72
k =  20   RPCholesky trace error 2.34e-01   best rank-k 1.03e-01   ratio 2.27
k =  40   RPCholesky trace error 4.93e-02   best rank-k 1.32e-02   ratio 3.73
k =  60   RPCholesky trace error 7.02e-03   best rank-k 1.86e-03   ratio 3.77
k =  80   RPCholesky trace error 1.59e-03   best rank-k 2.85e-04   ratio 5.56
k = 120   RPCholesky trace error 5.50e-05   best rank-k 7.64e-06   ratio 7.20

How it's checked

lib/randla/check_ref.py runs the same two experiments in NumPy: Hutch++'s fitted slope must be steeper than -0.9 on the j^-2 spectrum, and RPCholesky's trace error at k = 40 must be below the best rank-20 error on a 500-point kernel.

print(
    "check():",
    vp.randla.trace_hutchpp(A=A, m=48, seed=0).check().passed,
    vp.randla.rp_cholesky(A=K[:300, :300], k=40, seed=7).check().passed,
)

check(): True True

Source

gallery/math/plot_randomized_trace_and_lowrank.py, run as written when these pages were built. The output and the figures above are from that run.

The Viv program it calls is lib/randla/gallery_trace.viv: its atlas, the defs and the calls between them, or open it in the workbench.