Transfer learning of individualized treatment rules from experimental to real-world data

A personalized recommendation, tailored to individual characteristics, is becoming increasingly popular in real life such as healthcare, education, e-commerce, etc. An optimal individualized treatment rule (ITR) is defined as maximizing the mean value function of an outcome of interest over the target population if all individuals in the whole population follows the given rule. Many machine learning approaches to estimating the optimal ITR are available, such as outcome-weighted learning with support vector machine (Zhao et al., 2012; Zhou et al., 2017), regression-based methods with Adaboost (Kang et al., 2014), regularized linear basis function (Qian and Murphy, 2011), K-nearest neighbor (Zhou and Kosorok, 2017) and generalized additive model (Moodie et al., 2014). However, the ITRs derived by machine learning methods can be too complex to extract practically meaningful insights for policymakers. For example, in healthcare, clinicians need to scrutinize the estimated ITRs for scientific validity but the black-box rules conceal the clear relationship between patient characteristics and treatment recommendations. In these cases, parsimonious and interpretable ITRs are desirable.

It is well known that randomized experiments are the gold standard for learning the optimal ITRs because randomization of treatments in the design stage can ensure no unmeasured confounding (Greenland, 1990). However, randomized experiments might lack external validity or representativeness (Rothwell, 2005) because of the selective criteria for individual eligibility in the experiments. Consequently, the distribution of the covariates in the randomized experiments might differ from that in the target population. Thus, a parsimonious ITR constructed from randomized experiments cannot directly generalize to a target population (Zhao et al., 2019). Alternatively, real-world data (RWD) such as population survey data usually involve a large and representative sample to the target population. A large body of literature in statistics has focused on transporting or generalizing the average treatment effects from the experiments to a target population or a different study population (e.g., Cole and Stuart, 2010; Hartman et al., 2015). One exception for the ITRs is Zhao et al. (2019); Mo et al. (2020) which estimates the optimal linear ITR from the experimental data by optimizing the worst-case quality assessment among all covariate distributions in the target population satisfying some moment conditions or distributional closeness, while we address the problem in a different situation where we already have a representative sample at hand.

In this article, we propose a general framework of estimating a generalizable interpretable optimal ITR from randomized experiments to a target population represented by the RWD. We propose transfer weighting that shifts the covariate distribution in the randomized experiments to that in the RWD. To correctly estimate the population optimal ITR, transfer weights are used in two places. First, the empirical estimator of the population value function is biased when considering only the experimental sample, and hence the transfer weights are used to correct the bias of the estimator of the value function. Second, the nuisance functions in the value function such as the propensity score and the outcome mean also require transfer weighting in order to estimate the population counterparts. Then we can obtain the weighted interpretable optimal ITR by reformulating the problem of maximizing the weighted value to that of minimizing the weighted classification error with flexible loss function choices. We consider different parametric and nonparametric approaches to learning transfer weights. Moreover, a practical issue arises regarding the choice of transfer weights. Weighting can correct for the bias but also reduce the effective sample size and increase the variance of the weighted value estimator. To address the bias-variance tradeoff, we propose a cross-validation procedure that maximizes the estimated population value function to choose among the weighting schemes. Lastly, we provide the theoretical guarantee of the risk consistency of the proposed transfer weighted ITRs.

The remaining of the article is organized as follows. We introduce our proposed method in Section 3. In Section 4, we study the asymptotic properties of the proposed estimator to guarantee its performance in the target population theoretically. In Section 5, we use various simulation settings to illustrate how this approach works. In Section 6, we apply the proposed method to the National Supported Work program (LaLonde, 1986) for improving the decisions of whether to recommend the job training program based on the individual characteristics over the target population. We conclude the article in Section 7.

In this section, we provide a theoretical guarantee that the estimated ITR within a certain class $\mathcal{F}$ is consistent for the true population optimal ITR within $\mathcal{F}$. Toward this end, we show that the expected risk of the estimated ITR is consistent for that of the true optimal ITR within $\mathcal{F}$. In particular, we consider the ITRs learned by the nonparametric transfer weights, weighted AIPW estimator $\hat{\tau}_{\text{aw}}^{w}$ with the logistic regression model $\pi_{A}(\bm{X};\gamma)$ and the linear $Q$ function. The extensions to other cases with parametric transfer weights and regression and IPW estimators are straightforward.

We introduce more notation. Let $\theta=(\mathbf{\gamma^{\intercal},}\beta^{\intercal})^{\intercal}$ and let the contrast function

