# Amortized Inference for Causal Structure Learning

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

## Key facts

- Cited by 90 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 109 references that are in the map.
- First posted on arXiv 25 May 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [A simple neural network module for relational reasoning](https://mlsanity.com/paper/1706.01427/) (2017)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [Causality for Machine Learning](https://mlsanity.com/paper/1911.10500/) (2019)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [Learning Bayesian Networks: The Combination of Knowledge and Statistical Data](https://mlsanity.com/paper/1302.6815/) (2013)
- [Causal Discovery from Heterogeneous/Nonstationary Data with Independent Changes](https://mlsanity.com/paper/1903.01672/) (2019)
- [Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/1906.04477/) (2019)
- [Causal Consistency of Structural Equation Models](https://mlsanity.com/paper/1707.00819/) (2017)
- [Learning Neural Causal Models from Unknown Interventions](https://mlsanity.com/paper/1910.01075/) (2019)
- [Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning](https://mlsanity.com/paper/2106.02584/) (2021)
- [DiBS: Differentiable Bayesian Structure Learning](https://mlsanity.com/paper/2105.11839/) (2021)
- [Amortized Causal Discovery: Learning to Infer Causal Graphs from Time-Series Data](https://mlsanity.com/paper/2006.10833/) (2020)
- [Deep Generalized Method of Moments for Instrumental Variable Analysis](https://mlsanity.com/paper/1905.12495/) (2019)
- [Causal Structure Learning](https://mlsanity.com/paper/1706.09141/) (2017)
- [Permutation-based Causal Inference Algorithms with Interventions](https://mlsanity.com/paper/1705.10220/) (2017)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [Learning Gaussian Networks](https://mlsanity.com/paper/1302.6808/) (2013)
- [Differentiable Causal Discovery Under Unmeasured Confounding](https://mlsanity.com/paper/2010.06978/) (2020)
- [Causal Discovery from Changes](https://mlsanity.com/paper/1301.2312/) (2013)
- [From Ordinary Differential Equations to Structural Causal Models: the deterministic case](https://mlsanity.com/paper/1408.2063/) (2014)
- [Scaling structural learning with NO-BEARS to infer causal transcriptome networks](https://mlsanity.com/paper/1911.00081/) (2019)
- [Data Analysis with Bayesian Networks: A Bootstrap Approach](https://mlsanity.com/paper/1301.6695/) (2013)

## Built on by (most cited first, 30 of 90)

- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [GFlowNets for AI-Driven Scientific Discovery](https://mlsanity.com/paper/2302.00615/) (2023)
- [Do-PFN: In-Context Learning for Causal Effect Estimation](https://mlsanity.com/paper/2506.06039/) (2025)
- [Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network](https://mlsanity.com/paper/2305.19366/) (2023)
- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [The Essential Role of Causality in Foundation World Models for Embodied AI](https://mlsanity.com/paper/2402.06665/) (2024)
- [Stable Differentiable Causal Discovery](https://mlsanity.com/paper/2311.10263/) (2023)
- [CausalBench: A Large-scale Benchmark for Network Inference from Single-cell Perturbation Data](https://mlsanity.com/paper/2210.17283/) (2022)
- [A Scale-Invariant Sorting Criterion to Find a Causal Order in Additive Noise Models](https://mlsanity.com/paper/2303.18211/) (2023)
- [A Meta-Learning Approach to Bayesian Causal Discovery](https://mlsanity.com/paper/2412.16577/) (2024)
- [Sample, estimate, aggregate: A recipe for causal discovery foundation models](https://mlsanity.com/paper/2402.01929/) (2024)
- [Position: The Future of Bayesian Prediction Is Prior-Fitted](https://mlsanity.com/paper/2505.23947/) (2025)
- [Standardizing Structural Causal Models](https://mlsanity.com/paper/2406.11601/) (2024)
- [Federated Causal Discovery from Heterogeneous Data](https://mlsanity.com/paper/2402.13241/) (2024)
- [Amortized In-Context Bayesian Posterior Estimation](https://mlsanity.com/paper/2502.06601/) (2025)
- [Estimating Interventional Distributions with Uncertain Causal Graphs through Meta-Learning](https://mlsanity.com/paper/2507.05526/) (2025)
- [Deep Causal Learning: Representation, Discovery and Inference](https://mlsanity.com/paper/2211.03374/) (2022)
- [Demystifying amortized causal discovery with transformers](https://mlsanity.com/paper/2405.16924/) (2024)
- [Constraint-Free Structure Learning with Smooth Acyclic Orientations](https://mlsanity.com/paper/2309.08406/) (2023)
- [Discovering Mixtures of Structural Causal Models from Time Series Data](https://mlsanity.com/paper/2310.06312/) (2023)
- [Amortized Active Causal Induction with Deep Reinforcement Learning](https://mlsanity.com/paper/2405.16718/) (2024)
- [Variational Sequential Optimal Experimental Design using Reinforcement Learning](https://mlsanity.com/paper/2306.10430/) (2023)
- [LimiX-2: A Contextual Mechanism Network Towards General Structured-Data Intelligence](https://mlsanity.com/paper/2609.17488/) (2026)
- [A Fixed-Point Approach for Causal Generative Modeling](https://mlsanity.com/paper/2404.06969/) (2024)
- [CORE: Towards Scalable and Efficient Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/2401.16974/) (2024)
- [Optimal Transport for Structure Learning Under Missing Data](https://mlsanity.com/paper/2402.15255/) (2024)
- [Scalable Differentiable Causal Discovery in the Presence of Latent Confounders with Skeleton Posterior (Extended Version)](https://mlsanity.com/paper/2406.10537/) (2024)
- [Large-Scale Targeted Cause Discovery via Learning from Simulated Data](https://mlsanity.com/paper/2408.16218/) (2024)
- [End-to-End Compression for Tabular Foundation Models](https://mlsanity.com/paper/2602.05649/) (2026)
- [CauScale: Neural Causal Discovery at Scale](https://mlsanity.com/paper/2602.08629/) (2026)

## Abstract (arXiv)

Inferring causal structure poses a combinatorial search problem that typically involves evaluating structures with a score or independence test. The resulting search is costly, and designing suitable scores or tests that capture prior knowledge is difficult. In this work, we propose to amortize causal structure learning. Rather than searching over structures, we train a variational inference model to directly predict the causal structure from observational or interventional data. This allows our inference model to acquire domain-specific inductive biases for causal discovery solely from data generated by a simulator, bypassing both the hand-engineering of suitable score functions and the search over graphs. The architecture of our inference model emulates permutation invariances that are crucial for statistical efficiency in structure learning, which facilitates generalization to significantly larger problem instances than seen during training. On synthetic data and semisynthetic gene expression data, our models exhibit robust generalization capabilities when subject to substantial distribution shifts and significantly outperform existing algorithms, especially in the challenging genomics domain. Our code and models are publicly available at: https://github.com/larslorch/avici.