# From Mice to Trains: Amortized Bayesian Inference on Graph Data

arXiv: https://arxiv.org/abs/2601.02241 · HTML: https://mlsanity.com/paper/2601.02241/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 79 references that are in the map.
- First posted on arXiv 5 Jan 2026.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Recipe for a General, Powerful, Scalable Graph Transformer](https://mlsanity.com/paper/2205.12454/) (2022)
- [A Generalization of Transformer Networks to Graphs](https://mlsanity.com/paper/2012.09699/) (2020)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [Graph Neural Network for Traffic Forecasting: A Survey](https://mlsanity.com/paper/2101.11174/) (2021)
- [A Note on Over-Smoothing for Graph Neural Networks](https://mlsanity.com/paper/2006.13318/) (2020)
- [Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows](https://mlsanity.com/paper/1805.07226/) (2018)
- [Benchmarking Simulation-Based Inference](https://mlsanity.com/paper/2101.04653/) (2021)
- [Inference Suboptimality in Variational Autoencoders](https://mlsanity.com/paper/1801.03558/) (2018)
- [Validating Bayesian Inference Algorithms with Simulation-Based Calibration](https://mlsanity.com/paper/1804.06788/) (2018)
- [All-in-one simulation-based inference](https://mlsanity.com/paper/2404.09636/) (2024)
- [A Survey of Graph Neural Networks for Social Recommender Systems](https://mlsanity.com/paper/2212.04481/) (2022)
- [Talking-Heads Attention](https://mlsanity.com/paper/2003.02436/) (2020)
- [Fast ε-free Inference of Simulation Models with Bayesian Conditional Density Estimation](https://mlsanity.com/paper/1605.06376/) (2016)
- [JANA: Jointly Amortized Neural Approximation of Complex Bayesian Models](https://mlsanity.com/paper/2302.09125/) (2023)

## Built on by (most cited first, 2 of 2)

- [High-Dimensional Simulation-Based Inference in Latent Spaces](https://mlsanity.com/paper/2609.37381/) (2026)
- [Mitigating Representation Gaps in Amortized Bayesian Inference with Auxiliary Supervision](https://mlsanity.com/paper/2609.39525/) (2026)

## Abstract (arXiv)

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.