Estimation of Partially Conditional Average Treatment Effect by Hybrid Kernel-covariate Balancing

Causal inference often concerns not only the average effect of the treatment on the outcome but also the conditional average treatment effect (CATE) given a set of individual characteristics, when treatment effect heterogeneity is expected or of interest. Specifically, let $T\in\{0,1\}$ be the treatment assignment, $0$ for control and $1$ for active treatment, $X\in\mathcal{X}\subset\mathbb{R}^{d}$ a vector of all pre-treatment confounders, and $Y$ the outcome of interest. Following the potential outcomes framework, let $Y(t)$ be the potential outcome, possibly contrary to fact, had the unit received treatment $t\in\{0,1\}$. Then, the individual treatment effect is $Y(1)-Y(0)$, and the (fully) CATE can be characterized through $\gamma(x)=\mathbb{E}\{Y(1)-Y(0)\mid X=x\}$, $x\in\mathcal{X}$. Due to the fundamental problem in causal inference that the potential outcomes are not jointly observable, identification and estimation of the CATE in observational studies require further assumptions. A common assumption is the no unmeasured confounding (UNC) assumption, requiring $X$ to capture all confounding variables that affect the treatment assignment and outcome. This often results in a multidimensional $X$. Given the UNC assumption, many methods have been proposed to estimate $\gamma(x)$ (Nie and Wager, 2017; Wager and Athey, 2018; Kennedy, 2020). However, in clinical settings, researchers may only concern the variation of treatment effect over the change of a small subset of covariates $V\in\mathcal{V}\subseteq\mathcal{X}$, not necessarily the full set $X$. For example, researchers are interested in estimating the CATE of smoking (treatment) on birth weight (outcome) given mother’s age but not mother’s educational attainment, although this variable can be a confounder. In this article, we focus on estimating $\tau(v)=\mathbb{E}\{\gamma(X)\mid V=v\}$ for $v\in\mathcal{V}$, which we refer to as the partially conditional average treatment effect (PCATE). When $V$ is taken to be $X$, $\tau(v)$ becomes the fully conditional average treatment effect (FCATE) $\gamma(x)$. Despite our major focus on cases when $\mathcal{V}$ is a proper subset of $\mathcal{X}$, the proposed method in this paper does not exclude the setting with $\mathcal{V}=\mathcal{X}$, which results in the FCATE.

When $V$ contains discrete covariates, one can divide the whole sample into different groups by constricting the same values of discrete covariates of $V$ in the same group. Then, as long as there are enough samples in such stratum, $\tau(v)$ can be obtained by estimating the PCATE over the remaining continuous covariates in $V$ separately for every stratum. Therefore, for simplicity, we focus on the setups with continuous $V$ (Abrevaya et al., 2015; Lee et al., 2017; Fan et al., 2020; Zimmert and Lechner, 2019; Semenova and Chernozhukov, 2020) while keeping in mind that the proposed method can be used to handle $V$ that consists of continuous and discrete variables. The typical estimation strategy involves two steps. The first step is to estimate nuisance parameters including the propensity score function and the outcome mean functions for the construction of adjusted responses (through weighting and augmentation) that are (asymptotically) unbiased for $\gamma(x)$ given $X=x$. The nuisance parameters can be estimated by parametric, nonparametric, or even machine learning models. This step serves to adjust for confounding biases. In the second step, existing methods typically adopt nonparametric regression over $V$ using the adjusted responses obtained from the first step. However, these methods suffer from many drawbacks. Firstly, all parametric methods are potentially sensitive to model misspecification especially when the CATE is complex. On the other hand, although nonparametric and machine learning methods are flexible, the first-step estimator of $\gamma(X)$ with high-dimensional $X$ requires stringent assumptions for the possibly low-dimensional PCATE estimation to achieve the optimal convergence rate. For example, Abrevaya et al. (2015), Zimmert and Lechner (2019), Fan et al. (2020) and Semenova and Chernozhukov (2020) specify restrictive requirements for the convergence rate of the estimators of the nuisance parameters. Detailed discussions are provided in Remarks 3 and 6.

