Moderately-Balanced Representation Learning for Treatment Effects with Orthogonality Information

TL;DR

This paper introduces a novel representation learning framework that balances treatment and control groups while leveraging orthogonality information to improve the accuracy and robustness of treatment effect estimation from observational data.

Contribution

It proposes a moderately-balanced representation learning method that avoids over-balance and fully utilizes orthogonality information for better ATE estimation.

Findings

Outperforms existing methods on benchmark datasets

Achieves more robust treatment effect estimates

Effectively balances treatment and control groups

Abstract

Estimating the average treatment effect (ATE) from observational data is challenging due to selection bias. Existing works mainly tackle this challenge in two ways. Some researchers propose constructing a score function that satisfies the orthogonal condition, which guarantees that the established ATE estimator is "orthogonal" to be more robust. The others explore representation learning models to achieve a balanced representation between the treated and the controlled groups. However, existing studies fail to 1) discriminate treated units from controlled ones in the representation space to avoid the over-balanced issue; 2) fully utilize the "orthogonality information". In this paper, we propose a moderately-balanced representation learning (MBRL) framework based on recent covariates balanced representation learning methods and orthogonal machine learning theory. This framework protects the representation from being over-balanced via multi-task learning. Simultaneously, MBRL incorporates the noise orthogonality information in the training and validation stages to achieve a better ATE estimation. The comprehensive experiments on benchmark and simulated datasets show the superiority and robustness of our method on treatment effect estimations compared with existing state-of-the-art methods.