Causal Schrödinger Bridges: Constrained Optimal Transport on Structural Manifolds
- DiffMOTCML
Generative modeling typically seeks the path of least action via deterministic flows (ODE). While effective for in-distribution tasks, we argue that these deterministic paths become brittle under causal interventions, which often require transporting probability mass across low-density regions (``off-manifold'') where the vector field is ill-defined. This leads to numerical instability and spurious correlations. In this work, we introduce the Causal Schrödinger Bridge (CSB), a framework that reformulates counterfactual inference as Entropic Optimal Transport. Unlike deterministic approaches that require strict invertibility or rely on low-rank approximations, CSB leverages diffusion processes (SDEs) to robustly ``tunnel'' through support mismatches while strictly enforcing structural admissibility constraints. We prove the Structural Decomposition Theorem, showing that the global high-dimensional bridge factorizes exactly into local, robust transitions. Crucially, we demonstrate that CSB breaks the Curse of Dimensionality in regimes of high intrinsic dimension. We empirically validate this on a full-rank causal system (, intrinsic rank ), completing the transport in 26.48 seconds on a single GPU (RTX 3090). This stands in stark contrast to structure-agnostic baselines, which are estimated to require over 6 years for dense computations of this scale regardless of the data's intrinsic rank. Empirical validation on Morpho-MNIST and -D extremal stress tests demonstrates that CSB significantly outperforms deterministic baselines in structural consistency and distribution coverage, capturing the underlying manifold with high fidelity (MSE 0.04).
View on arXiv