Discussion of Kallus (2020) and Mo et al (2020)

Abstract

We discuss the results on improving the generalizability of individualized treatment rule following the work in Kallus [1] and Mo et al. [5]. We note that the advocated weights in Kallus [1] are connected to the efficient score of the contrast function. We further propose a likelihood-ratio-based method (LR-ITR) to accommodate covariate shifts, and compare it to the CTE-DR-ITR method proposed in Mo et al. [5]. We provide the upper-bound on the risk function of the target population when both the covariate shift and the contrast function shift are present. Numerical studies show that LR-ITR can outperform CTE-DR-ITR when there is only covariate shift.

1. Introduction

The problem of constructing individualized treatment rules (ITRs), a function that maps patient characteristics to an available treatment, has recently received significant attention among statistical researchers [6, 11, 8, 12, 4, 9]. These methods typically assume that the population of the training samples and the target population where the ITR will be implemented in the future are identical. However, when these two populations are different from each other, the estimated ITR may perform poorly on the target population [10]. We congratulate Kallus (2020) and Mo, Qi, and Liu (2020) on their contributions in proposing robust ITR estimation approaches where there are covariate changes between the training and target populations, also known as covariate shift.

Let (X, A,Y) be the triplet, where X denotes patients covariates, A denotes the treatments, and Y denotes the outcome. Suppose that the treatment space is $\mathcal{A}=\left\{1,-1\right\}$. Let Y(1) and Y(−1) be the potential outcomes given treatment A = 1 and −1, respectively. Let $\mathbb{P}$ be the distribution of the training population and ${\mathbb{P}}_{\text{test}}$ be the distribution of the testing population. Covariate shift assumes that the conditional distribution of (Y(1),Y(−1)) | X on the training and target populations are the same, but the covariate distributions may be different. Consequently, given X, the contrast functions satisfy ${\mathbb{E}}_{\mathbb{P}}[Y(1)-Y(-1)\mid X]={\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[Y(1)-Y(-1)\mid X]$. With the presence of covariate shifts, Kallus [1] defines the retargeted policy value and proposes a weighting approach that minimizes the efficient variance in estimating the finite-sample objective. Mo et al. [5] maximizes the worst value function over a set of possible weights.

We point out that the weights advocated in Kallus [1] to adjust treatment-control overlapping is related to the efficient score of estimating the contrast function. Therefore, his work can be considered as improving the efficiency of estimating ITR. Different from Kallus [1], Mo et al. [5] focus on controlling the possible bias when covariate distribution shifts. We propose an alternative method, termed as likelihood-ratio weighted ITR (LR-ITR), which re-weights the learning objective with the directly estimated likelihood ratio from the training and target data. Numerical results show that the LR-ITR method can improve the performance of the standard ITR approach, and outperform CTE-DR-ITR proposed in Mo et al. [5] when there is only covariate shift. On the other hand, CTE-DR-ITR method can be more generalizable when there is contrast shift in addition to covariate shift, i.e., ${\mathbb{E}}_{\mathbb{P}}[Y(1)-Y(-1)\mid X]\ne {\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[Y(1)-Y(-1)\mid X]$.

This article is organized as follows. In Section 2, we discuss the connection between the weights in Kallus [1] and the efficient estimation of the contrast function [3]. In Section 3, we introduce the LR-ITR approach, along with its theoretical properties, and conduct simulation studies. In Section 4, we briefly summarize our discussion.

2. Retargeting weights and efficient estimation

The weighting approach for adjusting treatment-control overlapping is closely connected with the efficient score of the contrast functions. Given a policy π(a | X), which is a distribution on the treatment space $\mathcal{A}$ given X, the retargeted learning objective in Kallus [1] is defined as

$$R(\pi ;w,\rho )=\mathbb{E}\left[w(X)\sum _{a\in A}(\pi (a\mid X)-\rho (a\mid X))\mu (a\mid X)\right],$$ (2.1)

where $\mu (a\mid X)=\mathbb{E}[Y(a)\mid X]$. Kallus [1] advocates to use ρ(a | X) = 1/2 and w(X) ∝ [Var{Y(1) | X}/P(A = 1 | X) + Var{Y(−1) | X} / P(A = −1 | X)]⁻¹. The proposed choice of ρ(a | X) = 1/2 compares the value function under the policy π with the value function under the pure randomization.

We highlight that the retargeting weight w(X) is connected with the efficient estimation of the contrast function, Δ(X) = μ(1 | X) – μ(−1 | X). In Liang and Yu [3], it is shown that assuming Δ(X) = g(X^⊤ β), where g is an unknown function, the efficient score for β is

$$\tilde{w}(X)\frac{A}{P(A\mid X)}\nabla g\left(X^{\top }\beta \right)\left\{X-\frac{\mathbb{E}\left[\tilde{w}(X)X\mid X^{\top }\beta \right]}{\mathbb{E}\left[\tilde{w}(X)\mid X^{\top }\beta \right]}\right\}\left\{ϵ-\mathbb{E}[ϵ\mid X]\right\},$$

where ϵ = Y − Ag(X^⊤ β) / 2 and $\tilde{w}(X)\propto {\left(\mathbb{E}\left[{ϵ}^{2}P{(A\mid X)}^{-2}\mid X\right]-\mathbb{E}{[ϵ\mid X]}^2\mathbb{E}\left[P{(A\mid X)}^{-2}\mid X\right]\right)}^{-1}$. Since ϵ can be written as $ϵ=S+\widetilde{ϵ}$, where $[\widetilde{ϵ}\mid X,A]=0$, we have

$$\begin{gathered} \mathbb{E}\left[\frac{{ϵ}^2}{{[P(A\mid X)]}^2}\mid X\right]-\mathbb{E}{\left[ϵ|X\right]}^2\mathbb{E}\left[\frac{1}{{[P(A\mid X)]}^2}\mid X\right] \\ =\mathbb{E}\left[\frac{{(S+\widetilde{ϵ})}^2}{{[P(A\mid X)]}^2}\mid X\right]-S^2\left\{\frac{1}{P(A=1\mid X)}+\frac{1}{P(A=-1\mid X)}\right\} \\ =\mathbb{E}\left[\frac{{\widetilde{ϵ}}^2}{{[P(A\mid X)]}^2}\mid X\right] \\ =\frac{\mathbb{E}\left[{\widetilde{ϵ}}^2\mid A=1,X\right]}{P(A=1\mid X)}+\frac{\mathbb{E}\left[{\widetilde{ϵ}}^2\mid A=-1,X\right]}{P(A=-1\mid X)}. \end{gathered}$$

Given that $\mathbb{E}\left[{\widetilde{ϵ}}^2\mid A=1,X\right]=\text{Var}(Y(1)\mid X)$ and $\mathbb{E}\left[{\widetilde{ϵ}}^2\mid A=-1,X\right]=\text{Var}(Y(-1)\mid X)$, $\tilde{w}(X)$ is also proportional to [Var{Y(1) | X} / P(A = 1 | X) + Var{Y(−1) | X} / P(A = −1 | X)]⁻¹. As such, the adjustment for treatment-control overlapping by the weights in Kallus [1] can also benefit the estimation of the contrast function.

3. Likelihood-ratio weighted ITR

Mo et al. [5] propose targeting a class of populations rather than a specific population to learn a distributionally robust individualized treatment rule (ITR). We propose an alternative likelihood-ratio-weighted approach. Following Mo et al. [5], we consider the value function $V(d)={\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[C(X)d(X)]$ given an ITR d(X). Under Assumption 1 (Covariate Changes) in Mo et al. [5], we have

$$\frac{\text{d}{\mathbb{P}}_{\text{test}}}{\text{d}\mathbb{P}}=\frac{q_x(X)}{p_x(X)},$$

where q_x and p_x are the density functions of the target population and the training population, respectively. Let w*(X) = q_x(X) / p_x(X). The target value function, V(d), can be reformulated as

$$V(f)={\mathbb{E}}_{\mathbb{P}}\left[w^{*}(X)C(X)d(X)\right].$$

3.1. Algorithm

We estimate w*(X) by using the covariates information from the training population and the target population. Unconstrained least-squares importance fitting (uLSIF) is an efficient algorithm to estimate the likelihood ratio [2], which has a closed-form solution and can be calculated by solving a linear system. Let $\gamma (X)=\sum _{l=1}^b{\alpha }_l{\varphi }_l(X)$, where ϕ_l(X)’s are pre-specified basis functions and b is the number of the basis. Notice that ${\mathbb{E}}_{\mathbb{P}}[\gamma (X)]=1$. The uLSIF aims to minimize ${\mathbb{E}}_{\mathbb{P}}\left[{\left(\gamma (X)-w^{*}(X)\right)}^2\right]$, and

$${\mathbb{E}}_{\mathbb{P}}\left[{\left(\gamma (X)-w^{*}(X)\right)}^2\right]={\mathbb{E}}_{\mathbb{P}}\left[{\gamma }^2(X)\right]-2{\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[\gamma (X)]+C,$$ (3.1)

where C is a constant irrelevant to ${{\alpha }^{\prime }}_l\text{s}$. Let α = (α₁,⋯, α_b)^⊤. (3.1) can be written as

$$\underset{\alpha }{\text{min}}{\alpha }^{\top }\widehat{\Psi }\alpha -2{\alpha }^{\top }h+\lambda \Vert \alpha {\Vert }_2^2,$$ (3.2)

where $\widehat{\Psi }$ is a b × b matrix with its (i, j)th coordinate, ${\widehat{\Psi }}_{i,j}={\mathbb{E}}_{\mathbb{P}}\left[{\varphi }_i(X){\varphi }_j(X)\right]$, and h is a b-dimensional vector with its ith coordinate, ${\widehat{h}}_i={\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}\left[{\varphi }_i(X)\right]$. The ${\mathbb{E}}_{\mathbb{P}}[\cdot ]$ and test ${\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[\cdot ]$ are the empirical expectations defined by the samples from the training and target populations. Typically, we only need a small calibration set from the target population to calculate any ${\widehat{h}}_i$’s. The uLSIF estimator is

$$\widehat{\gamma }(X)=\sum _{l=1}^b\text{max}\left\{0,{\widehat{\alpha }}_l\right\}{\varphi }_l(X),$$

where $\alpha ={\left({\widehat{\alpha }}_1,\cdots ,{\widehat{\alpha }}_b\right)}^{\top }$ minimizes (3.2). To stabilize the estimation of $\widehat{\gamma }(X)$, we cap $\widehat{\gamma }(X)$ at 10. This choice of cap corresponds to the bound of w*(X) assumed in the Mo et al. [5] simulations. The implementation procedure is outlined in Algorithm 1. We call this method as the likelihood-ratio weighted ITR (LR-ITR).

Algorithm 1: Likelihood-ratio weighted ITR.

Input: n samples from the training population (X, A,Y) and n_calib samples from the target population including only covariates.

Output: An ITR $\widehat{d}(X)=\text{sgn}\left\{X^{\top }\beta \right\}$.

1. Estimate $\widehat{\gamma }(X)=\sum _{l=1}^b\text{max}\left\{0,{\widehat{\alpha }}_l\right\}{\varphi }_l(X)$ where α is the minimizer of

   $$\underset{\alpha }{\text{min}}{\alpha }^{\top }\widehat{\Psi }\alpha -2{\alpha }^{\top }h+\lambda \Vert \alpha {\Vert }_2^2,$$

   where ϕ_l (·)’s are chosen as Gaussian kernels with kernel bandwidth σ. The parameter σ and λ are tuned by cross-validation. The n_calib samples from the target population are used to calculate h;

2. Obtain an estimator $\widehat{C}(X)$ by fitting a causal forest [7] using the training data;

3. Obtain the linear decision rule $\widehat{d}(X)=\text{sgn}\left\{X^{\top }\beta \right\}$, where β minimizes

   $${\mathbb{E}}_{\mathbb{P}}\left[\widehat{\gamma }(X)\widehat{C}(X)\left\{\psi \left(X^{\top }\beta \right)-1\right\}\right],$$

   and ψ is the robust smoothed ramp loss [12].

3.2. Risk Bounds on the Target Population

In this section, we provide risk bounds for the risk functions on the target population via different methods. Define the risk function on the target population as

$$L_{{\mathbb{P}}_{\text{test}}}(d(X))={\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}\left[C_{\text{test}}(X)(-d(X))\right],$$

where $C_{\text{test}}(X)={\mathbb{E}}_{{\mathbb{P}}_{\text{test}}}[Y(1)-Y(-1)\mid X]$. The proposed LR-ITR method minimizes the LR-risk function, defined as

$$L_{LR}(d(X))={\mathbb{E}}_{\mathbb{P}}\left[w^{*}(X)C(X)(-d(X))\right],$$

where $C(X)={\mathbb{E}}_{\mathbb{P}}[Y(1)-Y(-1)\mid X]$. The CTE-DR-ITR approach proposed in Mo et al. [5] minimizes the DR-risk function associated with parameters (c, k) on the training population, defined as

$$L_{DR}(d(X))=\underset{w(X)\in \mathcal{W}(c,k)}{\text{sup}}{\mathbb{E}}_{\mathbb{P}}[w(X)C(X)(-d(X))],$$

where $\mathcal{W}(c,k)=\left\{w(X):{\mathbb{E}}_{\mathbb{P}}[w(X)]=1,{\mathbb{E}}_{\mathbb{P}}\left[w^k(X)\right]\le c^k,w(X)\in {ℝ}_{+}\right\}$. The weight function w(X) represents a general density ratio in $\mathcal{W}(c,k)$, and we assume that $w^{*}(X)\in \mathcal{W}(c,k)$. Given a pre-specified class of ITRs, $\mathcal{D}$, let the minimizer of L_LR(d(X)) in $\mathcal{D}$ be $d_{LR}^{*}(X)$ and the minimizer of L_DR(d(X)) in $\mathcal{D}$ be $d_{DR}^{*}(X)$. We provide upper bounds on $L_{{\mathbb{P}}_{\text{test}}}\left(d_{LR}^{*}(X)\right)$ and $L_{{\mathbb{P}}_{\text{test}}}\left(d_{DR}^{*}(X)\right)$.

Let $v_{LR}^{*}$ be the minimizer (or the sequence converging to the minimum) of ${\Vert C_{\text{test}}(X)-vC(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}$ for v > 0, and ${\delta }_{LR}^{*}(X)=C_{\text{test}}(X)-v_{LR}^{*}C(X)$. Also, let $v_{DR}^{*}(X)$ be the minimizer (or the sequence converging to the minimum) of ${\Vert C_{\text{test}}(X)-v(X)C(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}$ over $\mathcal{V}$, where $\mathcal{V}=\left\{v(X)\in L_{\infty }(\mathbb{P}):v(X)\in {ℝ}_{+}\right\}$, and ${\delta }_{DR}^{*}(X)=C_{\text{test}}(X)-v_{DR}^{*}(X)C(X)$. Given k ≥ 2, we have the following inequalities.

Theorem 3.1.

1. For the LR-ITR approach, we have

   $$L_{{\mathbb{P}}_{\text{test}}}\left(d_{LR}^{*}(X)\right)\le v_{LR}^{*}{L}_{LR}\left(d_{LR}^{*}(X)\right)+{\Vert {\delta }_{LR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}.$$
2. For the CTE-DR-ITR approach, we have

   $$L_{{\mathbb{P}}_{\text{test}}}\left(d_{DR}^{*}(X)\right)\le v_{DR}^{*}{L}_{DR}\left(d_{DR}^{*}(X)\right)+{\Vert {\delta }_{DR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)},$$

   where $v_{DR}^{*}={\mathbb{E}}_{\mathbb{P}}\left[w^{*}(X)v_{DR}^{*}(X)\right]$, and

   $$L_{DR}\left(d_{DR}^{*}(X)\right)=\underset{d}{\text{inf}}\left\{\underset{w(X)\in \mathcal{W}(\tilde{c},k)}{\text{sup}}{\mathbb{E}}_{\mathbb{P}}[w(X)C(X)(-d(X))]\right\}$$

   with $\tilde{c}=c{\Vert v_{DR}^{*}(X)\Vert }_{L_{\infty }(\mathbb{P})}/v_{DR}^{*}$.

When there is a shift in the contrast function, i.e., C_test (X) ≠ C(X), we have ${\Vert {\delta }_{LR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}\ge {\Vert {\delta }_{DR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}$.

If the difference between ${\Vert {\delta }_{LR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}$ and ${\Vert {\delta }_{DR}^{*}(X)\Vert }_{L_2\left({\mathbb{P}}_{\text{test}}\right)}$ dominates the difference between $v_{LR}^{*}{L}_{LR}\left(d_{LR}^{*}(X)\right)$ and $v_{DR}^{*}{L}_{DR}\left(d_{DR}^{*}(X)\right)$, then $d_{LR}^{*}(X)$ leads to a larger upper bound compared with $d_{DR}^{*}(X)$. When C_test (X) = C(X), we have ${\delta }_{LR}^{*}(X)={\delta }_{DR}^{*}(X)=0$. We also have $v_{LR}^{*}=1$ and $v_{DR}^{*}(X)\equiv 1$. Further, we have $\tilde{c}=c$. Notice that L_DR(d(X)) ≥ L_LR(d(X)) when $\tilde{c}=c$. The $d_{LR}^{*}(X)$ can lead to a lower upper bound compared with $d_{DR}^{*}(X)$.

3.3. Simulations

We compare the performance of the LR-ITR approach with the CTE-DR-ITR approach. Both methods only require a calibration dataset with the covariate vector from the targeted distribution.

3.3.1. Mixture of Subgroups with only Covariate Shifts

We first consider the simulation settings with only covariate shifts present. We modify the simulation setup on the mixture of subgroups in Mo et al. [5] such that the two subgroups share the same contrast function between the training and the target populations, but have different covariate distributions.

We set n = 1000 and p = 10. We generate the covariate vector as: $X\mid \xi ~\xi {\mathcal{N}}_p\left({\mu }_1,I_p\right)+(1-\xi ){\mathcal{N}}_p\left({\mu }_2,I_p\right)$, where ξ ~ Bernoulli(p_mix) is the unobserved indicator of the mixture with p_mix determining the proportion of the two subgroups, μ₁ = (0,0,0,⋯,0)^⊤ and μ₂ = (1.958,1.958,0,⋯,0)^⊤. We consider different mixture proportions on the training and the target populations. We fix p_mix = 0.75 on the training population. For the target population, we change p_mix ∈{0.1,0.25,0.5,0.75,0.9}. We then generate A | X ~ Bernoulli(1/2) and $Y\mid (X,A)=m(X)+(A-1/2)C(X)+\mathcal{N}(0,1)$, where $m(x)=1+\sum _{j=1}^p{x}_j/p$ and $C(X)=x_2-\left(x_1^3-2x_1\right)$.

We generate a calibrating dataset from the targeted population with n_calib = 50, which only contains the covariate information. While CTE-DR-ITR approach utilizes this calibrating dataset to select the tuning parameter, we use this calibrating dataset to estimate the likelihood ratio. We compare the performances of the Standard ITR, the CTE-DR-ITR, and the LR-ITR approaches. To evaluate the estimated ITRs, we generate a testing dataset with n_test = 10⁶. We calculate the value function as ${\mathbb{E}}_{n_{\text{test}}}[C(X)\widehat{d}(X)]$, where ${\mathbb{E}}_{n_{\text{test}}}[\cdot ]$ represents the empirical mean on the testing dataset.

Table 1 presents the simulation results summarized over 500 replicates. The results show that the LR-ITR could outperform the CTE-DR-ITR with the presence of only covariate shifts.

Table 1.

Simulation results with n_calib = 50 and mixture of subgroups. Mean (Standard error) over 500 repeats are reported.

	The pmix on the target population
Type	0.1	0.25	0.5	0.75	0.9
Standard ITR	7.827 (0.0170)	6.607 (0.0145)	4.582 (0.0105)	2.560 (0.00654)	1.347 (0.00470)
CTE-DR-ITR	7.868 (0.0185)	6.641 (0.0158)	4.602 (0.0113)	2.557 (0.00738)	1.337 (0.00534)
LR-ITR	8.074 (0.0165)	6.846 (0.0133)	4.730 (0.0101)	2.616 (0.00608)	1.417 (0.00403)

3.3.2. Mixture of Subgroups with Contrast Shifts

The simulation setting is the same as Section 4.2 in Mo et al. [5], which involves contrast shifts in addition to covariate shifts. Specifically, on the CTE function, we follow the same generative model of Y | (A, X), but replace the C(X) with C(ξ, X) = −1.5 × (2ξ − 1) − 2x₁ + x₂. The optimal decision rule given X is $\text{sgn}\left\{\mathbb{E}[C(\xi ,X)\mid X]\right\}$, which involves the unobserved ξ. Consequently, the distributions of (Y(1),Y(−1)) | X also shift on the target population compared with those on the training population. We choose μ₁ = (−1/2,1/2,0,⋯,0)^⊤ and μ₂ = μ₁.

Table 2 shows that the LR-ITR does not improve over the standard ITR when the proportion of the mixture changes. Nonetheless, the CTE-DR-ITR performs better than the Standard ITR.

Table 2.

Simulation results with n_calib = 50 and mixture of subgroups. Mean (Standard error) over 500 repeats are reported.

	The pmix on the target population
Type	0.1	0.25	0.5	0.75	0.9
Standard ITR	1.143 (0.00434)	1.232 (0.00329)	1.383 (0.0015)	1.535 (0.000543)	1.632 (0.00142)
CTE-DR-ITR	1.16 (0.00409)	1.247 (0.00323)	1.388 (0.00137)	1.534 (0.00055)	1.628 (0.00149)
LR-ITR	1.111 (0.00416)	1.216 (0.00319)	1.379 (0.00151)	1.530 (0.00279)	1.616(0.00147)

4. Discussion

Kallus [1] and Mo et al. [5] provide different methods to accommodate the covariate shift. We point out the connection of the advocated weight in Kallus [1] with the efficient score of the contrast function estimation. We also propose a likelihood-ratio weighted approach to address the covariate shift. Upper bounds on the risk function under the target population are provided when distributional shifts exist. There is strength in the performance of the LR-ITR when only covariate shift exists. When other distributional shifts exist, such as shifts in conditional distributions of (Y(1),Y(−1)) | X, CTE-DR-ITR can perform better than LR-ITR even if the likelihood ratio between two populations is known.