Source-linked AI summary
Semi-Amortized Bayesian Inference with Reusable Local Posteriors
Niek Stevenson, Paul-Christian Bürkner, Andrew Heathcote, Stefan T. Radev
TL;DR
High-dimensional models with complex dependencies make full-joint posterior amortization challenging. SABI learns conditional posteriors for selected parameter blocks and alternates their direct draws with conventional MCMC updates. Across realistic applications, it reports close agreement with reference samplers, faster inference than NUTS on a latent-variable model, and better calibration than compositional inference under the same simulation budget.
Problem
Amortizing a full joint posterior is challenging as model dimension and dependencies grow across groups, time points, or locations.
Source: §1 · p. 1Method
SABI trains neural estimators for selected block conditionals and combines their direct draws with conventional MCMC updates for remaining variables.
Source: §1 · p. 2, §2 · p. 2, §2.1 · p. 2Results
Takeaways & Limitations
The same trained estimator can be reused across different hierarchical structures at scale, without approximating the entire joint posterior.
Source: §1 · p. 2, §6 · p. 8SABI is not fully amortized: each new dataset requires iterative sampling, warm-up, and convergence assessment, and uncorrected conditional approximation errors can bias the posterior.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Amortized Bayesian inference uses simulations to perform Bayesian estimation when likelihoods are expensive or analytically unavailable. However, learning the full joint posterior is challenging in high-dimensional models with complex dependencies. We introduce a semi-amortized method that divides parameters into smaller blocks, and uses conditional independence to identify the information needed to infer each block. We train neural estimators for blocks whose posterior updates are computationally expensive or frequently repeated. Each estimator learns from simulations of the relevant model components to approximate the block's full conditional posterior, given the relevant observations and the values of the variables on which it depends. A Markov chain then alternates direct draws from these estimators with conventional updates of the remaining variables. Each neural update uses the current conditioning values and requires neither likelihood evaluation nor an acceptance-rejection step. Where the same conditional structure recurs, one estimator can be reused across blocks and datasets, with an output dimension determined by the block rather than the full model. We illustrate the approach in high-dimensional and computationally expensive hierarchical and spatiotemporal models.
1 INTRODUCTION
Simulation-based inference can handle expensive or unavailable likelihoods, but amortizing a full joint posterior becomes difficult as model dimension and dependencies grow. SABI learns selected conditional updates and combines them with conventional MCMC, with reported agreement against reference samplers across realistic applications.
1 INTRODUCTION
- Full-posterior amortization becomes challenging as parameter counts and dependencies grow in hierarchical, state-space, and spatial models.The posterior dimension grows with groups, time points, or locations while dependencies must still be represented jointly.Source: §1 · p. 1
- SABI combines learned conditional updates for selected blocks with conventional MCMC updates for the remaining parameters.The method focuses neural estimation on difficult inverse factors rather than approximating the full joint posterior.Source: §1 · pp. 1–2, Figure 1 · p. 2
- Across three realistic applications with tens of thousands of parameters, SABI's posterior summaries closely agree with reference samplers.In a simulation-based application, it outperforms alternative SBI methods and reuses a trained estimator across hierarchical structures at scale.Source: §1 · p. 2
2 SEMI-AMORTIZED INFERENCE
SABI uses each learned block's Markov blanket to identify its conditional posterior inputs. It trains estimators from simulations over conditioning contexts, then alternates their conditional draws with conventional MCMC updates in a blockwise sampler.
2 SEMI-AMORTIZED INFERENCE
- A learned block's full conditional depends on other model variables through its Markov blanket, represented by relevant observations, conditioning parameters, and fixed inputs.For example, hierarchical blocks use population parameters as conditions, while spatial blocks may use neighboring states and shared parameters.Source: §2 · p. 2
2.1 Training the Neural Updates
- Training simulates block parameters and observations across sampled conditioning parameters and fixed inputs to approximate each block's conditional posterior.The training distribution p(ci, xi) controls context coverage; varying population means, scales, covariates, and observation counts supports reuse across groups and datasets.Source: §2.1 · pp. 2–3
2.2 Semi-Amortized Inference
- Each sweep updates learned blocks by drawing from their trained conditional estimators and updates remaining blocks with Gibbs draws or MCMC transitions preserving their full conditionals.Learned blocks use the latest conditioning values, including values updated earlier in the sweep; multiple blocks may share an estimator.Source: §2.2 · p. 3
3 ACCURACY OF THE PROPOSED SAMPLER
SABI preserves the intended joint posterior when its learned conditionals are exact, while approximate updates introduce finite-sweep and stationary-accuracy risks. The analysis bounds finite-sweep error and shows why small posterior-averaged conditional error alone is not enough to guarantee accurate stationary behavior.
3.1 Posterior Invariance
- With exact learned conditionals, SABI preserves the intended joint posterior, including dependencies between parameter blocks.The exact sweep kernel preserves Π.Source: §3.1 · p. 3
3.2 Finite-Sweep Error
- After t sweeps, the total-variation error is bounded by min{1, tDy/2}, where Dy is posterior-averaged conditional error across learned blocks.The finite-sweep bound concerns the approximate sampler started from the exact posterior.Source: §3.2 · p. 4
- Small conditional error averaged under the exact posterior does not guarantee an accurate stationary distribution, because rare inaccurate contexts can trap the chain.The analysis also notes that conditional errors can propagate through conditioning inputs supplied to subsequent updates.Source: §3.2 · p. 4
- In real-data experiments, conditional errors were less concerning in practice and could be reduced with longer training runs.The effect of conditional errors on posterior expectations depends on their persistence and cancellation across updates.Source: §3.2 · p. 4
4 RELATED WORK
Related SBI work uses structured posterior factorizations, hierarchical estimators, or neural Monte Carlo proposals with correction. SABI instead uses learned full-conditional approximations as direct sampling updates and extends likelihood-free approximate Gibbs with modern SBI tools and analysis.
4 RELATED WORK
- Structured posterior methods construct inference factorizations, but some variational approaches retain prior dependencies or impose mean-field independence and target full amortization.Faithful inversion methods instead seek to avoid spurious conditional independencies (Webb et al., 2018; Habermann et al., 2026).Source: §4 · p. 4
- Hierarchical neural posterior estimators share networks across groups, while variants differ in whether they require complete hierarchical simulations, fit population parameters separately, or compose exchangeable-factor scores.Compositional approaches still pursue full amortization and may accumulate error (Arruda et al., 2026).Source: §4 · p. 4
- Neural Monte Carlo proposals retain importance weighting or Metropolis–Hastings correction, whereas SABI uses learned full-conditional approximations directly without acceptance tests or importance weights.Related approaches include inference compilation (Le et al., 2017; Baydin et al., 2019), corrected block proposals (Wang et al., 2018; Liang et al., 2021), and amortized population Gibbs sampling (Wu et al., 2020).Source: §4 · pp. 4–5
- SABI extends likelihood-free approximate Gibbs sampling with modern SBI estimators and summaries, correctness and error analysis, and benchmarks on large hierarchical and spatiotemporal models.The reported applications include a simulation-based model with over 30,000 parameters and a spatiotemporal model with nearly two million observations.Source: §4 · p. 5
5 EXPERIMENTS
The experiments evaluate SABI in hierarchical, spatiotemporal, and simulation-based models with complex dependencies. Comparisons show close posterior agreement in some applications, higher sampling efficiency for Citi Bike at larger station counts, and lower RMSE than alternative SBI methods in the DMC simulation study.
5.1 TIMSS: Three-level hierarchical model
- In the 21,757-parameter TIMSS model, SABI showed close marginal agreement with full Gibbs, which was 10.4 times faster.Global C2ST accuracy was 0.508, while the full joint including local parameters remained distinguishable at 0.764.Source: §5.1 · p. 5, Figure 3 · p. 6
5.2 Citi Bike: Spatial and Temporal Dependence
- The Citi Bike model has 23,006 free parameters after marginalizing latent log rates; explicitly sampling those rates exceeds Stan’s available memory, and numerical marginalization was stopped after one day.Four SABI chains with 500 burn-in sweeps and 1,250 retained draws took 157 minutes.Source: §5.2 · p. 6
- For 200 Citi Bike stations, global posterior means differed from Stan by at most 0.26 Stan posterior standard deviations, but the full joint including locals was perfectly discriminable.The authors attribute the full-joint mismatch likely to a small, consistent inverse-dispersion bias.Source: §5.2 · p. 6, Figure 5 · p. 7
- At higher station counts, SABI exceeded NUTS by more than 3× in ESS per second, with timing including warmup.The comparison uses median bulk ESS/s across shared parameters.Source: §5.2 · p. 6, Figure 6 · p. 7
5.3 DMC: Simulation-Based Inference
- SABI reused the participant-level network without retraining in a richer ACDC hierarchy with 30,402 parameters across 5,015 participants.The model spans three conflict tasks, 34 datasets, and ten publications.Source: §5.3 · pp. 7–8
- On the ACDC model, three chains retained 2,500 draws in 10.79 minutes, with maximum bR of 1.037.Across global parameters, median bulk and tail ESS were 6,540 and 6,699, with minima of 71 and 23.Source: §5.3 · p. 8
- Across 100 paired datasets of 500 participants, SABI achieved lower RMSE than both alternatives; methods recovered means similarly, but SABI vastly outperformed them on variance.Table 1 reports RMSE and expected calibration error averaged over six population means and variances.Source: §5.3 · p. 7, Table 1 · p. 8, Figure 8 · p. 8
6 DISCUSSION
SABI combines learned conditional updates with conventional MCMC, with experiments showing close agreement to reference hyperposteriors and benefits over alternatives. The discussion also describes transformations, refinement diagnostics, posterior invariance, and limitations of uncorrected learned updates.
6 DISCUSSION
- Across three realistic experiments, SABI closely matched Gibbs and NUTS hyperposteriors, ran faster than NUTS with latent variables, and calibrated markedly better than compositional inference at the same simulation budget.The authors note that each new dataset still requires iterative sampling, warm-up, and convergence assessment; approximation errors can bias results and are not removed by longer chains.Source: §6 · p. 8
- A proposed future direction is to add Metropolis–Hastings correction using a reference posterior, or learn local marginal likelihoods to update conditioning parameters without drawing local parameters.The reference posterior could come from a separately trained estimator or teacher network.Source: §6 · p. 8
A PARAMETER TRANSFORMATIONS
- A learned block can be trained in transformed coordinates when its parameter map is a differentiable bijection with nonsingular Jacobian, then mapped back before later updates.Applying the same change of coordinates preserves KL and total variation distances; Gaussian conditional priors can use prior-standardized coordinates.Source: A PARAMETER TRANSFORMATIONS · p. 11
B REFINEMENT DIAGNOSTICS
- Conditional calibration checks compare generating parameters with independent neural posterior draws across conditioning regions, while sweep checks detect changes in joint parameter-observation summaries after inference sweeps.The checks can guide targeted simulation coverage in problematic regions.Source: B REFINEMENT DIAGNOSTICS · p. 11
C.1 Posterior invariance
- Exact full-conditional draws preserve the joint posterior, and composing block updates preserves it for any fixed update order; standard recurrence and aperiodicity conditions give convergence in total variation.Conventional MCMC transitions also preserve the posterior when they preserve each block’s full conditional.Source: C PROOFS AND SUPPLEMENTARY DERIVATIONS · p. 11, §C.1 · p. 11
C.2 Proof of Proposition 1
The proof bounds divergence between exact and approximate sampler paths by accumulating errors from learned block updates. The resulting bound captures errors caused by information discarded by a summary network.
C.2 Proof of Proposition 1
- Conventional transitions contribute zero pathwise KL, while each learned block contributes its conditional-update divergence, averaged under the exact path distribution.Each exact block update preserves the posterior, so every intermediate exact-chain state has distribution Π.Source: §C.2 · pp. 11–12
- Marginalizing the paths cannot increase KL, yielding a bound for the distributions after t sweeps; Pinsker’s inequality then bounds total variation.The bound accounts for learned-update errors, including information discarded by the summary network.Source: §C.2 · p. 12
- A summary-based estimator targets the posterior conditional on its summary, which equals the full conditional only when the summary retains all information relevant to that block.Source: §C.2 · p. 12
C.3 Proof of Proposition 2
The proof constructs an example where an approximate update can yield a stationary distribution far from the target even when the exact sampler mixes rapidly. It also analyzes how inherited errors in Gaussian conditional prior means affect learned posterior means.
C.3 Proof of Proposition 2
- As ε tends to zero, the approximate stationary distribution can approach total variation distance 1 from the target, even when the exact sampler mixes rapidly.The construction replaces one exact θ update with a conditional draw whose probability depends on ϑ; the approximate sweep induces transitions in ϑ of order ε.Source: §C.3 · p. 12
C.4 Gaussian conditional structure
- For a scalar Gaussian block, inherited mean error is approximately the prior-mean error multiplied by the posterior-to-prior variance ratio; ratios below one reduce it and ratios above one amplify it.The inherited error arises when preceding updates alter the prior mean supplied to the learned update.Source: §C.4 · p. 13
C.5 Stationary error and the Poisson equation
The stationary bias of the approximate sampler can be expressed as signed contributions from learned conditional errors, weighted by their effects after subsequent exact updates. The section also gives convergence conditions for the Poisson-series representation and describes the experiments and TIMSS model setup.
C.5 Stationary error and the Poisson equation
- Stationary bias is a sum of signed learned-update errors: preceding approximate updates determine the states where errors are evaluated, and subsequent exact updates determine their effect on the Poisson solution.For a learned block, the update difference integrates the downstream function against the difference between the learned and exact conditional densities.Source: §C.5 · pp. 13–14
- The Poisson-series solution applies when the series converges and the exact transition kernel can be applied term by term; slower decay lets conditional errors persist, while opposite signs can cancel.Each term is the expected difference between f and its posterior mean after t exact sweeps.Source: §C.5 · p. 14
D EXPERIMENTAL DETAILS
- Experiments used high-dimensional models on a MacBook Pro with an Apple M4 Pro, 14 CPU cores, 20 GPU cores, and 24 GB RAM.Source: D EXPERIMENTAL DETAILS · p. 14
D.1.1 Model and priors
- The TIMSS model uses 256,416 students in 10,820 schools across 54 countries or systems, with standardized achievement and home-resources variables and no survey weights.The model specifies school- and country-level intercept–slope covariance matrices and independent global priors.Source: §D.1.1 · p. 14
D.1.2 Training and inference
These sections specify blockwise training and inference for the school, country, and station–week parameters, including how spatial and temporal graph information reaches a local network. The station–week estimator is reused across station counts, while a separate hourly-profile estimator is trained for each count.
D.1.2 Training and inference
- Inference alternates school, country, and conjugate global updates, using separate bivariate affine coupling flows for school and country coefficients.The country network trains for 600 epochs; the school network trains for 8,700 epochs followed by 200 diagnostic-based refinement epochs.Source: §D.1.2 · p. 14
D.2.1 Model and priors
- A Gaussian graph prior links neighboring stations and consecutive weeks; each station–week network needs only the resulting conditional-prior mean η_it and covariance Σ/d_it.The conditional prior uses spatial and temporal neighbor counts and summed parameter vectors; the covariance is represented through graph precision and a shared covariance matrix.Source: §D.2.1 · p. 15
D.2.2 Training and inference
- The station–week network receives hourly count/profile pairs and conditional-prior parameters, while the hourly-profile network updates disjoint hour pairs.The station–week network is reused across station counts; the hourly-profile network is trained separately for each count. A 100-station pilot also supplies conditioning values and profiles for refinement training.Source: §D.2.2 · pp. 15–16
D.2.3 Posterior predictives
Posterior predictive checks assess daily Citi Bike departures and predictions from the conflict-task DMC model. The checks identify undercoverage at selected busy bike stations and specific mismatches in predicted Stroop behavior.
D.2.3 Posterior predictives
- Daily predictive intervals cover 94.6%, 87.5%, 80.4%, and 83.9% of observations at selected stations near the 3rd, 50th, 75th, and 90th activity percentiles, understating variation at busier stations.The predictive intervals include parameter and observation variation; the fit uses all 56 days.Source: §D.2.3 · p. 16, Figure 9 · p. 16
D.3.1 Model and priors
- The DMC models conflict-task decisions through controlled evidence accumulation and a transient automatic response, with response times combining boundary-crossing time and a nondecision draw.The model uses six positive participant parameters; the starting point, pulse shape, and diffusion noise are fixed.Source: §D.3.1 · p. 16
- The empirical analysis retains 5,015 participant units and 601,800 trials across 34 datasets, ten publications, and three tasks after applying participant and RT filters.It selects 60 trials per congruency condition and removes a narrow spike of incorrect Rey-Mermet trials before selection.Source: §D.3.1 · p. 17
D.3.2 Training and inference
The DMC application evaluates global-parameter recovery, participant-level inference, and task-level posterior predictions. A separate section considers learning local marginal likelihood factors as an alternative to SABI’s local posterior updates.
D.3.2 Training and inference
- The participant network estimates six prior-standardized log parameters from conditional-prior means and standard deviations plus a summary of simulated trials.It trains for 600 epochs and then undergoes 200 refinement epochs based on sweep checks and simulation-based calibration.Source: §D.3.2 · p. 17
- SABI and compositional inference each use 26,214,400 participant simulations, while direct inference uses 26,200,000; Table 3 reports coverage, prior-SD-scaled RMSE, and calibration error.Evaluation uses 100 paired datasets; the table averages RMSE and calibration error over six coordinates in each parameter family.Source: §D.3.2 · pp. 17–18, Table 3 · p. 17
D.3.3 Posterior predictives
- Predictions reproduce the increasing Stroop delta curve but underestimate the slowest-quantile RT difference in all three tasks; incongruent Stroop accuracy rises with RT while observed accuracy is nearly flat or decreasing.Figure 10 shows accuracy across within-participant RT bins and incongruent-minus-congruent correct-RT quantiles against average RT.Source: Figure 10 · p. 18, §D.3.3 · p. 18
E LEARNING LOCAL MARGINAL LIKELIHOODS
- The alternative trains conditional density estimators for local data given conditioning inputs, then evaluates learned likelihood factors during MCMC instead of drawing local parameters.This reduces the sampled parameter space but introduces approximation error through the learned factors.Source: E LEARNING LOCAL MARGINAL LIKELIHOODS · pp. 18–19