Full Probabilistic Inference Schema

This note lays out a unified probabilistic schema for state and parameter estimation, organised along two axes:

1. Two model tracks — a simulator (no internal latent variable) and an emulator (a latent-variable model trained on simulator outputs).

2. Three inference regimes — exact posteriors, per-observation variational inference, and amortized inference that generalises over observations.

The most general generative model factorises as

Notation¶

The full probabilistic graphical model implied by (1):

Track 1 — Simulator¶

Generative Model¶

With no internal latent $\boldsymbol{z}$, the simulator maps $\boldsymbol{\theta}, \boldsymbol{x}$ directly to the state $\boldsymbol{u}$. The state is the only latent object besides $\boldsymbol{\theta}$.

The marginal likelihood integrates out both the state and the parameters,

1A · Exact Inference¶

The exact target posteriors are:

The asymmetry is deliberate: the params-only posterior must marginalise out the still-unknown state $\boldsymbol{u}$ (hence the integral), whereas the state-only posterior conditions on $\boldsymbol{\theta}$, leaving $\boldsymbol{u}$ as the only unknown. It is therefore a plain Bayes posterior with no second latent to integrate — only the evidence $p(\boldsymbol{y} \mid \boldsymbol{x}, \boldsymbol{\theta}) = \int p(\boldsymbol{y} \mid \boldsymbol{u}, \boldsymbol{\theta}, \boldsymbol{x}) \, p(\boldsymbol{u} \mid \boldsymbol{\theta}, \boldsymbol{x}) \, \mathrm{d}\boldsymbol{u}$ appears in its normaliser.

1B · Variational Inference¶

Introduce a variational distribution $q$ with parameters $\boldsymbol{\psi}$ optimised per observation. Here $\boldsymbol{\psi}$ does not generalise across different $\boldsymbol{y}$ or $\boldsymbol{x}$ — a fresh $\boldsymbol{\psi}$ is optimised for each observation.

Params only — $q(\boldsymbol{\theta} \mid \boldsymbol{\psi}) \approx p(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x})$

State only — $q(\boldsymbol{u} \mid \boldsymbol{\psi}) \approx p(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

Joint — two ways to factor the variational family:

1C · Amortized Inference¶

The variational family now conditions on both $\boldsymbol{y}$ and $\boldsymbol{x}$. A single forward pass replaces per-observation optimisation — train once, generalise over $\boldsymbol{y}$ and $\boldsymbol{x}$.

Params only — $q(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\psi}) \approx p(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x})$

State only — $q(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\psi}) \approx p(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

Joint

Track 2 — Emulator¶

Generative Model¶

The emulator introduces an internal latent variable $\boldsymbol{z}$ that compresses the full state $\boldsymbol{u}$. It is itself a generative model, trained on simulator outputs before inference is performed.

The marginal likelihood now also integrates out $\boldsymbol{z}$,

2.0 · Emulator Training¶

Before inference, train the emulator on simulator outputs $\{\boldsymbol{u}, \boldsymbol{x}, \boldsymbol{\theta}\}$ using its own internal ELBO. This introduces emulator inference parameters $\boldsymbol{\psi}_{\mathrm{em}}$ and an encoder $q(\boldsymbol{z} \mid \boldsymbol{u}, \boldsymbol{x}, \boldsymbol{\theta}, \boldsymbol{\psi}_{\mathrm{em}})$.

2.1 · Emulator Uncertainty Characterization¶

Characterise the discrepancy between the emulator and the true simulator before using the emulator for inference,

2A · Exact Inference¶

The targets now involve three unknowns: $\boldsymbol{u}, \boldsymbol{z}, \boldsymbol{\theta}$.

Here every conditional target sits in the chain $\boldsymbol{z} \to \boldsymbol{u} \to \boldsymbol{y}$ and marginalises out the other latent: conditioning on $\boldsymbol{\theta}$, the state-only posterior integrates out the deeper latent $\boldsymbol{z}$; the latent-only posterior — the emulator’s compressed state given observations — integrates out the intervening state $\boldsymbol{u}$; and the params-only posterior integrates out both. Only the joint, which targets all unknowns at once, is a single Bayes ratio with no marginalisation.

2B · Variational Inference¶

Per-observation $q$ with parameters $\boldsymbol{\psi}$. The emulator’s internal latent $\boldsymbol{z}$ is now an additional object the variational family must handle.

Params only — $q(\boldsymbol{\theta} \mid \boldsymbol{\psi}) \approx p(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x})$

with $p(\boldsymbol{y} \mid \boldsymbol{\theta}, \boldsymbol{x}) = \iint p(\boldsymbol{y} \mid \boldsymbol{u}, \boldsymbol{\theta}, \boldsymbol{x}) \, p(\boldsymbol{u} \mid \boldsymbol{z}, \boldsymbol{\theta}, \boldsymbol{x}) \, p(\boldsymbol{z} \mid \boldsymbol{\theta}, \boldsymbol{x}) \, \mathrm{d}\boldsymbol{u}\, \mathrm{d}\boldsymbol{z}$ requiring a further inner approximation over both $\boldsymbol{u}$ and $\boldsymbol{z}$.

State only — $q(\boldsymbol{u} \mid \boldsymbol{\psi}) \approx p(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

where $p(\boldsymbol{u} \mid \boldsymbol{\theta}, \boldsymbol{x}) = \int p(\boldsymbol{u} \mid \boldsymbol{z}, \boldsymbol{\theta}, \boldsymbol{x}) \, p(\boldsymbol{z} \mid \boldsymbol{\theta}, \boldsymbol{x}) \, \mathrm{d}\boldsymbol{z}$ is itself intractable and approximated via the emulator ELBO (19).

Latent only — $q(\boldsymbol{z} \mid \boldsymbol{\psi}) \approx p(\boldsymbol{z} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

with $p(\boldsymbol{y} \mid \boldsymbol{z}, \boldsymbol{\theta}, \boldsymbol{x}) = \int p(\boldsymbol{y} \mid \boldsymbol{u}, \boldsymbol{\theta}, \boldsymbol{x}) \, p(\boldsymbol{u} \mid \boldsymbol{z}, \boldsymbol{\theta}, \boldsymbol{x}) \, \mathrm{d}\boldsymbol{u}$.

Joint

2C · Amortized Inference¶

The variational family conditions on both $\boldsymbol{y}$ and $\boldsymbol{x}$; a single forward pass replaces per-observation optimisation.

Params only — $q(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\psi}) \approx p(\boldsymbol{\theta} \mid \boldsymbol{y}, \boldsymbol{x})$

State only — $q(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\psi}) \approx p(\boldsymbol{u} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

Latent only — $q(\boldsymbol{z} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\psi}) \approx p(\boldsymbol{z} \mid \boldsymbol{y}, \boldsymbol{x}, \boldsymbol{\theta})$

Joint

Summary Table¶