Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure
Authors: Daniel Habermann, Andreas Bulling, Stefan T. Radev, Paul-Christian Bürkner
Organizations: 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
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.
Figures & tables
Figure 1: Graph processing pipeline. Graph description : Directed acyclic graph of a two-level model with population mean μ , population standard deviation τ , group-level means {λj} , observations yji , and observation-level standard deviation ω . The dashed box denotes that the group-level means λj and observations yji are sampled independently. Graph expansion : The interior node λj is split into two conditionally independent instances λ1 and λ2 , each receiving the same parent nodes μ and τ . This makes the exchangeability of groups explicit in the graph structure, allowing the inversion to determine that each instance conditions only on its own data and the global parameters. Graph inversion : Inversion of the expanded two-level model using outer-nodes-first ordering (μ,τ,ω,λ1,λ2) . In the inverse graph, each λj conditions only on its own group data yj={yji} and the global parameters μ , τ , ω , so the per-group inference becomes amortizable: a single inference network can be used for all J groups. Network architecture : Architecture derived from the inverse graph. Each inference network has its own summary network. The local summary network encodes the observation of each group into a group-level summary, from which the local inference network infers λj . The global summary network pools all observations into a single representation of the dataset, from which the global inference network infers μ , τ , and ω . The dashed box denotes that the local summary and inference networks are applied independently for each group.
Figure 2: Comparison of our method against Stan on the three case studies. Left: marginal posterior intervals of the global parameters, showing the posterior mean (point) and 90% credible intervals (line) under our approximator (blue) and Stan (red). Right : The posterior mean (grey) and posterior standard deviation (gold) of each group-level parameter under our method, plotted against the corresponding Stan estimate. The dashed line marks exact agreement and N is the number of groups.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Annotated two-level model graph (left) with its function specification (right). Each node is annotated with a sampling function f and a sample size function g . The root nodes μ , τ and ω have fixed sample sizes of 100. λj draws a sample size uniformly from [1,20] and observations draw a sample size uniformly from [1,10] . The right columns the output shape for each node, where the trailing dimension is the data dimension (always 1 for scalar parameters).
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