|
| 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. |
0 commit comments