fix: mean-field coupling compiled to an O(H²) IR#444
Merged
Conversation
`Autodiff.simplify` matched BinOp/UnOp/Cond and fell through `| _ -> e`
on everything else, so `Reduce` (the n-ary sum) was an opaque leaf.
Differentiating a global sum `Σ_h N[h]` w.r.t. any parameter yields
`Reduce [Const 0.0; ...]` — one zero per term, not `Const 0.0`. Left
unfolded, the quotient rule's `-f·g'` term never meets `Mul, _, Const
0.0`, so `Div, Const 0.0, _` cannot fire and the entire O(H) denominator
rides into EVERY parameter's gradient — including parameters absent from
the rate. On a mean-field model that is an O(H²) rate_grad.
The committed goldens carried the symptom directly: seir_age's infection
transition had a rate_grad for `sigma` (E→I) and `gamma` (recovery),
parameters that occur in no infection rate, emitted as `beta·S·reduce[
0/(N·N), 0/(N·N) ]` — structurally present, zero-valued.
Add the missing `Reduce` arm: drop `Const 0.0` summands (the additive
identity of a sum, exactly as the existing `Add, x, Const 0.0 -> x`
rule assumes), collapse a singleton, and empty folds to `Const 0.0`.
Value-preserving: removing an exact 0.0 from a sum changes no value.
`simplify_fixpoint` is reached only from derivative expressions and
lineage weights, never the forward rate.
Corpus effect: 12 stratified goldens shed gradient bytes 2,389,443 →
651,307 (3.67×), forward rate byte-identical. Because BindingRef is
memoized per propensity-vector evaluation, this also cuts per-step work
on gradient-based inference.
Red → green (ocaml, test_compiler trig_primitives):
before: FAIL ∂(beta*I/Σ N_h)/∂k must fold to a genuine zero and be
dropped; a Reduce of zero derivatives kept it alive
after: Test Successful in 0.004s. 19 tests run.
Ground truth (rust, gradient_check on sir_coupling, the mean-field
shape): the folded gradient matches finite differences —
d(ll)/d(beta) = 697.5910 vs 697.5915 rel_err 7.6e-7
d(ll)/d(gamma) = -280.3777 vs -280.3776 rel_err 2.3e-7
See docs/dev/notes/2026-07-16-mean-field-coupling-quadratic-ir.md.
`body_refs_param_or_let` rejected a `let` from hoisting if its body
named ANY other `let`, commented "conservatively excluded; it may
transitively carry a parameter." The contract it guards
(validate.ml's ParamInBinding) is only param-freeness — autodiff
differentiates BindingRef to 0 under WrtParam, so a param reachable
from a hoisted body would be a silent zero gradient. Param-freeness is
closed under let-reference, so "names a let" is strictly coarser than
the property.
That over-rejection hit the textbook mean-field denominator:
let N[h in hh] = S[h] + ... # compartments only
let N_glob = sum(h in hh, N[h]) # names N -> rejected
N_glob is param-free (N is), so it is hoistable — but was refused and
its O(H) sum inlined into all H force-of-infection rates: an O(H²) IR
in the number of coupled strata. N[h] and I_glob already hoisted, so
only the aggregate-of-an-aggregate was missing.
Recurse into a referenced let instead of rejecting on sight; disqualify
only if that let transitively reaches a parameter. A `visited` set
guards a self-referential let (least fixpoint = "no param here"). The
ParamInBinding backstop is unchanged, so a genuine param leak is still
caught.
On the TB household probe this drops global-coupling IR from O(H²) to
O(H) (40.2 MB → 1.27 MB at H=160) and forward sim from superlinear to
linear (1.93s → 0.59s at H=160, 3.27×), via the per-eval BindingRef
cache. Forward sim is byte-identical across seeds; the one seed that
errors (household depopulates, DivByZero) errors identically before and
after. Among the goldens only seir_cross_dim carries this shape; its
rate becomes `.../binding_ref N_total_north` with N_total_* defined
once — a value-preserving substitution.
Red → green (ocaml, test_compiler trig_primitives):
before: FAIL 'N_glob' is param-free and must hoist into
model.bindings; got [N_h0; I_glob; N_h1; N_h2]
after: Test Successful in 0.005s. 22 tests run.
See docs/dev/notes/2026-07-16-mean-field-coupling-quadratic-ir.md.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two independent compiler bugs made dense mean-field (metapopulation)
coupling compile to an O(H²) IR in the number of coupled strata
H—in IR size, compile time, memory, and per-simulation cost. Both are fixed;
scaling drops to O(H).
The models affected are the textbook two-scale force of infection —
a per-stratum total, a global total summing it, and a FoI mixing a
within-stratum and a between-stratum term:
The two bugs
1.
Autodiff.simplifycould not fold aReduceof zeros (commit 1).It matched
BinOp/UnOp/Condand fell through| _ -> e, so the n-arysum was an opaque leaf.
d(Σ_h N[h])/dp = Reduce[Const 0.0; …]— a sum ofzeros, not
Const 0.0— soDiv, Const 0.0, _never fired and the wholeO(H) denominator rode into every parameter's gradient, including
parameters absent from the rate.
2. The hoist predicate over-rejected (commit 2).
body_refs_param_or_letrefused anyletnaming anotherlet. Thecontract it guards (
validate.ml'sParamInBinding) is onlyparam-freeness, which is closed under let-reference.
N_globnamesN;Nis compartments-only, soN_globis param-free and hoistable — but wasrefused, inlining its O(H) sum into all H rates.
They compose: bug 2 puts an O(H) sum in each rate; bug 1 copies it into
every gradient column of each transition.
Why this is not cosmetic
The committed goldens carried the symptom.
seir_age's infectiontransition had a gradient for
sigma(E→I progression) andgamma(recovery) — parameters absent from any infection rate:
Zero-valued but structurally present, differentiated w.r.t. a parameter
the expression does not contain.
Measured (TB household probe,
camdl simulate, chain_binomial, 10y daily)IR size, O(H²) → O(H):
Forward-sim wall clock (via the per-eval
BindingRefcache), superlinear→ linear:
Verification
Forward sim byte-identical across seeds (the one failing seed errors
identically before and after — a household depopulates,
DivByZero).sir_coupling(the mean-fieldshape) — folded gradient matches finite differences to
rel_err ~1e-7.make testgreen.Golden changes
12 stratified goldens change; gradient bytes across the corpus
2,389,443 → 651,307 (3.67×). 11 of 12 have a byte-identical forward rate
(gradients only);
seir_cross_dim(the one aggregate-of-an-aggregate)substitutes
reduce[…]→binding_ref N_total_north. Reviewed andapproved by the maintainer.
Not addressed
payoff of the fold — is unmeasured; it needs a fit run.
camdlc render ocaml/golden/sir_basic.ir.jsonfails with a parse error(pre-existing; the pre-change
camdlcfails identically). Separate issue.Details:
docs/dev/notes/2026-07-16-mean-field-coupling-quadratic-ir.md.