# Optimal Transport for Treatment Effect Estimation

arXiv: https://arxiv.org/abs/2310.18286 · HTML: https://mlsanity.com/paper/2310.18286/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 89 references that are in the map.
- First posted on arXiv 27 Oct 2023.

## Builds on (most cited first, 30 of 44)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](https://mlsanity.com/paper/2103.03606/) (2021)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [Computational Optimal Transport: Complexity by Accelerated Gradient Descent Is Better Than by Sinkhorn's Algorithm](https://mlsanity.com/paper/1802.04367/) (2018)
- [Theoretical Analysis of Domain Adaptation with Optimal Transport](https://mlsanity.com/paper/1610.04420/) (2016)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [AutoDebias: Learning to Debias for Recommendation](https://mlsanity.com/paper/2105.04170/) (2021)
- [On Efficient Optimal Transport: An Analysis of Greedy and Accelerated Mirror Descent Algorithms](https://mlsanity.com/paper/1901.06482/) (2019)
- [On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm](https://mlsanity.com/paper/2002.03293/) (2020)
- [Sinkhorn Divergences for Unbalanced Optimal Transport](https://mlsanity.com/paper/1910.12958/) (2019)
- [ESCM2: Entire Space Counterfactual Multi-Task Model for Post-Click Conversion Rate Estimation](https://mlsanity.com/paper/2204.05125/) (2022)
- [DeepMatch: Balancing Deep Covariate Representations for Causal Inference Using Adversarial Training](https://mlsanity.com/paper/1802.05664/) (2018)
- [Large-scale Causal Approaches to Debiasing Post-click Conversion Rate Estimation with Multi-task Learning](https://mlsanity.com/paper/1910.09337/) (2019)
- [Unbalanced Optimal Transport through Non-negative Penalized Linear Regression](https://mlsanity.com/paper/2106.04145/) (2021)
- [Learning Causal Effects on Hypergraphs](https://mlsanity.com/paper/2207.04049/) (2022)
- [Enhanced Doubly Robust Learning for Debiasing Post-click Conversion Rate Estimation](https://mlsanity.com/paper/2105.13623/) (2021)
- [On the Opportunity of Causal Learning in Recommendation Systems: Foundation, Estimation, Prediction and Challenges](https://mlsanity.com/paper/2201.06716/) (2022)
- [StableDR: Stabilized Doubly Robust Learning for Recommendation on Data Missing Not at Random](https://mlsanity.com/paper/2205.04701/) (2022)

## Built on by (most cited first, 30 of 32)

- [FreDF: Learning to Forecast in the Frequency Domain](https://mlsanity.com/paper/2402.02399/) (2024)
- [BELM: Bidirectional Explicit Linear Multi-step Sampler for Exact Inversion in Diffusion Models](https://mlsanity.com/paper/2410.07273/) (2024)
- [Time-o1: Time-Series Forecasting Needs Transformed Label Alignment](https://mlsanity.com/paper/2505.17847/) (2025)
- [Bounds on Representation-Induced Confounding Bias for Treatment Effect Estimation](https://mlsanity.com/paper/2311.11321/) (2023)
- [DistDF: Time-Series Forecasting Needs Joint-Distribution Wasserstein Alignment](https://mlsanity.com/paper/2510.24574/) (2025)
- [PRANCE: Joint Token-Optimization and Structural Channel-Pruning for Adaptive ViT Inference](https://mlsanity.com/paper/2407.05010/) (2024)
- [Quadratic Direct Forecast for Training Multi-Step Time-Series Forecast Models](https://mlsanity.com/paper/2511.00053/) (2025)
- [GAD-PVI: A General Accelerated Dynamic-Weight Particle-Based Variational Inference Framework](https://mlsanity.com/paper/2312.16429/) (2023)
- [CausalRM: Causal-Theoretic Reward Modeling for RLHF from Observational User Feedbacks](https://mlsanity.com/paper/2603.18736/) (2026)
- [Debiased Collaborative Filtering with Kernel-Based Causal Balancing](https://mlsanity.com/paper/2404.19596/) (2024)
- [MetaCoCo: A New Few-Shot Classification Benchmark with Spurious Correlation](https://mlsanity.com/paper/2404.19644/) (2024)
- [Orthogonal Representation Learning for Estimating Causal Quantities](https://mlsanity.com/paper/2502.04274/) (2025)
- [Revisiting Counterfactual Regression through the Lens of Gromov-Wasserstein Information Bottleneck](https://mlsanity.com/paper/2405.15505/) (2024)
- [Rethinking Guidance Information to Utilize Unlabeled Samples:A Label Encoding Perspective](https://mlsanity.com/paper/2406.02862/) (2024)
- [Estimating Long-term Heterogeneous Dose-response Curve: Generalization Bound Leveraging Optimal Transport Weights](https://mlsanity.com/paper/2406.19195/) (2024)
- [Fine-tuning Flow Matching Generative Models with Intermediate Feedback](https://mlsanity.com/paper/2510.18072/) (2025)
- [PairNet: Training with Observed Pairs to Estimate Individual Treatment Effect](https://mlsanity.com/paper/2406.03864/) (2024)
- [Enhancing Treatment Effect Estimation via Active Learning: A Counterfactual Covering Perspective](https://mlsanity.com/paper/2505.05242/) (2025)
- [Optimal Transport for LLM Reward Modeling from Noisy Preference](https://mlsanity.com/paper/2605.06036/) (2026)
- [Be Aware of the Neighborhood Effect: Modeling Selection Bias under Interference](https://mlsanity.com/paper/2404.19620/) (2024)
- [Double-Bounded Nonlinear Optimal Transport for Size Constrained Min Cut Clusterin](https://mlsanity.com/paper/2501.18143/) (2025)
- [Treatment response as a latent variable](https://mlsanity.com/paper/2502.08776/) (2025)
- [A Relative Error-Based Evaluation Framework of Heterogeneous Treatment Effect Estimators](https://mlsanity.com/paper/2510.16419/) (2025)
- [Evaluating Uplift Modeling under Structural Biases: Insights into Metric Stability and Model Robustness](https://mlsanity.com/paper/2603.20775/) (2026)
- [Leveraging a Simulator for Learning Causal Representations from Post-Treatment Covariates for CATE](https://mlsanity.com/paper/2502.05037/) (2025)
- [Learning Counterfactual Outcomes Under Rank Preservation](https://mlsanity.com/paper/2502.06398/) (2025)
- [Minimax Optimal Two-Stage Algorithm For Moment Estimation Under Covariate Shift](https://mlsanity.com/paper/2506.23453/) (2025)
- [Mitigating Hidden Confounding by Progressive Confounder Imputation via Large Language Models](https://mlsanity.com/paper/2507.02928/) (2025)
- [PIPCFR: Pseudo-outcome Imputation with Post-treatment Variables for Individual Treatment Effect Estimation](https://mlsanity.com/paper/2512.18737/) (2025)
- [Deep Time-Series Forecasting in 10 Years: A Survey](https://mlsanity.com/paper/2603.19899/) (2026)

## Abstract (arXiv)

Estimating conditional average treatment effect from observational data is highly challenging due to the existence of treatment selection bias. Prevalent methods mitigate this issue by aligning distributions of different treatment groups in the latent space. However, there are two critical problems that these methods fail to address: (1) mini-batch sampling effects (MSE), which causes misalignment in non-ideal mini-batches with outcome imbalance and outliers; (2) unobserved confounder effects (UCE), which results in inaccurate discrepancy calculation due to the neglect of unobserved confounders. To tackle these problems, we propose a principled approach named Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport in the context of causality. Specifically, based on the framework of stochastic optimal transport, we propose a relaxed mass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.