Adversarially Balanced Representation for Continuous Treatment Effect Estimation

Estimating the individual treatment effect (ITE) from observational datasets has important applications in domains such as personalized medicine (Kent, Steyerberg, and van Klaveren 2018; Prosperi et al. 2020), economics (Varian 2016), and recommendation systems (Wang et al. 2020). The observational dataset contains units with their covariates $X$, the assigned treatment $T$, and outcome after intervention $Y$ (also known as the factual outcome). Since the treatment assignment policy is unknown, there typically exists an inherent treatment-selection bias stemming from confounding covariates that influence both the treatment assignment and the outcome. Consequently, a causal ITE estimator requires adjusting for the covariate shift among different treatment populations (Bareinboim, Tian, and Pearl 2014; Bareinboim and Pearl 2012).

In recent years, the representation learning approach (Johansson, Shalit, and Sontag 2016; Shalit, Johansson, and Sontag 2017) has demonstrated remarkable success in adjusting for covariate shift. Briefly, the idea is to learn a balanced representation of covariates (rather than balancing the covariates themselves) using an encoder and then predicting the outcomes from the representation using an outcome prediction network. However, the majority of the representation learning methods have been proposed for binary treatments (Hassanpour and Greiner 2019; Zhang, Liu, and Li 2021; Yao et al. 2018), and incorporating continuous treatments into their architectures is challenging (Chu et al. 2023). We elaborate on the associated challenges with continuous treatments in the following.

i) Balancing representation for continuous treatments
To achieve a balanced representation $Z$, most existing methods minimize the shift in terms of Integral Probability Metric (IPM) distance ¹¹1$IPM_{G}(p,q)=\sup_{g\in G}\int_{S}g(s)\big{(}p(s)-q(s)\big{)}ds$, where specifying $G$ leads to different distributional distance. (Sriperumbudur et al. 2009) between the distributions of $P(Z,T)$ and $P(Z)\,P(T)$ as it allows bounding the counterfactual prediction error (Shalit, Johansson, and Sontag 2017). Given that the marginal distribution $P(Z)$ is unknown, (Bellot, Dhir, and Prando 2022) proposed minimizing the IPM distance between $P(Z,T=t)$ and $P(Z,T=\lnot t)$ where $t$ denotes a treatment value in the data and $\lnot t$ denotes all treatments except $t$. Similarly, (Wang et al. 2022a) suggested discretizing the treatment range into intervals and minimized the maximum IPM distance between the distributions of the two intervals. Despite minimizing an upper bound, these methods involve non-parametric approximations of several IPM distances in practice, which may be inaccurate for high-dimensional representation and small training data (Liang 2019). Furthermore, IPMs are by definition worst-case distances and and obtaining a treatment-invariant representation through IPM might be overly restrictive, potentially excluding important confounding factors for outcome prediction (Zhang, Bellot, and van der Schaar 2020).

There are also non-IPM representation learning methods. For instance, (Du et al. 2021; Berrevoets et al. 2020) adopt an adversarial discriminator in order to balance the distributions in the latent representation, (Yao et al. 2018) proposed preserving the local similarity in the representation space, (Hassanpour and Greiner 2019; Wu et al. 2022) learn a disentangled representation to distinguish different latent factors, (Zhang, Bellot, and van der Schaar 2020) enforce invertibility of the encoder function to prevent the loss of covariate information, and (Alaa and van der Schaar 2017) introduced a regularization scheme in order to generalize to counterfactual outcomes. However, these methods lack theoretical justification (Du et al. 2021; Hassanpour and Greiner 2019; Yao et al. 2018) or their guarantees are limited to binary treatments (Zhang, Bellot, and van der Schaar 2020; Alaa and van der Schaar 2017).

