research: EP review phase 1 — statistics audit of autofit/graphical

[Jammy2211]

Overview

Phase 1 of the EP framework review (umbrella: PyAutoMind/research/graphical_ep/ep_framework_review.md; Phase 0 verdicts: #1330; fix decisions: #1331). A full statistical audit of the expectation-propagation machinery in autofit/graphical/, in the same style as the priors/messages audit: re-derive the math each component claims to implement, verify the code against it, and record findings with minimal reproductions.

Phase 0's key lesson sets the priority: bugs hide on the generic exponential-family path while direct paths stay correct (e.g. the TruncatedNormal log-partition bug) — and that generic path is exactly what EP consumes.

Plan

• Audit the core EP statistics in autofit/graphical/: message algebra (natural parameters, sufficient statistics, log-normalisers, products/quotients), cavity distribution computation, tilted-distribution moment matching / KL projection direction, damping, the mean-field approximation, EPHistory convergence criteria (KL evaluation), and the Laplace optimiser step.
• Ground the audit in the three real use cases: HowToFit/scripts/chapter_3_graphical_models, the IC50 cancer workspace, and the cosmology scripts — confirming the code paths they actually exercise.
• Output: a findings census (confirmed bugs / suspect math / verified-correct inventory), each finding with a minimal reproduction, posted here.
• Write the formal statistical description (equations as implemented, not textbook ideals) — this seeds the Phase 2 package documentation.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary — read-only: audit against clean main)
• PyAutoMind (findings census + verdicts)

Suggested branch