Instead of separating confounding adjustment and kernel smoothing in two steps, we propose a new framework that unifies the confounding adjustment and kernel smoothing in the weighting step. In particular, we generalize the idea of covariate balancing weighting in the average treatment effect (ATE) estimation literature (Qin and Zhang, 2007; Hainmueller, 2012; Imai and Ratkovic, 2014; Zubizarreta, 2015; Wong and Chan, 2018). This generalization, however, is non-trivial because we require covariate balancing in terms of flexible outcome models between the two treatment groups given all possible values of $v$. We assume that the outcome models lie in the reproducing kernel Hilbert space (RKHS, Wahba, 1990), a fairly flexible class of functions of $X$. We then propose covariate function balancing (CFB) weights that are capable of controlling the balancing error with respect to the $L_{2}$-norm of any function with a bounded norm over the RKHS after kernel smoothing. The construction of the proposed weights specifically involves two kernels — the reproducing kernel of the RKHS and the kernel of the kernel smoothing — and the goal of these weights can be understood as to balance covariate functions generated by the hybrid of these two kernels. Our method does not require any smoothness assumptions on the propensity score model, in sharp contrast to existing methods, and only require mild smoothness assumptions for the outcome models. Invoking the well-known representer theorem, a finite-dimensional representation form of optimization objective can be derived and it can be solved by a gradient-based algorithm. Asymptotic properties of the proposed estimator are derived under the complex dependency structure of weights and kernel smoothing. In addition, our proposed weighting estimator can be slightly modified to incorporate the estimation of the outcome mean functions, similar to the augmented inverse probability weighting (AIPW) estimator. We show that the augmentation of the outcome models relaxes the selection of tuning parameters theoretically.

The rest of paper is organized as follows. Section 2 provides the basic setup for the CATE estimation. Section 3 introduces our proposed CFB weighting estimator, together with the computation techniques. Section 4 introduces an augmented version of our proposed estimator. In Section 5, the asymptotic properties of the proposed estimators are developed. A simulation study and a real data application are presented in Sections 6 and 7, respectively.

Suppose $\{(T_{i},Y_{i}(1),Y_{i}(0),X_{i}):i=1,\ldots,N\}$ are $N$ independent and identically distributed copies of $\{T,Y(1),Y(0),X\}$. We assume that the observed outcome is $Y_{i}=T_{i}Y_{i}(1)+(1-T_{i})Y_{i}(0)$ for $i=1,\ldots,N$. Thus, the observed data $\{(T_{i},Y_{i},X_{i}):i=1,\dots,N\}$ are also independent and identically distributed. For simplicity, we drop the subscript $i$ when no confusion arises.

We focus on the setting satisfying treatment ignorability in observational studies (Rosenbaum and Rubin, 1983).

Assumption 1 rules out latent confounding between the treatment and outcome. In observational studies, its plausibility relies on whether or not the observed covariates $X$ include all the confounders that affect the treatment as well as the outcome.

Most of the existing works (Nie and Wager, 2017; Wager and Athey, 2018; Kennedy, 2020; Semenova and Chernozhukov, 2020) focus on estimating the CATE given the full set of $X$, i.e., $\gamma(x):=\mathbb{E}\{Y(1)-Y(0)\mid X=x\}$, $x\in\mathcal{X}$, which we refer to as the FCATE. However, to ensure Assumption 1 holds, $X$ is often multidimensional, leading to a multidimensional CATE function $\gamma(x)$ that is challenging to estimate. Indeed, it is common that some covariates in $X$ are simply confounders but not treatment effect modifiers of interest. Therefore, a more sensible way is to allow the conditioning variables to be a subset of confounders (Abrevaya et al., 2015; Zimmert and Lechner, 2019; Fan et al., 2020). Instead of $\gamma(x)$, we focus on estimating the PCATE

where $V$ is a subset of $X$. It is worth noting that $V=X$ is also allowed, and therefore $\gamma(x)$ can be estimated by our framework. For simplicity, we assume $V$ is a continuous random vector for the rest of the paper. When $V$ contains discrete random variables, one can divide the sample into different strata, of which the units have the same level of discrete covariates. Then $\tau(v)$ can be estimated by estimating the PCATE at every strata.

In addition to Assumption 1, we require sufficient overlap between the treatment groups. Let $\pi(x)=\mathbb{P}(T=1\mid X=x)$ be the propensity score. Throughout this paper, we also assume that the propensity score is strictly bounded above zero and below one to ensure overlap.

Under Assumptions 1 and 2, $\tau(v)$ is identifiable based on the following formula

First suppose $\pi(X_{i}),i=1,\dots,N,$ are known. Common procedures construct adjusted responses $Z_{i}=T_{i}Y_{i}/\pi(X_{i})-(1-T_{i})Y_{i}/\{1-\pi(X_{i})\}$ and apply kernel smoother to the data $\{(V_{i},Z_{i}),i=1,\dots,N\}$. Specifically, let $K(v)$ be a kernel function and $h>0$ be a bandwidth parameter (with technical conditions specified in Section 5.1). The above strategy leads to the following estimator for $\tau(v)$:

where

