PaperlayerView exact revision on arXiv

Source-linked AI summary

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure

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

arXiv:2609.40024v1stat.MLcs.LGstat.CO

TL;DR

Multilevel posterior dimensions vary with group count, challenging fixed-width approximators and making group-level amortization depend on posterior factorization. The paper derives factorizations and neural architectures from generative DAGs; across three case studies, its approximators closely matched Stan on models with more than 6,500 parameters while reducing inference to a near-instant forward pass.

  • Problem

    Multilevel posterior dimension varies with group count, challenging fixed-width approximators, while posterior factorization determines whether inference amortizes across groups.

    Source: §1 · pp. 1–2

  • Method

    Graph expansion and inversion derive posterior factorizations and matching inference architectures from generative DAGs, selecting amortizable factors while preserving conditional independence and exchangeability assumptions.

    Source: §1 · p. 2, §4 · p. 12

  • Results

    Across three case studies, approximators closely matched Stan on models with more than 6,500 parameters, reducing inference to a near-instant forward pass.

    Source: §1 · p. 2, §4 · p. 12

  • Takeaways & Limitations

    The method retains the generative model's conditional independencies and exchangeability assumptions; for crossed designs, autoregressive inference keeps approximation amortized over group count.

    Source: §4 · p. 12

    Amortization over group count typically holds only within the range seen during training, so models with many groups require training on many groups.

    Source: §4 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Multilevel posterior dimensions vary with group counts, creating a challenge for fixed-width approximators and making posterior factorization important to amortization. The paper derives inference architectures from generative graphs and reports close agreement with Stan across three case studies.

1 Introduction

  • Amortized Bayesian inference trains neural networks on simulated parameter–dataset pairs, requiring simulation rather than likelihood evaluation.Its computational cost is incurred during training, enabling inference across repeated datasets.Source: §1 · p. 1
  • Multilevel posterior dimensions vary with group counts, conflicting with the fixed-width inputs and outputs assumed by standard neural approximators.Earlier work derived architectures by hand for particular model classes, and the chosen posterior factorization determines whether inference amortizes across groups.Source: §1 · pp. 1–2
  • Transformer and meta-amortized approaches do not address multilevel dependency structure: masks may not guarantee conditional independence, while exchangeable parameters remain an open problem.Transformer attention masks must be pre-specified for each model; meta-amortized approaches had been restricted to exchangeable observations.Source: §1 · pp. 1–2
  • Graph expansion and inversion identify valid posterior factorizations, including choices that amortize over groups when possible.When no such factorization exists, the method selects one that leaves the fewest group-level parameters conditioned on other groups, then infers remaining parameters autoregressively.Source: §1 · p. 2
  • The method automatically determines posterior factorization and the inference and summary networks, extending amortized inference to DAG models including crossed designs.The paper notes that no amortized method had previously been available for crossed designs.Source: §1 · p. 2
  • Across three case studies, derived approximators closely matched Stan, including on a four-level model with more than 6,500 parameters.The studies cover eight schools, crossed image and annotator effects, and the UK Breeding Bird Survey.Source: §1 · p. 2

2 Methods

The methods transform a generative DAG into an inference architecture and use its conditional structure to represent simulation, including variable numbers of groups and observations.

2 Methods

  • The method transforms a generative graph into a neural architecture through graph expansion, graph inversion, amortization analysis, and architecture derivation.The resulting design specifies which nodes share inference networks, how networks condition on one another, and how observations are summarized.Source: §2 · p. 2

2.1 Generative graphs and ancestral sampling

  • Generative models also support posterior predictive checks and prior elicitation by simulating data under model assumptions or prior specifications.These are described as uses of generative models in a Bayesian workflow.Source: §2.1 · pp. 2–3
  • A generative DAG represents random quantities as nodes and conditional dependencies as edges, with latent nodes forming parameters and observed nodes forming data.Its joint distribution factorizes according to the graph, and data can be simulated by visiting nodes in topological order from roots to leaves.Source: §2.1 · p. 3
  • In the two-level example, population parameters inform group-level means, which inform observations alongside a shared observation-level standard deviation.The model's group count and the number of observations within groups may vary across datasets.Source: §2.1 · p. 3
  • Each node carries a conditional distribution and a distribution over how many samples it draws for each parent-value combination, allowing groups to differ in size.The resulting networks encounter varying group and dataset sizes during training and can be applied to any size at inference.Source: §2.1 · p. 3

2.2 Graph inversion

Graph inversion derives posterior factorizations from a generative graph; the choice of factorization determines whether group-level parameters can be inferred independently and amortized.

2.2 Graph inversion

  • Graph inversion starts from observed data and reasons back toward latent parameters, producing an inverse factorization of the posterior.It is not sufficient simply to reverse every edge in the generative graph.Source: §2.2 · p. 3
  • Reversing a collider's edges can incorrectly assert independence between its parents given their common child.The example notes that larger population variability in τ also implies larger uncertainty about the population mean µ.Source: §2.2 · p. 3
  • Multiple valid inverse factorizations may exist, and their choice determines whether group-level parameters can be inferred independently.For the example with µ and τ inferred jointly, five factorizations exist; two factorize group parameters using group-specific data and shared parameters.Source: §2.2 · pp. 3–4

2.2.1 Algorithm

  • The method expands exchangeable nodes before inversion and uses the inverse graph to design neural approximators for amortized posterior inference.Inversion cost is data-independent and on the order of milliseconds, allowing even complex models' factorizations to be enumerated without meaningful overhead.Source: §2.2.1 · p. 4
  • The inversion algorithm adds latent nodes in order, choosing a minimal set of existing nodes that d-separates each node from the rest of the inverse graph.Varying the node order enumerates possible inverse factorizations.Source: §2.2.1 · p. 4

2.3 Exchangeable nodes and graph expansion

Graph expansion makes exchangeability explicit before inversion so the algorithm can test whether groups have the same dependencies.

2.3 Exchangeable nodes and graph expansion

  • Without expansion, an inversion algorithm sees an exchangeable node as one node and can force group-level parameters to depend jointly on all groups' data.That dependency breaks amortization over groups.Source: §2.3 · p. 4
  • Expansion splits each interior node that can draw multiple samples per parent-value combination into two instances, allowing d-separation to test whether each depends only on its own data and global parameters.By symmetry, two instances suffice to reveal whether any instance depends on another.Source: §2.3 · p. 4
  • Algorithm 1 describes the expansion procedure, and Figure 1's second panel shows the expanded running-example graph.Source: §2.3 · p. 4

2.4 Graph inversion with expanded graphs

Graph expansion makes exchangeable instances explicit so graph inversion can test their dependencies. The ordering of latent nodes in the inverse graph determines whether group-level parameters can be inferred independently.

2.4 Graph inversion with expanded graphs

  • Graph expansion splits an exchangeable interior node into two instances with the same parents, making group exchangeability explicit.The expanded graph can then be tested for dependencies between instances using d-separation.Source: §2.4 · pp. 4–5, Figure 1 · p. 5
  • A latent node is independently amortizable when none of its instances conditions on another instance of that node.Its instances share an inverse conditional, allowing one inference network to be reused across groups.Source: §2.4 · p. 6
  • With outer-nodes-first ordering, each group mean λj conditions on its own data and global parameters, so one network can infer all J group means.Merging µ, τ, and ω reduces the required inference networks from four to two.Source: Figure 1 · p. 5, §2.4 · p. 6
  • With inner-nodes-first ordering, λ2 conditions on λ1, so group means cannot be inferred independently by one reusable network.The inversion algorithm can be applied directly to the expanded graph to test such dependencies.Source: §2.4 · pp. 4, 6

2.5 Non-independent amortization

In crossed designs, no posterior factorization makes every grouping factor independently amortizable. The method instead infers one factor autoregressively and selects the factorization that minimizes the number of instances requiring sequential inference.

2.5 Non-independent amortization

  • Conditioning on observations opens collider paths between instances, so the factor added first remains dependent while the second can be d-separated given the first.Thus, either factor can be independently amortizable, but not both.Source: §2.5 · p. 6
  • In crossed designs, no factorization makes both grouping factors independently amortizable, so one factor must be inferred autoregressively.Inferring its instances jointly would fix output dimension and prevent amortization over the number of groups.Source: §2.5 · p. 6
  • The chain rule expresses a non-independent node as single-instance conditionals whose output dimension does not depend on the number of instances K.Their conditions vary across datasets and grow with the preceding instances.Source: §2.5 · p. 6
  • A permutation-invariant set network summarizes observations from all groups together with preceding instances, giving each autoregressive step fixed-size conditions.Because training uses simulator-provided preceding instances, all K outputs can be computed in one pass; sampling remains sequential.Source: §2.5 · p. 6
  • The method selects the factorization with the fewest sequentially inferred instances, amortizing over the larger grouping factor in a crossed design.Training summaries scale quadratically with K for the sequential node, and sampling requires K inference-network passes.Source: §2.5 · pp. 6–7

2.6 Summary networks

Summary networks must reflect that observations are exchangeable within groups and that groups themselves are exchangeable. The generative graph determines the pooling structure and the conditions summarized for each inference network.

2.6 Summary networks

  • Ordinary summary networks require modification because multilevel observations form nested exchangeable sets, not one exchangeable set.Encodings for crossed and nested structures pool observations over groups specified by the model structure.Source: §2.6 · p. 7
  • For arbitrary generative graphs, the graph defines whether grouping factors are nested or crossed and therefore which groups the encoding pools over.Nested factors are exchangeable only within their enclosing grouping factors.Source: §2.6 · p. 7
  • Each inference network gets its own summary network because its conditioning parameters and required pooling level differ.Observations are extended with the parameter values relevant to the inference node and encoded with exchangeable layers.Source: §2.6 · p. 7
  • Exchangeable layers update each observation using averages over every group to which it belongs.For the breeding bird survey, the encoding combines counts with averages over the dataset, region, square, year, and region-year groups.Source: §2.6 · p. 7
  • Encoded observations are pooled to each target instance, with the log observation count and conditioned ancestor parameters appended to produce fixed-size conditions.Appending parameters also permits inference for groups without observations.Source: §2.6 · p. 7

2.7 Network architecture

The selected inverse graph specifies each inference network’s conditions and how nodes are grouped into stages. Summary and inference networks are trained jointly on simulated parameter–dataset pairs, while sampling follows the order in which conditions become available.

2.7 Network architecture

  • For each non-observation node, its parents in the inverse graph specify the conditions supplied to its inference network.Observation nodes are known and therefore require no network.Source: §2.7 · p. 7
  • Nodes are assigned to stages according to when their conditions become available, with one inference network per stage.Each stage’s network is conditioned on the union of its nodes’ parents.Source: §2.7 · p. 8
  • Summary and inference networks are trained jointly on simulated parameter–dataset pairs by minimizing the negative log density of the simulated parameters.With a sufficiently large simulation budget, minimizing the loss minimizes expected forward KL divergence between approximate and true posteriors.Source: §2.7 · p. 8
  • During sampling, inference networks are evaluated in the order their conditions become available.During training, they may be evaluated in any order because parameters and observations are known.Source: §2.7 · p. 8

3 Case studies

The case studies test the method on increasingly complex multilevel models, from eight schools to crossed skin-lesion effects and a four-level bird survey. Their posterior estimates generally agree well with Stan, though some bird-model standard deviations are slightly underestimated.

3 Case studies

  • The evaluation covers three real-world applications of increasing structural complexity and assesses calibration, predictive accuracy, agreement with Stan, and partial pooling.The applications are an eight-schools model, a three-level skin-lesion model with crossed image and annotator effects, and a four-level bird-count model.Source: §3 · p. 8

3.1 Eight schools

  • The eight-schools model treats school effects as draws from a shared population, with known standard errors simulated over the range reported in the original study.Its approximator uses a pooled summary network and separate inference networks for the global parameters and school effects; training varies the number of schools from 1 to 100.Source: §3.1 · p. 8, §3.1.1 · p. 9, §3.1.2 · p. 9

3.1.1 Network architecture

  • In the schools case, a SetTransformer summarizes school observations, while affine coupling flows infer global and school-level parameters.The networks are trained jointly on 100,000 simulated datasets.Source: §3.1.1 · p. 9, §3.1.2 · p. 9

3.2 Skin Lesions

  • The skin-lesion model predicts pairwise annotator agreement using crossed image and annotator effects, with covariates distinguishing automated masks from mixed-tool pairs.The factorization amortizes across 2,394 images and infers 15 annotator effects autoregressively; its architecture uses three inference flows and five summary networks, trained with varying numbers of images and annotators.Source: §3.2 · pp. 9–10, §3.2.1 · p. 10, §3.2.2 · p. 10

3.2.1 Network architecture

  • For skin lesions, the selected factorization leaves the group-level factor with the fewest instances to infer autoregressively while amortizing across images.The architecture uses three affine coupling flows and five summary networks; training uses 200,000 simulated datasets with varying image and annotator counts.Source: §3.2 · pp. 9–10, §3.2.1 · p. 10, §3.2.2 · p. 10

3.3.3 Results

  • Across three studies, posterior estimates generally agree well with Stan; in the bird model, σr and σs are slightly too low.Training took about 1 minute for the schools model, 1.5 hours for the skin-lesion model, and 12 hours for the bird model. Drawing 4,000 samples took about 0.04 s versus about 0.04 s for Stan to reach effective sample size 4,000 in schools; 2 s versus about 13 min in skin lesions; and about 4 s versus about 35 min in the bird model.Source: §3.1.3 · p. 9, Figure 2 · p. 10, §3.2.3 · p. 11

3.3 UK Breeding Bird Survey

  • The bird analysis models European robin counts across 6,444 squares in 82 regions and 31 years, using a negative-binomial model for overdispersed counts.The data comprise 89,865 observations, and 45% of square-year combinations are observed.Source: §3.3 · p. 11

3.3.1 Network architecture

  • The bird factorization amortizes across regions and squares, leaving 31 year effects for sequential inference; four affine coupling flows and eight DeepSets summarize conditions.Training uses one million simulated datasets with up to 100 regions and 250 squares per region.Source: §3.3 · p. 11, §3.3.1 · p. 11, §3.3.2 · p. 12

4 Discussion

The method derives inference architectures automatically for arbitrary multilevel models while retaining the generative model’s conditional independencies and exchangeability assumptions. Across three case studies it closely matched Stan on models with more than 6,500 parameters, and the discussion notes limits to group-count amortization and robustness under model misspecification.

4 Discussion

  • Graph expansion and inversion derive architectures automatically while retaining the generative model’s conditional independence and exchangeability assumptions.For crossed designs, autoregressive inference keeps the approximator amortized over the number of groups when no factorization makes every grouping factor independently amortizable.Source: §4 · p. 12
  • Across three case studies, the derived approximators closely matched Stan on models with more than 6,500 parameters while reducing inference to a near-instant forward pass.The models had tractable likelihoods for comparison with Stan, but the method does not require a likelihood and also applies to purely simulation-based multilevel models.Source: §4 · p. 12
  • Group-count amortization typically holds only within the range seen during training, so models with many groups require training on many groups.Approximators trained on simulations may also become inaccurate under model misspecification; self-consistency losses may improve robustness and can in principle be combined with the architecture.Source: §4 · p. 12
  • The results establish graph-based derivation as a principled, general route to amortized inference for multilevel models of arbitrary structure.Source: §4 · p. 12

A Enumerating inverse factorizations

The enumerated posterior factorizations differ in whether group-level effects are inferred jointly or independently. Factoring them jointly breaks amortization over the number of groups, whereas independent group-level estimation preserves that property.

A Enumerating inverse factorizations

B Full implementation of a two-level hierarchical model

The implementation uses user-provided sampling and sample-size functions to generate draws for a two-level model. It stores draws in long format with group indices, joining parent nodes by index columns as keys.

B Full implementation of a two-level hierarchical model

  • For each internal node, the implementation joins parent nodes using index columns as keys, then generates a sample size for each row and runs the node’s sampling function that many times.Draws receive a new index and are appended with the indices of their inputs.Source: B Full implementation of a two-level… · p. 16
  • The model graph and function specification annotate each node with sampling function f and sample-size function g, which the user supplies.Root nodes µ, τ, and ω have fixed sample sizes of 100; λj sample sizes are uniform on [1, 20], and observation sample sizes are uniform on [1, 10].Source: Figure 3 · p. 16, B Full implementation of a two-level… · p. 16

Related papers

Loading 2609.40024v1…