📊 statistics

Causal Partial Identification via Conditional Optimal Transport

This paper proposes a direct, consistent, non-parametric estimator for conditional optimal transport in causal partial identification by establishing the functional's continuity under the adapted Wasserstein distance, thereby avoiding the need for nuisance parameter estimation and overcoming the inconsistency of traditional plug-in methods.

Original authors: Sirui Lin, Zijun Gao, Jose Blanchet, Peter Glynn

Published 2026-02-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sirui Lin, Zijun Gao, Jose Blanchet, Peter Glynn

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery, but you only have half the clues.

In the world of Causal Inference (figuring out cause and effect), this is the standard problem. You want to know: If I give a patient a new drug, will they get better?

The Reality: You can only see one outcome. If you give the drug, you see them get better (or not). You can never go back in time and see what would have happened if you hadn't given them the drug.
The Missing Clue: The "counterfactual" outcome (what would have happened otherwise) is invisible.

Because you can't see both outcomes at once, you can't calculate the exact truth. Instead, statisticians have to settle for a Partial Identification (PI) Set. Think of this not as a single answer, but as a range of possibilities. "The drug likely improves health by somewhere between 5% and 15%." The goal is to make this range as narrow (precise) as possible.

The Problem: The "Ghost" in the Machine

Usually, to narrow down this range, we look at covariates (background details like age, weight, or blood pressure). If we know two people are exactly the same age and weight, we can compare them more fairly.

However, there's a catch. In real life, you rarely find two people who are exactly identical in every single detail (especially with continuous numbers like height or blood sugar).

The Old Way: Previous methods tried to force a match between people who were "close enough." But mathematically, this is like trying to build a bridge with mismatched bricks. The structure is unstable. If you add a tiny bit of data, the whole bridge (your estimate) can collapse or jump wildly. This is called discontinuity.

The Solution: A New Kind of Ruler

The authors of this paper, Sirui Lin and colleagues, invented a new mathematical tool to fix this instability. They call it Conditional Optimal Transport (COT).

Here is the analogy:

1. The Old Ruler (Weak Topology)

Imagine trying to match two piles of sand (Treatment Group vs. Control Group) by looking at the total shape of the piles. If you move a single grain of sand, the shape looks almost the same, but the "perfect match" might jump to a completely different grain. This makes the math "jumpy" and unreliable.

2. The New Ruler (Adapted Wasserstein Distance)

The authors realized that to match these piles correctly, you don't just look at the total shape. You have to look at how the sand is layered.

The Metaphor: Imagine the covariates (age, weight) are the floor plan of a building, and the outcomes (health) are the furniture inside.
The old method tried to match the whole building at once.
The new method says: "First, match the floor plans perfectly. Then, inside each specific room (e.g., '30-year-old males'), match the furniture."

This new ruler, called the Adapted Wasserstein Distance, forces the math to respect the structure of the data. It ensures that if you change the data slightly, your answer changes only slightly. It makes the bridge stable.

The Magic Trick: Discretization (Pixelating the World)

Even with the new ruler, calculating the perfect match is hard because there are infinite possibilities for continuous numbers (like 30.0001 years old).

The authors' solution is Discretization.

The Analogy: Imagine you have a high-definition photo of a crowd. It's too detailed to process perfectly. So, you turn it into a pixelated image (like an 8-bit video game).
Instead of trying to match "30.0001 years old" with "30.0002 years old," you group everyone into "30-year-olds."
You create "bins" or "cells" for your data.
Inside each bin, you calculate the average outcome.
Then, you match the bins.

Because you are matching groups (bins) rather than individual points, the math becomes smooth and stable. The authors proved that as you get more data, you can make the pixels smaller and smaller, and your answer gets closer and closer to the true truth without ever becoming "jumpy."

Why This Matters

