Part A: Randomized Kaczmarz

Randomized Kaczmarz (RK) solves a consistent linear system Ax=b, where A\in\mathbb{R}^{m\times n} has rows a_1^\top,\dots,a_m^\top and x^\star denotes a solution. Starting from an initial guess x_0 (for instance x_0=0), at each iteration k the algorithm:

1. Samples a row. Draw a row index I_k\in\{1,\dots,m\} at random from a fixed PMF p=(p_1,\dots,p_m), independently of all previous draws (sampling with replacement).
2. Projects onto that equation. Move the current iterate to the closest point on the hyperplane \{x : a_{I_k}^\top x=b_{I_k}\}: x_{k+1}=x_k+\frac{b_{I_k}-a_{I_k}^\top x_k}{\|a_{I_k}\|^2}\,a_{I_k}.

Collision probability in randomized Kaczmarz

Suppose two consecutive iterations sample rows I_1,I_2 independently, each \sim p (with replacement). What’s the probability the algorithm samples the same row twice in a row – a collision?

LOTUS: why the row-norm PMF?

A popular distribution over the rows is the row-norm PMF p_I(i) = \|a_i\|^2/\|A\|_F^2. We’ll see why now!

Linearity + LOTUS: average one-step progress of RK

Recall randomized Kaczmarz’s update, and note that the Pythagorean theorem (the update is an orthogonal projection) provides, \|x_{k+1}-x^\star\|^2 = \|v\|^2 - \frac{(a_I^\top v)^2}{\|a_I\|^2}. where v = x_k-x^\star.

Watching the bound play out over many steps

The one-step bound says the expected squared error shrinks by a fixed factor every step. Let’s actually run RK and watch \|x_k-x^\star\|^2 decrease and compare it to the bound.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(151)
m, n = 20, 8
A = rng.standard_normal((m, n))
x_star = rng.standard_normal(n)
b = A @ x_star
row_norm_sq = (A ** 2).sum(axis=1)
p_row_norm = row_norm_sq / row_norm_sq.sum()
sigma_min = np.linalg.svd(A, compute_uv=False)[-1]
frob_sq = row_norm_sq.sum()

num_steps = 60
num_runs = 300
errors_sq = np.zeros((num_runs, num_steps + 1))

for r in range(num_runs):
    x = np.zeros(n)
    errors_sq[r, 0] = np.linalg.norm(x - x_star) ** 2
    idx_sequence = rng.choice(m, size=num_steps, p=p_row_norm)
    for t, i in enumerate(idx_sequence, start=1):
        a_i, b_i = A[i], b[i]
        x = x + (b_i - a_i @ x) / row_norm_sq[i] * a_i
        errors_sq[r, t] = np.linalg.norm(x - x_star) ** 2

mean_error_sq = errors_sq.mean(axis=0)
bound_curve = errors_sq[:, 0].mean() * (1 - sigma_min**2 / frob_sq) ** np.arange(num_steps + 1)

fig, ax = plt.subplots(figsize=(6, 3.5))
ax.semilogy(mean_error_sq, label="average of $\\|x_k-x^\\star\\|^2$ over runs")
ax.semilogy(bound_curve, "--", label="one-step bound, applied repeatedly")
ax.set_xlabel("step $k$")
ax.set_ylabel("squared error (log scale)")
ax.legend()
plt.tight_layout()
plt.show()

Recursing the randomized Kaczmarz bound

Recall our previous one-step bound, valid conditional on the current iterate x_k: E\big[\|x_{k+1}-x^\star\|^2 \mid x_k\big] \le (1-c)\,\|x_k-x^\star\|^2, \qquad c = \frac{\sigma_{\min}(A)^2}{\|A\|_F^2}.

Uniform vs. row-norm sampling: the distribution controls the rate

One can prove a similar (but worse) bound for the error for uniform sampling. Let’s run both samplers on a matrix whose row norms vary by orders of magnitude, and compare each empirical curve with its bound.

rng = np.random.default_rng(151)
m, n = 40, 8
A = rng.standard_normal((m, n)) * (10 ** rng.uniform(-1, 1, size=(m, 1)))  # very unequal row norms
x_star = rng.standard_normal(n)
b = A @ x_star
row_norm_sq = (A ** 2).sum(axis=1)
frob_sq = row_norm_sq.sum()
sigma_min = np.linalg.svd(A, compute_uv=False)[-1]

samplers = {
    "row-norm": row_norm_sq / frob_sq,
    "uniform": np.ones(m) / m,
}
rates = {
    "row-norm": sigma_min**2 / frob_sq,
    "uniform": sigma_min**2 / (m * row_norm_sq.max()),
}

num_steps, num_runs = 400, 300
mean_err = {}
for name, p in samplers.items():
    errs = np.zeros((num_runs, num_steps + 1))
    for r in range(num_runs):
        x = np.zeros(n)
        errs[r, 0] = np.linalg.norm(x - x_star) ** 2
        for t, i in enumerate(rng.choice(m, size=num_steps, p=p), start=1):
            x = x + (b[i] - A[i] @ x) / row_norm_sq[i] * A[i]
            errs[r, t] = np.linalg.norm(x - x_star) ** 2
    mean_err[name] = errs.mean(axis=0)

e0 = np.linalg.norm(x_star) ** 2
ks = np.arange(num_steps + 1)
fig, ax = plt.subplots(figsize=(7, 4))
for name, color in [("row-norm", "C0"), ("uniform", "C1")]:
    ax.semilogy(mean_err[name], color=color, label=f"{name}: empirical")
    ax.semilogy(e0 * (1 - rates[name]) ** ks, "--", color=color, label=f"{name}: bound")