Denote the general smooth ramp loss function $l^{\zeta}(u)=l(\zeta_{1}u)/\zeta_{2}$, where $\zeta_{1},\zeta_{2}>0$, $l$ is defined in (8), so it is easy to see that when $\zeta=(\zeta_{1},\zeta_{2})\rightarrow\zeta^{*}\overset{\mathrm{\Delta}}{=}(+\infty,2)$, the general smooth ramp loss function will converge to the 0-1 loss function, i.e., $l^{\zeta}(u)\rightarrow l_{0-1}(u),\forall u$. Note that the general smooth ramp loss function can also be written as the difference of two convex functions, $l^{\zeta}(u)=l_{1}(\zeta_{1}u)/\zeta_{2}-l_{0}(\zeta_{1}u)/\zeta_{2}$, thus the d.c. algorithm in Algorithm 1 can be easily generalized to $l^{\zeta}(u)$. Recalling the loss function $\mathcal{\mathcal{L}}(\tau;\eta,l_{0-1})$ and the corresponding $\mathcal{F}$-risk $\mathcal{R}_{\mathcal{F}}(\tau;\eta,l_{0-1})$ defined in (3), we replace $\tau$ with $\tau_{\text{aw}}^{\theta}$ and obtain:

denoted as $\mathcal{R}_{\mathcal{F}}(\eta)$ and $\mathcal{L}(\eta,\theta)$ respectively for simplicity. We define the minimal $\mathcal{F}$-risk as $\mathcal{R}_{\mathcal{F}}^{*}=\inf_{\eta\in\mathcal{F}}\mathcal{R}_{\mathcal{F}}(\eta)$. Similarly, replacing $l_{0-1}$ in $\mathcal{R}_{\mathcal{F}}(\eta)$ and $\mathcal{L}(\eta,\theta)$ with the general smooth ramp loss $l^{\zeta}$, we can obtain $\mathcal{R}_{\mathcal{F}}(\tau_{\text{aw}}^{\theta};\eta,l^{\zeta})$ and $\mathcal{\mathcal{L}}\left(\tau_{\text{aw}}^{\theta};\eta,l^{\zeta}\right)$, denoted as as $\mathcal{R}_{\zeta\mathcal{F}}(\eta)$ and $L^{\zeta}(\eta,\theta)$, respectively. Likewise, the minimal $\mathcal{F}$-risk with $l^{\zeta}$ is defined as $\mathcal{R}_{\zeta\mathcal{F}}^{*}=\inf_{\eta\in\mathcal{F}}\mathcal{R}_{\zeta\mathcal{F}}(\eta)$.

The goal is to show that the decision rule is $\mathcal{F}$-consistency (Bartlett et al., 2006), i.e.,

where $\hat{\eta}=\arg\min_{\eta\in\mathcal{F}}\sum_{i=1}^{N}S_{i}\hat{w}_{i}L^{\zeta}(\eta,\hat{\theta}_{w})$, where $\hat{\theta}_{w}=(\hat{\gamma}_{w}^{\intercal},\hat{\beta}_{w}^{\intercal})^{\intercal}$, the estimated logistic regression coefficients $\hat{\gamma}_{w}=\arg\max_{\gamma}\sum_{i=1}^{N}\hat{w}_{i}S_{i}[A_{i}\log\pi_{A}(\bm{X}_{i};\gamma)+(1-A_{i})\log\left\{1-\pi_{A}(\bm{X}_{i};\gamma)\right\}]$ which is converge in probability to $\arg\max_{\gamma}\mathbb{E}[A\log\pi_{A}(\bm{X};\gamma)+(1-A_{i})\log\left\{1-\pi_{A}(\bm{X};\gamma)\right\}]$ $\overset{\mathrm{\Delta}}{=}\gamma^{*}$ as $N\rightarrow\infty$ under regular conditions, and the estimated outcome regression coefficients $\hat{\beta}_{w}=\arg\min_{\beta}\sum_{i=1}^{N}\hat{w}_{i}S_{i}(Y_{i}-\beta^{\intercal}L_{i})^{2}$, where $L_{i}=(1,\bm{X}_{i}^{\intercal},A_{i},A_{i}\bm{X}_{i}^{\intercal})$, thus we have

