Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation

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.

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.