None — research task; fixes graduate to their own tasks (coordinating with #1331's batch).

Audit order (generic-path-first, per the Phase 0 lesson)

1. autofit/messages/abstract.py message interface algebra: __mul__/__truediv__ natural-parameter arithmetic, sum_natural_parameters, log_normalisation, project, from_mode — the operations every EP update composes.
2. Factor approximation / cavity: how q^{\i} is formed and how the tilted distribution is projected back (autofit/graphical/expectation_propagation/, mean_field.py).
3. KL machinery: which direction KL(p||q) vs KL(q||p) is computed where; moment-matching equivalence assumptions; EPHistory convergence checks.
4. Damping / step control and the LaplaceOptimiser (autofit/graphical/laplace/) — gradients, Hessians, line search.
5. Deterministic variables plumbing (factor_graphs/), as groundwork for Phase 5.
6. Cross-check against the three use cases; note which components each exercises.

Key Files

• autofit/graphical/expectation_propagation/ — EP optimiser, history, factor optimiser
• autofit/graphical/mean_field.py, factor_graphs/ — approximation structure
• autofit/graphical/laplace/ — Laplace/Newton machinery
• autofit/messages/ — the algebra underneath (post-Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331 fixes where relevant)

[Jammy2211]

Phase 1 — initial findings (core EP machinery pass)

First audit pass over mean_field.py, ep_mean_field.py, optimiser.py, history.py, factor_optimiser.py and the message algebra in messages/abstract.py + interface.py. Verdicts: 2 confirmed bugs, 2 code defects, 1 doc defect — and the central EP update logic verifies correct.

Confirmed bugs (with reproductions)

F1. MeanField.__truediv__ and __pow__ pass log_norm into the plates argument — mean_field.py:322-336 calls type(self)({...}, self.log_norm - other.log_norm), but __init__ is (dists, plates=None, log_norm=0.0). Verified: dividing two mean fields with log_norm 3.0 and 1.0 yields log_norm = 0.0 (silently dropped) and _plates = 2.0 (a float in the plates slot). Same shape in __pow__ (_plates = 1.5). log_norm feeds factor_evidence → log_evidence, so EP evidence bookkeeping through these operators is silently wrong; the corrupted plates is a latent landmine for any plate-aware path. (__pow__'s MeanField-exponent branch also computes self.log_norm * other.log_norm, which is not a meaningful quantity even in the right slot.)

F2. GammaMessage.kl computes the reverse KL direction vs NormalMessage.kl — verified numerically: n1.kl(n2) = 0.443147 = KL(n1‖n2) ✓, but g1.kl(g2) = 0.106770 = KL(g2‖g1) (forward would be 0.173602). gamma.py sets P, Q = dist, self (the source even carries a # TODO check this is correct). MeanField.kl sums per-variable KLs, and EPHistory.kl_divergence uses that sum as the convergence criterion — a graph mixing Normal and Gamma variables sums KLs measured in opposite directions. Fix: align Gamma (and audit Beta/others) to self.kl(other) = KL(self‖other).

Code defects

F3. ParallelEPOptimiser.run crashes at completion when paths=None — unlike the serial EPOptimiser.run (which guards with if self.paths:), the parallel variant calls self.visualiser() / self._output_results(...) unconditionally at the end (optimiser.py:498-499), and self.visualiser is None without paths. Also worth documenting: the parallel variant computes all factor approximations from the same (stale) mean field per sweep — correct "parallel EP" semantics, but a behavioural difference from serial that users should know.

F4. AbstractMessage.update_invalid scalar branch is self-flagged broken — abstract.py:336 carries # TODO: Fairly certain this would not work on the non-array path; this is the invalid-projection recovery path used by update_factor_mean_field under BAD_PROJECTION.

Doc defect

F5. FactorApproximation.project_mean_field docstring writes the damped update as qʳₐ = (qᶠₐ)ᵟ * (qᶠₐ)¹⁻ᵟ — the second term should be the previous factor dist (qₐ)¹⁻ᵟ.

Verified correct (inventory so far)

• Damped factor update (update_factor_mean_field): (new**δ * last**(1-δ)) / cavity**δ = exponential moving average on natural parameters — algebraically correct damped EP.
• Cavity/tilted construction (EPMeanField.factor_approximation): cavity = product of other factors' mean fields; textbook.
• AbstractMessage.project: importance-weighted sufficient-statistics projection with correct max-shift numerical stabilisation.
• log_normalisation: product-of-exponential-family normaliser Σ(logh−A) − (logh−A)_prod; structurally correct.
• NormalMessage.kl: exact.
• Evidence decomposition (EPMeanField.log_evidence): matches the stated Z = ∏Zᵢ∏Zₐ factorisation.
• Exact-update infrastructure exists (ExactFactorFit, calc_exact_update/calc_exact_projection/has_exact_projection) — Phase 6 (analytic likelihood updates) has a hook to build on rather than green-field.

Cross-link to Phase 0

Phase 0's finding 04 (TruncatedNormalMessage.log_partition missing the Gaussian term, #1331) directly poisons log_normalisation → variable_evidence → log_evidence for any graph containing truncated-normal messages — the two fixes belong in the same batch.

Remaining Phase 1 scope

laplace/ (Newton step, line search), declarative/ (AnalysisFactor plumbing, hierarchical factors), factor_graphs/ (jacobians, deterministic-variable propagation — groundwork for Phase 5), Beta/other message KL audits, and grounding against the three use cases (HowToFit ch3, IC50, cosmology). Findings census will be updated here as the audit proceeds.

[Jammy2211]

Phase 1 — findings update 2 (message KLs, Laplace machinery, declarative path, use-case grounding)

New confirmed findings

F2 (extended). BetaMessage.kl is also reverse-direction — verified numerically: b1.kl(b2) = 1.266291 = KL(b2‖b1) (forward is 1.750376), and its docstring explicitly claims KL(self || dist). Both Beta and Gamma set P = dist, Q = self; Normal and TruncatedNormal compute KL(self‖dist). The KL direction contract must be stated once and enforced family-wide (property test).

F6. TruncatedNormalMessage.kl uses the untruncated Gaussian formula — the docstring admits it's approximate, but even with identical supports the true truncated-normal KL involves truncated moments; the formula degrades as posterior mass approaches the bounds. This is live in the flagship tutorial: HowToFit chapter 3 tutorial 5 uses af.TruncatedGaussianPrior, so its EP convergence metric runs on this approximation (compounded by the Phase 0 log-partition bug #1331).

F7. Evidence bookkeeping is systemically broken across three independent legs:
(a) MeanField.__truediv__/__pow__ silently drop log_norm (F1);
(b) MeanField.from_priors — the projection used by every search-driven factor update (AbstractSearch.optimise → result.projected_model.priors) — constructs with default log_norm = 0, so the sampler's per-factor evidence never enters the mean field;
(c) the Phase 0 TruncatedNormalMessage.log_partition bug poisons log_normalisation → variable_evidence.
Bottom line: EPMeanField.log_evidence should not be trusted for model comparison in sampler-driven EP fits on current main. Fixing (a) and (c) is mechanical; (b) needs a design decision on where factor evidence should be recorded.

F8. Dead/suspect quasi-Newton code in laplace/newton.py — diag_sr1_bfgs_update computes locals then implicitly returns None (would nuke the state if ever wired in); bfgs1_update disagrees with the exported bfgs_update on the sign of the B·d term (they can't both be right); diag_sr1_update_ is unused. Only full_bfgs_update / full_sr1_update / full_diag_update are exported. Recommend deleting the dead variants.

F9 (robustness). laplace/line_search.py imports scipy.optimize._linesearch (private API) — works today, breaks on a scipy reorganisation. Also noted: TransformedMessage.kl delegates to the base messages without checking the two messages share a transform stack — correct inside the EP loop (same variable ⇒ same transforms), silently wrong as a public API on differently-transformed messages.

Verified correct (additions)

• The search-driven EP factor update is textbook: cavity messages become the model's priors (prior.with_message), the search samples the tilted distribution, and result.projected_model moment-matches the posterior back onto the prior families via the importance-weighted projection.
• Quasi-Newton core: exported bfgs_update/sr1_update algebra checks out (SR1 with correct tolerance guard; the deterministic-variable curvature-transfer update satisfies its secant-like condition).
• Exact updates (ExactFactorFit): implemented via message-as-factor (AbstractMessage.as_factor + calc_exact_update returning the conjugate same-family multiply-in), auto-selected by EPOptimiser.from_meanfield when has_exact_projection — this is the Phase 6 extension point (e.g. linear-Gaussian factors could join the analytic path).

Use-case grounding

Use case	EP API exercised	Notes
HowToFit ch3 (tut 5 + optional hierarchical EP)	FactorGraphModel, AnalysisFactor, Dynesty + LaplaceOptimiser, EPHistory, TruncatedGaussianPrior	exercises truncated messages → F6 + #1331-04 live here
Cosmology (concr/scripts/cosmology/ep.py)	FactorGraphModel, Nautilus + LaplaceOptimiser, EPHistory, UniformPrior	transformed (uniform) messages in EP
IC50 cancer (ep_sim/ep_real)	FactorGraphModel with a free shared hill_coef; the global linear relation latents @ coef_matrix + coef_mean is evaluated inside the likelihood	does not use the low-level deterministic_variable/factor_out API — the Phase 5 tension confirmed: three composition mechanisms exist (compound prior arithmetic in mapper/prior/arithmetic/, free shared variables, low-level deterministic factors) and the real use case picks model-space composition

Remaining

Hierarchical factor math (declarative/factor/hierarchical.py), a quantitative pass on the truncated-KL approximation error, and the formal statistical write-up (equations as implemented) that seeds Phase 2 docs. Findings census then final verdict table.

[Jammy2211]

Phase 1 complete — consolidated verdict

The statistics audit of autofit/graphical is done (core EP machinery, message algebra + KLs, Laplace/quasi-Newton, declarative/search path, hierarchical factors, use-case grounding). Final additions since update 2:

• F6 quantified: the truncated-normal KL approximation error is 1.5% (mass central) → 72% (near a bound) → 140% (mode at the bound) against numerically-exact truncated KL — directly biasing the default kl_tol = 0.1 termination for truncated priors.
• Hierarchical factors reviewed — sound: one factor per drawn variable (deliberate dimensionality choice, author's open design question noted in the docstring); no math defects found at this level.

Verdict table

#	Finding	Verdict	Action
F1	MeanField.__truediv__/__pow__ put log_norm in the plates slot	confirmed (repro)	fix now — mechanical
F2	Gamma and Beta kl compute KL(dist‖self); Normal/TruncNormal compute KL(self‖dist)	confirmed (repro)	fix now — align to KL(self‖other) + family-wide property test
F3	ParallelEPOptimiser.run unguarded visualiser() at completion	confirmed	fix now — trivial
F4	update_invalid scalar branch self-flagged broken	confirmed (author TODO)	fix with F1 batch
F5	Damped-update docstring typo	confirmed	fix in Phase 2 docs
F6	Truncated-normal KL uses untruncated formula (up to 140% error)	confirmed + quantified	decision: exact truncated KL vs documented approximation
F7	Evidence bookkeeping broken in 3 legs (F1 + from_priors log_norm + #1331-04)	confirmed	(a),(c) mechanical; (b) needs design: where sampler evidence enters the mean field
F8	Dead/suspect quasi-Newton variants (diag_sr1_bfgs_update returns None; bfgs1_update sign-inconsistent)	confirmed	delete dead code
F9	scipy private _linesearch import; TransformedMessage.kl doesn't check transform equality	confirmed	robustness fixes
—	Core EP update algebra (cavity/tilted/projection/damping-EMA), evidence decomposition structure, importance-weighted projection, exported quasi-Newton updates, search-driven factor loop, exact-update plumbing	verified correct	document in Phase 2

Outputs & hand-offs

• Statistical write-up done (equations as implemented: mean-field factorisation, cavity/tilted/moment-matching, damping as natural-parameter EMA, evidence decomposition, exact conjugate updates, the three deterministic-composition mechanisms). This seeds the Phase 2 public documentation, which must additionally state two contracts the code currently leaves implicit: the KL direction (m.kl(other) ≡ KL(m‖other)) and the log-prior normalisation convention (Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331 decision 4).
• Phase 5 groundwork: confirmed the IC50 use case composes deterministic relations in model space (compound prior arithmetic / free shared variables), not via the low-level factor_out deterministic API — three mechanisms to reconcile.
• Phase 6 groundwork: analytic updates already exist for conjugate message-factors (ExactFactorFit); scoping = which non-trivial factor forms (linear-Gaussian first) can join that path.
• Recommended EP fix batch (alongside the Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331 priors batch): F1 + F2 + F3 + F4 + F8, all mechanical with reproductions in hand; F6/F7b join after their decisions.

Phase 1 of the umbrella review is complete pending guidance on #1331 (which also settles F6/F7's adjacent decisions). Phase 2 (formal package documentation + doc-level fixes) is ready to start.

[Jammy2211]

Phase 3 addendum — F10: undamped Laplace-only EP collapses the mean field (behavioural)

Found while building the Phase 3 integration scripts (autofit_workspace#81), reproduced independently twice:

• Running factor_graph.optimise with LaplaceOptimiser as every factor's optimiser and the default undamped updater (SimplerUpdater(delta=1.0)) for the default max_steps=100: every posterior sigma collapses to ~1e-15 around the starting means — repeated cycles over-count shared-variable information and nothing converges to the data.
• Evidence this is a known sharp edge: the multi-factor pure-Laplace unit test test_autofit/graphical/gaussian/test_declarative.py::_test_optimise_factor_model is disabled (leading-underscore) in the unit suite.
• The canonical pattern (per-AnalysisFactor sampler + Laplace only as the PriorFactor default) does not exhibit this; nor do short runs (max_steps<=3).

Implication for the fix batch / Phase 4 tooling: there is no guardrail — no warning when sigmas collapse below any floor, no damping default that prevents it, and EPHistory happily iterates to max_steps while it happens. Suggest (a) considering a default delta < 1 or a DynamicUpdater default, and (b) a Phase 4 diagnostic that flags monotone sigma collapse across cycles. Related: F2/F6 (the KL convergence metric this pathology sails past).

Phase 3 outputs (step-by-step feature example + ep_parity/ep_deterministic/ep_exact integration scripts) are at ship sign-off on autofit_workspace#81.