Estimating Heterogeneous Treatment Effects for General Responses

We adopt the Neyman-Rubin potential outcome model. Each unit is associated with a covariate vector $X$, a treatment assignment indicator $W$, and two potential outcomes $Y(0)$, $Y(1)$. We observe the response $Y=Y(1)$ if the unit is under treatment, i.e., $W=1$, and $Y=Y(0)$ if the unit is under control, i.e., $W=0$. We make the standard assumptions in causal inference (Imbens and Rubin, 2015).

To facilitate asymptotic analyses, we introduce the super-population model

where $f_{X}$ denotes the distribution of covariate, $e(x)$ denotes the propensity score, and $f_{1}$ and $f_{0}$ denote the distributions of potential outcomes.

For survival analysis, we let $C(0)$, $C(1)\in[0,\infty]$ be the counterfactual censoring times and $C$ be the observed censoring time. Then $C=C(1)$ if $W=1$ and $C=C(0)$ if $W=0$. Let $\Delta:=\mathbbm{1}_{\{C\geq Y\}}$ be the observed censoring indicator. The observed responses are pairs $(Y^{c}:=\min\{Y,C\},\Delta)$. We make the following assumption on the counterfactual censoring times.

Assumption 4 implies that the counterfactual censoring times do not directly depend on the treatment assignment or survival times.

Let $\lambda_{0}(y;x)$ be the conditional control hazard rate at time $y$, and similarly we define the conditional treatment hazard rate $\lambda_{1}(y;x)$. We follow the Cox model (Cox, 1972) and make the proportional hazards assumption.

In this paper, we aim to estimate the heterogeneous (conditional) treatment effects, denoted by $\tau(x)$. The exact form of $\tau(x)$ depends on the type of the response.

1. 1.

   For continuous data, we use the difference of the conditional means

   $\displaystyle\tau(x):=\mathbb{E}[Y(1)\mid X=x]-\mathbb{E}[Y(0)\mid X=x].$

   (4)

2. 2.

   For binary responses, we use the difference of log conditional odds, i.e., the log of conditional odds ratio,

   $\displaystyle\tau(x):=\log\left(\frac{\mathbb{P}(Y(1)=1\mid X=x)}{\mathbb{P}(Y(1)=0\mid X=x)}\right)-\log\left(\frac{\mathbb{P}(Y(0)=1\mid X=x)}{\mathbb{P}(Y(0)=0\mid X=x)}\right).$

   (5)

3. 3.

   For count data, we use the difference of the log of conditional means, i.e., the log of conditional mean ratio,

   $\displaystyle\tau(x):=\log\left(\mathbb{E}[Y(1)\mid X=x]\right)-\log\left(\mathbb{E}[Y(0)\mid X=x]\right).$

   (6)

4. 4.

   For survival data, we use the difference of the log of conditional hazards, i.e., the log of hazards ratio,

   $\displaystyle\begin{split}\quad\tau(y;x)&:=\log\left(\lambda_{1}(y;x)\right)-\log\left(\lambda_{0}(y;x)\right)\\ &~{}=\log\left(\lim_{\delta\to 0^{+}}\frac{\mathbb{P}(Y(1)\in[y,y+\delta]\mid X=x)}{\delta\mathbb{P}(Y(1)\geq y\mid X=x)}\right)\\ &\quad~{}-\log\left(\lim_{\delta\to 0^{+}}\frac{\mathbb{P}(Y(0)\in[y,y+\delta]\mid X=x)}{\delta\mathbb{P}(Y(0)\geq y\mid X=x)}\right).\end{split}$

   (7)

   Under Assumption 5, $\tau(y;x)$ does not depend on $y$. Without further specification, we will omit $y$ and use $\tau(x)$.

If the responses indeed follow Gaussian, Bernoulli, Poisson distributions, or Cox model, then the above estimands are the differences in the treatment and control group natural parameter functions (DINA). While the conditional means are usually supported on intervals for non-continuous responses, natural parameter functions often take values over the entire real axis and are more appropriate for modeling the dependence on covariates.

Under Assumptions 1, 2, 3, the above causal estimands are identifiable. In fact, for non-survival responses,

where $\mathbb{E}[Y\mid X=x,W=1]$ is the conditional mean in the treatment group. Estimands (4), (5), and (6) are functions of $\mathbb{E}[Y(1)\mid X=x]$ and $\mathbb{E}[Y(0)\mid X=x]$, and thus estimable. For survival responses, since

