Skip to content

perf(ode-grad): hoist state-Jacobian eval out of per-parameter loop (gh#400)#401

Merged
vsbuffalo merged 2 commits into
mainfrom
ode-grad-perf-1-2
Jul 17, 2026
Merged

perf(ode-grad): hoist state-Jacobian eval out of per-parameter loop (gh#400)#401
vsbuffalo merged 2 commits into
mainfrom
ode-grad-perf-1-2

Conversation

@vsbuffalo

Copy link
Copy Markdown
Owner

Cuts the per-step cost of the ODE forward-sensitivity gradient (the NUTS-on-ODE wall for long-burn-in fits) by removing a redundant expression evaluation. First of the two cost-per-step levers in gh#400; the second (kill per-step allocation) is deferred — see below.

What

In sensitivity_derivs, the state-Jacobian term ∂rate_r/∂x_j (eval_deriv_entry) was re-evaluated inside the for p in 0..d loop over estimated parameters, even though its value is independent of the parameter column p — only the S[j,p] multiplier varies. This evaluates each transition's J_x entries once (during the first column) and reuses them across the remaining columns, cutting the J_x expression-eval work from P·nnz to nnz per transition per RK4 stage. The force-of-infection rates (dense J_x, ~8–9 nonzeros) benefit most.

Byte-neutral

The gradient is bit-identical before/after. The mechanism matters: column 0 keeps the exact original J_θ-then-J_x evaluation order, so the shared binding cache is populated identically; columns p>0 previously re-called eval_deriv_entry against the warm cache and got exactly the values now stashed. A fully-hoisted variant that evaluated all J_x before J_θ reorders cache population and perturbs one gradient component by 1 ULP — the column-0 interleave avoids that.

A/B gate (det_grad on the seir_observations FD-oracle fixture, f64::to_bits):

baseline  ll=c0598ed08da0c428
          grad=c05a4bbe00ed4375 4070de40d6693ee5 bffbd780d1241668 405e2ba1ec971515
after     identical — all four grad words + loglik bit-for-bit

Verification

  • FD gradient oracles green (3/3)
  • make test green (all phases: Rust 907, inference, OCaml, integration)

Deferred: lever 2 (per-step allocation)

The companion optimization — threading a reusable Rk4Workspace through rk4_step to eliminate per-step heap allocation — was implemented and reverted. It is not byte-identical: pure storage relocation drifts one gradient component by 1–2 ULP (value trajectory provably unchanged; only the dense-binding sensitivity chain moves), revealing that the forward-sensitivity accumulation is ULP-fragile to RK4 buffer structure via FP-order sensitivity. Held pending a decision on whether to accept a numerically-equivalent (FD-oracle-validated) change or root-cause the fragility first. Tracked in gh#400.

vsbuffalo added 2 commits July 8, 2026 10:03
…gh#400)

In `sensitivity_derivs`, the state-Jacobian term ∂rate_r/∂x_j
(`eval_deriv_entry`) was re-evaluated inside the `for p in 0..d` loop
over estimated parameters, but its value is independent of the
parameter column `p` — only the `S[j,p]` multiplier varies. Evaluate
each transition's J_x entries once during the first column and reuse
them for the remaining columns, cutting the J_x expression-eval work
from P·nnz to nnz per transition per RK4 stage. The FOI-type rates
(dense J_x, ~8-9 nonzeros) benefit most.

Byte-neutral: column 0 keeps the exact original J_θ-then-J_x
evaluation order, so the shared binding cache is populated identically;
columns p>0 previously re-called `eval_deriv_entry` against the warm
cache and got exactly the values now stashed in `jx`. An earlier
fully-hoisted variant that evaluated all J_x before J_θ reordered
cache population and perturbed one gradient component by 1 ULP — the
column-0 interleave avoids that.

A/B gate (det_grad on the seir_observations FD-oracle fixture,
f64::to_bits):
  baseline  ll=c0598ed08da0c428
            grad=c05a4bbe00ed4375 4070de40d6693ee5
                 bffbd780d1241668 405e2ba1ec971515
  after     identical — all four grad words + loglik bit-for-bit.
FD oracles green (3/3); sim --lib green (323/0).
@vsbuffalo
vsbuffalo merged commit d132af9 into main Jul 17, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant