Amortized Inference

Every method up to now — OI, 3DVar, strong / weak / incremental 4DVar — solves a fresh minimisation per event. A new observation triggers a new analysis: fine for retrospective studies, too slow for real-time alerting. Amortized inference instead learns a function $q_{\boldsymbol{\phi}}(\mathbf{x} \mid \mathbf{y})$ that maps observations directly to posteriors in a single forward pass. The training cost amortises over many events; the per-event cost collapses to one network evaluation.

This is not a departure from variational data assimilation — it is its endpoint. The variational cost we have minimised all along is the posterior’s energy; amortized inference learns to produce that posterior for all observations at once.

From Per-Event Optimisation to a Learned Posterior Map¶

Recall the variational cost shared by 3DVar and 4DVar, $J(\mathbf{x}) = \tfrac{1}{2}\|\mathbf{x}-\mathbf{x}_b\|^2_{\mathbf{B}^{-1}} + \tfrac{1}{2}\sum_t \|\mathbf{y}_t - H_t(M_t(\mathbf{x}))\|^2_{\mathbf{R}_t^{-1}}$. It is exactly the negative log of the joint, so the posterior is a Gibbs density in $J$:

The methods so far all chase the mode of (1):

one optimisation per observation. Variational inference instead approximates the whole posterior with a tractable $q$, and amortized variational inference trains a single $q_{\boldsymbol{\phi}}(\mathbf{x}\mid\mathbf{y})$ that does this for every $\mathbf{y}$ — the amortized regime of the inference schema, with the conditioning network playing the role of the variational parameters.

The amortization spectrum¶

The DA methods form a ladder of how much is moved from per-event solve-time into trained-once weights:

Read top-to-bottom, amortized inference is the $K \to 0$ limit of 4DVarNet: there is no inner optimisation left at all — the network emits the posterior directly.

The Variational Objective¶

Train the conditional density to match the exact posterior in expectation over observations:

For each $\mathbf{y}$, that KL is one ELBO away from the evidence,

so minimising the KL maximises the ELBO. Substituting $\log p(\mathbf{x},\mathbf{y}) = -J(\mathbf{x}) + \text{const}$ from (1) gives the clean variational-DA reading:

Simulation-Based Training¶

When the prior $p(\mathbf{x})$ and the forward model are samplable — exactly the simulator track of the generative model — the training set comes from simulation: draw a truth, push it through the physics, add noise, and learn to invert it.

This is simulation-based inference Cranmer et al., 2020: every pair has known ground truth and the training distribution is fully controlled — but the deployment distribution must match it, or $q_{\boldsymbol{\phi}}$ is miscalibrated (the central risk, below).

import jax

def simulate_pair(key, prior_sample, forward, obs_noise):
    """One synthetic (x, y) pair from the generative model."""
    kx, ke = jax.random.split(key)
    x = prior_sample(kx)                       # x ~ p(x)
    y = forward(x) + obs_noise(ke)             # y = H(M(x)) + ε
    return x, y

The Posterior Family¶

The conditioning network $c_{\boldsymbol{\psi}}(\mathbf{y})$ (an encoder over the gappy observations) feeds a head that defines $q_{\boldsymbol{\phi}}$. Three choices trade expressiveness against cost:

Training Objective¶

Flow and regression heads admit a tractable density, so train by maximum likelihood on simulated pairs; score heads use denoising score matching:

Either way it is a standard minibatch loop over simulated data — the same optax training machinery as 4DVarNet:

import jax
import jax.numpy as jnp
import optax

optimizer = optax.adam(1e-3)
opt_state = optimizer.init(params)            # params = (φ, ψ): head + encoder

@jax.jit
def train_step(params, opt_state, x, y):
    # maximum likelihood:  L = -E log q_φ(x | y)
    loss, grads = jax.value_and_grad(lambda p: -jnp.mean(log_q(p, x, y)))(params)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

Calibration: the Hard Part¶

Amortized inference is dangerous precisely because it looks confident. A tight $q_{\boldsymbol{\phi}}$ centred in the wrong place gives a confident wrong answer with no warning — the opposite failure mode to a slow-but-honest solver. Promote a head to operational use only after three gates pass:

The adjoint-calibration gate is the one that keeps amortized inference inside DA rather than outside it: the learned map’s Jacobian $\partial \mathbf{x}^\star / \partial \mathbf{y}$ must match the physics-based tangent-linear / adjoint of the true inverse problem. Matching the oracle MAP on the training set is not enough — without matching gradients, the head extrapolates badly the moment the observation distribution shifts.

The Oracle Pattern¶

This is what ties amortized inference back to the rest of the DA hierarchy. The right design is both, in sequence:

1. Run a solver-based method (incremental / strong 4DVar) to produce oracle posteriors on a representative sample.

2. Train the amortized head against the oracle (and/or by simulation).

3. Validate with the three gates.

4. Deploy the head for real-time, high-throughput work.

The physics-based solvers are the teachers; the amortized head is the student that runs in milliseconds once it has earned trust. Amortized inference does not replace variational DA — it compiles it.

When It Helps, and the Trade-offs¶