the causal estimand (7) can be simplified so as not to depend on the counterfactuals and is identifiable.

In this paper, we work under the partially linear assumption (Robinson, 1988; Chernozhukov et al., 2018).

The partially linear assumption allows non-parametric nuisance functions (the natural parameter functions and the propensity score) and assumes that the heterogeneous treatment effect follows the linear model. Predictors $x$ in model (8) can be replaced by any known functions of the covariates. For instance, if the treatment effect is believed to be homogeneous, we will use $x=1$; if we are interested in the treatment effect in some sub-populations, we will design categorical predictors to specify the desired subgroups. Our method with non-parametric $\tau(x)$ is discussed in Section 9.

The semi-parametric model encodes the common belief that the natural parameter functions are more complicated than the treatment effect (Hansen, 2008; Künzel et al., 2019; Gao and Han, 2020). Consider a motivating example of hypertension: the blood pressure of a patient could be determined by multiple factors over a long period, such as the income level and living habits, while the effect of an anti-hypertensive drug is likely to interact with only a few covariates, such as age and gender, over a short period.

There is a rich literature on using flexible modeling techniques to estimate heterogeneous treatment effects. Tian et al. (2012); Imai et al. (2013) formulate the estimation of heterogeneous treatment effects as a variable selection problem and consider a LASSO-type approach. Athey and Imbens (2015); Su et al. (2009); Foster et al. (2011) design recursive partitioning methods for causal inference and (Hill, 2011) adapts the Bayesian Additive Regression Trees (BART). Ensemble learners, such as random forests (Wager and Athey, 2018) and boosting (Powers et al., 2018), have been investigated under the counterfactual framework. In addition, neural network-based causal estimators have also been proposed (Künzel et al., 2018; Shalit et al., 2017).

More recently, meta-learners for heterogeneous treatment effect estimation are of increasing popularity. Meta-learners decompose the estimation task into sub-problems that can be solved by off-the-shelf machine learning tools (base learners) (Künzel et al., 2019). One common meta-algorithm, which we call separate estimation (SE) later, applies base learners to the treatment and control groups separately and then takes the difference (Foster et al., 2011; Lu et al., 2018; Hu et al., 2021; Foster et al., 2011). Another approach regards the treatment assignment as a covariate and uses base learners to learn the dependence on the enriched set of covariates. Künzel et al. (2019) propose X-learner that first estimates the control group mean function, subtracts the predicted control counterfactuals from the observed treated responses, and finally estimates treatment effects from the differences. X-learner is effective if the control group and the treatment group are unbalanced in sample size. A different approach R-learner (Nie and Wager, 2017), motivated by Robinson’s method (Robinson, 1988), estimates the propensity score and the marginal mean function (nuisance functions) using arbitrary machine learning algorithms and then minimizes a designed loss with the estimators of nuisance function plugged in to learn treatment effects. R-learner is able to produce accurate estimators of treatment effects from less accurate estimators of nuisance functions (more details are discussed in Section 3).

In this paper, we aim to employ available predictive machine learning algorithms to quantify causal effects on the natural parameter scale. In particular, we extend the framework of R-learner which focuses on the conditional mean difference to estimate differences in natural parameters and hazard ratios. Inherited from R-learner, our proposed method uses black-box predictors and is robust to observed confounders. Beyond R-learner, our proposed method provides quantifications of causal effects that may be of more practical use and protects against non-collapsibility under non-linear link functions.

Confounders are covariates that affect both the treatment assignment and potential outcomes. Confounders are common in observational studies. We show that R-learner is insensitive to confounding while separate estimation will produce biased estimators.

We consider an example of the additive error model (9) with

and $\tau(x)=0$. We adopt the propensity score $e(x)=e^{x_{1}}/(1+e^{x_{1}})$ and assume it is known. In the example, $x_{1}$ influences both the treatment assignment and potential outcomes, and is thus a confounder.

Separate estimation fits separate models in the control and treatment groups to get $\hat{\eta}_{0}(x)$ and $\hat{\eta}_{1}(x)$, and then estimates the HTE by $\hat{\tau}_{\text{SE}}(x)=\hat{\eta}_{1}(x)-\hat{\eta}_{0}(x)$. For illustration, we estimate nuisance functions $\eta_{0}(x)$, $\eta_{1}(x)$ by linear regression. By the design of propensity score, the distributions of $x_{1}$ are different in the control and treatment groups, and the best linear approximations $\hat{\eta}_{0}(x)$, $\hat{\eta}_{1}(x)$ to $x_{1}^{2}$ are also different. Therefore, the estimator $\hat{\eta}_{1}(x)-\hat{\eta}_{0}(x)$ is far from zero and inconsistent (Figure 1 panel (a)).

