fit is one loop: optimizer/callback hooks, calibrate(), init="glorot" — the exact reference protocol

[MArpogaus]

What this PR is

CausalFlowDAG.fit is one minibatch Adam loop with two hooks, and the paper replication runs the reference's exact rule with the reference's init. Stacked on #43 (feat/followups); the diff shown is this PR's own.

Three independent reviews of flow.py — user code or framework?, does it duplicate torch/zuko?, which behaviours are strategies tuned on our own DGPs? — agreed on the cut, and a fourth reviewer verified the result. Every training strategy that left had a default the paper replication or the tests had to switch off.

1. fit (breaking)

flow.fit(train_df, *, epochs, learning_rate=1e-2, batch_size=512,
         seed=None, optimizer=None, callback=None, vc_ehat=None)

• One optimizer over all parameters (the per-node NLLs have independent gradients, so this is per-node training); optimizer= takes any torch optimizer — how a torch.optim.lr_scheduler attaches; callback(flow, epoch, optimizer) after every epoch, True stops. flow.history["train"] is all the loop records.
• Gone, with the replacement: val_df (compute flow.nll(val) in the callback); schedule="plateau" + plateau_* + min_delta (torch's ReduceLROnPlateau on the optimizer you pass); freeze_patience and the per-node freeze (experiments/benchmarks/bench_training.py::_PerNodePlateau, the recipe that benchmark measures); restore_best (a six-line snapshot callback, docs/fitting.md); verbose and the logger; epoch_callback (now callback, receives the optimizer); marginal_init (→ calibrate); vc_warm_start and its hidden classical proxy fit (two lines of user code; measured on vc_hetero: beta0 0.16 from the truth without it, recovery corr 0.99 either way); vc_oof_fit, the hidden five-fold stage-1 fits and VC(center_folds=) (→ fit(vc_ehat=), below).
• flow.calibrate(train_df, marginal_init=True) — the data-dependent state once (transform ranges, network-input min-max, calibrated start); the first fit/fit_classical calls it; one calibrated buffer replaces three per-module latches that load re-closed by hand.
• fit(vc_ehat={node: {t: array}}) — a centered VC term needs its out-of-fold propensities from the caller; fit refuses a centered spec without them or a mismatch; the spec refuses chained centering (the guard used to live in the deleted stage).
• fit_classical lost verbose, uses torch.nn.utils.get_total_norm. save no longer records the machine (machine_info → experiments/benchmarks/perf_machine.py). transforms: the inverse is zuko's own (Transform.inv, closed-form tail) — the 85-line expanding bisection duplicated it (identical residuals, spline inverse 166 ms → 0.3 ms on 4000 rows); icdf is torch.logit.
• flow.py 2089 → 1606 lines. All notebooks, validate_ls, the benchmark and the paper scripts run on the hooks; tests/test_fit_hooks.py carries the exact-MLE guard through a torch scheduler.

2. CausalFlowDAG(init="glorot") — and why VACA needed it

With fit on the hooks, helpers.fit_paper runs the reference's global ReduceLROnPlateau on the summed validation NLL (patience 49 / threshold 0 abs / factor 0.1 / min 1e-7 — verified against torch's source to match update_learning_rate's wait >= patience, strict <, reset after reduction). Under that exact rule VACA fell apart with torch's default init — and one deviation at a time, under the exact protocol:

VACA, exact R protocol	do(x2) error at −3 / −2 / 0
per-node plateau, torch init (#43)	0.026 / 0.034 / 0.077
global plateau (R), torch init	0.523 / 0.334 / 0.129
… no plateau at all	0.562 / 0.354 / 0.142
… Adam eps 1e-7	0.523 / 0.334 / 0.129
… Bernstein on train min/max	0.154 / 0.012 / 0.062
… glorot init	0.035 / 0.006 / 0.007

Keras' Dense default (glorot-uniform, zero bias) is now CausalFlowDAG(init="glorot"), set in the VACA/CAREFL configs (the triangle configs use init="normal", see §3); the framework default stays torch's. The glorot row above is one init draw; at the config's seed 7 the errors are 0.098 / 0.159 / 0.026 (do(x2) = −3 / −2 / 0 in this diagnostic; the committed configs intervene at −3 / −1 / 0 per the R code, §3). Under the exact protocol VACA is seed-sensitive at the off-manifold point do(x2 = −3) — 0.03–0.27 over four init draws, ≤ 0.026 at do(x2 = 0) — and CAREFL is stable (± 0.03) with x4 1.5–2× worse than the per-node approximation gave. All of it, per experiment with every hyperparameter and its source, is in docs/paper-replication.md; the ground truths are re-pinned under the exact protocol.

3. Second verification round (four fresh small-context reviewers)

Ground truth vs paper text + R code — four corrections, all applied:

• The triangle scripts' LinearMasked layers use Keras random_normal (sd 0.05), not glorot → CausalFlowDAG(init="normal"), set in the two triangle configs (glorot stays for VACA/CAREFL's layer_dense).
• The triangle scripts take the Bernstein domain from the 5 %/95 % quantiles as we do — no deviation there; min/max (scale_df) is the comparison scripts' choice only.
• Fig. 5 and vaca_triangle.r intervene at do(x2) = −3, −1, 0 (the text says −2) → DO_X2_VALUES = (-3, -1, 0), matching the frozen truth.json.
• carefl_fig5.r trains on CAREFL's own X.csv with val = train and sd-standardized x3/x4; our draw, validation split and raw units are stated as repo choices. The paper's own fitted numbers (1.98 / −0.21 / 0.26; 2.07 / −0.203; OR 7.74 [7.16, 8.38]) now sit in the report's "paper" column.

Implementation — init="glorot" had overwritten the VC head's deliberate zero output layer (bug; fixed + tested); fit validates vc_ehat (range, alignment) before calibrating; ordinal levels beyond the spec raise clearly; calibrate's second-dataset semantics documented and tested.

Experiments — bench_training stroke-ls crashed on val=None (fixed), per-phase plateau settings, freeze documented as rate 0; the torch plateau mapping was simulated against R's update_learning_rate on 500 random loss sequences: identical in 500/500.

Docs — stale CHANGELOG entries that still described intermediate 0.4 designs rewritten; numbers quoted at the config's seed; tests/README completed.

4. CI runtime — one documented deviation, chosen by measurement

A 24-run grid (batch 128/256/512 × lr 1e-3…1e-2 × 250/500 epochs, four triangle variants, each scored against the ground-truth bands):

batch / lr / epochs	linear-cs	atan-cs	mixed exp-cs	steps
32 / 0.001 / 500 (paper)	0.110	0.080	0.071	1×
128 / 0.001 / 500	0.170	0.041	0.126	1/4
256 / 0.004 / 500	0.132	0.073	0.072	1/8
256 / 0.008 / 500	0.163	0.113	0.073	1/8
512 / 0.010 / 500	0.183	0.104	0.119	1/16

Batch 256 / lr 0.004 is the only row within 0.02 of the paper protocol on every cs error; the triangle configs use it (paper values kept beside it in the YAML), 500 epochs stay (100 → cs errors 2–3×, 250 → mixed exp-cs +60 %). Early stopping at min lr gains nothing for VACA/CAREFL. Ground truth re-pinned from the CI config with init: normal: cs errors 0.088 / 0.077 / 1.024 / 0.107, β's within one SE of truth.

5. Reviews applied earlier

• Cut verifier: chained-centering guard restored (spec level, tested); code-map/fitting/training-speed/varying-coefficients rewritten for the hooks; README import torch; tracked Colab .ipynb regenerated.
• Notebooks executed end to end after the change (four of four).

6. Revision round (three more reviewers: deck, framework, configs)

A CI blocker found: pyyaml was missing from the test dependency group, while testpaths collects experiments/paper/tests/test_configs.py, which imports it — all 15 CI test jobs would have aborted at collection on main/dev-* (this branch never triggers that workflow's push filter). Verified in an isolated uv sync --no-default-groups --group test env, fixed.

Framework

• Every query and read-out is now (df, node, *, ...): varying_coef and intercept_contributions took the node first and called the frame data, unlike pmf/density/scores/design_matrix/effect_modifier_scan. do=, seed=, t=, candidates=, calibrate(marginal_init=) are keyword-only.
• Degenerate data fails loudly instead of training into NaN: coinciding 5 %/95 % quantiles (a constant or 95 %-constant column) and ordinal columns outside 0..levels-1 (previously truncated silently); unknown activation=/init= raise with the valid choices.
• tramdag.utils is gone — config_section moved to experiments/common.py (its only caller), so the package is modelling code only. pmf/density share one prelude; abduct keeps the caller's index; fit_classical records its report in history["classical"]; two one-expression helpers became one.
• New known-truth tests: an affine SI at θ = 0 is the standard logistic including the pre-map Jacobian (the continuous path had no closed-form anchor before), history across two fit calls, vc_ehat outside [0, 1], init="normal" on the biases, both degenerate-column errors.

Configs / CI — experiments.yaml timeout-minutes 300 → 60 (jobs measure ≤ 11 min), the plan job gets one, the variants lookup tolerates a config without them, the result-dir name dashes only the script part (as make_output_dir does); test_configs.py now covers every <area>/<script>.yaml including validate_ls; per-file-ignores pruned to the rules the directories actually raise; the commitizen-branch pre-push hook removed (pre-commit exports PRE_COMMIT_FROM_REF/TO_REF only when both refs exist, so a branch's first push hands commitizen a literal string and it aborts — the commit-msg hook already checks every message).

Docs / YAML headers — the VACA do-values are the R code's −3/−1/0 everywhere (the results table still had −2 rows and pre-change numbers), the batch/lr grid is labelled as glorot-measured, the "now" hyperparameter columns carry the CI deviation, the triangle init claim is scoped to the linear layers.

Checks

ruff (extended set), complexipy 15/10, full pre-commit run --all-files, fast suite 182 passed at 97 % coverage, all four notebooks run, all eight paper variants re-run and green against the re-pinned ground truth (CI run 32974872751), plus three variants re-run after the signature change — check green on all three.

🤖 Generated with Claude Code