Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect Estimation

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.

💡 Research Summary

The paper tackles the fundamental problem of estimating heterogeneous treatment effects (HTE) from observational data, where treatment selection bias and the lack of counterfactual outcomes make causal inference difficult. Existing representation‑learning approaches, such as Counterfactual Regression (CFR), mitigate bias by aligning the treated and control groups in a latent space using a global distribution discrepancy (e.g., Wasserstein distance). However, these methods ignore two crucial aspects: (1) local proximity—the intuition that units with similar covariates tend to have similar outcomes—and (2) the curse of dimensionality, which renders distance‑based discrepancy measures unreliable when data are high‑dimensional and scarce.

To address these gaps, the authors propose CFR‑Pro, a framework that integrates (i) a pair‑wise proximity regularizer (PPR) and (ii) an informative subspace projector (ISP) within an optimal transport (OT) formulation.

Pair‑wise Proximity Regularizer (PPR).
Standard OT minimizes the total transport cost ⟨π, D⟩ where D contains Euclidean distances between treated and control representations. PPR augments this objective with a Gromov‑Wasserstein‑style term that penalizes mismatches in the internal geometry of each group. Formally, the regularized cost is
  F = min₍π∈Π₎ κ⟨π,D⟩ + (1‑κ)∑₍i,j,k,l₎ ‖D⁰_{ik} − D¹_{jl}‖² π_{ij} π_{kl},
where D⁰ and D¹ are intra‑group distance matrices and κ∈