In contrast, R-learner produces consistent estimators even if the estimator of the nuisance function $m(x)$ is incorrect. Continuing from the example above, the true marginal mean function is $m(x)=\eta_{0}(x)=x_{1}^{2}$. In the first step of R-learner, we also estimate the nuisance function $m(x)$ by linear regression to make the comparison fair. In the second step, we regress the residuals $Y-\hat{m}(X)$ on $(W-e(X))X$ with the true propensity score $e(x)$ provided. The residuals $Y-\hat{m}(X)$ consist of three parts: the component of treatment effect $(W-e(X))\tau(X)$, the bias of the nuisance-function estimator $m(X)-\hat{m}(X)$, and an uncorrelated noise $\varepsilon$. Since there is no treatment effect in the example, the residuals are essentially realizations of a function of covariates plus a zero-mean noise,

Notice that for any function of covariates $g(x)$,

Therefore, the residuals $Y-\hat{m}(X)$ are uncorrelated with the regressors $(W-e(X))X$ (Figure 1 panel (b)), and the regression coefficients of R-learner are asymptotically zero and consistent.

Non-collapsibility refers to the phenomenon that the mean of conditional measures does not equal the marginal counterpart. Non-collapsibility (Fienberg, 2007; Greenland et al., 1999) and confounding are two distinct issues. In other words, even if there is no confounder, the non-collapsibility issue may arise when we use non-collapsible association measures. For example, we consider randomized experiments ($e(x)=0.5$) and the exponential family model (12) with the logistic link. We assume the true natural parameter functions are

where $Z$ is an unobserved binary random variable. For constant $\tau(x)$, the true treatment effect is the log odds ratio conditional on $Z=0$ or $Z=1$, but may not equal the log of the marginal odds ratio. More generally, for natural parameter functions (11) and constant $\tau(x)$, the separate estimation method that estimates $\eta_{0}(x,z)$, $\eta_{1}(x,z)$ by generalized linear regression is only consistent for Gaussian and Poisson responses (Gail et al., 1984). For non-constant $\tau(x)$, the separate estimation method is only consistent for Gaussian (Claim 1 in the appendix). Figure 2 verifies the non-collapsibility issue numerically.

R-learner is designed to estimate the difference in conditional means but not natural parameter functions. Direct extension of R-learner to the natural parameter scale is infeasible: the response distribution conditional on the covariates $X$ is often not a member of the exponential family, and thus the marginal natural parameter function analogous to the marginal mean function $m(x)$ is not well-defined. In Section 4 and 5, we introduce our extension that inherits the robustness to confounders from R-learner and provides protection against non-collapsibility on the natural parameter scale.

We provide insights into Algorithm 1 by comparing it to R-learner. As discussed above, Algorithm 1 is an extension of R-learner, and shares its motivation as well as the method skeleton. The major difference lies in the nuisance functions (14). In R-learner, $a(x)$ equals the propensity function $e(x)$; in the proposed DINA estimator, $a(x)$ not only depends on $e(x)$ but also the variance functions $V_{0}(x)$ and $V_{1}(x)$ for the particular exponential family. Since $a(x)$ lies between 0 and 1, it can be decomposed as

$a(x)$ is large if a unit is likely to be treated or their response under treatment has higher variance than under no treatment. As a result, the $a(x)$ not only balances the treatment group and the control group, but also reduces the impact of differential variance between the two groups. For more comments on the adjustment based on $V_{w}(x)=\frac{d\mu}{d\eta}(\eta_{w}(x))$, see the multi-valued treatment setting in Section 8.

Similar to R-learner, we estimate the DINA in (15) in two steps. We estimate nuisance functions (14) and the DINA using independent subsamples so that the estimation of the nuisance functions does not interfere with that of DINA. We employ cross-fitting (Chernozhukov et al., 2018) to boost data efficiency. We now give details.

In the first step, we estimate the nuisance functions $\nu(x)$ and $a(x)$ in (14). These in turn depend on the propensity score $e(x)$, and separate estimators of $\eta_{0}(x)$ and $\eta_{1}(x)$. Given these latter two, the variance functions $V_{0}(x)$ and $V_{1}(x)$ are immediately available from the corresponding exponential family. For example for the binomial family, we have that $V_{0}(x)=\mu_{0}(x)(1-\mu_{0}(x))$, where $\mu_{0}(x)=e^{\eta_{0}(x)}/(1+e^{\eta_{0}(x)})$. The propensity scores can be estimated by any classification methods that provide probability estimates. Likewise the functions $\eta_{0}(x)$ and $\eta_{1}(x)$ can be separately estimated by any suitable method that operates within the particular exponential family.

