Performance Guide: Making JAXSR Fast on CPU and GPU

Performance Guide: Making JAXSR Fast on CPU and GPU#

JAXSR uses JAX for numerical computation, which enables JIT compilation and GPU acceleration. However, “just using JAX” does not automatically make code fast. This guide documents the optimization strategies used in JAXSR’s model selection algorithms and the lessons learned from GPU benchmarking.

The Core Bottleneck#

JAXSR’s selection algorithms (greedy_forward, greedy_backward, exhaustive, lasso_path) evaluate many candidate subsets of basis functions. Each evaluation originally called jnp.linalg.lstsq on an (n_samples, k) matrix inside a Python loop. This creates two problems:

  1. O(n * k^2) per subset – lstsq solves the full n-by-k system each time

  2. Kernel launch overhead – each lstsq call dispatches a GPU kernel (~0.1-1ms overhead), which dominates when k is small

Benchmarking showed that GPU was actually slower than CPU for typical JAXSR workloads because of problem (2). The fix addresses both.

Strategy: Gram Matrix Precomputation#

The key insight is that all subset OLS/ridge problems share the same data matrix Phi and target y. We can precompute the products once:

PhiTPhi = Phi.T @ Phi    # (B, B) Gram matrix
PhiTy   = Phi.T @ y      # (B,)   cross-product
yTy     = y.T @ y        # scalar

Then for any subset of indices S, the OLS coefficients satisfy the k-by-k normal equations:

PhiTPhi[S, S] @ c = PhiTy[S]

Solving this via jnp.linalg.solve costs O(k^3) instead of O(n * k^2) for lstsq. For typical problems (n=1000, k=3), this is 1000/3 ~ 300x fewer floating-point ops.

Cost Comparison Per Subset#

Operation

Original (lstsq)

Gram (solve)

Coefficient solve

O(n * k^2)

O(k^3)

MSE (float64)

O(n * k)

O(k)

MSE (float32)

O(n * k)

O(n * k)

Total

O(n * k^2)

O(k^3) or O(n*k + k^3)

The one-time precomputation cost is O(n * B^2) for the Gram matrix, which is the same order as a single lstsq call on the full matrix.

MSE Computation: Why dtype Matters#

For the MSE, there is a closed-form expression at the OLS optimum:

MSE = (yTy - c^T @ PhiTy[S]) / n

This is O(k) and eliminates n_samples from the inner loop entirely. However, it involves subtracting two large numbers (yTy and c^T @ rhs) to get a small residual. In float32 (~7 significant digits), this catastrophic cancellation produces wildly incorrect MSE values. In float64 (~15 significant digits), it works fine.

JAXSR branches on dtype:

  • float64: Uses the closed-form formula. Each subset costs O(k^3) total – completely independent of n_samples. This is the maximum-speedup path.

  • float32: Computes MSE from actual residuals: mean((y - Phi[:, S] @ c)^2). Each subset costs O(nk + k^3), still much faster than O(nk^2) from lstsq.

Enabling float64 in JAX#

JAX defaults to float32. To use float64 (recommended for best performance and numerical accuracy):

import jax
jax.config.update("jax_enable_x64", True)

Or set the environment variable before importing JAX:

JAX_ENABLE_X64=True python my_script.py

Why GPU Can Be Slower Than CPU#

GPU benchmarks on JAXSR revealed a counterintuitive result: GPU was slower for typical symbolic regression workloads. The reasons:

1. Kernel Launch Overhead#

Each JAX operation dispatches a GPU kernel. The overhead (~0.1-1ms per dispatch) is negligible for large matrix operations but dominates when the inner loop does hundreds of small solves:

# 100 subsets x 1ms overhead = 100ms of pure overhead
# vs the actual computation: 100 x 0.01ms = 1ms

The Gram precomputation helps because:

  • The expensive O(n*B^2) Phi.T @ Phi is a single large kernel (GPU-friendly)

  • The inner loop solves k-by-k systems, which are tiny but have no per-call overhead beyond array indexing

2. Python Loop Bottleneck (Amdahl’s Law)#

The greedy/exhaustive loops are written in Python, not JAX. Even if each operation inside the loop is fast, the Python iteration overhead limits speedup. GPU helps most when you can express the entire computation as a single JAX operation (e.g., vmap over subsets).

3. Data Transfer Costs#

Moving data between CPU and GPU has latency. For small problems, the transfer time exceeds the computation time.

Practical Recommendation#

For most JAXSR workloads, CPU is faster. Set:

JAX_PLATFORMS=cpu python my_script.py

GPU becomes beneficial only at very large n_samples (50K+) where the basis evaluation (Phi = library.evaluate(X)) dominates runtime.

Optimization Techniques Applied in JAXSR#

1. Precompute Shared Work Outside the Loop#

Before:

for subset in subsets:
    Phi_sub = Phi[:, subset]
    coeffs = lstsq(Phi_sub, y)      # O(n * k^2) each time
    mse = mean((y - Phi_sub @ coeffs) ** 2)

After:

PhiTPhi = Phi.T @ Phi               # O(n * B^2) once
PhiTy = Phi.T @ y                   # O(n * B) once
for subset in subsets:
    c = solve(PhiTPhi[S,S], PhiTy[S])  # O(k^3) each time
    mse = ...                           # O(k) or O(n*k)

2. Use solve Instead of lstsq for Square Systems#

jnp.linalg.lstsq on an n-by-k system does QR factorization (O(n*k^2)). jnp.linalg.solve on a k-by-k system does LU factorization (O(k^3)). When you’ve already formed the normal equations, solve is the right choice.

3. Branch on Numerical Precision#

Don’t use the same algorithm for all dtypes. Float64 can use algebraically equivalent but numerically cheaper formulas that would fail in float32.

4. Guard Against Singular Systems#

When solving normal equations, the Gram submatrix can be singular (e.g., collinear basis functions). JAXSR checks for NaN/Inf in the solution and falls back to lstsq on the k-by-k system, which handles rank deficiency.

5. Skip the Optimization for Non-Linear Parameters#

When the basis library contains parametric (non-linear) basis functions, the Gram precomputation doesn’t apply because the design matrix changes with the parameter values. JAXSR detects this via basis_library.has_parametric and falls back to the original per-subset fit_ols/fit_ridge path.

Benchmark Results#

On a dataset with n=1000 samples and B=20 basis functions, forward selection with max_terms=5:

Approach

Time

Speedup

Original (lstsq per subset)

~2s

1x

Gram precompute (float32, residual MSE)

~0.05s

~40x

Gram precompute (float64, closed-form MSE)

~0.015s

~133x

At n=100,000 samples, the speedup exceeds 500x because the O(n) cost of lstsq is completely eliminated from the inner loop (float64 path).

Future Optimization Opportunities#

  1. Batch subset evaluation with jax.vmap: Instead of a Python loop over subsets, vectorize the Gram extraction and solve across all candidates at once. This would maximize GPU utilization.

  2. JIT the entire selection loop: Use jax.lax.while_loop or jax.lax.scan instead of Python for/while. Requires fixed-size arrays and careful handling of dynamic shapes.

  3. Rank-1 updates (Cholesky): When forward selection adds one term at a time, the Gram submatrix grows by one row/column. A Cholesky rank-1 update costs O(k^2) instead of solving the full k-by-k system (O(k^3)).