In observational studies, the propensity scores $\pi(X_{i}),i=1,\dots,N$, are often unknown. Abrevaya et al. (2015) propose to estimate these scores using another kernel smoother, and construct the adjusted responses based on the estimated propensity scores. There are two drawbacks with this approach. First, it is well known that inverting the estimated propensity scores can result in instability, especially when some of the estimated propensity scores are close to zero or one. Second, this procedure relies on the propensity score model to be correctly specified or sufficiently smooth to approximate well.

To overcome these issues, instead of obtaining the weights by inverting the estimated propensity scores, we focus on estimating the proper weights directly. In the next section, we adopt the idea of covariate balancing weighting, which has been recently studied in the context of average treatment effect (ATE) estimation (e.g., Hainmueller, 2012; Imai and Ratkovic, 2014; Zubizarreta, 2015; Chan et al., 2016; Wong and Chan, 2018; Zhao et al., 2019; Kallus, 2020; Wang and Zubizarreta, 2020).

Inspired by the augmented inverse propensity weighting (AIPW) estimators in the ATE literature, we also propose an augmented estimator that directly adjusts for the outcome models $m_{1}(\cdot)$ and $m_{0}(\cdot)$.

Recall that the outcome regression functions $m_{1}(\cdot)$ and $m_{0}(\cdot)$ are assumed to be in an RKHS ${\mathcal{H}}$, kernel-based estimators $\hat{m}_{1}(\cdot)$ and $\hat{m}_{0}(\cdot)$ can be employed. We can then perform augmentation and obtain the adjusted response $Z_{i}$ in (2) as

Correspondingly, the decomposition in (4) becomes

Now, our goal is to control the difference between $N^{-1}\sum_{i=1}^{N}T_{i}\hat{w}_{i}\tilde{K}_{h}(V_{i},v)\{m_{1}(X_{i})-\hat{m}_{1}(X_{i})\}$ and $N^{-1}\sum_{i=1}^{N}\tilde{K}_{h}(V_{i},v)\{m_{1}(X_{i})-\hat{m}_{1}(X_{i})\}$. The weight estimators in Section 3.2 can be adopted similarly to control this difference. It can be shown that the term $S_{N,h}(\hat{w},m_{1}-\hat{m}_{1}):=\|\sum_{i=1}^{N}(T_{i}\hat{w}_{i}-1)\tilde{K}_{h}(V_{i},\cdot)\{m_{1}(X_{i})-\hat{m}_{1}(X_{i})/N\}\|_{2}^{2}$ can achieve a faster rate of convergence than $S_{N,h}(\hat{w},m_{1})$ does with the same estimated weights $\hat{w}$ as long as $\hat{m}_{1}$ is a consistent estimator. However, this property does not improve the final convergence rate of the PCATE estimation. This is because the term $\|{N}^{-1}\sum_{i=1}^{N}\tilde{K}_{h}(V_{i},\cdot)m_{1}(X_{i})-\mu_{1}(\cdot)\|_{2}^{2}$ dominates other terms, and thus the final rate can never be faster than the optimal non-parametric rate. See Remark 4 for more details. Our theoretical results reveal that the benefit of using the augmentations lies in the relaxed order requirement of the tuning parameters to achieve the optimal convergence rate. Therefore, the performance of the augmented estimator is expected to be more robust to the tuning parameter selection.

Unlike other AIPW-type estimators (Lee et al., 2017; Fan et al., 2020; Zimmert and Lechner, 2019; Semenova and Chernozhukov, 2020) which often rely on data splitting for estimating the propensity score and outcome mean functions to relax technical conditions, our estimator does not require data splitting to facilitate the convergence with augmentation. See also Remark 7. We defer the theoretical comparison between our estimator and the existing AIPW-type estimators in Remark 6 in Section 5.

Last, we note that there are existing work using weights to balance the residuals (e.g. Athey et al., 2016; Wong and Chan, 2018), which is similar to what we consider here for the proposed augmented estimator. These estimators are designed for ATE estimation and the balancing weights cannot be directly adopted here with theoretical guarantee.

In this section, we conduct an asymptotic analysis for the proposed estimator. For simplicity, we assume $\mathcal{X}=[0,1]^{d}$. To facilitate our theoretical discussion in terms of smoothness, we assume the RKHS $\mathcal{H}$ is contained in a Sobolev space (see Assumption 3). Our results can be extended to other choices of ${\mathcal{H}}$ if the corresponding entropy result and boundedness condition for the unit ball $\{u\in{\mathcal{H}}:\|u\|_{\mathcal{H}}\leq 1\}$ are provided. Recall that we focus on $\mathbb{E}\{Y(1)\mid V=v\}$. Similar analysis can be applied to $\mathbb{E}\{Y(0)\mid V=v\}$ and finally the PCATE.