Semi-Amortized Bayesian Inference with Reusable Local Posteriors

Niek Stevenson; Paul-Christian Bürkner; Andrew Heathcote; Stefan T. Radev

stat.MLOct 13, 2026

Semi-Amortized Bayesian Inference with Reusable Local Posteriors

Authors: Niek Stevenson, Paul-Christian Bürkner, Andrew Heathcote, Stefan T. Radev

Organizations: University of Amsterdam · TU Dortmund University · Rensselaer Polytechnic Institute

Abstract

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.

Figures & tables

Figure 1: Comparison of sampling methods for the same graphical model, where shared parameters

\theta_{c}

and

\theta_{b}

govern interconnected local parameters

\theta_{a_{i}}

, each generating observations

y_{i}

. (a) Conventional MCMC updates all parameter blocks in repeated sweeps. (b) Amortized neural posterior estimation (NPE) tries to approximate the full joint posterior directly through a suitable inverse factorization (see text) without resorting to MCMC. (c) Semi-amortized inference (ours) combines conventional MCMC updates of

\theta_{c}

with draws from amortized conditional posteriors for the remaining

\theta_{-c}

. Estimator A is shared across the

\theta_{a_{i}}

blocks, while estimator B updates

\theta_{b}

. During each sweep, a conditional update uses the current values of the other blocks.

Panel of Figure 2 — Figure 2: TIMSS hierarchy fit in Experiment 1 . Students nested within schools and countries/systems. Separate neural estimators update school coefficients $\theta_{cs}=(\alpha_{cs},\beta_{cs})$ (green) and country coefficients $b_{c}=(\alpha_{c},\beta_{c})$ (purple). Population means $\mu=(\mu_{\alpha},\mu_{\beta})$ , covariance matrices $\Sigma_{C}$ and $\Sigma_{S}$ , and residual variance $\sigma^{2}$ receive Gibbs updates. Gray nodes denote the achievement scores ( $y$ ), and the home-resources scores ( $x$ ).

Panel of Figure 3 — Figure 3: TIMSS posterior comparison in Experiment 1 . Top: global posterior medians (dots) and 95% credible intervals (bars), using Gibbs (blue) or SABI (red); $C$ and $S$ denote country and school levels. Bottom: paired posterior medians for every school and country intercept and slope, with Gibbs estimates on the x axis and SABI estimates on the y axis. We observe high agreement across all parameters.

Panel of Figure 4 — Figure 4: Citi Bike Gaussian field fit in Experiment 2 . Departures $y_{ith}$ at station $i$ in week $t$ and hour $h$ . Neural estimators update station–week parameters $\theta_{it}=(a_{it},b_{it})^{\top}$ , controlling activity and overdispersion (green), and the shared hourly profile $u$ (purple), that together determine latent activity rates. Population mean $\mu$ , covariance $\Sigma$ , and spatial and temporal coupling strengths $\rho_{S}$ and $\rho_{T}$ receive conventional updates.

Panel of Figure 5 — Figure 5: Citi Bike posterior comparison for 200 stations in Experiment 2 . Top: global posterior medians (dots) and 95% credible intervals using Stan (blue) or SABI (red). $\tau_{b}$ is divided by two and $\rho_{T}$ by 100. Bottom: paired posterior medians of $a_{it}$ and $b_{it}$ for all 1,600 station–weeks, with Stan on the x-axis and SABI on the y-axis. We find high agreement across all parameters, with slight overestimation of local inverse dispersion that did not translate to global bias.

Figure 6: Citi Bike sampling efficiency in Experiment 2 : median bulk ESS/s across shared parameters for SABI relative to NUTS. The dashed line marks equal efficiency; timing includes warmup.

Panel of Figure 7 — Figure 7: ACDC hierarchy of Experiment 3 . Overlapping plates represent publications and tasks, with $D_{pt}$ datasets in each combination. Parameter nodes collect six coordinates; gray $y_{i}$ contains a participant’s response times and choices. Fixed trial congruencies are omitted. Light green indicates updates from the participant network; hierarchy parameters receive conventional updates.

Figure 8: DMC estimation and ACDC inference of Experiment 3 . (a,b) Posterior means and 95% credible intervals for the mean

\mu_{A}

and log standard deviation

\log\sigma_{A}

\log A

across 100 paired datasets. The three methods show similar recovery for the means, but SABI vastly outperforms the other two on the variance. (c) Posterior medians and 95% intervals for task- and dataset-level population medians of

A

; dots are participant posterior medians.

N

denotes the number of participants and

L

the number of trials.

Method	Train	ESS/s	ECE		RMSE
	(min)		$\mu$	$\sigma^{2}$	$\mu$	$\sigma^{2}$
SABI (ours)	179	32	0.042	0.059	0.120	0.307
Compositional	292	12 ∗	0.157	0.308	0.246	0.738
Direct amortized	116	3853 ∗	0.031	0.030	0.377	0.764

Table 1: Reference comparisons at

N=500

