Counterfactual Identifiability via Dynamic Optimal Transport

AI summary

Overview

Research area: Causal machine learning — specifically counterfactual identifiability, optimal transport theory, and continuous-time generative flow models.

Technical level: Advanced. The paper is built on measure-theoretic optimal transport (Brenier's theorem, Benamou–Brenier dynamic formulation, Caffarelli regularity), structural causal models, and the mathematics of flow matching. Readers need a background in causal inference and probability to follow the theoretical arguments.

Scope: The paper establishes when high-dimensional, multivariate counterfactuals can be provably recovered from observational data alone, and prescribes a flow-matching estimator that respects those conditions.

What This Paper Is About

Pearl argues that counterfactual claims are only justified if they are identifiable — that is, if there is a unique mapping from the observed data distribution to the counterfactual of interest. Prior deep-learning work on counterfactual inference produces promising-looking samples but offers no such guarantee, because the monotonicity conditions known to secure identification were only defined for scalar variables. This paper generalises monotonicity to multi-dimensional variables (e.g. images) using dynamic optimal transport, proving that a particular class of continuous-time flows yields a unique, monotone, rank-preserving counterfactual transport map.

Key Contributions

1. A foundation for multivariate counterfactual identification from observational data. The authors establish identification for high-dimensional variables with dim(X) = dim(U) = d > 1 in Markovian SCMs, a case that the prior state of the art (Nasr-Esfahany et al., 2023) explicitly flagged as unresolved, noting "it is not clear how to generalise the monotonicity condition to multi-dimensional variables."

2. A dynamic optimal transport characterisation of counterfactual mechanisms. Lemma 4.6 shows that under absolute continuity and bounded, strictly positive densities on bounded, open, convex domains, the time-1 map of a dynamic OT flow is the gradient of a convex function T(u; pa) = ∇_u φ(u; pa), is monotone, bijective almost everywhere, and uniquely determined by the pair (P_U, P^ℭ_{X|PA}). Theorem 4.12 then proves the induced counterfactual map T* is strictly monotone in x, extending the classical monotone-quantile notion to the multivariate setting.

3. Extension to non-Markovian settings. The paper proves ∼_{L3} identifiability beyond the Backdoor Criterion case that prior work covered — adding multivariate Instrumental Variable results and a Frontdoor Criterion proof — using the canonical graphs shown in Figure 1 (Markovian, IV, Backdoor, Frontdoor).

4. A Markovian Batch-OT coupling that fixes a consistency bug in standard flow matching. The authors show that naively applying Batch-OT flow matching to a Markovian structure implicitly entangles the parents PA with the exogenous noise U, violating the independence requirement U ⟂⟂ PA and invalidating abduction. Their fix samples batched data from the conditional distribution P^ℭ_{X|PA=pa} and solves the OT problem separately for each fixed parent value.

Main Findings

• Multivariate monotonicity is achievable without an arbitrary coordinate order. Proposition 4.4 shows that if a mechanism f(pa, u) is monotone in u in the sense of monotone operators, then the counterfactual transport map T*(pa*, pa, x) is monotone in x. This preserves the rank order of factual outcomes under intervention, preventing rank inversions across individuals — a property the authors connect to fairness.

• Prior-indeterminacy does not break identification. Because the true exogenous prior P_U is typically unknown, choosing one induces indeterminacy. Definition 4.8 formalises when two transport maps are "counterfactually equivalent" (T^(1) ∼_{L3} T^(2)), and Lemma 4.11 shows that between any two absolutely continuous priors with finite second moments there exists a transport map that is unique P_U^(1)-a.e., monotone, and a.e. bijective. Under Theorem 4.12 the exogenous prior is taken to be the continuous uniform measure on [0,1]^d.

• Bijectivity alone suffices in one non-Markovian regime. The paper's front-door experiment with non-linear unobserved confounding validates the theory's claim that bijectivity by itself is enough for consistent counterfactual inference in that case.

• Naive Batch-OT flow matching is incorrect for Markovian settings. The paper identifies the specific failure mode: solving OT over a batch reassigns each u^(i) to a pair (x^(j), pa^(j)), implicitly coupling u^(i) to pa^(j), so U is "no longer strictly exogenous, and the counterfactuals will be incorrect."

• On the counterfactual ellipse benchmark, OT-coupled flows reach near-exact ground truth with very few function evaluations. The authors report that their flows "produce near-exact ground-truth counterfactuals, and using just two function evaluations with OT." Counterfactual reversibility is also improved by straighter transport paths (Figure 2). For calibration, the prior spline-flow model of Nasr-Esfahany et al. (2023) failed in the Markovian case with μ_APE = 607%, and succeeded only under the Backdoor scheme BC_Z with a μ_APE of 1% (reproduced at .98%). Attempts to learn P(X | PA) or P(X | PA, Z) also failed when Z is a confounder (Cond_PA, Cond_{PA,Z}).

• Axiomatic soundness improves on real images, but the authors are explicit about the limits. On a real-world medical imaging dataset widely used for counterfactual inference, the method improves the counterfactual soundness axioms of composition, effectiveness, and reversibility (following Monteiro et al., 2023 and Ribeiro et al., 2023). The authors state plainly that "while useful, these metrics alone do not imply identification and should not be construed as evidence of causal validity." Specific numerical values from the real-image experiments are not reported in the available content.

Methodology in Plain English

The authors start from a structural causal model — a set of functional assignments X_i := f_i(PA_i, U_i) with independent exogenous noise in the Markovian case. A counterfactual query ("what would X have been had its parents been pa*?") can be answered by composing the mechanism with its own inverse: T*(pa*, pa, x) = f(pa*, f^{-1}(pa, x)). The hard part is that this map can never be learned directly, because paired observational and counterfactual samples essentially do not exist.

Their insight is to import a classical result from optimal transport. Brenier's theorem says that, under general conditions, the optimal transport map between two distributions is unique and is the gradient of a convex function — which means it is monotone in the vector sense. The authors show that a time-1 map of a dynamic optimal transport flow (the Benamou–Brenier formulation, where a velocity field moves mass along an ODE) supplies exactly the mechanism they need. Caffarelli's regularity results are invoked to additionally guarantee bijectivity, since Brenier's map does not guarantee it by default; in practice, mild continuous noise (uniform or Gaussian) added to empirical targets ensures the smoothed distribution admits a smooth density.

Practically, the estimator is a continuous-time flow trained by flow matching, where counterfactuals are computed by solving the ODE backwards for abduction (recovering u) and forwards for prediction (producing x*). Uniqueness of both initial value problems follows from the Picard–Lindelöf theorem when the velocity field is Lipschitz continuous.

The one place where off-the-shelf practice breaks is the coupling. Standard Batch-OT flow matching pairs source and target samples across an entire batch, which for a Markovian SCM silently correlates the noise with the parents. The authors instead fix the parent value and solve the OT problem within that conditional slice, solving a family of OT couplings rather than one global one.

Why This Matters

Impact on research. A large body of recent work parameterises structural causal models with deep generative models (Pawlowski et al., 2020; Sanchez and Tsaftaris, 2021; Ribeiro et al., 2023; Komanduri et al., 2024; Kumar et al., 2025; Rasal et al., 2025; Xia et al., 2025). The authors argue this practice carries operational risk: without identification, there can exist observationally equivalent models that give different answers, so the counterfactuals cannot support causal claims. This paper supplies the missing identification guarantees for the multivariate case and, in doing so, gives the field a principled stopping criterion for when a learned counterfactual model can be trusted. It also connects two previously separate literatures — dynamic optimal transport and graphical causality — and shows that the multivariate-monotonicity problem flagged by Nasr-Esfahany et al. (2023) has a natural solution.

Real-world applications (these are the application areas the paper itself names or implies):

• Evaluating interventions — as in the cited work on individual treatment effects and algorithmic recourse (Kusner et al., 2017; Tsirtsis and Rodriguez, 2023).
• Characterising causal relationships — quantifying how variables influence one another (Karimi et al., 2020; Budhathoki et al., 2022).
• Generating targeted synthetic data for downstream tasks (Pitis et al., 2022; Roschewitz et al., 2024; Mehta et al., 2025).
• Fairness and ranking — monotonicity of T* in x means the rank order of outcomes is preserved under a given intervention, preventing rank inversions across individuals.
• Medical imaging — the authors validate on a real-world medical imaging dataset widely used for counterfactual inference.

Industry relevance. Any setting that makes consequential decisions from observational data — treatment recommendation, lending, hiring, insurance, or clinical imaging pipelines — inherits the risk that a model's counterfactual prediction is one of several equally plausible answers consistent with the data. Identification guarantees turn those predictions from statistical associations into defensible causal statements. The paper's Markovian Batch-OT coupling also matters practically: it shows that simply adopting a popular off-the-shelf flow-matching recipe will silently violate the assumptions and produce wrong counterfactuals.

Future Directions

• Scaling the OT coupling. The proposed approach is described as nontrivial to operationalise at scale. Batch-OT flow matching asymptotically recovers the global OT map but needs large batches in practice, especially in high dimensions (Klein et al., 2025). The authors point to amortising OT pairing via semidiscrete couplings (Mousavi-Hosseini et al., 2025) as a promising lower-cost route to improved sample quality.

• Relaxing the regularity and prior assumptions. Theorem 4.12 assumes the exogenous prior is continuous uniform on [0,1]^d, and the identification arguments rely on absolute continuity, bounded and strictly positive densities on bounded, open, convex domains, and Caffarelli-style regularity. How far these can be weakened — and whether the prior-transition result (Lemma 4.11) can be made constructive in practice — remains open.

• Ground-truth evaluation beyond controlled settings. The ellipse benchmark has known counterfactuals, but the authors note that on real data, axiomatic metrics (composition, effectiveness, reversibility) "alone do not imply identification." Building evaluation protocols that can actually test identification on real, high-dimensional data would substantially strengthen empirical claims across the field.

• Broadening the causal-graph coverage. The paper proves results for Markovian, IV, Backdoor, and Frontdoor structures. Other identification criteria, nested counterfactual queries, and fairness-aware criteria (Zhang and Bareinboim, 2018; Correa et al., 2021) may or may not admit the same dynamic-OT treatment.

Target Audience

This paper is aimed at causal machine learning researchers and graduate students working on counterfactual inference, identifiable representation learning, or generative models of causal structure — particularly those who have used flow-based or diffusion-based counterfactual estimators and want to know what guarantees, if any, those methods carry. It is also relevant to statisticians and applied mathematicians interested in the interface between optimal transport theory and causality, and to practitioners in high-stakes domains (medical imaging, fairness-sensitive decision systems) who need counterfactual predictions that can survive scrutiny. Readers without a background in measure-theoretic optimal transport will find Sections 4 and 5 demanding; the experimental section and the causal-graph framing in Figure 1 are considerably more accessible.

Authors’ abstract

We address the open question of counterfactual identification for high-dimensional multivariate outcomes from observational data. Pearl (2000) argues that counterfactuals must be identifiable (i.e., recoverable from the observed data distribution) to justify causal claims. A recent line of work on counterfactual inference shows promising results but lacks identification, undermining the causal validity of its estimates. To address this, we establish a foundation for multivariate counterfactual identification using continuous-time flows, including non-Markovian settings under standard criteria. We characterise the conditions under which flow matching yields a unique, monotone, and rank-preserving counterfactual transport map with tools from dynamic optimal transport, ensuring consistent inference. Building on this, we validate the theory in controlled scenarios with counterfactual ground-truth and demonstrate improvements in axiomatic counterfactual soundness on real images.