in probability as $N\rightarrow\infty$. The convergences of $\hat{\gamma}_{w}$ and $\hat{\beta}_{w}$ are implied from the weak law of large numbers and the result $\max_{i}\mid N\hat{w}_{i}-1/\pi_{S}(\bm{X}_{i})\mid=o_{p}(1)$ proved in Theorem 2 in Wang and Zubizarreta (2020). Let $\theta^{*}=(\gamma^{*\intercal},\beta^{*}{}^{\intercal})^{\intercal}$, thus we have

It is easy to see that $\mathcal{R}_{\zeta\mathcal{F}}(\eta)$ will converge to $\mathcal{R}_{\mathcal{F}}(\eta)$ when $\zeta\rightarrow\zeta^{*},$ therefore, to obtain $\mathcal{F}$-consistency, it is suffice to show that $\lim_{N\rightarrow\infty}\mathcal{R}_{\zeta\mathcal{F}}(\hat{\eta})=\mathcal{R}_{\zeta\mathcal{F}}^{*}$ which is stated in the theorem below.

We evaluate the finite-sample performances of the proposed estimators with a set of simulation studies. We first generate a target population of a size $N=10^{6}$. The covariates are generated by $\bm{X}_{i}=(X_{i1},X_{i2})^{\intercal}\sim N(\mathbf{1}_{2\times 1},I_{2\times 2})$, and the potential outcome is generated by $Y_{i}^{*}(a)\mid\bm{X}_{i}=1+2X_{i1}+3X_{i2}+a\tau(\bm{X}_{i})+\epsilon_{i}(a),$ where $\epsilon_{i}(a)$ i.i.d. follows $N(0,0.5^{2})$ for $a=0,1.$ We consider three contrast functions: (I). $\tau(\bm{X}_{i})=\arctan\{\exp(1+X_{i1})-3X_{i2}-5\};$ (II). $\tau(\bm{X}_{i})=\cos(1)+\cos(X_{i1})+\cos(X_{i2})-3/2;$ (III). $\tau(\bm{X}_{i})=1+2X_{i1}+3X_{i2}$.

In order to evaluate the necessity and performances of weighting methods, we consider two categories of data generating models. One is the scenarios where the true optimal policies are not in linear forms , i.e., the settings (I) and (II), while the other one is considering the true optimal policy is in a linear form as the setting (III). We start with the first two scenarios. From the data visualization, we observed that the covariates shifted between the experimental data and RWD to a varying extent in the two generative models: the covariates generated by (I) have similar distributions between the experimental data and RWD, while those generated by (II) have very different ones. Then we generate the RWD by randomly sample $m=5000$ subjects from the target population $\{\bm{X}_{i}\}_{i=1}^{N}$, denoted as $\{\bm{X}_{i}\}_{i\in\mathcal{I}_{m}}$. To form the experimental sample, we generate the indicator of selection according to $S_{i}\mid\bm{X}_{i}\sim\text{Bernoulli}\left[\exp(\alpha_{0}+\alpha_{1}^{{}^{\intercal}}\bm{X}_{i})/\{1+\exp(\alpha_{0}+\alpha_{1}^{\intercal}\bm{X}_{i})\}\right]$, where $\alpha_{1}=(1,-2)^{\intercal}$ and $\alpha_{0}=-8$ which gives the average sample sizes of the experimental data over the simulation replicates round 1386, respectively. In the experimental data denoted as $\mathcal{I}_{n}$, the treatments are randomly assigned according to $A_{i}\sim\mathrm{Bernoulli}(0.5)$, and then the actual observed outcomes $Y_{i}=Y_{i}^{*}(A_{i})$, $i\in\mathcal{I}_{n}$.

For each setting, we compare the following estimators:

1. “$w_{1}$”: our proposed weighted estimator with parametric weights learned by the modified maximum likelihood introduced in Section 3.2.1.

2. “$w_{2}$”: our proposed weighted estimator with parametric weights learned by the estimating equations introduced in Section 3.2.1.

3. “cv”: our proposed weighted estimator using cross-validation procedure introduced in Section 3.4, and we set the number of sample splits as 10 in all the simulation studies.

4. “np”: our proposed weighted estimator with nonparametric weights.

5. “unweight”: only using the experimental data to learn the optimal linear ITRs.

6. “bm”: using the weighted estimator but we replace the estimated transfer weights $\hat{w}_{i}$ with the normalized true weights $w_{i}=\{\pi_{S}(\bm{X}_{i})\}^{-1}/\sum_{i\in\mathcal{I}_{n}}\{\pi_{S}(\bm{X}_{i})\}^{-1}$, and replace the estimated contract functions $\hat{\tau}_{i}$ with the true contrast values $Y_{i}^{*}(1)-Y_{i}^{*}(0)$.

