Searching optimal adjustment features for treatment effect estimation

Haotian Wang, Kun Kuang, Mingyang Geng, Haoang Chi, Ruoxuan Xiong, Longqi Yang, Wanrong Huang, Wenjing Yang

Keywords: treatment effect estimation, covariate separation, confounder balancing

TL;DR: We construct a reinforcement-learning based framework to search the optimal adjustment features for more precise treatment effect estimation.

Abstract: Most efforts devoted to causal inference focus on controlling the adjustment features to further alleviate the confounding effect. In realistic applications, the collected covariates often contain variables correlating to only one of the treatment (e.g., instrumental variables) and the outcome (e.g., precision variables). Due to the absence of prior knowledge, the brute-force approach for the practitioner is to include every covariate for adjustment. However, previous literature shows that adjusting the former covariates (treatment-only) hurts the treatment effect estimation, while adjusting the latter covariates (outcome-only) brings benefits. Consequently, it is meaningful to find an optimal adjustment set rather than the brute-force approach for more efficient treatment effect estimation. To this end, we establish a variance metric which is computationally tractable to measure the optimality of the adjustment set. From the non-parametric viewpoint, we theoretically show that our metric is minimized if and only if the adjustment features contain the confounders and the outcome-only variables. As optimizing the proposed variance metric is a combinational optimization problem, we incorporate the Reinforcement Learning (RL) to search the corresponding optimal adjustment set. More specifically, we adopt the encoder-decoder model as the actor to generate the binary feature mask on the original covariates, which serves as the differentiable policy. Meanwhile, the proposed variance metric serves as the reward to guide the policy update. Empirical results on synthetic and real-world datasets demonstrate that ~(a) our method successfully searches the optimal adjustment sets and (b) the searched adjustment features achieve more precise treatment effect estimation.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Probabilistic Methods (eg, variational inference, causal inference, Gaussian processes)