Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Causal Schrödinger Bridges: Constrained Optimal Transport on Structural Manifolds

Rui Wu, Li YongJun

TL;DR

The Causal Schr\"odinger Bridge (CSB) is introduced, a framework that reformulates counterfactual inference as Entropic Optimal Transport and proves the Structural Decomposition Theorem, showing that the global high-dimensional bridge factorizes into local, robust transitions.

Abstract

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, 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 into local, robust transitions. Empirical validation on high-dimensional interventions (Morpho-MNIST) demonstrates that CSB significantly outperforms deterministic baselines in structural consistency, particularly in regimes of strong, out-of-distribution treatments.

Causal Schrödinger Bridges: Constrained Optimal Transport on Structural Manifolds

TL;DR

The Causal Schr\"odinger Bridge (CSB) is introduced, a framework that reformulates counterfactual inference as Entropic Optimal Transport and proves the Structural Decomposition Theorem, showing that the global high-dimensional bridge factorizes into local, robust transitions.

Abstract

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, 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 into local, robust transitions. Empirical validation on high-dimensional interventions (Morpho-MNIST) demonstrates that CSB significantly outperforms deterministic baselines in structural consistency, particularly in regimes of strong, out-of-distribution treatments.
Paper Structure (25 sections, 2 theorems, 16 equations, 9 figures, 2 tables, 1 algorithm)

This paper contains 25 sections, 2 theorems, 16 equations, 9 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Schrödinger Bridges (SB)
  4. Structural Causal Models (SCMs)
  5. Causally Constrained Optimal Transport
  6. The Factorized Reference Process
  7. Causal Admissibility
  8. The Causal Schrödinger Bridge Problem
  9. The Structural Decomposition Theorem
  10. Theoretical Justification: The Geometry of Tunneling (SDE vs. ODE)
  11. Algorithm: Causal Sequential Fitting (CSF)
  12. Why SDEs and not ODEs?
  13. Counterfactual Inference: Structural Abduction
  14. Experiments
  15. Experimental Setup
  16. ...and 10 more sections

Key Result

Theorem 1

Let the reference $\mathbb{Q}$ and boundary marginals $\mu_0, \mu_1$ factorize according to $\mathcal{G}$. The solution $\mathbb{P}^*$ to the constrained problem (eq:csb_prob) factorizes structurally: where each $\mathbb{P}^*_i$ is the solution to a local, conditional SB problem: subject to matching the conditional marginals $p_0(x_i|\mathbf{Pa}_i)$ and $p_1(x_i|\mathbf{Pa}_i)$.

Figures (9)

  • Figure 1: The Geometry of Causal Transport: Manifold Adherence vs. Structural Adherence.(a) Conceptual illustration. Standard generative models (e.g., Optimal Transport, orange dashed) minimize transport energy by adhering to the statistical manifold $\mathcal{M}_{data}$, leading to spurious correlations. Our Causal Schrödinger Bridge (CSB, green solid) respects the structural constraints, transporting probability mass "off-manifold" to the causally valid region. (b) Empirical validation on a confounder structure $Y \leftarrow X \rightarrow Z$. Under a strong intervention $do(Y=3)$, Standard OT incorrectly increases $Z$ due to observed correlation, whereas CSB correctly maintains $Z$'s value, demonstrating structural independence.
  • Figure 1: Quantitative comparison on Confounder Test.
  • Figure 2: Why Stochastic Bridges? Robust Transport Across Disjoint Supports. Deterministic flows (ODE, red dashed) struggle with the "void" of low-density regions between source and target, often leading to numerical instability or unnatural detours. In contrast, Causal SB (SDE, green solid) utilizes entropic regularization (diffusion) to "tunnel" through support mismatch regions, ensuring robust and smooth transport even when the observational support is disconnected from the counterfactual target.
  • Figure 3: The Cost of Determinism in High Dimensions ($D=50$).(a) Projected trajectories (PCA). Standard Flow Matching (ODE, pink) minimizes kinetic energy by taking rigid linear paths, failing to cover the target manifold's non-convex geometry. In contrast, CSB (SDE, teal) leverages entropic regularization to "tunnel" through the void, adapting to the target shape. (b) Density estimation on the principal component. The ODE baseline exhibits severe mode collapse, converging to the conditional mean. CSB accurately recovers the distributional envelope, demonstrating that stochasticity is essential for capturing aleatoric uncertainty in scientific tasks.
  • Figure 4: Disentangling Mechanism from Style.Left: Factual thin digit '3'. Middle: Structural abduction recovers the latent "skeleton" ($U_X$), stripping away thickness. Right: CSB generates a thick digit that rigorously preserves the original topology and tilt. The "halo" effect highlights the minimal action principle: the model adds thickness to the existing skeleton rather than regenerating a new digit.
  • ...and 4 more figures

Theorems & Definitions (10)

  • Definition 1: Causal Reference Process
  • Definition 2: Causal Admissibility
  • Definition 3: CSB Problem
  • Theorem 1: Structural Decomposition
  • proof
  • Remark 1: Geometric Interpretation: Transport on Fiber Bundles
  • Remark 2: Error Propagation vs. Structural Correctness
  • Proposition 2: Structural Abduction
  • proof
  • proof