ii) Treatment impact in outcome prediction network
In order to predict outcomes of different treatments, the outcome prediction network on top of the representation needs to incorporate the treatment value. Considering the treatment variable as an input of the outcome prediction network along with the representation causes overfitting the much higher-dimensional representation and largely limits the impact of the treatment value. Also, a distinct prediction head for each treatment value (as in the case of binary treatments) is not practical. Instead, (Schwab et al. 2020) proposed dose response network (DRNet) which divides the treatment range into intervals and consider a distinct network for the prediction of each interval. (Nie et al. 2021) proposed a varying coefficient network (VCNet) that involves the treatment value in the network parameters through spline functions of treatment. Nonetheless, both networks can not capture the dependency between representation and treatment effectively since the choice of spline functions in VCNet (or the intervals in DRNet) are made irrespective to representation values (Zhang et al. 2022).

In this paper, we propose a representation learning method to accurately predict potential outcomes of a continuous treatment. We address the above challenges through the following contributions:

1. 1.

   We prove that under certain assumptions the counterfactual error is bounded by the factual error and the KL divergence between $P(Z)P(T)$ and $P(Z,T)$. Unlike the IPM distance, the KL divergence can be estimated parametrically, leading to a more reliable bound.

2. 2.

   We propose Adversarial Balanced Counterfactual Regression (ACFR) network. ACFR minimizes the KL divergence using an adversarial game extracting a balanced representation for continuous treatments. ACFR also minimizes the factual prediction error by a cross-attention network that captures the complex dependency between treatment and the representation.

3. 3.

   We conduct an experimental comparison of ACFR against state-of-the-art methods on semi-synthetic datasets, News and TCGA, and analyze the robustness to varying-levels of treatment-selection bias for the methods.

We assume a dataset of the form $D=\{x_{i},t_{i},y_{i}\}_{i=1}^{N}$, where $x_{i}\in\mathcal{X}\subseteq\mathbb{R}^{d}$ denotes the covariates of the $i$th unit, $t_{i}\in[0,1]$ is the continuous treatment that unit $i$ received, and $y_{i}\in\mathcal{Y}\subseteq\mathbb{R}$ denotes the outcome of interest for unit $i$ after receiving treatment $t_{i}$. $N$ is the total number of units, and $d$ is the dimension size of covariates. We are interested in learning a machine learning model to predict the causal quantity $\mu(x,t)=E_{\mathcal{Y}}[Y(t)|X=x]$, which is the potential expected outcome under treatment $t$ for the individual with covariates $x$. Note that, unlike binary ITE, the goal is to predict all potential outcomes, not just the difference between them. Similar to previous works, we rely on the following standard assumptions to make treatment effects identifiable from an observational dataset.

{restatable*}[ Assumption 1 - Unconfoundedness]thm ($\{Y_{t}\}_{t\in T}\perp T|X$). In words, given covariates, treatment and potential outcomes are independent.

{restatable*}[ Assumption 2 - Overlap]thm ($P(T=t|X=x)>0,\forall t\in[0,1],\forall x\in X$). In words, every unit receives treatment level $t$ with a probability greater than zero.

With these assumptions, $\mu(x,t)$ can be rewritten as follows, and we can estimate it:

We analyze the properties of encoder function $\phi:\mathcal{X}\rightarrow\mathcal{Z}$, where $\mathcal{Z}$ is the representation space, outcome prediction function $h:\mathcal{Z}\times[0,1]\rightarrow\mathcal{Y}$ and loss function $L:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$.

{restatable*}[Definition 1]thm Define $\ell_{L,h,\phi}(x,t)=L(h(\phi(x),t),y):\mathcal{X}\times[0,1]\rightarrow\mathbb{R}^{+}$ to be the unit-loss for a unit with covariate $x$ that is intervened with treatment $t$. Unit-loss $\ell_{L,h,\phi}(x,t)$ measures loss $L$ between the predicted outcome $\hat{y}=h(\phi(x),t)$ and the ground-truth outcome $y=\mu(x,t)$.