7. “Imai”: the method proposed in Imai et al. (2013). They construct the transfer weights by fitting a Bayesian additive regression trees (BART) model using covariates as predictors and whether to be in the experimental data as outcomes, and then use a variant of SVM to estimate the conditional treatment effect (implemented in the R package FindIt). Then given $\bm{X}$, we can assign treatment $1$ if its estimated conditional treatment effect is positive, and $0$ otherwise.

8. “DBN”: the method proposed in Mo et al. (2020), where they propose to use distributionally robust ITRs, that is, value functions are evaluated under all testing covariate distributions that are ”close” to the training distribution, and the worse-case takes a minimal one, and then the ITRs are learned by maximizing the defined worst-case value function.

For the estimators 1. – 5. in the above, we estimate the contrast function using $\tau_{\text{aw}}$ with linear $Q$ functions and constant propensity scores estimated via the proportion of $A=1$ in the experimental data, $n^{-1}\sum_{i\in\mathcal{I}_{n}}A_{i}$.

To study the impact of model misspecification of $\pi_{S}(\bm{X})$, we consider to estimate $\alpha$ by fitting logistic regression wrongly: $\pi_{S}(\bm{X};\alpha)=\exp(\alpha_{0}+\alpha_{1}^{\intercal}\widetilde{\bm{X}})/\{1+\exp(\alpha_{0}+\alpha_{1}^{\intercal}\widetilde{\bm{X}})\}$, where $\widetilde{\bm{X}}=(X_{1}^{2},X_{2}^{2})$. And as described in the proposed methods, only estimators 1. – 3. need to specify the sampling model. To measure the performances of the different approaches, we use the mean value MSE defined as $N^{-1}\sum_{i=1}^{N}\left(Y_{i}^{*}\{d_{\eta}(\bm{X}_{i})\}-Y_{i}^{*}[\mathbf{1}\{\tau(\bm{X}_{i})>0\}]\right)^{2}$. Thus the lower the MSE, the better performances the approach has.

The results are summarized in Figure 1, the dark red point in each boxplot is the mean value MSE averaged over 200 Monte Carlo replicates. For the generative model (I) shown in Figure 1, unweighted and DBN rules perform worst in all the cases. The second parametric weighting method is relatively robust compare with the first parametric one when the sampling model is misspecified; this is reasonable since the first one uses maximum likelihood estimation which is more reliable on the correctness of the parametric model specification. When the training data and testing data are similarly distributed, the performances of parametric, nonparametric weighting method and the cross-validation procedure are similar; in this case, sampling model misspecification does not have much influence on the rule learning.

For the generative model (II) shown in Figure 1, among the performances of weighted and unweighted estimators, we can see the weighted estimators outperform the unweighted estimators when model correctly specified; the nonparametric weighted estimator outperforms other estimators and the performance of the parametric weighted estimators have dramatically decreased when the sampling model misspecified. These findings make sense, because when the more different the training and testing data are, the more crucial the bias correction is. Thus sampling model misspecification might have huge influence on the learned rules.

From the both generative models, we can see that the nonparametric weighting and cross-validation procedure are more robust when the sampling model is unknown and recommended in practice use. As for DBN estimator, it performs badly in terms of MSE magnitude or variability, one possible reason might be that the tuning parameters when measuring the ”closeness” of two distributions need to be chosen more carefully in practice Mo et al. (2020). As for the comparison with Imai also has relative good performance among all estimators and has relative small variability.

As for setting (III), we consider the correctly specified sample score model. As shown in Figure 2, unweighted method performances the best, close to the benchmark. This demonstrates that when the true optimal polices are in linear forms, then learning with only the experimental data performances the best. Thus weighting is not a necessity, and it even will induce the variance of the estimators. The other competitive estimator is using cross-validation. Therefore, when we have no idea of the true data generating models, using CV is a good choice in practice.

