Tightening Causal Bounds via Covariate-Aware Optimal Transport

Sirui Lin, Zijun Gao, Jose Blanchet, Peter Glynn

Abstract: Causal estimands can vary significantly depending on the relationship between outcomes in treatment and control groups, leading to wide partial identification (PI) intervals that impede decision making. Incorporating covariates can substantially tighten these bounds, but requires determining the range of PI over probability models consistent with the joint distributions of observed covariates and outcomes in treatment and control groups. This problem is known to be equivalent to a conditional optimal transport (COT) optimization task, which is more challenging than standard optimal transport (OT) due to the additional conditioning constraints. In this work, we study a tight relaxation of COT that effectively reduces it to standard OT, leveraging its well-established computational and theoretical foundations. Our relaxation incorporates covariate information and ensures narrower PI intervals for any value of the penalty parameter, while becoming asymptotically exact as a penalty increases to infinity. This approach preserves the benefits of covariate adjustment in PI and results in a data-driven estimator for the PI set that is easy to implement using existing OT packages. We analyze the convergence rate of our estimator and demonstrate the effectiveness of our approach through extensive simulations, highlighting its practical use and superior performance compared to existing methods.

Lay Summary:

Causal conclusions can vary a lot depending on how outcomes from treated and untreated groups relate to each other. This uncertainty results in wide ranges—called partial identification (PI) intervals—that make it hard to draw firm conclusions. However, by taking into account additional information about individuals, known as covariates, we can significantly narrow these ranges and improve the reliability of decisions.

Using covariates in a statistically sound way is not straightforward. The problem turns into a complex mathematical task known as conditional optimal transport (COT), which is much harder to solve than the standard tools researchers typically use.

In this work, we introduce a simpler and more practical method that approximates this difficult problem using well-known tools from optimal transport. Our method still uses covariate information, provides tighter and more useful results, and becomes nearly exact as a tuning parameter increases. It’s also easy to apply using existing software. We prove that our method works well in theory and show through simulations that it outperforms current approaches.

Primary Area: General Machine Learning->Causality

Keywords: Partial identification, causal inference, conditional optimal transport

Submission Number: 5226