MULTI-TREATMENT EFFECT ESTIMATION WITH PROXY: CONTRASTIVE LEARNING AND RANK WEIGHTING

Abstract

We study the treatment effect estimation problem for continuous and multidimensional treatments, in the setting with unobserved confounders, but highdimension proxy variables for unobserved confounders are available. Existing methods either directly adjust the relationship between observed covariates and treatments or recover the hidden confounders by probabilistic models. However, they either rely on a correctly specified treatment assignment model or require strong prior of the unobserved confounder distribution. To relax these requirements, we propose a Contrastive regularizer (Cr) to learn the proxy representation that contains all the relevant information in unobserved confounders. Based on the Cr, we propose a novel Rank weighting method (Rw) to de-bias the treatment assignment. Combining Cr and Rw, we propose a neural network framework named CRNet to estimate the effects of multiple continuous treatments under unobserved confounders, evaluated by the Average Dose-Response Function. Empirically, we demonstrate that CRNet achieves state-of-the-art performance on both synthetic and semi-synthetic datasets.



Causal inference is widely applied for explanatory analysis and decision making, e.g., Precision Medicine (Raita et al., 2021) , Advertisement (Lada et al., 2019) , Education (Johansson et al., 2016) and Digital Economy (Nazarov, 2020) . With accessible observation data, many existing algorithms accurately estimate the effect of binary treatment by adjusting the confounders (i.e., the common causes of treatments and outcomes) which rely on unconfoundedness assumption that all confounders are observed. However, continuous and multi-dimensional treatments and unmeasured confounders are common in practice. For instance, practitioners seek to develop precise medicine by studying the response of multiple drug dosages (i.e., treatment) on patient health state (i.e., outcome) (Shi et al., 2020) . Besides, due to technique and manipulation issues, some key variables, associated with the treatments and outcomes, like patient's immunity maybe missing in the historical data, which are referred to as unmeasured confounders. To detect and adjust unmeasured confounders, practitioners would record some proxy variables (noised unobserved confounders, e.g., antibodies) which don't have a direct effect on treatments and outcome of interest but has a spurious association through shared common confounders (Fig. 1(a) ). In continuous treatments setting, under unconfoundedness assumption, recent works discretize the continuous treatment into multi-valued treatment (Hill, 2011; Wager & Athey, 2018) to traditional models, or develop generalize balancing methods for continuous scenario (Hirano & Imbens, 2004; Vegetabile et al., 2021; Huling et al., 2021) . Among them, state-of-the-art works (Wu & Fukumizu, 2021; Schwab et al., 2020; Nie et al., 2021) learn a low-dimensional representation for raw data and balance it using minimizing mutual information, which discard the imbalance part of raw data and lose most information for predictive task in practice. In fact, the technique implements a trade-off decreasing the estimator variance at the price of increasing the bias. Furthermore, with unobserved confounders, if we control the proxy rather than unobserved variables, the effect estimation will induce additional bias, referred as recovery bias. To deal with this bias, instead of balancing representations and discarding information to block the relationship between observed covariates and treatments, we propose a novel Contrastive regularizer (Cr) to learn a proxy representation for capturing all the relevant information in unobserved confounders with contrastive learning (He et al., 2020; Chen et al., 2020; Grill et al., 2020) which regularize representation space by positive and negative pairs. In Cr, we define the positive pair is the pair of treatments and proxies from the same sample, and the negative pair is the pair of treatment from one sample and proxies from different samples. And with an ideally representation for confounders, we would adopt a balancing methods to eliminate confounding bias, such as generalized propensity score (Hirano & Imbens, 2004 ). However, one limitation is that the covariate balancing methods rely on the correct specified models. If we don't have any prior for the models of propensity score, i.e., the conditional distribution of treatment conditioning on the covariates, the effect estimation would still be biased, especially for high-dimensional data and continuous treatment. Besides, balancing methods still suffer from extreme values problem. Although recent methods (Fong et al., 2018; Vegetabile et al., 2021; Huling et al., 2021) propose to clip the score value or optimize balancing weights directly, they still fail in complex data, especially, under multi-continuous treatment setting. So a balancing method that have no extreme values and adapted to unobserved confounders is urgently needed. Therefore, to control for bias from treatment assignment, we propose to rank the weights obtained from inverse propensity score for more effective balancing weighting. Based on the proxy representation learned above, we sort the propensity score based weights in descending order and record their rank (the order in sorted data) as rank weights (Rw), which is an effective and robust weights for treatment effect estimation, theoretically. Combining Contrastive regularizer (Cr) and Rank weighting (Rw) methods, we propose a neural network framework CRNet to alleviate the outcome approximate bias in estimating the Average Dose-Response Function (ADRF). CRNet can accurately estimate the effects of multiple continuous treatments with high-dimension proxy variables. Empirically, we demonstrate that CRNet achieves state-of-the-art performance on both synthetic and semi-synthetic datasets.

2. RELATED WORK

Causal effect identification with proxy methods Proxy (Guo et al., 2020) et al., 2018) for estimating the causal effect. The setting of this paper is similar to the proxy. But our method need no data distribution prior and outperforms others in performance. Estimation methods for continuous treatments For estimating the continuous treatment effect, a branch of methods include spline (Imai & Van Dyk, 2004 ), kernel methods (Flores et al., 2012) , ensemble methods (Hill, 2011; Wager & Athey, 2018 ), representation-based methods (Schwab et al., 2020; Nie et al., 2021; Bica et al., 2020) model the relationship between treatments and outcomes. There is also a branch of methods (Hirano & Imbens, 2004; Imai & Van Dyk, 2004; Robins et al., 2000; Vegetabile et al., 2021; Arbour et al., 2021; Huling et al., 2021) aim at balancing the covariates shifts. Few previous works take into account of unobserved variables with continuous treatment assignment bias. In this paper, we propose the contrastive regularizer to gain the balancing methods with the presence of unobserved confounders. Also, we propose a new rank weighting method which have no extreme values and is not much sensitive to model misspecified. Combining Cr and Rw, we design a framework CRNet to estimate continuous treatment with proxy.



Figure 1: (a) Causal Structure of Raw Data, i.e., Y ⊥ T | U; (b) Target Relationship from proxy representation, i.e., Y (T) ⊥ U | E(X).

assumes that the unobserved confounders can be recovered from the observed covariates. CEVAE(Louizos et al., 2017), intact-VAE (Wu & Fukumizu, 2021) recover unobserved confounders with VAE (Kingma et al., 2019) constraint. Negative controls (Lipsitch et al., 2010) assume that there exist two negative control variables: one is related to treatments and confounders, and another is related to outcomes and confounders. DFPV (Xu et al., 2021) introduces neural networks to model the bridge function (Miao