Upon obtaining $\hat{e}(x)$, $\hat{\eta}_{0}(x)$, and $\hat{\eta}_{1}(x)$, we plug them into (14) to get $\hat{a}(x)$ and $\hat{\nu}(x)$. We remark that though we can directly use the difference of the two estimators $\hat{\eta}_{1}(x)$, $\hat{\eta}_{0}(x)$ as a valid DINA estimator — separate estimation — we show in Proposition 1 that our method yields more accurate DINA estimators. Note that in the case of the Gaussian distribution, $a(x)=e(x)$, $\nu(x)=\mathbb{E}[Y\mid X=x]$, and we can directly estimate $\nu(x)$ by any conditional mean estimator of $Y$ given $X$ and avoid having to separately estimate $\hat{\eta}_{0}(x)$ and $\hat{\eta}_{1}(x)$.

In the second step, we maximize the log-likelihood corresponding to (15) with the functions $\hat{\nu}(x)$ and $\hat{a}(x)$ fixed from the first step¹¹1If there are collinearity or sparsity patterns, penalties such as ridge and LASSO can be directly added to the loss function.. The function $\nu(x)$ is regarded as an offset, and the function $a(x)$ is used to construct the predictors $(W-a(X))X$. The second step can be implemented, for example, using the glm function in R.

From the practical perspective, Algorithm 1 decouples the estimation of DINA from that of the nuisance functions and allows for flexible estimation of nuisance functions. Though the nuisance functions $a(x)$ and $\nu(x)$ are complicated, various methods are applicable since there are no missing data in the nuisance-function estimation. By modularization, in step one, we are free to use any off-the-shelf methods, and in step two, we solve the specially designed MLE problem with the robustness protection.

The motivation behind the method is to gain robustness to nuisance functions, and we make our idea rigorous in the following proposition.

Proposition 1 states that for $\hat{\tau}(x)=x^{\top}\hat{\beta}$ to achieve a certain accuracy, the conditions on $\hat{a}(x)$ and $\hat{\nu}(x)$ are relatively weak. This implies the method is locally insensitive to the nuisance functions, and thus is robust to noisy plugged-in nuisance-function estimators. In other words, we can view Algorithm 1 as an accelerator: we input crude estimators $\hat{\eta}_{1}(x)$, $\hat{\eta}_{0}(x)$, and $\hat{e}(x)$, and Algorithm 1 outputs twice accurate DINA estimators compared to the simple difference $\hat{\eta}_{1}(x)-\hat{\eta}_{0}(x)$. As a corollary of Proposition 1, if both $\hat{a}(x)$ and $\hat{\nu}(x)$ can be estimated at the rate $n^{-1/4}$, then the DINA can be estimated at the parametric rate $n^{-1/2}$.

The following proposition illustrates a connection between the Cox model and the exponential distribution.

If there is no censoring and the hazard function $\lambda(y)$ is known, Claim 2 implies the Cox model estimation can be simplified to the case based on exponential distribution with responses $\Lambda(Y)$ and the log-link.

Now we discuss the Cox model with censoring under Assumption 4. Similar to Section 4, in Proposition 5 in the appendix, we derived a result analogous to (14),

under the reparametrization (15) to provide protection to misspecified nuisance functions. The $a(x)$ depends on the propensity score as well as the probability of being censored. Plugging the $a(x)$, $\nu(x)$ in (19) into Algorithm 1 leads to the hazard ratio estimator of the Cox model (Algorithm 2).

In Figure 4, we illustrate the efficacy of the proposed method applied to the Cox model with known baseline hazard. We consider the baseline hazard function $\lambda(y)=y$, and uniform censoring with $5\%$, $50\%$, $95\%$ censored units. The tilting functions $\eta_{0}(x)$, $\eta_{1}(x)$ are non-linear, and we approximate them by linear functions. We observe that regardless of the magnitude of censoring, the proposed method makes significant improvements over the separate estimation method. As the censoring gets heavier, $a(x)$ in (19) differs more from $e(x)$, and thus the proposed method outperforms the direct extension of R-learner by larger margins.

