Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation

Jiaxuan Zhang, Emadeldeen Eldele, Fuyuan CAO, Yang Wang, Xiaoli Li, Jiye Liang

TL;DR: We propose a method for ITE estimation that leverages a derived error bound to ensure fine-grained alignment and robust performance, even with individual heterogeneity.

Abstract: Estimating Individual Treatment Effects (ITE) from observational data is challenging due to covariate shift and counterfactual absence. While existing methods attempt to balance distributions globally, they often lack fine-grained sample-level alignment, especially in scenarios with significant individual heterogeneity. To address these issues, we reconsider counterfactual as a proxy to emulate balanced randomization. Furthermore, we derive a theoretical bound that links the expected ITE estimation error to both factual prediction errors and representation distances between factuals and counterfactuals. Building on this theoretical foundation, we propose FCCL, a novel method designed to effectively capture the nuances of potential outcomes under different treatments by (i) generating diffeomorphic counterfactuals that adhere to the data manifold while maintaining high semantic similarity to their factual counterparts, and (ii) mitigating distribution shift via sample-level alignment grounded in our derived generalization-error bound, which considers factual-counterfactual similarity and category consistency. Extensive evaluations on benchmark datasets demonstrate that FCCL outperforms 13 state-of-the-art methods, particularly in capturing individual-level heterogeneity and handling sparse boundary samples.

Lay Summary:

In many real-world scenarios, making personalized decisions is essential: digital marketing for tailoring customer strategies, social sciences for policy evaluation, and healthcare for personalized treatment planning. Estimating individual treatment effects (ITE) supports these needs by evaluating how a specific treatment would affect a specific individual. However, this task is challenging because we can only observe factual outcomes, not the counterfactual outcomes—i.e., what would have happened under a different treatment scenario.

To address this, we emulate the randomized controlled trials (RCTs) conditions to infer the missing counterfactual outcomes. Our approach generates counterfactuals that are both realistic and meaningful, and it leverages the relationship between factual and counterfactual samples to learn consistent representations for estimating potential outcomes under different treatments.

Notably, our method improves the accuracy of ITE estimation, particularly for boundary cases and heterogeneous individuals.

Primary Area: General Machine Learning->Causality

Keywords: Causal inference, individual treatment effect, counterfactual generation, contrastive learning

Submission Number: 5420