We apply the proposed methods to personalized recommendations of a job training program on post-market earning. Our analysis is based on two data sources (downloaded from https://users.nber.org/~rdehejia/data/.nswdata2.html): an experimental dataset from the National Supported Work (NSW) program (LaLonde, 1986) and a non-experimental dataset from the Current Population Survey (CPS; LaLonde (1986)). The NSW dataset includes $185$ individuals who received the job training program ($A=1$) and $260$ individuals who did not receive the job training program ($A=0$), and the CPS dataset includes $15992$ individuals without the job training program. Let $\bm{X}$ be a vector including an intercept $1$ and the pre-intervention covariates: age, education, Black (1 if black, 0 otherwise), Hispanic (1 if Hispanic, 0 otherwise), married (1 if married, 0 otherwise), whether having high school degree (1 if no degree, 0 otherwise), earning in 1974, and earning in 1975. The outcome variable is the earning in 1978. We transfer the earnings in the year 1974, 1975, and 1978 by taking the logarithm of the earnings plus one, denoted as “log.Re74”, “log.Re75” and “log.Re78”. We aim to learn a linear rule $d(\bm{X};\eta)=\mathbf{1}(\bm{X}^{\intercal}\eta>0)$ tailoring the recommendation of the job training program to an individual with characteristics $\bm{X}$.

The CPS sample is generally representative of the target population; while the participants to the NSW program may not represent the general population. Table 1 contrasts sample averages of each pre-intervention covariate in the NSW and CPS samples and shows that the covariates are highly imbalanced between the two samples. On average, the NSW sample is younger, less educated, more black people, fewer married, more without high school degree, fewer earnings in 1974 and 1975. Because of imbalance, ITRs cannot be learned based only on the NSW sample, thus we are motivated to use the proposed weighted transfer learning approaches to learn optimal linear ITRs which benefit the target population most.

We consider the unweighted and weighted transfer learning methods. For the weighted transfer learning method, we consider the nonparametric weighting approach, because the parametric approach requires the target population size $N$ which is unknown in this application. To estimate the weighted and unweighted decision rules, we use the AIPW estimator $\hat{\tau}_{\text{aw}}$ with linear Q functions since it outperforms other competitors in the simulation study in Section 5. We also implement the cross-validation procedure to choose between the unweighted and weighted transfer learning methods.

Table 1 presents the estimated linear coefficients of each corresponding variable in the linear decision rule from the weighted transfer and unweighted learning methods, respectively. Moreover, from the cross-validation procedure, the weighted transfer learning method produces a smaller population risk than the unweighted learning method and therefore is chosen as the final result. Compared with the learned decision rule with only the experiment data NSW which oversamples the low-income individuals, individuals without high school degrees (i.e. Nodegr = 1) are more likely to have higher outcomes if offered the job training program in the target population.

To assess the performances of the learned rules for the target population, following Kallus (2017), we create a test sample that is representative of the target population. We sample $100$ individuals at random from the CPS sample to form a control group and then find their closest matches in the NSW treated sample to form the treatment group (use pairmatch in R package optmatch). The test sample thus contains the $100$ matched pairs, mimicking a experimental sample where binary actions are randomly assigned with equal probabilities. We use random forest on the test data to estimate $Q$ function $\mathbb{E}(Y\mid\bm{X},a)$, denoted the estimator as $\widehat{Q}_{\text{test}}(\cdot,a),$ $a=0,1$, then we can use $\sum_{i\in\text{CPS}}\widehat{Q}_{\text{test}}\left\{\bm{X}_{i},d(\bm{X}_{i};\eta)\right\}/\mid\text{CPS}\mid$ to estimate the value of an ITR $d(\cdot;\eta)$, where $\mid\text{CPS}\mid$ is the sample size of CPS. Besides the weighted linear rule $d(\cdot;\hat{\eta}_{w})$ and the unweighted linear rule $d(\cdot;\hat{\eta}_{u})$ learned from the above, we also compare the Imai method introduced in Imai et al. (2013): first to estimate the conditional probability of being selected into NSW given $\bm{X}$ from the total sample which consists of NSW and CSP using BART model, and then to estimate the conditional treatment effect weighted by the inverse of the sampling probability. Then the Imai rule will recommend the job training program if the conditional treatment effect is positive and will not recommend it otherwise. We replicate this sampling test data procedure 100 times, estimate values of the three ITRs for each replicate, and get the averaged estimated values (standard error) as follows: 8.395 (0.009), 6.106 (0.011), and 8.049 (0.006) for $d(\cdot;\hat{\eta}_{w})$, $d(\cdot;\hat{\eta}_{u})$, and Imai respectively, thus the transfer weighted rule has the highest value estimate. Although this evaluation method needs assumption of no unmeasured confounders for the non-experimental data CPS, which is not needed in our learning procedure, it can still bring us confidence, to some degree, that the weighted linear rule has better generalizability to the target population.