From Mice to Trains: Amortized Bayesian Inference on Graph Data

TL;DR

The paper extends Amortized Bayesian Inference (ABI) to graph-structured data by pairing permutation-invariant graph encoders with flexible neural posterior estimators, enabling fast, likelihood-free posterior inference on graph parameters. It systematically evaluates multiple graph-aware summary networks—notably the Set Transformer—against alternatives like Graph Convolutional Networks and Graph Transformers, using simulation-based calibration, posterior contraction, and recovery as quality metrics. Across toy, biological (mice interaction), and logistics (train scheduling) case studies, the Set Transformer with attention pooling consistently delivers strong parameter recovery and sharp posteriors, though calibration can be challenging for some parameters and architectures. The work highlights practical ABI for graphs, identifies a strong default architecture, and discusses limitations and avenues for extending ABI to more complex graph types and dynamics.

Abstract

Graphs arise across diverse domains, from biology and chemistry to social and information networks, as well as in transportation and logistics. Inference on graph-structured data requires methods that are permutation-invariant, scalable across varying sizes and sparsities, and capable of capturing complex long-range dependencies, making posterior estimation on graph parameters particularly challenging. Amortized Bayesian Inference (ABI) is a simulation-based framework that employs generative neural networks to enable fast, likelihood-free posterior inference. We adapt ABI to graph data to address these challenges to perform inference on node-, edge-, and graph-level parameters. Our approach couples permutation-invariant graph encoders with flexible neural posterior estimators in a two-module pipeline: a summary network maps attributed graphs to fixed-length representations, and an inference network approximates the posterior over parameters. In this setting, several neural architectures can serve as the summary network. In this work we evaluate multiple architectures and assess their performance on controlled synthetic settings and two real-world domains - biology and logistics - in terms of recovery and calibration.

Figures (8)

• Figure 1: Overview of the proposed framework, which comprises two phases. In the training phase (left), parameters are sampled from the prior and passed to a simulator to generate graph-structured datasets. These simulated graphs are then used to jointly train the summary network (encoder) and the inference network (posterior estimator). In the inference phase (right), an observed graph is processed by the trained summary and inference networks to obtain near instant posterior inference on the parameters.
• Figure 2: Results for the toy example across four summary-network architectures with each three different aggregation layers. We report parameter recovery (higher is better), posterior contraction (higher is better), and calibration ($\ell_\gamma$; values > 0 indicate good calibration) for the parameters $\pi = \{\pi_{AA}, \pi_{BB}, \pi_{AB} \}$ and $\lambda$. Markers (circles/triangles/squares) indicate the median across five runs, and error bars span the minimum to maximum values.
• Figure 3: Recovery (median and 95% credible interval) and Calibration (ECDF Difference plots) of parameters for one run of the Graph Convolutional Network (with mean pooling) and Set Transformer (with multi-head attention pooling).
• Figure 4: Recovery parameters $\pi_{AB}$ and $\lambda$, shown as posterior median and $95\%$ credible intervals, for graphs with $N=15$ nodes (left) and $N=45$ nodes (right). Posterior samples are drawn from the same model (Set Transformer with PMA as the summary network), trained on graphs with $N \in \{10,\dots, 50\}$.
• Figure 5: Results for the mice interaction network case study across four summary-network architectures, each evaluated for three different observation horizons. We report parameter recovery (higher is better), posterior contraction (higher is better), and calibration ($\ell_\gamma$; values > 0 indicate good calibration) for the parameters network density $\delta$ and exchange factor $\alpha$. Markers (circles/triangles/squares) indicate the median across five runs, and error bars span the minimum to maximum values.