# Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect Estimation

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 61 references that are in the map.
- First posted on arXiv 1 Jul 2024.

## Builds on (most cited first, 24 of 24)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Scalable Gromov-Wasserstein Learning for Graph Partitioning and Matching](https://mlsanity.com/paper/1905.07645/) (2019)
- [Fused Gromov-Wasserstein distance for structured objects: theoretical foundations and mathematical properties](https://mlsanity.com/paper/1811.02834/) (2018)
- [Learning Counterfactual Representations for Estimating Individual Dose-Response Curves](https://mlsanity.com/paper/1902.00981/) (2019)
- [Faster Wasserstein Distance Estimation with the Sinkhorn Divergence](https://mlsanity.com/paper/2006.08172/) (2020)
- [AutoDebias: Learning to Debias for Recommendation](https://mlsanity.com/paper/2105.04170/) (2021)
- [ESCM2: Entire Space Counterfactual Multi-Task Model for Post-Click Conversion Rate Estimation](https://mlsanity.com/paper/2204.05125/) (2022)
- [Projection Robust Wasserstein Distance and Riemannian Optimization](https://mlsanity.com/paper/2006.07458/) (2020)
- [Large-scale Causal Approaches to Debiasing Post-click Conversion Rate Estimation with Multi-task Learning](https://mlsanity.com/paper/1910.09337/) (2019)
- [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)
- [Learning Disentangled Representations for Counterfactual Regression via Mutual Information Minimization](https://mlsanity.com/paper/2206.01022/) (2022)
- [A Generalized Doubly Robust Learning Framework for Debiasing Post-Click Conversion Rate Prediction](https://mlsanity.com/paper/2211.06684/) (2022)
- [Balancing Unobserved Confounding with a Few Unbiased Ratings in Debiased Recommendations](https://mlsanity.com/paper/2304.09085/) (2023)
- [DESCN: Deep Entire Space Cross Networks for Individual Treatment Effect Estimation](https://mlsanity.com/paper/2207.09920/) (2022)
- [Local Similarity-Aware Deep Feature Embedding](https://mlsanity.com/paper/1610.08904/) (2016)
- [Entire Space Counterfactual Learning for Reliable Content Recommendations](https://mlsanity.com/paper/2210.11039/) (2022)
- [Optimal transport weights for causal inference](https://mlsanity.com/paper/2109.01991/) (2021)
- [Deep Treatment-Adaptive Network for Causal Inference](https://mlsanity.com/paper/2112.13502/) (2021)
- [Revisiting Counterfactual Regression through the Lens of Gromov-Wasserstein Information Bottleneck](https://mlsanity.com/paper/2405.15505/) (2024)

## Built on by (most cited first, 7 of 7)

- [Time-o1: Time-Series Forecasting Needs Transformed Label Alignment](https://mlsanity.com/paper/2505.17847/) (2025)
- [DistDF: Time-Series Forecasting Needs Joint-Distribution Wasserstein Alignment](https://mlsanity.com/paper/2510.24574/) (2025)
- [Quadratic Direct Forecast for Training Multi-Step Time-Series Forecast Models](https://mlsanity.com/paper/2511.00053/) (2025)
- [CausalRM: Causal-Theoretic Reward Modeling for RLHF from Observational User Feedbacks](https://mlsanity.com/paper/2603.18736/) (2026)
- [Learning Counterfactual Outcomes Under Rank Preservation](https://mlsanity.com/paper/2502.06398/) (2025)
- [Deep Time-Series Forecasting in 10 Years: A Survey](https://mlsanity.com/paper/2603.19899/) (2026)
- [Disentangled Double Machine Learning for Accurate Causal Effect Estimation](https://mlsanity.com/paper/2605.24808/) (2026)

## Abstract (arXiv)

Heterogeneous treatment effect (HTE) estimation from observational data poses significant challenges due to treatment selection bias. Existing methods address this bias by minimizing distribution discrepancies between treatment groups in latent space, focusing on global alignment. However, the fruitful aspect of local proximity, where similar units exhibit similar outcomes, is often overlooked. In this study, we propose Proximity-enhanced CounterFactual Regression (CFR-Pro) to exploit proximity for enhancing representation balancing within the HTE estimation context. Specifically, we introduce a pair-wise proximity regularizer based on optimal transport to incorporate the local proximity in discrepancy calculation. However, the curse of dimensionality renders the proximity measure and discrepancy estimation ineffective -- exacerbated by limited data availability for HTE estimation. To handle this problem, we further develop an informative subspace projector, which trades off minimal distance precision for improved sample complexity. Extensive experiments demonstrate that CFR-Pro accurately matches units across different treatment groups, effectively mitigates treatment selection bias, and significantly outperforms competitors. Code is available at https://github.com/HowardZJU/CFR-Pro.