Leveraging a Simulator for Learning Causal Representations from Post-Treatment Covariates for CATE

TL;DR

This paper explores how simulators can be used to learn causal representations from post-treatment covariates for estimating CATE, proposing a new method called SimPONet that adapts to simulator relevance.

Contribution

It provides a theoretical analysis of CATE error with simulators, introduces SimPONet, and demonstrates its effectiveness across various data scenarios.

Findings

SimPONet outperforms state-of-the-art baselines in diverse settings.

The generalization bound links CATE error to real-simulator distribution mismatch.

Adjusting simulator influence improves causal representation learning.

Abstract

Treatment effect estimation involves assessing the impact of different treatments on individual outcomes. Current methods estimate Conditional Average Treatment Effect (CATE) using observational datasets where covariates are collected before treatment assignment and outcomes are observed afterward, under assumptions like positivity and unconfoundedness. In this paper, we address a scenario where both covariates and outcomes are gathered after treatment. We show that post-treatment covariates render CATE unidentifiable, and recovering CATE requires learning treatment-independent causal representations. Prior work shows that such representations can be learned through contrastive learning if counterfactual supervision is available in observational data. However, since counterfactuals are rare, other works have explored using simulators that offer synthetic counterfactual supervision. Our goal in this paper is to systematically analyze the role of simulators in estimating CATE. We analyze the CATE error of several baselines and highlight their limitations. We then establish a generalization bound that characterizes the CATE error from jointly training on real and simulated distributions, as a function of the real-simulator mismatch. Finally, we introduce SimPONet, a novel method whose loss function is inspired from our generalization bound. We further show how SimPONet adjusts the simulator's influence on the learning objective based on the simulator's relevance to the CATE task. We experiment with various DGPs, by systematically varying the real-simulator distribution gap to evaluate SimPONet's efficacy against state-of-the-art CATE baselines.