Geometry-Aware Normalizing Wasserstein Flows for Optimal Causal Inference

TL;DR

This work addresses finite-sample causal inference under distribution shift by marrying continuous normalizing flows (CNFs) with Wasserstein gradient flows to create geometry-aware normalizing Wasserstein flows. The approach refines parametric submodels used in TMLE by guiding perturbations along optimal transport-geometric paths from a prior $p_0$ to a data-driven $p_1$, aiming to minimize the Cramér-Rao bound through $\mathcal{W}_2$-gradient dynamics. It introduces velocity-field alignment to avoid PDE gradient computations, proposes variance-regularized and velocity-aligned objectives, and provides OT and diffusion-interpretations to justify the framework. Preliminary experiments on toy distributions (e.g., 8gaussians, Pinwheel) show reduced mean-squared error and improved efficiency over naïve flows, indicating potential for more robust, finite-sample causal estimations.

Abstract

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ér-Rao bound and transitioning from a predefined distribution $p_0$ to a data-driven distribution $p_1$. 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 $p_t$. 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.

Figures (5)

• Figure 1: Illustration of Trajectories in Continuous Normalizing Flows: This figure presents the evolution of continuous normalizing flows from a 2-dimensional standard Gaussian prior to distinct data distributions. Specifically, the trajectories are shown for transformations leading to the data distributions generated by two datasets: the 8Gaussians dataset and the Pinwheel dataset, respectively.
• Figure 2: Illustrative Transition from Initial Noise to Biased Sample Distribution: This figure depicts the evolution of the first dimension in the 8gaussians dataset, starting from pure noise and gradually transitioning towards a biased sample distribution.
• Figure 3: Analysis of Normalizing Wasserstein Flows on 8gaussians
• Figure 4: Consistent Outcomes in Normalizing Wasserstein Flows for pinwheel
• Figure :

Theorems & Definitions (13)

• Theorem 2.1: Dual representation; benamou2000computational
• Theorem 2.2: jordan1998variationalvillani2021topics
• Theorem 3.1
• Theorem 3.2
• Theorem 3.3
• proof : Proof of Theorem \ref{['thm:W2_dual']}
• proof : Proof of Theorem \ref{['thm:var_reg']}
• proof : Proof of Theorem \ref{['thm:general_formulation']}
• proof : Proof of Theorem \ref{['thm:lower_bound_reg_loss']}
• Remark 2.1