ax.set_ylim(1e-20, 1e2)
ax.set_xlabel("step $k$")
ax.set_ylabel("squared error (log scale)")
ax.legend()
plt.tight_layout()
plt.show()

Part B: Randomized trace estimation

Recall the randomized trace estimator: for A\in\mathbb R^{n\times n}, draw z\in\{-1,+1\}^n with i.i.d. Rademacher entries and set \tau = z^\top A z = \sum_{i,j} z_iz_j A_{ij}.

Example: unbiasedness of randomized trace estimation

We’ll now prove E[\tau]=\mathrm{tr}(A) – recall this is called unbiasedness.

Example: continuous RV in randomized trace estimation

The unbiasedness proof only ever used E[z_i^2] = 1 \qquad\text{and}\qquad E[z_iz_j] = 0 \ (i\ne j), which followed from mean 0, variance 1, and independence – nothing about z_i being discrete. So a continuous z_i works exactly as well.

rng = np.random.default_rng(151)
d = 50
A = rng.standard_normal((d, d)) + 5 * np.eye(d)
true_trace = np.trace(A)

num_samples = 20_000
Z_rad = rng.choice([-1.0, 1.0], size=(num_samples, d))  # discrete: Rademacher
Z_gauss = rng.standard_normal((num_samples, d))          # continuous: standard Gaussian

tau_rad = np.einsum('si,ij,sj->s', Z_rad, A, Z_rad)
tau_gauss = np.einsum('si,ij,sj->s', Z_gauss, A, Z_gauss)

print(f"true trace(A)                        = {true_trace:.4f}")
print(f"mean of tau, Rademacher z (discrete)  = {tau_rad.mean():.4f}")
print(f"mean of tau, Gaussian z (continuous)  = {tau_gauss.mean():.4f}")

true trace(A)                        = 250.0783
mean of tau, Rademacher z (discrete)  = 249.9342
mean of tau, Gaussian z (continuous)  = 249.5031

How much variance? The Rademacher estimator

Unbiased is not enough – a single sample of \tau could be far from \mathrm{tr}(A). Let A be symmetric, so A_{ij}=A_{ji} and \tau=\sum_i A_{ii}z_i^2+2\sum_{i<j}A_{ij}z_iz_j=\mathrm{tr}(A)+2\sum_{i<j}A_{ij}z_iz_j for Rademacher z (since z_i^2=1). The first term is a constant, so only the cross terms contribute to the variance.

Comparing the two estimators

Since \|A\|_F^2=\sum_i A_{ii}^2+\sum_{i\ne j}A_{ij}^2, the two variances differ by \mathrm{Var}(\tau_{\text{Gauss}})-\mathrm{Var}(\tau_{\text{Rad}})=2\sum_i A_{ii}^2\;\ge 0. Rademacher always has smaller variance, and it wins by the most when A has a heavy diagonal. If A has zero diagonal, the two are identical. Let’s check on a diagonal-dominant matrix and on one with zero diagonal.

rng = np.random.default_rng(151)
d = 30
G = rng.standard_normal((d, d))
S = (G + G.T) / 2                                   # symmetric, generic diagonal
A_diag_heavy = S + 5 * np.eye(d)                    # diagonal-dominant
A_zero_diag = S - np.diag(np.diag(S))               # zero diagonal

num_samples = 200_000
for name, M in [("diagonal-dominant", A_diag_heavy), ("zero diagonal", A_zero_diag)]:
    Z_rad = rng.choice([-1.0, 1.0], size=(num_samples, d))
    Z_gauss = rng.standard_normal((num_samples, d))
    tau_rad = np.einsum('si,ij,sj->s', Z_rad, M, Z_rad)
    tau_gauss = np.einsum('si,ij,sj->s', Z_gauss, M, Z_gauss)
    off = (M ** 2).sum() - (np.diag(M) ** 2).sum()
    print(name)
    print(f"  Rademacher: empirical var {tau_rad.var():9.2f}   formula {2*off:9.2f}")
    print(f"  Gaussian:   empirical var {tau_gauss.var():9.2f}   formula {2*(M**2).sum():9.2f}")

diagonal-dominant
  Rademacher: empirical var    978.64   formula    983.70
  Gaussian:   empirical var   2578.98   formula   2595.05
zero diagonal
  Rademacher: empirical var    984.84   formula    983.70
  Gaussian:   empirical var    981.17   formula    983.70

Part C: Randomized matrix multiplication

Let A\in\mathbb R^{m\times n} and B\in\mathbb R^{n\times q}. Write a_j for the j-th column of A and b_j^\top for the j-th row of B. Then AB=\sum_{j=1}^n a_j b_j^\top, a sum of n rank-one outer products.

Unbiasedness of randomized matrix multiplication

X is a matrix-valued function h(J) of the discrete RV J, and by LOTUS: E[X]=\sum_{j=1}^n \frac{1}{p_j}\,a_jb_j^\top\;p_j=\sum_{j=1}^n a_jb_j^\top=AB.

The best sampling distribution

Part (b) of the exercise above can be summed to show: E\|AB-\hat C\|_F^2=\frac1c\left(\sum_{j}\frac{w_j^2}{p_j}-\|AB\|_F^2\right), where w_j=\|a_j\|\|b_j\|. Only the sum depends on p, so minimize \sum_j w_j^2/p_j.