We discuss several special instances of (19). In particular, we consider the subcases that the censoring times satisfy $C(0)\mid X\stackrel{{\scriptstyle d}}{{=}}C(1)\mid X$, i.e., the conditional distributions of the censoring times are the same across the treatment and control arms.

1. 1.

   No treatment effect. If there is no treatment effect, then the distributions of $Y$ and $C$ are conditionally independent of the treatment assignment indicator $W$. As a result, the probability ratio of not being censored is one and $a(x)=e(x)$ — R-learner.

2. 2.

   Light censoring. If the proportion of censored units is small, the censoring time will be longer than the survival time despite the treatment effect, suggesting the ratio of not being censored to be close to one. Consequently, the case approximately reduces to R-learner.

3. 3.

   Heavy censoring. If the proportion of censored units is high, then as argued by Lin et al. (2013), almost all the information is contained in the indicator $\Delta$, and we can directly model the rare event — “not censored” by Poisson distribution. In this case,

   $\displaystyle\begin{split}\frac{\mathbb{P}(C\geq Y\mid W=1,X=x)}{\mathbb{P}(C\geq Y\mid W=0,X=x)}\stackrel{{\scriptstyle\text{c.p.}\to 1}}{{\longrightarrow}}e^{\tau(x)},\end{split}$

   (21)

   and plug (21) in (19) gives the $a(x)$ and $\nu(x)$ corresponding to a Poisson distribution.

To provide intuition for Algorithm 2, we compare the nuisance-function construction (19) with that of the direct extension of R-learner and the DINA estimator for the exponential family (Algorithm 1). Different from R-learner, the $a(x)$ in (19) also relies on the censoring probabilities under treatment and control. For a unit with covariate value $x$, the multiplier $(w-a(x))$ associated with the HTE in (15) is smaller for $w=1$ compared to $w=0$ if the unit tends to be censored under treatment. Consequently, the weight $a(x)$ emphasizes more the units not censored, which agrees with the common sense that censored units contain less information. Compared with (14), (19) does not include the derivatives $\frac{d\mu}{d\eta}(\eta_{w}(x))$ due to different parameter scales (variances): in the Cox model, we focus on the tilting functions $\eta_{w}(x)$ instead of the natural parameters $e^{\eta_{w}(x)}$ in the exponential family.

We now probe into the implementation details in Algorithm 2. Algorithm 2 is reminiscent of Algorithm 1 except in the estimation of $a(x)$, and in particular, the estimation of the censoring probabilities. We list several common random censoring mechanisms and the associated probabilities of “not censored” in Table 1. If the censoring mechanism is known, for example, all the units are enrolled simultaneously and operate for the same amount of time (singly-censored), then we can compute the censoring probabilities in closed form or via numerical integration given rough nuisance-function estimators $\hat{\eta}_{0}(x)$, $\hat{\eta}_{1}(x)$. If the censoring mechanism is unknown, we can directly estimate the censoring probabilities from the data: solving the classification problem with the response $\Delta$, predictors $(X,W)$ by any classifier that comes with predicted class probabilities. If the censoring is extremely low or high, we can turn to the aforementioned special instances.

In the following proposition, we extend the robustness property of Algorithm 1 in Proposition 1 to Algorithm 2.

To end the section, we discuss how to carry out the estimation when the baseline hazard $\Lambda(y)$ is inaccessible. One option is to start with estimating the baseline hazard, and plug it in for subsequent procedures. Another option is using the partial likelihood that cleverly avoids the baseline hazard (see Section 5.2 for more details). We remark that the proposed method is not provably robust to the baseline hazard misspecification.

Instead of focusing on the full likelihood, Cox (1972) proposes to maximize the partial likelihood. The partial likelihood is prevalent in practice because it does not require the baseline hazard function and preserves promising statistical properties (Tsiatis, 1981; Andersen and Gill, 1982).

Van der Vaart (1998) shows that the partial likelihood $\text{pl}_{n}$ can be obtained from the full likelihood (58) by profiling out the baseline hazard

where $\Lambda_{i}$ denotes the cumulating hazard at $Y_{i}^{c}$ if the subject $i$ is not censored. Let $\hat{\Lambda}_{i}$ be the baseline hazard estimators associated with the partial likelihood maximization. As a corollary, the partial likelihood maximizer is equivalent to that of the full likelihood with $\hat{\Lambda}_{i}$. The connection motivates Algorithm 3 — an application of Algorithm 2 with the true hazard baseline function replaced by $\hat{\Lambda}_{i}$ from the partial likelihood maximization.