A two-stage disentangled and balanced representation learning method for counterfactual regression

Estimating ITEs is crucial for personalized decision-making. Prior research has sought to estimate ITEs using the causal representation balancing method, where all features are mapped to an embedding space for the tasks of outcome modeling and representation balancing. Nevertheless, two challenges remain: (i) the potential adverse interactions between the two tasks due to their sharing of the same representation of confounders and (ii) the negative impact of instrumental variables on outcome predictions. These observations motivate us to disentangle instrumental, confounding, and outcome-specific factors according to their causal relationships with the treatment and outcome within the representation balancing framework. We propose an ITE error bound using confounding and outcome-specific factors, showing that representation balancing through confounding factors alone is sufficient. Building upon this ITE error bound, we introduce a two-stage method, Disentangled and Balanced Representation Learning for Counterfactual Regression (DBR-CFR), to estimate ITEs. In the first stage, DBR-CFR effectively disentangles observed pre-treatment variables into three desired factors. In the second stage, DBR-CFR minimizes the objective function based on the empirical ITE error bound using disentangled factors. Extensive experimental results on synthetic and benchmark datasets verify the effectiveness and superiority of our proposed method. The source code is available at https://github.com/sssywang/DBR-CFR.