# Optimal Transport for Structure Learning Under Missing Data

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 104 references that are in the map.
- First posted on arXiv 23 Feb 2024.

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

- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [Causal Discovery with Continuous Additive Noise Models](https://mlsanity.com/paper/1309.6779/) (2013)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [Learning Sparse Nonparametric DAGs](https://mlsanity.com/paper/1909.13189/) (2019)
- [On the Role of Sparsity and DAG Constraints for Learning Linear DAGs](https://mlsanity.com/paper/2006.10201/) (2020)
- [Handling Incomplete Heterogeneous Data using VAEs](https://mlsanity.com/paper/1807.03653/) (2018)
- [Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/1906.04477/) (2019)
- [DAGMA: Learning DAGs via M-matrices and a Log-Determinant Acyclicity Characterization](https://mlsanity.com/paper/2209.08037/) (2022)
- [Handling Missing Data with Graph Representation Learning](https://mlsanity.com/paper/2010.16418/) (2020)
- [DiBS: Differentiable Bayesian Structure Learning](https://mlsanity.com/paper/2105.11839/) (2021)
- [Missing Data Imputation using Optimal Transport](https://mlsanity.com/paper/2002.03860/) (2020)
- [Deep End-to-end Causal Inference](https://mlsanity.com/paper/2202.02195/) (2022)
- [Amortized Inference for Causal Structure Learning](https://mlsanity.com/paper/2205.12934/) (2022)
- [Masked Gradient-Based Causal Structure Learning](https://mlsanity.com/paper/1910.08527/) (2019)
- [Tutorial on amortized optimization](https://mlsanity.com/paper/2202.00665/) (2022)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [DAGs with No Fears: A Closer Look at Continuous Optimization for Learning Bayesian Networks](https://mlsanity.com/paper/2010.09133/) (2020)
- [Graphical Models for Processing Missing Data](https://mlsanity.com/paper/1801.03583/) (2018)
- [Ordering-Based Causal Discovery with Reinforcement Learning](https://mlsanity.com/paper/2105.06631/) (2021)
- [MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms](https://mlsanity.com/paper/2111.03187/) (2021)
- [Neural Topic Model via Optimal Transport](https://mlsanity.com/paper/2008.13537/) (2020)
- [Causal Discovery in the Presence of Missing Data](https://mlsanity.com/paper/1807.04010/) (2018)
- [On the Convergence of Continuous Constrained Optimization for Structure Learning](https://mlsanity.com/paper/2011.11150/) (2020)
- [MCFlow: Monte Carlo Flow Models for Data Imputation](https://mlsanity.com/paper/2003.12628/) (2020)
- [A Unified Wasserstein Distributional Robustness Framework for Adversarial Training](https://mlsanity.com/paper/2202.13437/) (2022)

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

- [LLM Reading Tea Leaves: Automatically Evaluating Topic Models with Large Language Models](https://mlsanity.com/paper/2406.09008/) (2024)
- [Neural Topic Modeling with Large Language Models in the Loop](https://mlsanity.com/paper/2411.08534/) (2024)
- [Robust Simulation-Based Inference under Missing Data via Neural Processes](https://mlsanity.com/paper/2503.01287/) (2025)
- [Ordering-based Causal Discovery via Generalized Score Matching](https://mlsanity.com/paper/2601.16249/) (2026)
- [Nearly Optimal Bayesian Inference for Structural Missingness](https://mlsanity.com/paper/2601.18500/) (2026)
- [Generalizing Identifiability Results in Causal Discovery from Complete Data to Missing Data Settings](https://mlsanity.com/paper/2610.19211/) (2026)

## Abstract (arXiv)

Causal discovery in the presence of missing data introduces a chicken-and-egg dilemma. While the goal is to recover the true causal structure, robust imputation requires considering the dependencies or, preferably, causal relations among variables. Merely filling in missing values with existing imputation methods and subsequently applying structure learning on the complete data is empirically shown to be sub-optimal. To address this problem, we propose a score-based algorithm for learning causal structures from missing data based on optimal transport. This optimal transport viewpoint diverges from existing score-based approaches that are dominantly based on expectation maximization. We formulate structure learning as a density fitting problem, where the goal is to find the causal model that induces a distribution of minimum Wasserstein distance with the observed data distribution. Our framework is shown to recover the true causal graphs more effectively than competing methods in most simulations and real-data settings. Empirical evidence also shows the superior scalability of our approach, along with the flexibility to incorporate any off-the-shelf causal discovery methods for complete data.