Semi-Amortized Bayesian Inference with Reusable Local Posteriors

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

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…