Performance
This guide shows how to measure omnibias's performance wins on your own hardware, and how to read the numbers honestly.
The wins, recapped
| Win | What it means |
|---|---|
Laplacian is O(1) in D | the D-dependent term is computed once and reused |
68×–199× faster than dense-Hessian autodiff at D=240 | no Hessian graph to build or differentiate |
| 63×–108× less memory | no quadratic Hessian materialization |
Δᵏ flat in k | one tower per order; folx-nested OOMs at Δ⁴ |
All measured in float64 with answers identical to ≤ 1e-15. See
Complexity for the derivations.
Measure it yourself (CPU smoke tier)
The smoke tier runs in seconds on a laptop and verifies correctness and the scaling shape:
JAX_PLATFORMS=cpu python -m bench.laplacian_scaling.dimension_sweep \
--dims 3 12 30 --hidden 32 --batch 64 --repeats 3
You should see the per-sample Laplacian overhead stay roughly flat as D grows,
while a dense-Hessian baseline grows quadratically.
A minimal timing harness
import time, jax, jax.numpy as jnp
from omnibias.jax import neural_field_value_grad_laplacian
def bench(D, H=256, B=4096, reps=20):
x = jnp.ones((B, D))
W = jnp.ones((H, D)) / D ** 0.5
beta = jnp.zeros((H,)); c = jnp.ones((H,)); b = 0.0
f = jax.jit(lambda x: neural_field_value_grad_laplacian(x, W, beta, c, b, "tanh"))
f(x)[0].block_until_ready() # warm up / compile
t0 = time.perf_counter()
for _ in range(reps):
f(x)[2].block_until_ready()
return (time.perf_counter() - t0) / reps * 1e3 # ms
for D in (3, 30, 240):
print(D, f"{bench(D):.3f} ms")
Benchmark hygiene
- Warm up / compile first (the first call includes JIT time).
- Block on the result so you time compute, not dispatch.
- Compare identical answers. If your closed-form and autodiff results differ by more than float64 round-off, it is a bug, not a speedup.
- Report the hardware tier (memory class) and precision with every number.
Tuning notes
- Batch up. The Laplacian's
‖W‖²term is shared across the batch, so larger batches amortize it further. - Stay in float64 for correctness checks, then drop to float32 for throughput if your application tolerates it.
- Prefer the dedicated Laplacian kernel over building a full multivariate jet when you only need the trace.