Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect Estimation

This paper proposes CFR-Pro, a novel method for heterogeneous treatment effect estimation that enhances representation balancing by integrating a pairwise proximity regularizer based on optimal transport and an informative subspace projector to overcome the curse of dimensionality and mitigate treatment selection bias in observational data.

The Gist

Imagine you are a doctor trying to figure out if a new medicine works better for some people than others. You can't run a perfect experiment where you give the medicine to half the patients and a placebo to the other half (maybe it's too expensive or unethical). So, you have to look at observational data: records of people who chose to take the medicine and people who didn't.

The problem? People who choose the medicine are often different from those who don't. Maybe they are richer, healthier, or more health-conscious. This is called Selection Bias. It's like trying to judge if a sports car is faster than a minivan, but you only test the sports car on a race track and the minivan on a muddy dirt road. The results will be misleading.

This paper introduces a new method called CFR-Pro to fix this mess. Here is how it works, explained with simple analogies.

The Core Problem: The "Global" vs. "Local" Mistake

Previous methods tried to fix this bias by making the two groups (medicine-takers and non-takers) look exactly the same globally. Imagine you have two bags of marbles: one bag is mostly red, the other mostly blue. Old methods tried to mix them until the overall color of the pile looked identical.

The Flaw: Just because the overall pile looks mixed doesn't mean the neighbors are similar. You might have a red marble sitting right next to a blue one that is completely different in size and weight. In the real world, this means the method might compare a sick, elderly person taking the drug with a healthy, young person who didn't, just because they happen to be "close" in the big picture. This leads to wrong conclusions.

The Paper's Insight: We need to care about Local Proximity. If two people are very similar (neighbors), they should be compared to each other. If they are different, they shouldn't be matched, even if the groups look balanced on average.

The Solution: CFR-Pro (The "Smart Matchmaker")

The authors propose CFR-Pro, which uses two clever tricks to fix the matching process.

Trick 1: The "Neighborly Nudge" (Pair-wise Proximity Regularizer)

Imagine you are organizing a dance.

• Old Method: You just make sure there are equal numbers of men and women in the room.
• CFR-Pro Method: You add a rule: "If two people are already standing close together and look similar, you must pair them up."

In technical terms, the paper uses a math tool called Optimal Transport (think of it as a logistics planner moving boxes from one warehouse to another). They added a "Neighborly Nudge" to the planner. This nudge says: "Don't just move boxes to make the total weight equal; make sure you move a box to a spot that is geometrically similar to where it came from."

This ensures that when the computer compares a patient who took the drug to one who didn't, it's comparing "apples to apples" (similar neighbors), not "apples to oranges."

Trick 2: The "High-Definition Filter" (Informative Subspace Projector)

Here is the second problem: The Curse of Dimensionality.
Imagine you are trying to find similar people in a crowd. If you only look at 2 traits (height and weight), it's easy to find matches. But if you look at 1,000 traits (eye color, shoe size, favorite song, blood type, etc.), everything starts to look the same. In high-dimensional math, the distance between any two points becomes meaningless, like trying to find a specific grain of sand on a beach by looking at the whole beach at once.

This makes the "Neighborly Nudge" fail because the computer can't tell who is actually close to whom.

The Fix: CFR-Pro introduces a Smart Filter (the Informative Subspace Projector).
Imagine you have a giant, messy room full of 1,000 objects. Instead of trying to organize the whole room, you put on a special pair of glasses that only lets you see the 50 most important objects that actually matter for the task.

• The computer ignores the noise (the 950 useless traits).
• It focuses only on the "informative" traits.
• Now, finding similar neighbors is easy again, and the matching becomes accurate.

Why This Matters (The Results)

The authors tested this on real-world data (like medical studies on infant health).

• Old methods were like a blurry photo: they got the general idea right but missed the details, leading to biased results.
• CFR-Pro is like a high-definition photo: it correctly matched similar people, ignored the noise, and gave a much more accurate answer about whether the treatment actually worked.

Summary Analogy

Think of estimating treatment effects like trying to find the best route for a delivery truck.

• Old methods looked at the whole map and said, "The average traffic is fine," ignoring that one specific street is a dead end.
• CFR-Pro zooms in. It uses a Smart Filter to ignore irrelevant roads (curse of dimensionality) and a Neighborly Nudge to ensure the truck only compares routes that are actually similar (local proximity).

The result? A delivery truck that never gets lost and always finds the fastest path, leading to better decisions in healthcare, business, and policy.

Technical Summary

1. Problem Statement

The paper addresses the challenge of estimating Heterogeneous Treatment Effects (HTE) from observational data. The primary obstacles are:

• Treatment Selection Bias: Non-random treatment assignment causes covariate shifts between treated and untreated groups, leading to biased outcome estimators.
• Limitations of Existing Methods:
  • Global vs. Local: Current representation-based methods (e.g., CFR, TARNet) focus on minimizing global distribution discrepancies (e.g., via Maximum Mean Discrepancy or standard Wasserstein distance) between treatment groups. They often overlook local proximity, the principle that similar units should exhibit similar outcomes. Ignoring local structure can lead to erroneous matching and misleading gradient updates.
  • Curse of Dimensionality: Calculating distribution discrepancies in high-dimensional latent spaces requires massive sample sizes to be reliable. In real-world scenarios where treated units are scarce, standard distance metrics become ineffective due to the "concentration of measure" phenomenon, where Euclidean distances lose discriminative power.

