Skip to content

High-dimensional regression

A regression with 500 coefficients, fitted through the faer-backed MatVecMul path. The 500 coefficients become one contiguous vector parameter and one graph node, so a gradient evaluation costs two faer GEMV calls -- X @ beta forward and X.T @ adjoint backward -- instead of 500 scalar loops in each direction.

Cost note, which is arithmetic rather than a measurement: every leapfrog step streams the whole design matrix twice, so work grows with N_OBS * N_PARAMS. At 4,000 x 500 the matrix is 16 MB per GEMV. Raising the dimensions to 6,000 x 5,000 moves 240 MB per GEMV, fifteen times this script's traffic.

This script prints its own elapsed time when you run it. That number is not on this page, and no figure here is a performance claim: benchmarks/ and benchmarks/README.md hold the matched protocol, the retained raw output and the environment that a claim about speed needs.

Run it with python examples/large_linear_regression.py from the repository root. The blocks below are cells of that one file and share its state, so they assume the blocks above them have run.

import time

import numpy as np
import rustmc as rmc

Generate a 4,000 x 500 design

N_OBS = 4_000
N_PARAMS = 500

np.random.seed(42)
true_beta = np.random.randn(N_PARAMS) * 0.1
X = np.random.randn(N_OBS, N_PARAMS)  # row-major C order
y = X @ true_beta + np.random.randn(N_OBS) * 1.0

print(f"Dataset: {N_OBS:,} obs x {N_PARAMS:,} params")

Dataset: 4,000 obs x 500 params

What beta @ "X" does

Writing beta @ "X" instead of a sum of scalar products makes rustmc:

  1. detect that beta is used in a matrix multiply;
  2. infer the number of coefficients from the matrix's column count;
  3. promote beta to a contiguous vector parameter block;
  4. replace N scalar multiply-add nodes with a single MatVecMul op;
  5. compute the forward pass and the gradient through faer's GEMV, which uses Rayon once the matrix reaches 100,000 elements.

A 2D NumPy array in the data dict is detected and stored as a row-major matrix; passing "X": X where X.ndim == 2 is all that is needed. normal_prior combined with @ auto-promotes beta, so vector_normal_prior("beta", n=P) is only needed when you want to set the coefficient count yourself rather than infer it from the matrix.

Whether @ or scalar beta * "x" is faster for a given P is a question for benchmarks/, not for this page: it depends on the model, the data shape and the machine, and nothing here measures it. What the @ form does change is the graph -- one op instead of one per coefficient -- which is why it exists.

Build the model

t0 = time.time()
builder = rmc.ModelBuilder(data={"X": X, "y": y})
intercept = builder.normal_prior("intercept", mu=0.0, sigma=10.0)
beta = builder.normal_prior("beta", mu=0.0, sigma=1.0)
mu_expr = intercept + beta @ "X"
builder.normal_likelihood("obs", mu_expr=mu_expr, sigma=1.0, observed_key="y")
model = builder.build()
build_time = time.time() - t0
print(f"Model built in {build_time:.3f}s")

Model built in <elapsed>s

Sample

DRAWS, WARMUP = 200, 200
print(f"Sampling: NUTS, 1 chain, {WARMUP} warmup + {DRAWS} draws ...")
t0 = time.time()
result = rmc.sample(model, draws=DRAWS, warmup=WARMUP, chains=1, seed=42, show_progress=True)
elapsed = time.time() - t0

print(f"Elapsed : {elapsed:.2f}s")
print(f"Iters/s : {(DRAWS + WARMUP) / elapsed:.1f}")
print(f"Accept  : {result.accept_rates()[0]:.3f}")
print(f"Diverge : {sum(result.divergences())}")

Sampling: NUTS, 1 chain, 200 warmup + 200 draws ...
Elapsed : <elapsed>s
Iters/s : <rate>
Accept  : <rate>
Diverge : 0

Recover the coefficients

# A vector parameter is reported one entry per coordinate -- beta[0], beta[1], ... --
# not as a nested array.
samples = result.get_samples()
beta_means = np.array([samples[f"beta[{k}]"].mean() for k in range(N_PARAMS)])
rmse = np.sqrt(np.mean((beta_means - true_beta) ** 2))
# Reported against the scale the coefficients were drawn from, and at a precision
# this run can defend. One chain of 200 draws carries Monte Carlo error of order
# the fourth decimal on a single coordinate, and the faer GEMV path picks its SIMD
# kernel from the CPU, so the individual posterior means printed here to four
# places agreed on one machine and not on another. The ratio is the claim the
# example is making; a per-coordinate table would invite exactly the reading the
# Limitations section below warns against.
print(f"beta recovery RMSE : {rmse:.3f}  (generating coefficient sd {true_beta.std():.3f})")
print(f"  that is {rmse / true_beta.std():.0%} of the scale the coefficients were drawn from")

beta recovery RMSE : 0.017  (generating coefficient sd 0.098)
  that is 17% of the scale the coefficients were drawn from

Limitations of this run

One chain of 200 draws is enough to show the API and to recover coefficients whose generating scale is known. It is not enough for convergence diagnostics: R-hat compares chains and there is only one, and 200 draws is a small sample for any single coordinate. Raise chains and draws, and read the per-parameter ESS in result.summary(), before reading anything into an individual coefficient.

The elapsed time and rate printed above are replaced with placeholders on this page, because they depend on the machine and would otherwise change on every regeneration. Run the script to see real numbers for your hardware, and use benchmarks/ for a measurement with provenance.