# Geometry-Aware Normalizing Wasserstein Flows for Optimal Causal Inference

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 79 references that are in the map.
- First posted on arXiv 30 Nov 2023.

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

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Physics-informed neural networks (PINNs) for fluid mechanics: A review](https://mlsanity.com/paper/2105.09506/) (2021)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [Variance-based regularization with convex objectives](https://mlsanity.com/paper/1610.02581/) (2016)
- [Semiparametric doubly robust targeted double machine learning: a review](https://mlsanity.com/paper/2203.06469/) (2022)
- [Mode Regularized Generative Adversarial Networks](https://mlsanity.com/paper/1612.02136/) (2016)
- [Stochastic interpolants with data-dependent couplings](https://mlsanity.com/paper/2310.03725/) (2023)
- [Fused Gromov-Wasserstein distance for structured objects: theoretical foundations and mathematical properties](https://mlsanity.com/paper/1811.02834/) (2018)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [Distributionally Robust Optimization: A Review](https://mlsanity.com/paper/1908.05659/) (2019)
- [Semi-Supervised Learning with Normalizing Flows](https://mlsanity.com/paper/1912.13025/) (2019)
- [Distributionally Robust Losses for Latent Covariate Mixtures](https://mlsanity.com/paper/2007.13982/) (2020)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Smooth Normalizing Flows](https://mlsanity.com/paper/2110.00351/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Graphical Normalizing Flows](https://mlsanity.com/paper/2006.02548/) (2020)
- [Pathwise Derivatives Beyond the Reparameterization Trick](https://mlsanity.com/paper/1806.01851/) (2018)
- [Structured Optimal Transport](https://mlsanity.com/paper/1712.06199/) (2017)

## Abstract (arXiv)

This paper presents a groundbreaking approach to causal inference by integrating continuous normalizing flows (CNFs) with parametric submodels, enhancing their geometric sensitivity and improving upon traditional Targeted Maximum Likelihood Estimation (TMLE). Our method employs CNFs to refine TMLE, optimizing the Cram'er-Rao bound and transitioning from a predefined distribution p0 to a data-driven distribution p1. We innovate further by embedding Wasserstein gradient flows within Fokker-Planck equations, thus imposing geometric structures that boost the robustness of CNFs, particularly in optimal transport theory. Our approach addresses the disparity between sample and population distributions, a critical factor in parameter estimation bias. We leverage optimal transport and Wasserstein gradient flows to develop causal inference methodologies with minimal variance in finite-sample settings, outperforming traditional methods like TMLE and AIPW. This novel framework, centered on Wasserstein gradient flows, minimizes variance in efficient influence functions under distribution pt. Preliminary experiments showcase our method's superiority, yielding lower mean-squared errors compared to standard flows, thereby demonstrating the potential of geometry-aware normalizing Wasserstein flows in advancing statistical modeling and inference.