Amortised Bayesian inference for mechanistic models
Estimated time: 1.5-2 hours | Advanced session | Requires: particle MCMC, variational inference, and basic familiarity with neural networks (what layers and activation functions are)
In the particle MCMC session we fitted the stochastic SEIT4L model to the Tristan da Cunha outbreak. Every likelihood evaluation needed a full particle filter, every chain needed hundreds of thousands of those evaluations, and the resulting chain answered a question about exactly one dataset. Ask about a second dataset — a different island, a different reporting schedule, the same outbreak truncated at week three — and the whole cost is paid again.
The variational inference session already replaced sampling with optimisation. But it optimised an approximation to the posterior for one dataset: change the data and you re-run the optimisation from scratch.
This session goes one step further. We use the simulator to generate training data, and learn the mapping from data to posterior once. The learned estimator is amortised: afterwards a posterior costs milliseconds, for any dataset the model can produce, including observation schedules we invent after training.
The approach is neural posterior estimation (NPE), part of the wider family of simulation-based inference (Cranmer et al. 2020).
In R, the NeuralEstimators package on CRAN does neural posterior estimation by calling out to Julia, so the tool underneath is a Julia one. The established tools are Python’s sbi and BayesFlow (Radev et al. 2020), which package the whole simulation-based workflow rather than the flow alone, and in Julia NeuralEstimators.jl does the same, training on simulations generated as it goes. We build the flow ourselves with Lux.jl in about forty lines instead. For a course that is a feature: every part of the method stays visible in the session’s own code. The ideas transfer directly — everything below is ordinary Bayesian inference for the same mechanistic model you have been fitting all week.
In this session you will:
Everything here targets the same posterior as PMMH: the same SEIT4L model, the same priors, the same Poisson observation process. The neural network replaces the sampler, not the model. R_0 still means what it meant in the Structure and stochasticity session because it is still a parameter of a transmission model.
We build the flow with Lux.jl, the neural network library the universal differential equations session also uses, so this session needs nothing beyond the course environment.
using Distributions, LinearAlgebra, Random, Statistics
using CSV, DataFrames, DrWatson, Plots, StatsPlots
using Lux, Zygote, Optimisers
using MFIIDD
flu_tdc = CSV.read(datadir("flu_tdc_1971.csv"), DataFrame)
const YOBS = Vector{Float64}(flu_tdc.obs)
const T = length(YOBS)
const INIT = [279.0, 0.0, 2.0, 3.0, 0.0, 0.0, 0.0, 0.0] ## S, E, I, T1-T4, L8-element Vector{Float64}:
279.0
0.0
2.0
3.0
0.0
0.0
0.0
0.0Bayesian inference wants . NPE approximates it by training a conditional density estimator on pairs drawn from the joint distribution:
Minimising that expected negative log density over draws from has the true posterior as its minimiser (Papamakarios and Murray 2016). So this is the same inferential target you have been chasing with MCMC, reached by regression on simulations rather than by sampling.
The density estimator is a normalising flow: an invertible neural network that maps a simple base distribution (a standard normal) to a complicated one, conditioned on the data. Because it is invertible with a tractable Jacobian, we can both evaluate its density (needed for training) and sample from it (needed for inference).
If you have not used neural networks before, the ones here are doing something narrow. A network is a function with adjustable numbers inside it, built by alternating a linear map with a simple nonlinear function such as gelu. Given enough of those numbers it can approximate almost any smooth function, and training means nudging them with gradient descent until some loss is small.
In this session the networks never see a posterior and never output one directly. Each takes half the parameter vector plus a summary of the data, and returns a scale and a shift for the other half. The statistics live in the flow construction around them: the change-of-variables formula turns those scales and shifts into a density, and the invertibility lets us sample. You can follow every statistical claim below while treating each network as a black-box flexible function.
Both VI and NPE replace sampling with optimisation, so it is worth being precise about the difference.
| Variational inference | Neural posterior estimation | |
|---|---|---|
| What is optimised | for one dataset | over all datasets |
| Divergence | reverse KL, | forward KL, in expectation over |
| Behaviour at the optimum | mode-seeking, too narrow | mass-covering, too wide |
| Typical failure in practice | too narrow | too narrow; a flow that starts at the prior, like this one, errs wide if undertrained |
| Needs | the likelihood | only the simulator |
| New dataset | re-optimise from scratch | free |
The reverse KL that VI minimises is mode-seeking, which is why the mean-field posteriors in that session came out too narrow. NPE minimises a forward KL averaged over datasets, which is mass-covering, so at the optimum its errors tend towards posteriors that are too wide, and being conservative is the safer failure.
That argument holds at the optimum. It is not a guarantee for a network trained on a finite budget: benchmarks across simulation-based inference find overconfident posteriors to be common (Hermans et al. 2022), which is why the calibration check below is not optional.
The price of amortisation is the amortisation gap: the estimator must work for every dataset the prior can produce, so on any single dataset it will be worse than a bespoke fit. More training reduces the gap; it does not close it.
NPE can be seen as amortised variational inference with the KL reversed. If you found the VI session’s framing natural — pick a family, optimise until it fits — this is the same move, applied once to the whole joint distribution rather than repeatedly to individual posteriors.
We reuse the priors and the jump simulator from the PMMH session, which fitted the same SEIT4L model.
const PNAMES = [:R_0, :D_lat, :D_inf, :α, :D_imm, :ρ]
const PRIORS = [
truncated(Normal(3.0, 2.0), lower=1.0),
truncated(Normal(2.0, 1.0), lower=0.5),
truncated(Normal(3.0, 2.0), lower=0.5),
Beta(2, 2),
truncated(Normal(15.0, 10.0), lower=1.0),
Beta(2, 2),
]
const D = length(PNAMES)
"Simulate one SEIT4L outbreak and its Poisson-observed daily incidence."
function simulate_obs(rng, θ)
d = Dict{Symbol,Float64}(PNAMES[i] => Float64(θ[i]) for i in 1:D)
state = copy(INIT)
problem = seitl_jump_problem(d, state; rng = rng)
y = Vector{Float64}(undef, T)
for t in 1:T
inc = seitl_jump_step!(problem, state, 1.0)
y[t] = rand(rng, Poisson(max(θ[6] * inc, 1e-10)))
end
return y
endMain.Notebook.simulate_obsHow expensive is one simulation?
rng = Xoshiro(20)
θ_test = [rand(rng, p) for p in PRIORS]
simulate_obs(rng, θ_test) ## warm up
@elapsed for _ in 1:1000; simulate_obs(rng, θ_test); end0.110997081A single outbreak takes well under a millisecond. Compare that with the SEIT4L particle filter behind the PMMH session, where one likelihood evaluation runs 256 such outbreaks side by side. Streaming the 512,000 simulations this training run consumes takes a couple of minutes, while the committed PMMH chains needed about 650,000 likelihood evaluations between them, each roughly 256 times the work, and took hours. This is the trade NPE makes: pay a large one-off simulation and training cost, then get posteriors for free. In a slower language the simulation stage is usually the bottleneck.
The flow’s base distribution is a standard normal, so it helps enormously to give it parameters on a comparable scale. We use the probit of the prior CDF:
where is the prior CDF of parameter . Under the prior, is exactly standard normal. Two things follow: the flow starts out equal to the prior and learns the prior-to-posterior update, and inverting always lands back inside the prior support, so D_lat can never come out below its truncation bound.
const CLIP = 1e-6
unconstrain(θ) = [quantile(Normal(), clamp(cdf(PRIORS[i], θ[i]), CLIP, 1 - CLIP)) for i in 1:D]
constrain(u) = [quantile(PRIORS[i], clamp(cdf(Normal(), u[i]), CLIP, 1 - CLIP)) for i in 1:D]constrain (generic function with 1 method)The obvious choice, log(θ - lower), blows up as a parameter approaches its truncation bound. Over 100,000 prior draws you will get values near , which are extreme outliers once standardised, and they destabilise training badly. The probit transform has no such tail.
Real surveillance is ragged: locations start reporting at different times, some days are missed, some studies visit a site twice in seven years. A neural network needs a fixed-length input, so we put every record on a common grid of T days and add a binary mask saying which days were actually observed (Wang et al. 2024).
"Pack an observation vector and its mask into the conditioning channels."
pack(y, m) = vcat(Float32.(log1p.(y)) .* m, m)
"Random schedule: half the time fully observed, otherwise a random subset of days."
function draw_mask(rng)
m = ones(Float32, T)
if rand(rng) < 0.5
nobs = rand(rng, 8:T)
m[randperm(rng, T)[1:(T - nobs)]] .= 0.0f0
end
return m
endMain.Notebook.draw_maskThe zero-filling is only padding — it carries no information. All the information about missingness lives in the mask. Without it the network sees a zero on day 12 and cannot tell whether nobody reported that day or whether genuinely zero cases were reported, and those two imply very different things about .
Because a random mask is drawn for every training example, the network learns — a posterior conditioned on the observation schedule itself. That lets one trained network serve every schedule, including ones we invent later.
The network only learns about schedules it has seen. draw_mask is fully observed half the time and otherwise keeps a random 8 to T days, so a posterior conditioned on a schedule outside that range is an extrapolation, however confident it looks. Exercise 2 makes the point with a mask distribution that concentrates observations early.
function build_training_set(n; seed=1)
rng = Xoshiro(seed)
Θ = Matrix{Float32}(undef, D, n)
Y = Matrix{Float32}(undef, 2T, n)
for j in 1:n
θ = [rand(rng, PRIORS[i]) for i in 1:D]
Θ[:, j] = Float32.(unconstrain(θ))
Y[:, j] = pack(simulate_obs(rng, θ), draw_mask(rng))
end
return Θ, Y
end
@time Xva, Yva = build_training_set(5_000; seed=77) ## held out, for monitoring
size(Xva), size(Yva) 1.792699 seconds (8.44 M allocations: 534.510 MiB, 7.54% gc time, 52.23% compilation time)((6, 5000), (118, 5000))Θ is D × n and Y is 2T × n, the layout Lux’s Dense layers expect: one column per training pair.log1p — compresses the count scale and handles zeros, so a day with 40 cases does not dominate the network’s input.rng, threaded through — seitl_jump_problem takes the generator, so the parameter draws, the simulation and the mask all come off the same seeded stream and the function is reproducible on its own.pack returns 2T numbers: T (masked) observations followed by the T mask entries.We use a conditional coupling flow, the simplest invertible architecture that still has a tractable Jacobian (Dinh et al. 2017).
Split the six parameters into two halves. Leave one half untouched, and use it — together with the observed data — to choose an affine transformation, a scale and a shift, for the other half. Because the untouched half decides the transformation, the step inverts exactly: recompute the same scale and shift, then undo them. The Jacobian is triangular, so its log-determinant is simply the sum of the log-scales. Stack several such layers, alternating which half moves, and the composition is flexible enough to represent the posterior.
const FLOW = (couple=8, hidden=64, embed=64) ## layers, coupling width, embedding width
## Layer k transforms the entries where the mask is 0 and leaves the rest alone.
const MASKS = [Float32[(i + k) % 2 for i in 1:D] for k in 1:FLOW.couple]
function build_flow(rng)
embed = Chain(Dense(2T => FLOW.embed, gelu), Dense(FLOW.embed => FLOW.embed, gelu))
nets = Tuple(Chain(Dense(D + FLOW.embed => FLOW.hidden, gelu),
Dense(FLOW.hidden => FLOW.hidden, gelu),
Dense(FLOW.hidden => 2D; init_weight=zeros32, init_bias=zeros32))
for _ in 1:FLOW.couple)
model = (embed=embed, nets=nets)
return model, Lux.initialparameters(rng, model), Lux.initialstates(rng, model)
end
"Log-scale and shift for coupling layer `k`, both zero on the frozen half."
function coupling(model, ps, st, k, xm, h, mask)
out, _ = model.nets[k](vcat(xm, h), ps.nets[k], st.nets[k])
s = 2.0f0 .* tanh.(out[1:D, :]) .* (1 .- mask)
t = out[D+1:2D, :] .* (1 .- mask)
return s, t
end
"Parameters (probit scale) → latent, with the log-determinant of the Jacobian."
function flow_forward(model, ps, st, x, y)
h, _ = model.embed(y, ps.embed, st.embed)
z = x
logdet = zeros(Float32, size(x, 2))
for k in 1:FLOW.couple
mask = MASKS[k]
xm = z .* mask
s, t = coupling(model, ps, st, k, xm, h, mask)
z = xm .+ (1 .- mask) .* (z .* exp.(s) .+ t)
logdet = logdet .+ vec(sum(s, dims=1))
end
return z, logdet
end
"Latent → parameters (probit scale), conditional on the data."
function flow_inverse(model, ps, st, z, y)
h, _ = model.embed(y, ps.embed, st.embed)
x = z
for k in FLOW.couple:-1:1
mask = MASKS[k]
xm = x .* mask
s, t = coupling(model, ps, st, k, xm, h, mask)
x = xm .+ (1 .- mask) .* ((x .- t) .* exp.(-s))
end
return x
end
model, ps, st = build_flow(Xoshiro(2024))
Lux.parameterlength(ps) ## how many weights we are about to train87648Chain and Dense(in => out, activation) are the same Lux pieces the UDE session uses: Dense applies a linear map and then an element-wise activation, Chain runs several in sequence. Lux.initialparameters and Lux.initialstates walk the container and return the parameters and state separately from the architecture, which is why every call below passes ps and st explicitly.embed compresses the 2T conditioning inputs — the masked observations and the mask — into FLOW.embed numbers, which every coupling layer sees. The parameters and the data meet in vcat(xm, h).init_weight=zeros32 on the last layer of each coupling network means every scale starts at 1 and every shift at 0, so the untrained flow is exactly the identity. Combined with the probit transform, the network starts as the prior and learns the update from there.2 * tanh(...) bounds the log-scales to ±2. An unbounded scale is the usual way a coupling flow blows up in the first few hundred steps.logdet accumulates the log-scales of the entries that moved. That is the whole Jacobian: an affine map on one half, conditioned on the other.flow_inverse recomputes s and t from the frozen half and undoes the affine step. Nothing is approximated — it is the same arithmetic, backwards.We can check that claim rather than trust it — a density estimator whose inverse does not invert is not estimating anything. The check has to wait until after training: at initialisation every scale and shift is exactly zero, so the flow is the identity and the round trip would return zero for any inverse at all, including a wrong one.
A neural network is normally trained by cycling repeatedly over a fixed training set, and the usual failure is overfitting: validation loss bottoms out early and then climbs while training loss keeps falling. Density estimators are especially prone to it, because the model can win by placing a spike on each training point.
Here we do not have a fixed training set. A simulation costs a fraction of a millisecond, so we draw a fresh batch for every gradient step. No example is ever seen twice, and overfitting becomes impossible by construction.
Simulating a batch of 128 is the same order of cost as the gradient step it feeds, so the simulator takes a substantial share of the training time. At this model size that is affordable, and it is the price of the guarantee. With a simulator a hundred times slower it would not be: you would go back to a fixed training set, and back to watching for overfitting.
The training objective is the negative log density of the training parameters under the flow: half the squared norm of the latents (how far the flow pushed them from the standard normal) minus the log-determinant (how much it stretched the space). Averaging it over the batch and the six dimensions keeps the number readable.
function objective(model, ps, st, x, y)
z, logdet = flow_forward(model, ps, st, x, y)
return (sum(abs2, z) / 2 - sum(logdet)) / (size(x, 2) * D)
end
function train_streaming(model, ps, st, Xva, Yva; steps, batch, lr=1.0f-3,
val_every=25, sampler=build_training_set)
opt = Optimisers.setup(Optimisers.Adam(lr), ps)
history = NamedTuple[]
best = (loss=Inf, step=0, ps=deepcopy(ps))
for s in 1:steps
Xb, Yb = sampler(batch; seed=10_000 + s) ## fresh every step
loss, back = Zygote.pullback(p -> objective(model, p, st, Xb, Yb), ps)
grads = back(one(loss))[1]
opt, ps = Optimisers.update(opt, ps, grads)
if s == steps ÷ 2 ## step decay
Optimisers.adjust!(opt, lr / 2)
elseif s == (3 * steps) ÷ 4
Optimisers.adjust!(opt, lr / 4)
end
if s % val_every == 0
vl = objective(model, ps, st, Xva, Yva)
push!(history, (step=s, val=vl))
isfinite(vl) && vl < best.loss &&
(best = (loss=vl, step=s, ps=deepcopy(ps)))
end
end
return DataFrame(history), best ## keep the best network, not the last one
end
@time history, best = train_streaming(model, ps, st, Xva, Yva; steps=4_000, batch=128)
best.step, best.loss139.537297 seconds (746.12 M allocations: 76.724 GiB, 13.71% gc time, 19.78% compilation time)(3925, 0.036802832f0)plot(history.step, history.val, lw=2, legend=false,
xlabel="gradient step", ylabel="validation objective", title="Training")
vline!([best.step], ls=:dash, c=:grey)
Because every batch is new, the validation objective settles onto a plateau rather than turning back upwards. We still keep the best network rather than the last, since the plateau is noisy.
Now the invertibility check promised earlier is worth running, because the scales and shifts are no longer zero:
z_check, _ = flow_forward(model, best.ps, st, Xva[:, 1:64], Yva[:, 1:64])
maximum(abs, flow_inverse(model, best.ps, st, z_check, Yva[:, 1:64]) .- Xva[:, 1:64])7.1525574f-7A number at the level of floating-point noise means the two directions really are inverses of each other.
Normalising flows can train unstably. A validation objective that falls and then climbs steadily means the flow is diverging, and no amount of extra training will rescue it. If that happens, reduce the learning rate or the number of coupling layers.
Two choices in build_flow are doing the stabilising here, and both are worth copying into any flow you write: the zero-initialised output layer, which starts the network at the identity map, and the bounded log-scale, which stops a single coupling layer from stretching the space by an arbitrary factor.
Training is done. Every posterior from here on is essentially free.
function npe_sample(model, ps, st, y, m, nsamp; seed=7)
Yc = repeat(pack(y, m), 1, nsamp) ## same data for every draw
Zn = Float32.(randn(Xoshiro(seed), D, nsamp)) ## draws from the base normal
U = flow_inverse(model, ps, st, Zn, Yc)
return stack(constrain(Float64.(u)) for u in eachcol(U); dims=1)
end
mask_full = ones(Float32, T)
@time post = npe_sample(model, best.ps, st, YOBS, mask_full, 10_000)
posterior = DataFrame(post, PNAMES)
describe(posterior, :median, :std) 0.823600 seconds (1.43 M allocations: 153.005 MiB, 8.78% gc time, 78.35% compilation time)| Row | variable | median | std |
|---|---|---|---|
| Symbol | Float64 | Float64 | |
| 1 | R_0 | 5.58991 | 1.33513 |
| 2 | D_lat | 1.2182 | 0.45845 |
| 3 | D_inf | 2.16326 | 0.953009 |
| 4 | α | 0.453568 | 0.0684742 |
| 5 | D_imm | 10.8796 | 1.6334 |
| 6 | ρ | 0.650654 | 0.0750922 |
Ten thousand posterior draws, with no MCMC and no further simulation. The time printed above is a first call, so most of it is compilation.
pmmh = CSV.read(datadir("pmcmc_seit4l_chain.csv"), DataFrame)
plts = []
for p in PNAMES
pl = density(pmmh[!, p], lw=2.5, label="PMMH", title=string(p), titlefontsize=10)
density!(pl, posterior[!, p], lw=2.5, ls=:dash, label="NPE")
push!(plts, pl)
end
plot(plts..., layout=(2, 3), size=(950, 520))
Expect broad agreement in location, and expect the widths to disagree. On this run the NPE posterior is wider than PMMH for five of the six parameters, by 10% to 56%, and narrower for D_imm.
That is a property of this run. The mass-covering argument in the comparison with variational inference holds only at the optimum, and a flow trained on a finite budget can just as easily come out too narrow; benchmarks find that it often does (Hermans et al. 2022).
Widths are only half of it. R_0 shows why the two are worth reading separately: its width sits within 15% of PMMH’s, while its centre is about half a PMMH posterior standard deviation low, the NPE median being 5.59 against PMMH’s 6.25, and the two modes are visibly apart in that panel. ρ makes the same point from the other side, agreeing least well on centre while being the widest of the six. So a matching width tells you nothing about where an interval sits, and the next section flags R_0 as well.
Read the weakly identified parameters differently from the well identified ones. A single 59-day outbreak says little about D_lat and D_inf separately, because a shorter latent period and a longer infectious period reproduce much the same epidemic curve, so the two posteriors have most licence to disagree there.
Agreement with PMMH is reassuring but is not a check in itself. The standard check for a simulation-based posterior is simulation-based calibration (Talts et al. 2018): draw from the prior, simulate data, fit, and record the rank of the true among the posterior draws. If the posterior is calibrated, those ranks are uniform.
n_sbc, L_sbc = 200, 200
Xs, Ys = build_training_set(n_sbc; seed=999)
ranks = zeros(Int, n_sbc, D)
for j in 1:n_sbc
Yc = repeat(Ys[:, j:j], 1, L_sbc)
Zn = Float32.(randn(Xoshiro(1000 + j), D, L_sbc))
U = flow_inverse(model, best.ps, st, Zn, Yc)
for i in 1:D
ranks[j, i] = count(<(Xs[i, j]), U[i, :])
end
end
grid = range(0, 1, length=100)
## a rank takes L_sbc + 1 values, so under calibration P(u <= x) rises in steps
## of 1 / (L_sbc + 1) rather than along the diagonal
p0 = [(floor(Int, L_sbc * x) + 1) / (L_sbc + 1) for x in grid]
lo = [quantile(Binomial(n_sbc, p), 0.025) / n_sbc for p in p0]
hi = [quantile(Binomial(n_sbc, p), 0.975) / n_sbc for p in p0]
sbc_plots = []
for (i, p) in enumerate(PNAMES)
u = ranks[:, i] ./ L_sbc
ecdf_vals = [mean(u .<= x) for x in grid]
pl = plot(grid, hi .- p0, fillrange=lo .- p0, fillalpha=0.2, c=:grey, lw=0,
legend=false, title=string(p), titlefontsize=10)
plot!(pl, grid, ecdf_vals .- p0, lw=2)
hline!(pl, [0], c=:black, ls=:dot)
push!(sbc_plots, pl)
end
plot(sbc_plots..., layout=(2, 3), size=(950, 520))
A curve wandering outside the grey band means the posterior is miscalibrated for that parameter, and its shape says how. A curve that crosses zero in the middle, running above the line on one side and below it on the other, is a width problem: intervals systematically too narrow or too wide. A curve bowed to one side, staying above or below the line across the whole range, is bias.
Look at R_0. Its curve sits above zero across almost the whole range and leaves the band between roughly 0.24 and 0.35, with no sign change, so this is bias rather than width. Ranks skewed low mean the true value falls below the NPE draws more often than it should, which is the estimator centring R_0 too high. D_lat bows gently the other way but stays inside the band, and ρ bows the same way and dips below it over the lowest ranks. D_inf and D_imm stay inside throughout, and α touches the edge at a single point, which a pointwise band produces by chance. So both checks land on R_0, but they disagree about the direction. Against PMMH on the Tristan da Cunha data the NPE median came out low; averaged over datasets drawn from the prior, calibration says the estimator runs high. Both can hold at once, because they ask different questions: SBC describes the estimator across every dataset the prior can produce, while the comparison describes one fit to one outbreak. An estimator biased high on average can still land low on the particular dataset in front of you.
Two things to keep in mind when reading it. The band is pointwise, computed separately at each of the 100 grid points, so under perfect calibration a curve still strays outside it somewhere reasonably often; judge the shape of the excursion rather than the fact of one. (Simultaneous bands for this plot are given by Säilynoja et al. (2022).) And the ranks average over the distribution of schedules the estimator was trained on, so passing here does not certify calibration for one particular mask, including the fully observed one used above.
SBC checks the estimator against the simulator it was trained on. It cannot tell you that the simulator is a poor description of reality. A well-calibrated posterior for the wrong model is still the wrong answer (Hermans et al. 2022).
ppc_rng = Xoshiro(4242)
ppc = reduce(vcat, [simulate_obs(ppc_rng, post[j, :])' for j in 1:500])
tt = 1:T
lo95 = [quantile(ppc[:, t], 0.025) for t in tt]
hi95 = [quantile(ppc[:, t], 0.975) for t in tt]
med = [median(ppc[:, t]) for t in tt]
plot(tt, hi95, fillrange=lo95, fillalpha=0.25, lw=0, label="95%")
plot!(tt, med, lw=2, label="median")
scatter!(tt, YOBS, c=:black, ms=3, label="observed",
xlabel="day", ylabel="reported cases", title="Posterior predictive check")
MCMC on a badly specified model gives visibly bad fit and pathological diagnostics. A trained flow will happily map data that sits outside everything it was trained on to a tidy, narrow, confident and wrong posterior, with no convergence diagnostic to catch it. The posterior predictive check is your main guard.
The replicates above re-roll the epidemic from day 0, because the flow returns parameters and nothing else. No latent path appears anywhere in , so the estimator holds nothing that could be conditioned on the observations.
The median in that figure sits far below the observed second wave, which makes the model look unable to produce one. Count the replicates instead and most of them manage it; they disagree about when, with peaks landing anywhere across the second half of the outbreak. A median taken separately at each day averages over that disagreement and flattens a wave the individual replicates have.
wave = 25:45
totals = [sum(ppc[j, wave]) for j in 1:size(ppc, 1)]
peak_day = [wave[argmax(ppc[j, wave])] for j in 1:size(ppc, 1)]
println("replicates with at least 20 cases over days 25-45: ",
round(100 * mean(totals .>= 20)), "%")
println("their peak day, 90% range: ", quantile(peak_day, 0.05), " to ", quantile(peak_day, 0.95))
println("pointwise median on the observed peak day (32): ", median(ppc[:, 32]))replicates with at least 20 cases over days 25-45: 89.0%
their peak day, 90% range: 25.0 to 42.0
pointwise median on the observed peak day (32): 3.0A particle filter can hand back a path that was conditioned on the data, and the PMMH session does this for its own chain. The same helper works here. The parameter draws came from a neural network rather than from MCMC, which the filter neither knows nor needs to.
Random.seed!(2) ## the bootstrap filter inside `filtered_incidence` uses the global RNG
n_filt = 200
filtered = map(1:n_filt) do _
r = posterior[rand(1:nrow(posterior)), :]
θ = Dict(:R_0 => r.R_0, :D_lat => r.D_lat, :D_inf => r.D_inf,
:α => r.α, :D_imm => r.D_imm, :ρ => r.ρ)
inc = filtered_incidence(θ, flu_tdc.obs, 256)
[rand(Poisson(max(r.ρ * x, 1e-10))) for x in inc]
end
filt = reduce(hcat, filtered)'function band!(p, M, label, colour)
lo = [quantile(M[:, t], 0.025) for t in tt]
hi = [quantile(M[:, t], 0.975) for t in tt]
md = [median(M[:, t]) for t in tt]
plot!(p, tt, md, ribbon=(md .- lo, hi .- md), label=label,
colour=colour, fillalpha=0.2, lw=2)
end
pl = plot(xlabel="day", ylabel="reported cases",
title="Latent path re-simulated and conditioned on the data")
band!(pl, ppc, "re-simulated", :darkorange)
band!(pl, filt, "filtered", :seagreen)
scatter!(pl, tt, YOBS, c=:black, ms=3, label="observed")
Amortisation stops at the parameters. The flow returns in milliseconds and will keep doing so for any dataset, and it says nothing about the latent path: each of the two hundred trajectories above cost its own particle filter run, at the 256 particles the PMMH session settled on. That is still a bargain beside PMMH, which runs a filter at every iteration — about 650,000 of them across the committed chains — merely to arrive at . Here the filter runs once per trajectory you actually want, and only if you want one.
The filtered envelope tracks the observations because it was told what they were, and it would do the same for a model that generates nothing like this outbreak when left to run. Judge the fit on the re-simulated envelope, and read the filtered one as a picture of what the epidemic plausibly did. Each of the two hundred paths comes from its own filter run at its own posterior draw, and that keeps them independent: within one run the surviving particles share ever fewer distinct ancestors as you trace them back, so two hundred paths from a single run would understate the spread badly.
Here is the payoff. We never retrain — we just condition on a different mask.
widths = DataFrame(nobs=Int[], param=String[], width=Float64[])
for nobs in [10, 20, 30, 40, 50, T], rep in 1:5 ## 8 is the sparsest schedule trained on
m = zeros(Float32, T)
m[randperm(Xoshiro(rep), T)[1:nobs]] .= 1.0f0
p = npe_sample(model, best.ps, st, YOBS, m, 2_000; seed=100 + rep)
for i in 1:D
push!(widths, (nobs, string(PNAMES[i]),
quantile(p[:, i], 0.95) - quantile(p[:, i], 0.05)))
end
end
summary_widths = combine(groupby(widths, [:nobs, :param]), :width => mean => :w)
## The widths differ by more than an order of magnitude between parameters, from
## about 0.2 for α and ρ to six days or more for D_imm, so on a shared axis the
## narrow curves flatten along the bottom. Scale each parameter by its own width
## under the fully observed schedule instead, and every curve starts at 1 on the
## right.
pl = plot(xlabel="number of observed days",
ylabel="90% width, relative to fully observed",
title="Posterior width vs surveillance intensity", legend=:topright)
for p in string.(PNAMES)
s = sort(summary_widths[summary_widths.param .== p, :], :nobs)
plot!(pl, s.nobs, s.w ./ s.w[end], lw=2, marker=:circle, ms=3, label=p)
end
pl
Each of those posteriors took milliseconds. Doing the same with PMMH would mean 30 separate chains, and hours of compute.
In favour
Against
For a one-off fit to a single dataset, use PMMH. For the hundredth fit of the same model, amortisation wins.
Retrain with a much smaller budget and compare the posterior widths against the full run. Which parameters suffer most, and does the validation objective still reach a plateau?
Build a fresh flow with build_flow (reusing best.ps would continue from the trained weights) and call train_streaming with steps=400. Compare 90% interval widths.
model_s, ps_s, st_s = build_flow(Xoshiro(11))
hist_small, best_small = train_streaming(model_s, ps_s, st_s, Xva, Yva;
steps=400, batch=128)
post_small = npe_sample(model_s, best_small.ps, st_s, YOBS, mask_full, 5_000)
comparison = DataFrame(
parameter = string.(PNAMES),
width_400 = [quantile(post_small[:, i], 0.95) - quantile(post_small[:, i], 0.05) for i in 1:D],
width_4000 = [quantile(post[:, i], 0.95) - quantile(post[:, i], 0.05) for i in 1:D],
)
comparison.ratio = round.(comparison.width_400 ./ comparison.width_4000, digits=2)
comparison| Row | parameter | width_400 | width_4000 | ratio |
|---|---|---|---|---|
| String | Float64 | Float64 | Float64 | |
| 1 | R_0 | 5.08687 | 4.38929 | 1.16 |
| 2 | D_lat | 2.36554 | 1.48312 | 1.59 |
| 3 | D_inf | 2.80704 | 3.09704 | 0.91 |
| 4 | α | 0.30726 | 0.223383 | 1.38 |
| 5 | D_imm | 9.38511 | 5.34191 | 1.76 |
| 6 | ρ | 0.361755 | 0.246594 | 1.47 |
Every interval is wider, here by 13% to 101%. That direction follows from how this flow is built: it starts out equal to the prior, so with less training it has learned less of the prior-to-posterior update and falls back towards the prior, and the amortisation gap closes from above. An undertrained flow initialised this way is a cautious one; one initialised differently need not be.
The ordering is steadier than the ratios alone suggest. Repeat the exercise with other seeds in build_flow and D_inf and α stay the least affected, at around 1.1, while ρ stays the worst, at roughly twice the width its fully trained counterpart has. α narrows sharply under full training and yet loses least when training is cut short, because whether a second wave happens at all is written directly into the shape of the data, and even a short run captures most of that update.
draw_mask removes days completely at random, which is a strong assumption. Change it so missingness is structured — early days observed far more often than late ones — retrain briefly, and compare the posterior for the fully observed data.
Write a build_training_set variant that uses the new mask, then pass it to train_streaming through its sampler argument.
function build_structured(n; seed=1)
rng = Xoshiro(seed)
Θ = Matrix{Float32}(undef, D, n)
Y = Matrix{Float32}(undef, 2T, n)
for j in 1:n
θ = [rand(rng, PRIORS[i]) for i in 1:D]
m = zeros(Float32, T)
for t in 1:T ## early days far more likely
rand(rng) < (t <= 20 ? 0.95 : 0.4) && (m[t] = 1.0f0)
end
Θ[:, j] = Float32.(unconstrain(θ))
Y[:, j] = pack(simulate_obs(rng, θ), m)
end
return Θ, Y
end
Xv2, Yv2 = build_structured(2_000; seed=5)
model_2, ps_2, st_2 = build_flow(Xoshiro(12))
hist2, best2 = train_streaming(model_2, ps_2, st_2, Xv2, Yv2;
steps=400, batch=128, sampler=build_structured)
post_struct = npe_sample(model_2, best2.ps, st_2, YOBS, mask_full, 5_000)
DataFrame(parameter = string.(PNAMES),
random_masks = [median(post[:, i]) for i in 1:D],
structured_masks = [median(post_struct[:, i]) for i in 1:D])| Row | parameter | random_masks | structured_masks |
|---|---|---|---|
| String | Float64 | Float64 | |
| 1 | R_0 | 5.58991 | 5.72374 |
| 2 | D_lat | 1.2182 | 1.44426 |
| 3 | D_inf | 2.16326 | 1.70405 |
| 4 | α | 0.453568 | 0.351954 |
| 5 | D_imm | 10.8796 | 10.551 |
| 6 | ρ | 0.650654 | 0.656066 |
The two disagree, and that is the point. The estimator learns a posterior conditioned on the distribution of schedules it was trained on. Deploying it on a schedule far outside that distribution is extrapolation. In a real surveillance system the missingness is informative — reporting starts because the epidemic arrived — and neither mask distribution captures that.
Simulate 500 datasets from the prior and plot them against the Tristan da Cunha data. Does the observed outbreak look like something this model and prior can produce? What would it mean for the NPE posterior if it did not?
pp_rng = Xoshiro(11)
prior_pred = reduce(vcat, [simulate_obs(pp_rng, [rand(pp_rng, p) for p in PRIORS])' for _ in 1:500])
plo = [quantile(prior_pred[:, t], 0.025) for t in tt]
phi = [quantile(prior_pred[:, t], 0.975) for t in tt]
plot(tt, phi, fillrange=plo, fillalpha=0.2, c=:grey, lw=0, label="prior predictive 95%")
scatter!(tt, YOBS, c=:black, ms=3, label="observed",
xlabel="day", ylabel="reported cases", title="Prior predictive check")
If the observed data fell outside the prior predictive range, the trained network would never have seen anything like it during training. It would still return a posterior — confidently — but that posterior would be an extrapolation with no guarantees at all. This is the silent failure mode discussed above, and it is why the prior and posterior predictive checks matter more for NPE than for MCMC.
The thinning experiment used random schedules. Compare specific schedules with the same number of observations — the first 20 days, the last 20 days, and every third day. Which is most informative about R_0, and which about D_imm? Why might they differ?
schedules = Dict(
"first 20" => 1:20,
"last 20" => (T-19):T,
"every third" => 1:3:T,
)
rows = DataFrame(schedule=String[], parameter=String[], width=Float64[])
for (name, idx) in schedules
m = zeros(Float32, T); m[collect(idx)] .= 1.0f0
p = npe_sample(model, best.ps, st, YOBS, m, 3_000; seed=42)
for i in 1:D
push!(rows, (name, string(PNAMES[i]),
quantile(p[:, i], 0.95) - quantile(p[:, i], 0.05)))
end
end
unstack(rows, :parameter, :schedule, :width)| Row | parameter | first 20 | every third | last 20 |
|---|---|---|---|---|
| String | Float64? | Float64? | Float64? | |
| 1 | R_0 | 4.38407 | 5.38287 | 3.87886 |
| 2 | D_lat | 2.01168 | 2.17374 | 2.7936 |
| 3 | D_inf | 3.51807 | 3.50444 | 5.94425 |
| 4 | α | 0.577953 | 0.391745 | 0.585637 |
| 5 | D_imm | 20.5384 | 12.0063 | 30.3436 |
| 6 | ρ | 0.371073 | 0.484405 | 0.546254 |
The mechanistic expectation is that R_0 is driven by the exponential growth phase and so wants early observations, while D_imm governs the second wave and so wants late ones. Check that against the table you just printed: it is not what comes out. R_0 is narrowest under “last 20” and D_imm under “every third”, with “last 20” its widest.
Before concluding anything about surveillance design, look back at draw_mask. It produces either a fully observed series or a uniformly random subset of days, so a contiguous block like days 1-20 has essentially no chance of appearing in training. All three schedules here sit outside the distribution the estimator was trained on, which is the extrapolation Exercise 2 was built to demonstrate and the warning above tells you not to trust.
The lesson is the one the warning states rather than the one the table appears to offer: amortisation makes this calculation cheap, and that is exactly why it is easy to run it somewhere the estimator was never trained.
No code for this one. For each scenario, decide whether you would reach for NUTS, PMMH, VI, or NPE, and say why:
The framework extends well beyond what we have used it for. Sequential NPE focuses simulations near the observed data, buying accuracy at the cost of amortisation (Greenberg et al. 2019). Summary networks learn their own summary statistics rather than taking the raw series, which handles longer and more structured data, and they also enable misspecification detection by testing whether observed data falls outside the training distribution in summary space. The same machinery has been applied to compartmental models of COVID-19 (Radev et al. 2021) and to plant pathogen surveillance.
The masking idea used here is not specific to epidemiology: any surveillance system combining a mechanistic model with irregular longitudinal observations faces the same problem (Wang et al. 2024).
In Universal Differential Equations we combine mechanistic models with neural networks in a different way — instead of learning the inference, we learn part of the model itself.