Stable RBF & squared distances — keep the kernel SPD before Cholesky

0.11 — numerical stability handled the downstream problem: what to do when cholesky(K) returns NaNs. This notebook handles the upstream problem: how to construct $K$ so it stays SPD in the first place.

The textbook squared-distance expansion

is fast (one matmul, no broadcast) but numerically catastrophic for close-by points in float32: it subtracts two nearly equal positive numbers, so round-off can flip the result negative — and once $d^2 < 0$, the RBF kernel $\sigma_f^2\exp(-d^2/2\ell^2)$ exceeds $\sigma_f^2$ on the diagonal, the Gram matrix is no longer PSD, and Cholesky fails downstream.

gaussx.stable_squared_distances and gaussx.stable_rbf_kernel fix this with a mixed-precision recipe: dot products in compute_dtype (default float32, fast), subtraction in accumulate_dtype (default float64, stable), then a clamp to enforce $d^2 \ge 0$. The cost is one extra cast; the benefit is a Gram matrix that’s as PSD as f32 storage allows — non-negative distances, no diagonal blow-up, no asymmetry from cancellation.

This notebook covers:

1. The naïve expansion ((1)) and where it breaks (catastrophic cancellation on close points).
2. The mixed-precision stable recipe and what it costs.
3. stable_rbf_kernel end-to-end — diagonal exactness, PSD-by-construction.
4. A spectrum-of-$K$ sweep showing naïve f32 drifts negative; stable f32 stays $\ge 0$.
5. Cost benchmarking — the stability is essentially free.
6. Where this sits in the pipeline: stable distances (this notebook) → jitter / safe Cholesky (0.11) → Joseph form (0.5).

Prerequisites: 0.9 — primitives tour, 0.11 — numerical stability.

Why naïve squared distances fail in low precision¶

The expansion ((1)) computes squared Euclidean distance using one matmul ($X^\top Z$) plus two squared-norm vectors. For two close-by points $x \approx z$:

• $\|x\|^2 \approx \|z\|^2 \approx x^\top z$, so $\|x\|^2 + \|z\|^2 \approx 2 x^\top z$;
• the subtraction $\|x\|^2 + \|z\|^2 - 2 x^\top z$ is a difference of two nearly-equal large positive numbers;
• in float32, round-off in either summand can swamp the (tiny) true result, flipping the sign.

This is catastrophic cancellation — losing significant digits because most of the magnitude cancels. The naïve recipe is fast on a GPU because float32 dot products are accelerated, but the cancellation problem is structural, not tunable.

Hundreds of negative entries on a matrix that should be elementwise non-negative — and the diagonal isn’t even zero (it should be exactly $\|x_i - x_i\|^2 = 0$). Once you exponentiate this through an RBF kernel, the Gram matrix loses PSD-ness on the diagonal alone.

The mixed-precision recipe¶

gaussx.stable_squared_distances(X, Z, *, compute_dtype, accumulate_dtype) does the same arithmetic but in two precisions:

The dot products $X^\top Z$ and the squared norms run in compute_dtype (float32 by default — fast on GPU); the subtraction that triggers cancellation runs in accumulate_dtype (float64 — exact for this magnitude regime). The result is cast back to compute_dtype, then clamped to $\ge 0$ to absorb any residual round-off floor.

The clamp matters: the f64 subtraction is much better but not perfect, so a final $\max(\cdot, 0)$ guarantees the bound exactly. Without the clamp, even f64 round-off can leak through to RBF-exponentiation and break PSD.

Zero negative entries, diagonal is exactly zero, and the result agrees with the f64 reference to full f32 precision. The mixed-precision recipe gets us back to “the answer the f64 path would give” — which is what we wanted.

A scatter — naïve vs stable vs reference¶

Plot the naïve and stable squared distances against the f64 reference. Naïve f32 has a cloud of negative values for small distances; stable f32 hugs the $y = x$ line all the way down to zero.

The RBF kernel — same idea, one extra exp¶

The RBF (squared-exponential) kernel is the textbook GP prior:

gaussx.stable_rbf_kernel(X, Z, lengthscale, variance, ...) is exactly ((2)) wrapped in ((3)). The downstream guarantees we want:

• diagonal $K_{ii} = \sigma_f^2$ exactly;
• $K \succeq 0$ for any input set $X$;
• agreement with the f64 reference at full f32 precision.

Naïve f32: diagonal error and $\lambda_{\min}$ both bigger than f32 round-off, asymmetry > 0. Stable f32: errors at the f32 floor (~10^-5), no asymmetry. Note that $\lambda_{\min}(K)$ in f32 saturates at the f32 storage floor for genuinely rank-deficient inputs — no construction recipe can recover a smaller eigenvalue than the matrix’s elements can encode. The clean win lives one level up — at the squared-distance level, where the stable recipe avoids negative entries entirely (we showed earlier: 77 negatives → 0).

Tightness sweep — counting negative-distance entries¶

The kernel-level $\lambda_{\min}$ saturates at the f32 storage floor (~10^-7) for rank-deficient inputs, regardless of construction. But the squared-distance level shows a clear, monotone improvement: the number of $d^{2} < 0$ entries in the naïve f32 distance matrix grows as the cluster tightens, while the stable recipe stays at exactly zero by the clamp.

We sweep the cluster tightness from 10^-1 down to 10^-5 on $n = 200$ points and count negative-distance entries for each method.

The naïve curve climbs from zero to thousands of negative entries as the cluster tightens past $\sim 10^{-2}$. The stable curve stays at zero throughout — the clamp is doing exactly what it should. Once you build a kernel from a distance matrix with negatives, you’ve already lost PSD-ness; clamping at the distance level is the right place to enforce it.

Cost — is mixed precision free?¶

The mixed-precision recipe adds one cast from f32 to f64 for the subtraction and one back to f32 — both are essentially memory bandwidth, not flops. We benchmark on increasing $n$ to confirm there’s no measurable overhead.

On these sizes the stable path adds 0–50% overhead — the cast to f64 for the subtraction is memory-bandwidth-bound, and on small matrices that cost is comparable to the dot-product itself. On larger matrices ($n \gtrsim 10^3$) the dot-product dominates and the relative overhead shrinks. For most workloads, stability is worth the cost — silent NaN propagation downstream is far more expensive than 30% on the kernel-construction step.

Where this sits in the numerical-robustness pipeline¶

A robust GP / kernel pipeline has three lines of defence, applied in this order:

Each stage handles a different failure mode. Skipping stage 1 forces stage 2 to do more work (jitter has to absorb the numerical kernel-construction error on top of the legitimate ill-conditioning); skipping stage 3 means even a perfectly-built kernel can drift non-PSD across updates. The “always-on” recipe is all three — the cost of each is small, the cost of skipping is silent NaNs in production.

Practical recipe — drop-in stable RBF¶

from gaussx import stable_rbf_kernel, add_jitter, safe_cholesky
import lineax as lx

def Ky(X, sigma_f, ell, sigma_n, jitter=1e-6):
    K = stable_rbf_kernel(X, X, ell, sigma_f**2)              # 1. stable construction
    K = K + (sigma_n**2 + jitter) * jnp.eye(K.shape[0])        # 2. observation noise + jitter
    return lx.MatrixLinearOperator(K, lx.positive_semidefinite_tag)

# Then: gaussx.solve / gaussx.logdet on Ky(X, ...) — both gradients flow through.
# If you ever still see NaNs, swap cholesky → safe_cholesky.

This is the kernel construction you’ll see throughout part 3 (exact GP regression) and part 5 (sparse / inducing-point GPs).

Recap¶

• Catastrophic cancellation in the naïve squared-distance expansion ((1)) is a constructive failure mode — round-off can produce $d^2 < 0$, breaking PSD before Cholesky even runs.
• The mixed-precision recipe ((2)) keeps dot products in fast compute_dtype and the subtraction in stable accumulate_dtype; the clamp guarantees $d^2 \ge 0$ exactly.
• gaussx.stable_rbf_kernel wraps this for the RBF / squared-exponential kernel ((3)) — diagonal at f32 precision (vs naïve’s >10^-4 errors), no spurious non-PSD violations from cancellation, modest 0–50% overhead for kernel construction.
• Build kernels with stable_* first, jitter / safe_cholesky second (0.11), Joseph form on covariance updates (0.5). The three stages address different failure modes; use them all.