Using the definition of unit-loss, we can define the expected prediction error of some treatment $t$ by marginalizing over the covariate distribution. As a result of treatment-selection bias, covariate distribution of samples having received treatment $t$ (factual) and samples not having received treatment $t$ (counterfactual) are different. We define factual error $\varepsilon_{f}^{\ell}(t)$ by marginalizing over $p(x|t)$ and counterfactual error $\varepsilon_{cf}^{\ell}(t)$ by marginalizing over $p(x)$ as follows.

We also define the expected error of all treatments by marginalizing over their range $[0,1]$ as follows: $\varepsilon_{f}^{\ell}=\int_{[0,1]}\int_{\mathcal{X}}\ell_{L,h,\phi}(x,t)\,p(x,t)\,dx\,dt$ and $\varepsilon_{cf}^{\ell}=\int_{[0,1]}\int_{\mathcal{X}}\ell_{L,h,\phi}(x,t)\,p(x)\,p(t)\,dx\,dt$. Note that the expected factual error $\varepsilon_{f}^{\ell}$ integrates over joint distribution $p(x,t)$, and expected counterfactual error $\varepsilon_{cf}^{\ell}$ integrates over $p(x)\,p(t)$. We aim to reduce the distributional distance in representation space $Z$ to ensure that minimizing $\varepsilon_{f}^{\ell}$ results in minimizing $\varepsilon_{cf}^{\ell}$. We need the following assumption on encoder $\phi$ to ensure balancing properties generalize from the representation space to the covariate space.

{restatable*}[Assumption 3]thmass3 The encoder function $\phi$ is a twice-differentiable one-to-one mapping and the representation space $\mathcal{Z}$ is the image of $\mathcal{X}$ under $\phi$ with the induced distribution $p_{\phi}(z)$.

We also need the following assumption to ensure the unit-loss is not arbitrary large for any $(x,t)$ pair. Constraint on the unit-loss function is required for the IPM distance specification as well.

{restatable*}[Assumption 4]thmass4 Let $G$ be a class of functions with infinity norm less than 1, $G=\{g:\mathcal{Z}\times[0,1]\rightarrow\mathbb{R}^{+}\,|\,||g||_{\infty}\leq 1\}$. Then, there exist a constant $C>0$ such that $\frac{\ell_{L,h,\phi}(x,t)}{C}\in G$. This means for any $(x,t)$ we have $\frac{\ell_{L,h,\phi}(x,t)}{C}\leq 1$.

Note that Assumptions 3 and 4 are common in representation learning literature (Shalit, Johansson, and Sontag 2017; Bellot, Dhir, and Prando 2022; Hassanpour and Greiner 2019). Now we present our main theoretical results which demonstrates a bound on the expected counterfactual error $\varepsilon_{cf}^{\ell}$ consisting of the expected factual error $\varepsilon_{f}^{\ell}$ and the KL divergence between distributions $p_{\phi}(z,t)$ and $p_{\phi}(z)\,p(t)$.

{restatable*}[Proposition 1 - Counterfactual Generalization Bound]thmprop1 Given the one-to-one encoder function $\phi:\mathcal{X}\rightarrow\mathcal{Z}$, the outcome prediction function $h:\mathcal{Z}\times[0,1]\rightarrow\mathcal{Y}$, and the unit-loss function $\ell_{L,h,\phi}(x,t)$ that satisfies Assumption 4,

Note that the KL divergence is non-negative and becomes zero if and only if two distributions are the same. Therefore, $D_{KL}\bigg{(}p_{\phi}(z,t)||p_{\phi}(z)p(t)\bigg{)}=0$ implies $p_{\phi}(z,t)=p_{\phi}(z)\,p(t)$ and $\varepsilon_{f}^{\ell}=\varepsilon_{cf}^{\ell}$. We can also interpret the above bound from information theory perspective. Briefly, minimizing RHS of Proposition 1 results in representation $z$ which has maximum mutual information with outcome $y$ given treatment $t$ (equivalent to minimizing $\varepsilon_{f}^{\ell}$) and minimum mutual information with treatment $t$ (equivalent to minimizing KL divergence term).

