Skip to content

Commit 07072db

Browse files
Jammy2211claude
authored andcommitted
docs: formal EP specification for autofit.graphical (README + statistical docstrings)
- autofit/graphical/README.md (new): the mathematics as implemented — factor-graph model, exponential-family mean field, the EP step (cavity / tilted / moment-matching projection / damped division as natural-parameter EMA), the three factor-optimiser paths, convergence criterion + KL direction contract, evidence decomposition, deterministic-variable mechanisms, parallel-EP semantics; each equation anchored to its implementing class/method. - MeanField.update_factor_mean_field: damped-update equations, delta semantics, invalid-projection fallback. - MeanField.kl: KL direction contract (m.kl(other) == KL(m||other)). - FactorApproximation: fix damped-update docstring error ((q_a^f)^(1-d) -> (q_a)^(1-d)). - composed_transform.py: module docstring stating the transform composition-order convention (innermost-first storage; reversal asymmetry) with a worked UniformPrior trace; direction notes on _transform/_inverse_transform; validity condition on TransformedMessage.kl. - AbstractMessage.project: moment-matching projection equations. EP framework review phase 2 (#1333); audit findings context in #1332. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 0f26ff2 commit 07072db

4 files changed

Lines changed: 362 additions & 11 deletions

File tree

autofit/graphical/README.md

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
# `autofit.graphical` — the mathematics, exactly as implemented
2+
3+
This package fits **factor graphs** — joint models over many datasets
4+
and/or shared parameters — using **expectation propagation (EP)**. This
5+
document states the statistical machinery in formal equations, with each
6+
equation anchored to the class or method that implements it, so that the
7+
code can be verified against the math line-by-line (by a human or an AI
8+
agent). It documents what the code *does*; design intent beyond that is
9+
marked explicitly.
10+
11+
Notation: `a, b` index factors; `i, j` index variables; `η` are natural
12+
parameters; `T(x)` sufficient statistics; `A(η)` the log-partition
13+
function; `h(x)` the base measure.
14+
15+
---
16+
17+
## 1. The model
18+
19+
A factor graph over variables `x = (x₁, …, x_n)`:
20+
21+
p(x) ∝ ∏ₐ fₐ(xₐ) (1)
22+
23+
where `xₐ` is the subset of variables factor `a` touches. Factors are
24+
`Factor` objects (`factor_graphs/factor.py`) wrapping arbitrary
25+
callables; a `FactorGraph` (`factor_graphs/graph.py`) is their product.
26+
Priors participate as factors of their own (`declarative/factor/prior.py`,
27+
`MeanField.from_priors`): a prior is one more term in Eq. (1).
28+
29+
In the declarative interface (`declarative/`), each dataset contributes
30+
an `AnalysisFactor` whose value is `Analysis.log_likelihood_function`,
31+
and `FactorGraphModel` assembles Eq. (1) from those plus the prior
32+
factors.
33+
34+
## 2. The approximating family
35+
36+
The approximation is fully factorised over both factors and variables
37+
("mean field"):
38+
39+
q(x) = ∏ₐ qₐ(x), qₐ(x) = ∏ᵢ qₐᵢ(xᵢ) (2)
40+
41+
- `qₐᵢ` — one **message** per (factor, variable): an exponential-family
42+
distribution (`autofit/messages/`),
43+
44+
qₐᵢ(x) = h(x) exp( ηₐᵢ · T(x) − A(ηₐᵢ) ) (3)
45+
46+
- `MeanField` (`mean_field.py`) is one factor's `{Variable → message}`
47+
dictionary, i.e. one `qₐ`.
48+
- `EPMeanField` (`expectation_propagation/ep_mean_field.py`) is the full
49+
`{Factor → MeanField}` map, i.e. `q(x)`.
50+
51+
Because messages are exponential-family, products and quotients are
52+
sums and differences of natural parameters
53+
(`MessageInterface.sum_natural_parameters`, message `__mul__` /
54+
`__truediv__`; powers scale them, `AbstractMessage.__pow__`):
55+
56+
∏ₖ qₖ(x) ∝ h(x) exp( (Σₖ ηₖ) · T(x) ) (4)
57+
58+
The global approximation to the posterior of `xᵢ` is the product of
59+
every factor's message on it: `q(xᵢ) ∝ ∏ₐ qₐᵢ(xᵢ)`
60+
(`EPMeanField.mean_field`).
61+
62+
Non-Gaussian priors (Uniform, LogUniform, LogGaussian) are represented
63+
as a base `NormalMessage` composed with deterministic transforms
64+
(`TransformedMessage`, `messages/composed_transform.py`); the message
65+
algebra operates on the base-space natural parameters and the transform
66+
stack maps to physical space. See the module docstring of
67+
`composed_transform.py` for the composition-order convention.
68+
69+
## 3. One EP update
70+
71+
For each factor `a` in turn (`EPOptimiser.run`,
72+
`expectation_propagation/optimiser.py`):
73+
74+
### 3.1 Cavity — `EPMeanField.factor_approximation`
75+
76+
Everything except factor `a`:
77+
78+
q^{\a}(x) = ∏_{b ≠ a} q_b(x) η_cav,i = Σ_{b≠a} η_{bi} (5)
79+
80+
The code builds this as `MeanField({v: 1.0}).prod(*other_factor_fields)`
81+
and packages `(factor, cavity_dist, factor_dist, model_dist)` into a
82+
`FactorApproximation` (`mean_field.py`), where
83+
`model_dist = factor_dist × cavity_dist` is the current full
84+
approximation restricted to factor `a`'s variables.
85+
86+
### 3.2 Tilted distribution — the factor optimiser
87+
88+
The exact factor times the cavity:
89+
90+
p̂ₐ(x) = fₐ(x) q^{\a}(x) / Ẑₐ , Ẑₐ = ∫ fₐ(x) q^{\a}(x) dx (6)
91+
92+
`FactorApproximation.__call__` evaluates `log fₐ + log q^{\a}`. The
93+
tilted distribution is then *fitted* by the factor's optimiser:
94+
95+
- **Sampling path** (`AbstractSearch.optimise`,
96+
`non_linear/search/abstract_search.py`): the cavity messages are
97+
installed as the model's priors (`prior.with_message`), a non-linear
98+
search (Dynesty, Nautilus, Emcee, …) samples Eq. (6), and the result
99+
is projected per §3.3.
100+
- **Laplace path** (`LaplaceOptimiser`, `graphical/laplace/`):
101+
quasi-Newton ascent on `log p̂ₐ` (`newton.py`: BFGS/SR1 with Wolfe
102+
line search, `line_search.py`), then a Gaussian at the optimum with
103+
covariance from the (approximate) Hessian inverse
104+
(`MeanField.from_mode_covariance` via `from_mode` on each message).
105+
- **Exact path** (`ExactFactorFit`,
106+
`expectation_propagation/factor_optimiser.py`): if the factor is
107+
itself a message of the same family as the cavity
108+
(`Factor.has_exact_projection`), the tilted distribution is conjugate
109+
and the update is the closed-form natural-parameter sum — no sampler.
110+
`EPOptimiser.from_meanfield` auto-selects this path.
111+
112+
### 3.3 Projection (moment matching) — `AbstractMessage.project`
113+
114+
Find the family member closest to the tilted distribution in inclusive
115+
KL:
116+
117+
q* = argmin_q KL( p̂ₐ ‖ q ) (7)
118+
119+
For an exponential family this is **moment matching**:
120+
121+
E_{q*}[ T(x) ] = E_{p̂ₐ}[ T(x) ] (8)
122+
123+
Implemented as importance-weighted sample moments over the search's
124+
weighted samples `{x_s, log w_s}`:
125+
126+
E[T] ≈ Σ_s w̃_s T(x_s) / Σ_s w̃_s , w̃_s = exp(log w_s − max log w) (9)
127+
128+
(`AbstractMessage.project`; the max-log-weight shift is the standard
129+
numerical stabilisation, and the mean weight supplies the projection's
130+
`log_norm`). `from_sufficient_statistics` then inverts Eq. (8) to
131+
natural parameters per family. On the Laplace path the "projection" is
132+
the Gaussian mode/covariance construction instead.
133+
134+
### 3.4 Factor update with damping — `MeanField.update_factor_mean_field`
135+
136+
Divide out the cavity and damp with `δ ∈ (0, 1]`:
137+
138+
qₐ^new = (q*)^δ (qₐ^old)^{1−δ} / (q^{\a})^δ (10)
139+
140+
equivalently, an exponential moving average on natural parameters:
141+
142+
ηₐ ← (1 − δ) ηₐ + δ ( η_{q*} − η_cav ) (11)
143+
144+
`δ` is supplied by an `ApproxUpdater` (`optimiser.py`): `SimplerUpdater`
145+
(one fixed value), `FactorUpdater` (per factor), or `DynamicUpdater`
146+
(per variable, `δᵢ ∝ min_count / count(i)` — variables shared by more
147+
factors update more slowly). `δ` may therefore be a scalar or a
148+
per-variable `MeanField` of scalars.
149+
150+
**Invalid-projection fallback**: if the division produces an invalid
151+
message (e.g. negative variance — possible because Eq. (10)'s
152+
subtraction of natural parameters is not closed in the family),
153+
`update_factor_mean_field` reverts the invalid parameters to the
154+
previous message per-parameter (`update_invalid`) and flags
155+
`StatusFlag.BAD_PROJECTION`.
156+
157+
## 4. Convergence — `EPHistory` (`expectation_propagation/history.py`)
158+
159+
After each factor update the history records the new `EPMeanField`.
160+
Termination when either:
161+
162+
KL( q_t ‖ q_{t−1} ) = Σᵢ KL( q_t(xᵢ) ‖ q_{t−1}(xᵢ) ) < kl_tol (12)
163+
164+
(`FactorHistory.kl_divergence`, summed per-variable via `MeanField.kl`;
165+
default `kl_tol = 1e-1`), or the log-evidence change drops below
166+
`evidence_tol`, or a user callback fires.
167+
168+
**KL direction contract**: `m.kl(other)` means `KL(m ‖ other)`. (As of
169+
the 2026-07 audit, `NormalMessage` and `TruncatedNormalMessage` satisfy
170+
this; `GammaMessage` and `BetaMessage` compute the reverse direction,
171+
and `TruncatedNormalMessage` uses the untruncated formula — tracked in
172+
issue #1332, findings F2/F6.)
173+
174+
## 5. Evidence — `EPMeanField.log_evidence`
175+
176+
The EP approximation to the model evidence factorises as:
177+
178+
Zᵢ = ∫ ∏ₐ qₐᵢ(xᵢ) dxᵢ (per variable) (13)
179+
Zₐ = Ẑₐ / ∏_{i ∈ a} Zᵢ (per factor) (14)
180+
Z = ∏ᵢ Zᵢ ∏ₐ Zₐ (15)
181+
182+
`Zᵢ` is computed in closed form from log-partitions
183+
(`MessageInterface.log_normalisation`):
184+
185+
log ∫ ∏ₖ qₖ = Σₖ (log hₖ − A(ηₖ)) − ( log h − A(Σₖ ηₖ) ) (16)
186+
187+
and the per-factor `Ẑₐ` is carried on `MeanField.log_norm` by the
188+
projection. (Audit note: as of 2026-07 the `log_norm` bookkeeping is
189+
broken in three places — issue #1332 finding F7 — so Eq. (15) should
190+
not be used for model comparison until those fixes land; the
191+
decomposition itself is correct.)
192+
193+
## 6. Deterministic variables
194+
195+
Three composition mechanisms exist; they are **not** interchangeable
196+
and reconciling them is an open design item (EP review Phase 5):
197+
198+
1. **Graph-level deterministic variables**: `Factor(..., factor_out=v)`
199+
declares outputs computed by the factor
200+
(`FactorValue.deterministic_values`); the Laplace path maintains a
201+
separate curvature estimate for them
202+
(`quasi_deterministic_update`, `laplace/newton.py`), and their
203+
approximate distributions live in the cavity
204+
(`FactorApproximation.deterministic_dist`).
205+
2. **Compound prior arithmetic** (`mapper/prior/arithmetic/`):
206+
deterministic relations expressed on priors at model-composition
207+
time (e.g. `prior_a * matrix + prior_b`).
208+
3. **Free shared variables**: share a prior across factors and encode
209+
the relation inside the likelihood.
210+
211+
## 7. Parallel EP — `ParallelEPOptimiser`
212+
213+
All factor approximations for a sweep are built from the **same**
214+
mean field, factor optimisations run in a process pool, and the updates
215+
are applied sequentially afterwards. This is standard "parallel EP"
216+
semantics: cavities within a sweep are stale relative to serial EP, so
217+
serial and parallel runs converge along different trajectories (to the
218+
same fixed points when EP converges).
219+
220+
## 8. Reading
221+
222+
- T. Minka (2001), *Expectation Propagation for Approximate Bayesian
223+
Inference* — the algorithm of §3.
224+
- A. Vehtari et al. (2020), *Expectation propagation as a way of life*
225+
(JMLR 21(17)) — the data-partitioned framing this package follows
226+
(one factor per dataset, Eq. (1)).
227+
- M. Seeger (2005), *Expectation propagation for exponential families*
228+
— the natural-parameter algebra of §2–§3.

autofit/graphical/mean_field.py

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,17 @@ def sample(self, n_samples=None):
388388
return VariableData({v: dist.sample(n_samples) for v, dist in self.items()})
389389

390390
def kl(self, mean_field: "MeanField") -> np.ndarray:
391+
"""
392+
The KL divergence between two mean fields, summed over variables:
393+
394+
KL(self || mean_field) = Σᵢ KL( self[xᵢ] || mean_field[xᵢ] )
395+
396+
which is exact because both mean fields factorise over variables.
397+
398+
Direction contract: ``message.kl(other)`` means ``KL(message || other)``
399+
for every message family. ``EPHistory`` uses this sum between
400+
successive global mean fields as its convergence criterion.
401+
"""
391402
return sum(np.sum(dist.kl(mean_field[k])) for k, dist in self.items())
392403

393404
__hash__ = Factor.__hash__
@@ -405,7 +416,44 @@ def update_factor_mean_field(
405416
delta: Delta = 1.0,
406417
status: Status = Status(),
407418
) -> Tuple["MeanField", Status]:
419+
"""
420+
The EP factor update: given ``self`` = the newly-fitted model
421+
distribution q* (the moment-matched projection of the tilted
422+
distribution), divide out the cavity to recover the factor's new
423+
message, damped by ``delta``:
424+
425+
q_a_new = (q*)^δ · (q_a_old)^(1−δ) / (cavity)^δ
426+
427+
which on natural parameters is an exponential moving average,
428+
429+
η_a ← (1 − δ)·η_a + δ·(η_q* − η_cavity)
408430
431+
``delta`` may be a scalar or a per-variable ``MeanField`` of scalars
432+
(see ``DynamicUpdater``). At ``delta == 1`` this reduces to the
433+
undamped ``q* / cavity``.
434+
435+
The natural-parameter subtraction is not closed in the family (it can
436+
produce e.g. a negative variance). If the result is invalid, the
437+
invalid parameters are reverted per-parameter to ``last_dist`` via
438+
``update_invalid`` and the status is flagged ``BAD_PROJECTION``.
439+
440+
Parameters
441+
----------
442+
cavity_dist
443+
The cavity distribution: the product of every other factor's
444+
messages on this factor's variables.
445+
last_dist
446+
The factor's previous message distribution ``q_a_old`` (required
447+
when ``delta < 1`` and as the fallback for invalid projections).
448+
delta
449+
Damping coefficient in (0, 1]; scalar or per-variable MeanField.
450+
status
451+
Status carried through from the factor optimisation.
452+
453+
Returns
454+
-------
455+
The factor's new message distribution and the updated status.
456+
"""
409457
success, messages, _, flag = status
410458
updated = False
411459
try:
@@ -505,9 +553,9 @@ class FactorApproximation(AbstractNode):
505553
returns q⁺ᵃ(x₀), {xₐ: dq⁺ᵃ(x₀)/dxₐ}
506554
507555
project_mean_field(mean_field, delta=1., status=Status())
508-
for qᶠ = mean_field, finds qₐ such that qᶠₐ * q⁻ᵃ = qᶠ
509-
delta controls how much to change from the original factor qₐ
510-
so qʳₐ = (qᶠₐ)ᵟ * (qᶠₐ)¹⁻ᵟ
556+
for qᶠ = mean_field, finds qᶠₐ such that qᶠₐ * q⁻ᵃ = qᶠ
557+
delta controls how much to change from the previous factor
558+
distribution qₐ, so qʳₐ = (qᶠₐ)ᵟ * (qₐ)¹⁻ᵟ
511559
512560
returns qʳₐ, status
513561
"""

autofit/messages/abstract.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -266,9 +266,26 @@ def factor(self, x):
266266
def project(
267267
cls, samples: np.ndarray, log_weight_list: Optional[np.ndarray] = None, **kwargs
268268
) -> "AbstractMessage":
269-
"""Calculates the sufficient statistics of a set of samples
270-
and returns the distribution with the appropriate parameters
271-
that match the sufficient statistics
269+
"""
270+
Moment-matching projection: the exponential-family member closest
271+
(in inclusive KL) to the weighted sample distribution.
272+
273+
For an exponential family, argmin_q KL(p || q) is the member whose
274+
expected sufficient statistics match p's:
275+
276+
E_q[T(x)] = E_p[T(x)]
277+
278+
Estimated by importance-weighted sample moments over weighted
279+
samples {x_s, log w_s} (e.g. from a nested sampler run on the EP
280+
tilted distribution):
281+
282+
E[T] ≈ Σ_s w̃_s T(x_s) / Σ_s w̃_s , w̃_s = exp(log w_s − max log w)
283+
284+
The max-log-weight shift is the standard numerical stabilisation;
285+
the mean unshifted weight supplies the projection's ``log_norm``
286+
(the tilted distribution's normalisation estimate).
287+
``from_sufficient_statistics`` then inverts the moment equations to
288+
natural parameters for the concrete family.
272289
"""
273290
# if weight_list aren't passed then equally weight all samples
274291

0 commit comments

Comments
 (0)