No Guessing Games: Old methods often required estimating "nuisance parameters" (guessing the shape of the data distribution first). If your guess was slightly wrong, your final answer was wrong. This new method is direct. It doesn't need to guess the shape; it just matches the data as it is.
Robustness: If the data is slightly messy or shifted (e.g., the treatment group is slightly older on average), this method doesn't break. It absorbs the noise.
Better Bounds: In their tests, this method produced much tighter, more accurate ranges (PI sets) than existing methods, especially when the relationship between variables was complex and non-linear (like a curve rather than a straight line).

Summary

Think of this paper as inventing a new, ultra-stable measuring tape for causal inference.

The Problem: We can't see the future (or the past), so we have to estimate a range of possibilities.
The Flaw: Previous tapes were wobbly and broke if the data wasn't perfect.
The Fix: The authors created a tape that looks at the structure of the data (Adapted Wasserstein) and uses a "pixelated" approach to smooth out the rough edges.
The Result: We can now say, with much more confidence, exactly how effective a treatment is, even when we can't observe the "what if" scenario directly.

It turns a shaky, uncertain guess into a solid, mathematically proven range.

1. Problem Statement

The paper addresses the challenge of Partial Identification (PI) in causal inference. In standard potential outcome models, a unit has two potential outcomes, $Y(0)$ (control) and $Y(1)$ (treatment), but only one is observed. Consequently, the joint distribution of $(Y(0), Y(1))$ is unidentifiable, making causal estimands that depend on this joint distribution (e.g., variance of treatment effects, correlation between outcomes) impossible to estimate point-wise.

While Optimal Transport (OT) has been used to derive bounds (PI sets) for these estimands based on marginal distributions, the presence of pretreatment covariates ( $Z$ ) offers a way to tighten these bounds. By conditioning on $Z$ , one can utilize Conditional Optimal Transport (COT) to find the optimal joint distribution that preserves the outcome-covariate marginals $(Y(w), Z)$ for $w \in \{0, 1\}$ .

The Core Challenge:
Direct estimation of the COT value from finite samples is statistically inconsistent under the standard weak topology.

Discontinuity: The COT functional is discontinuous with respect to the weak topology.
Data Sparsity: In continuous covariate spaces, it is rare to observe the exact same covariate value in both treatment and control groups. A naive "plug-in" estimator (using empirical point masses) fails to converge because the conditional distributions are ill-defined at specific points, leading to a mismatch in supports.
Existing Limitations: Current methods rely on relaxations, indirect estimation of nuisance parameters (e.g., propensity scores or conditional densities), or heuristic procedures lacking theoretical guarantees.

2. Methodology

The authors propose a direct, consistent, non-parametric estimator for the COT value that avoids nuisance parameter estimation. The methodology relies on three key pillars:

A. Topological Shift: Adapted Wasserstein Distance

The authors establish that while the COT functional is discontinuous under the weak topology, it becomes continuous under a stronger topology induced by the Adapted Wasserstein Distance ( $W_a$ ).

Definition: $W_a$ measures the distance between joint distributions $(Y, Z)$ by optimally aligning the covariate marginals ( $Z$ ) while minimizing the transport cost of the conditional outcome distributions ( $Y|Z$ ).
Significance: This topology enforces the convergence of conditional distributions, which is necessary for COT consistency.

B. Adapted Empirical Distribution

To construct an estimator that converges under $W_a$ , the authors introduce an Adapted Empirical Distribution:

Discretization: The covariate space $Z$ is partitioned into a finite grid of cells (cubes) with edge length $n^{-r}$ .
Projection: Observed covariates are projected to the center of their respective cells.
Construction: The estimator aggregates outcomes within each cell, creating a discrete approximation of the conditional distribution.
- Crucially, only the covariate $Z$ is discretized; the outcome $Y$ remains continuous.
- This discretization ensures that the empirical conditional distributions converge to the true conditional distributions in the $W_a$ metric.

C. Estimation Algorithm

The proposed estimator, $\hat{V}_{n_0, n_1}$ , is computed as follows:

Discretize: Map observed covariates $Z_i$ to cell centers $\phi_r(Z_i)$ .
Coupling: Construct a coupling $\hat{\pi}$ between the discretized covariate distributions of the treatment and control groups. This coupling minimizes the distance between the discretized $Z$ marginals (effectively matching cells).
Integration: Calculate the expected cost over the coupled conditional distributions:
$\hat{V}_{n_0, n_1} = \int W_h(\hat{\mu}^{z_0}_{Y(0)}, \hat{\mu}^{z_1}_{Y(1)}) d\hat{\pi}(z_0, z_1)$
where $W_h$ is the standard OT distance between the conditional outcome distributions within matched cells.

D. Handling Covariate Shift

For observational studies where treatment assignment depends on covariates (non-constant propensity scores), the authors extend the method by:

Reweighting: Applying self-normalized inverse propensity weights to the samples to align the covariate distributions of the treatment and control groups with the population marginal.
Cross-Fitting: Using a two-fold cross-fitting procedure to estimate the propensity score and the adapted empirical distribution independently, ensuring statistical validity.

3. Key Contributions

Theoretical Continuity: Proved that the COT functional is continuous under the topology induced by the Adapted Wasserstein distance (Theorem 2), resolving the inconsistency of plug-in estimators.
Direct Consistent Estimator: Proposed a non-parametric estimator that does not require estimating nuisance functions (like conditional densities or propensity scores) directly, avoiding the "curse of dimensionality" associated with non-parametric regression.
Convergence Rates: Derived finite-sample convergence rates for the estimator. The rate depends on the dimensions of the covariate ( $d_Z$ $d_{Z}$ ) and outcome ( $d_Y$ $d_{Y}$ ) spaces:
- Rate: $O(N^{-\frac{1}{d_Z + 2 \vee d_Y}})$ (ignoring log factors).
- This rate is optimal for non-parametric estimation in this context.
Robustness: Established robustness guarantees against covariate distributional shifts (Theorem 5), showing the estimator remains stable under small perturbations.
Extension to Observational Data: Generalized the framework to unconfounded designs with covariate-dependent treatment assignment, providing a rigorous method for PI in real-world observational settings.

4. Results

Theoretical Results

Consistency: The estimator converges almost surely to the true COT value as sample size increases (Theorem 3).
Finite-Sample Bounds: The paper provides explicit error bounds (Theorem 4) showing the trade-off between bias (controlled by cell size) and variance.
Optimality: The derived PI sets are shown to be the tightest possible bounds given the observed outcome-covariate marginals (Theorem 1).

Empirical Results

Synthetic Data: The method was tested on linear, quadratic, and scale models with Gaussian noise.
- Performance: The proposed "Adapt" estimator consistently outperformed existing methods (DualBounds using Ridge Regression and KNN) in non-linear settings (Quadratic and Scale models) where parametric assumptions failed.
- Bernoulli vs. Observational: The method performed robustly in both randomized (Bernoulli) and observational (covariate-dependent) settings.
Real Data (STAR Project): Applied to a semi-synthetic dataset based on the Student Achievement and Retention (STAR) project.
- The estimator provided tighter and more accurate bounds for the correlation between potential outcomes compared to baseline methods.
Sensitivity Analysis: The method showed low sensitivity to the choice of the discretization constant $c$ and remained stable under covariate distributional shifts.

5. Significance and Impact

Bridging Theory and Practice: This work resolves a fundamental theoretical gap in causal inference by providing a consistent estimator for COT, a tool previously considered difficult to estimate due to topological discontinuities.
Model-Free Inference: By avoiding the estimation of nuisance parameters (densities/regressions), the method is more robust to model misspecification, a common issue in high-dimensional causal inference.
Tighter Bounds: The use of covariates via COT significantly shrinks the Partial Identification sets compared to standard OT, leading to more precise causal conclusions (e.g., narrower confidence intervals for treatment effect variance).
General Applicability: The framework is applicable to a broad class of causal estimands and handles both randomized and observational data, making it a versatile tool for policy evaluation and scientific research.

In summary, this paper provides a rigorous mathematical foundation and a practical algorithm for estimating causal bounds using conditional optimal transport, overcoming previous limitations related to data sparsity and topological discontinuity.