For some applications, one might be interested in the treatment effect between two different treatments rather than predicting all counterfactual outcomes. For instance, in binary treatment setting it is standard to report the model performance in terms of precision of estimating heterogeneous effect (PEHE) (Hill 2011) which measures the squared difference between ground-truth treatment effect $\tau(x)=\mu(x,1)-\mu(x,0)$ and predicted treatment effect $\hat{\tau}(x)=h(\phi(x),1)-h(\phi(x),0)$. We define the continuous counterpart $\varepsilon_{pehe}(t_{1},t_{2})$ between two treatment levels $t_{1}$ and $t_{2}$ and present an upper bound on it in Proposition 2.

{restatable*}[Definition 2]thmdef2 Define $\varepsilon_{pehe}(t_{1},t_{2})=\int_{\mathcal{X}}\big{[}\big{(}\mu(x,t_{1})-\mu(x,t_{2})\big{)}-\big{(}h(\phi(x),t_{1})-h(\phi(x),t_{2})\big{)}\big{]}^{2}p(x)dx$ to be the expected precision of estimating heterogeneous effect between treatment levels $t_{1}$ and $t_{2}$.

{restatable*}[Proposition 2 - Precision of Estimating Heterogeneous Effect Bound]thmprop2 Given the one-to-one encoder function $\phi$ and outcome prediction function $h$ as in Proposition 1, and a unit-loss function $\ell_{L,h,\phi}(x,t)$ that satisfies Assumption 4 and its associated $L$ is squared error $||.||^{2}$,

You can find the proof of Proposition 1 and 2 in the Appendix A.

Based on Proposition 1, we can control the counterfactual error by learning an encoder function $\phi$ and an outcome prediction function $h$ that jointly minimize distribution shift and factual outcome error. We propose our method that implements functions $\phi$ and $h$ using neural networks and is trained with an objective function inspired by Proposition 1. Figure 1 illustrates the architecture of the ACFR network, consisting of an encoder network $\phi$, an outcome prediction network $h$, and a treatment prediction network $\pi$. The key aspects of ACFR, distribution shift minimization and minimization of outcome prediction error, are presented as follows. ²²2Note that in the previous section we assumed $\phi$ is a one-to-one mapping, however, we obtained better results empirically for all methods using a neural network as encoder.

Treatment effect estimation methods have to be evaluated for predicting potential outcomes including counterfactuals which are unavailable in real-world observational datasets. Therefore, synthetic or semi-synthetic datasets are commonly used since their treatment assignment mechanism and outcome function are known and hence counterfactual outcomes can be generated. Note that this does not change the fact that only factual outcomes are accessible during training. In this section, we present our experimental results. The code for synthetic data generation and implementation of the methods can be found at here: https://github.com/amirrezakazemi/acfr

This paper has investigated the problem of continuous treatment effect estimation and introduced ACFR (Adversarial Counter-Factual Regression) method for predicting potential outcomes. We proved a new bound of the counterfactual error using the KL divergence instead of an IPM distance, which has the benefit that the KL divergence can be estimated parametrically and results in a more reliable bound. Based on the error bound, ACFR uses an adversarial neural network architecture to minimize the KL divergence of the representations and a cross-attention network to minimize the factual prediction error. It is worth mentioning that ACFR is not restricted to continuous treatments, and in future work we plan to extend and evaluate ACFR framework for structured and time series treatments. Nonetheless, we note that ACFR, similar to many treatment effect estimation methods, relies on strong ignorability assumption, which is not necessarily hold in real-world applications.

We would like to thank Oliver Schulte and Sharan Vaswani for feedback on the paper. This research was supported by the Natural Sciences and Engineering Research Council of Canada (NSERC) Discovery Grant.