2. Methodology: CFR-Pro

The authors propose Proximity-enhanced CounterFactual Regression (CFR-Pro), a unified framework based on Optimal Transport (OT) that simultaneously addresses local proximity preservation and the curse of dimensionality.

A. Pair-wise Proximity Regularizer (PPR)

To incorporate local proximity into the discrepancy calculation, the authors extend the standard Kantorovich OT problem using a Gromov-Wasserstein inspired approach.

• Mechanism: Instead of just minimizing the transport cost between two groups (Global Discrepancy), PPR adds a regularization term that penalizes the mismatch of internal geometric structures within each group.
• Formulation: The objective function $F$ minimizes:
  $$\kappa \langle \pi, D \rangle + (1-\kappa) \sum_{i,j,k,l} P_{i,j,k,l} \pi_{i,j} \pi_{k,l}$$
  • The first term is the standard global transport cost.
  • The second term measures the difference in pairwise distances between units within the treated group and units within the untreated group. If the local neighborhood structure is preserved (i.e., similar units in one group map to similar units in the other), the cost is minimized.
  • $\kappa$ balances global alignment vs. local structure preservation.
• Effect: This ensures that the transport plan matches units with similar neighbors, preventing "blurred" alignments that occur with standard OT.

B. Informative Subspace Projector (ISP)

To combat the curse of dimensionality without sacrificing the representational capacity of the neural network, CFR-Pro introduces an ISP module.

• Problem: Calculating OT in high dimensions ($d$) is computationally expensive and statistically unreliable with limited samples. Reducing the hidden layer size of the neural network directly limits its ability to learn complex patterns.
• Solution: The ISP projects the high-dimensional representations into a lower-dimensional informative subspace ($k < d$) specifically for the discrepancy calculation, while keeping the full representation mapping $\psi$ intact for outcome prediction.
• Optimization: The projector $U$ is learned to minimize the reconstruction error of the representations ($\min \| R - R U U^\top \|_2^2$), effectively performing Principal Component Analysis (PCA) on the fly to retain the most informative variance.
• Benefit: This reduces the sample complexity of the OT calculation from exponential to manageable levels while preserving the full expressive power of the feature extractor.

C. Overall Learning Objective

The total loss function combines factual outcome estimation error with the enhanced discrepancy term:
$$\mathcal{L} = \mathcal{L}^{(F)}(\psi, \phi) + \lambda \cdot \mathcal{P}_{\kappa, P}(\psi)$$
Where $\mathcal{L}^{(F)}$ is the standard factual loss, and $\mathcal{P}_{\kappa, P}$ is the PPR-enhanced OT discrepancy calculated in the informative subspace.

3. Key Contributions

1. Theoretical Insight: The paper identifies and formalizes the neglect of local proximity in existing HTE balancing methods and proposes a theoretical bound (Theorem 1) showing that minimizing PPR-enhanced discrepancy reduces the PEHE (Precision in Estimation of Heterogeneous Effect) error.
2. Novel Architecture (CFR-Pro):
   • Introduces PPR to enforce geometric consistency in the transport plan, ensuring similar units are matched.
   • Introduces ISP to mitigate the curse of dimensionality by calculating discrepancies in a learned low-dimensional subspace.
3. Empirical Superiority: Demonstrates that CFR-Pro significantly outperforms state-of-the-art baselines (including CFR-MMD, CFR-WASS, ESCFR, and SITE) on standard benchmarks.

4. Experimental Results

The method was evaluated on two semi-synthetic datasets: IHDP (Infant Health and Development Program) and ACIC (Atlantic Causal Inference Conference).

• Performance Metrics: Evaluated using PEHE (Heterogeneous Effect), ATE (Average Treatment Effect), and ATT (Average Treatment Effect on Treated).
• Key Findings:
  • CFR-Pro vs. Baselines: CFR-Pro achieved the lowest PEHE and ATE errors across both datasets. For example, on ACIC, CFR-Pro reduced PEHE to 2.04, significantly outperforming the next best representation-based method (ESCFR at 2.68) and the standard CFR-WASS (3.32).
  • Ablation Studies:
    • Removing PPR (keeping ISP) increased error, confirming the necessity of local proximity.
    • Removing ISP (keeping PPR) increased error, confirming the necessity of dimensionality reduction for reliable OT in high dimensions.
    • The combination of both yielded the best results.
  • Sensitivity Analysis: The model is robust to hyperparameters $\lambda$ (alignment strength) and $\kappa$ (proximity strength) within a reasonable range. The study also revealed an inverse relationship between $\kappa$ and the optimal dimensionality reduction ratio $P$: as the reliance on local geometry ($\kappa$) increases, the need for dimensionality reduction ($P$) becomes more critical to avoid noise.

5. Significance

• Advancing Causal Inference: CFR-Pro bridges the gap between global distribution alignment and local geometric structure, offering a more nuanced approach to debiasing observational data.
• Practical Applicability: By addressing the curse of dimensionality, the method is more viable for real-world applications where data is high-dimensional but sample sizes (especially for the treated group) are limited.
• Generalizability: The framework of integrating local proximity regularizers and informative subspaces into Optimal Transport can be extended to other domains beyond HTE, such as domain adaptation and fair machine learning.

In conclusion, the paper argues that "Proximity Matters"; accurate treatment effect estimation requires not just balancing the global distribution of covariates, but also preserving the local geometric relationships between units, all while managing the statistical challenges of high-dimensional data.