(100 paired datasets) for Experiment 3 . Expected calibration error (ECE) and RMSE for global means and variances; lower is better. SABI has substantially lower RMSE. Direct amortization also has low ECE, since its posteriors remain close to the prior.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Citi Bike daily departures at four stations, ordered by observed activity. Predictive medians and pointwise 95% intervals include parameter and observation variation. All 56 days are used for fitting.

Parameter	Meaning	Prior for $\mu_{k}$	Prior for $\log\sigma_{k}$
$A$	Automatic amplitude	$\mathcal{N}(\log 0.20,0.35^{2})$	$\mathcal{N}(\log 0.40,0.25^{2})$
$\tau$	Automatic peak time	$\mathcal{N}(\log 0.15,0.30^{2})$	$\mathcal{N}(\log 0.35,0.25^{2})$
$\mu_{c}$	Controlled drift	$\mathcal{N}(\log 3.00,0.35^{2})$	$\mathcal{N}(\log 0.35,0.25^{2})$
$\mu_{r}$	Nondecision location	$\mathcal{N}(\log 0.35,0.20^{2})$	$\mathcal{N}(\log 0.22,0.25^{2})$
$b$	Boundary distance	$\mathcal{N}(\log 0.60,0.35^{2})$	$\mathcal{N}(\log 0.35,0.25^{2})$
$\sigma_{r}$	Nondecision spread	$\mathcal{N}(\log 0.04,0.30^{2})$	$\mathcal{N}(\log 0.35,0.25^{2})$

Appendix

Table 2: DMC population priors. Normal distributions are parameterized by mean and variance; physical time parameters are measured in seconds.

	SABI			Compositional			Direct amortized
	Coverage	RMSE	ECE	Coverage	RMSE	ECE	Coverage	RMSE	ECE
$\mu_{A}$	91	0.209	0.065	83	0.194	0.117	97	0.216	0.057
$\mu_{\tau}$	89	0.256	0.051	59	0.489	0.279	96	0.519	0.015
$\mu_{\mu_{c}}$	95	0.052	0.025	78	0.123	0.131	97	0.236	0.041
$\mu_{\mu_{r}}$	93	0.061	0.061	77	0.138	0.086	99	0.232	0.039
$\mu_{b}$	94	0.062	0.022	87	0.095	0.095	97	0.148	0.017

Appendix

Table 3: DMC recovery: 95% coverage; RMSE in prior-SD units.

Figure 10: ACDC posterior predictives by task. Top: accuracy across five within-participant RT bins, separately for congruent and incongruent trials. Bottom: incongruent-minus-congruent correct-RT quantiles against their average RT. Observed summaries are compared with predictive medians and pointwise 95% intervals from 300 replications.

Bayesian Inference

Bayesian inference updates probability distributions over unknowns with evidence to quantify uncertainty and guide prediction and decisions.

See all 316 papers

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure
Sep 30, 2026Daniel Habermann, Andreas Bulling, Stefan T. Radev +1Amortized Inference Hierarchical Bayesian Modeling
Department of Statistics TU Dortmund University, Germany · Institute for Visualisation and Interactive Systems University of Stuttgart, Germany · Department of Cognitive Science Rensselaer Polytechnic Institute, USA
Neural Architectures for Amortized Bayesian Inference: Statistical Foundations and Empirical Assessments
Jan 12, 2026Roy Shivam Ram Shreshtth, Arnab Hazra, Gourab MukherjeeUncertainty Quantification Amortized Inference
Department of Statistics and Data Science, Indian Institute of Technology Kanpur, Kanpur, India 208016 · Marshall School of Business, University of Southern California, Los Angeles, United States 90089-0809
Amortized Factor Inference Networks for Posterior Inference
May 26, 2026Joohwan Ko, Justin DomkeAmortized Inference Bayesian Inference
Manning College of Information and Computer Sciences University of Massachusetts Amherst

Sep 30, 2026J/K move · Enter open · S save

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure

Daniel Habermann, Andreas Bulling, Stefan T. Radev, Paul-Christian Bürkner

Department of Statistics TU Dortmund University, Germany · Institute for Visualisation and Interactive Systems University of Stuttgart, Germany · Department of Cognitive Science Rensselaer Polytechnic Institute, USA

Amortized Inference Hierarchical Bayesian Modeling Probabilistic Graphical Models

Open paper AskPDF arXiv

We develop a general method for amortized Bayesian inference on multilevel models of arbitrary structure. Given a generative model specified as a directed acyclic graph, our method automatically derives valid factorizations of the joint posterior and matching neural network architectures. The key steps, graph expansion and graph inversion, yield an inverse graph that determines how inference networks are stacked and conditioned, producing factorizations that amortize over the number of groups and the number of observations within each group. Unlike approaches that simplify the dependency structure to speed up learning or inference, our method preserves all conditional independence and exchangeability assumptions of the generative model. Across three case studies, it closely matches gold-standard samplers on models with more than 6,500 parameters while reducing inference to a near-instant forward pass once trained.

Abstract

Bayesian Inference

Explore similar work

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure

Neural Architectures for Amortized Bayesian Inference: Statistical Foundations and Empirical Assessments

Amortized Factor Inference Networks for Posterior Inference