Online Testing of Subgroup Treatment Effects Based on Value Difference

The mixture sequential probability ratio test (mSPRT) [9] supposes that the independent and identically distributed (i.i.d.) random variables $Y_{1},Y_{2},\cdots$ have a probability density function $f_{\theta}(x)$ induced by parameter $\theta$, and aims to test

Its test statistics $\Lambda^{\pi}_{n}$ at sample size $n$ is a mixture of likelihood ratios weighted by a mixture density $\pi(\cdot)$ over the parameter space $\Theta$:

The mSPRT stops the sampling at stage

and rejects the null hypothesis $\mbox{H}_{0}$ in favor of $\mbox{H}_{1}$. If no such time exists, it continues the sampling indefinitely and accepts the $\mbox{H}_{0}$. Since the likelihood ratio under $\mbox{H}_{0}$ is a nonnegative martingale with the initial value equal to 1, and so is the mixture of such likelihood ratios $\Lambda_{n}^{\pi}$, the type I error of mSPRT can be proved to be always controlled at $\alpha$ by an application of Markov’s inequality and optional stopping theorem: $P_{H_{0}}(\Lambda_{n}^{\pi}\geq\alpha^{-1})\leq\mathbb{E}_{H_{0}}(\Lambda_{n}^{\pi})/\alpha^{-1}=\mathbb{E}_{H_{0}}(\Lambda_{0}^{\pi})/\alpha^{-1}=\alpha$. Besides, mSPRT is a test of power one [14], which means that any small deviation from $\theta_{0}$ can be detected as long as waiting long enough. It is also shown that mSPRT is almost optimal for data from an exponential family of distributions, with respect to the expected stopping time [15].

The mSPRT was brought to A/B tests by Johari et al. [12, 16], who assume that the observations in control and treatment groups arrive in pairs $(Y_{i}^{(0)},Y_{i}^{(1)})$, $i=1,2,\cdots$. They restricted their data model to the two most common cases in practice: normal distribution and Bernoulli distribution, with $\mu_{\text{A}}$ and $\mu_{\text{B}}$ denoting the mean for the control and treatment groups, respectively. They test the hypotheses as below:

by directly applying mSPRT to the distribution of the differences $Z_{i}=Y_{i}^{(1)}-Y_{i}^{(0)}$ (normal), or the joint distribution of data pairs $(Y_{i}^{(0)},Y_{i}^{(1)})$ (Bernoulli), $i=1,2,\cdots$. After making some approximations to the likelihood ratio and choosing a normal mixture density $\pi(\theta)\sim\mbox{N}(0,\tau^{2})$, the test statistic $\Lambda^{\pi}_{n}$ is able to have a closed form for both normal and Bernoulli observations.

However, the mSPRT does not work well in testing heterogeneous treatment effects due to the complexity of the likelihood induced by individual covariates. Specifically, a conjugate prior $\pi(\cdot)$ for the likelihood ratio may not exist anymore, making the computation for the test statistic challenging. The unknown baseline covariates effect also increases the difficulty in constructing and approximating the likelihood ratio [13].

The sequential score test (SST) [13] assumes a generalized linear model with a link function $g(\cdot)$ for the outcome $Y$:

where $A$ and ${\bm{\mathbf{{X}}}}$ respectively denote the binary treatment indicator and user covariates vector. It tests the multi-dimensional treatment-covariates interaction effect:

while accounting for the linear baseline covariates effect ${\bm{\mathbf{{\mu}}}}^{T}{\bm{\mathbf{{X}}}}$. For the test statistics $\Lambda_{n}^{\pi}$, instead of using a mixture of likelihood ratios as mSPRT, SST employs a mixture of asymptotic probability ratios of the score statistics. Since the probability ratio has the same martingale structure as the likelihood ratio, the type I error can still be controlled with the same decision rule as mSPRT (1). The asymptotic normality of the score statistics also guarantees a closed form of $\Lambda_{n}^{\pi}$ with a multivariate normal mixture density $\pi(\cdot)$. However, the considered parametric model (2) can only be used to test if there is a linear covariate-treatment interaction effect and may fail to detect the existence of a subgroup with enhanced treatment effects. In addition, the subgroup estimated based on the index ${\bm{\mathbf{{\theta}}}}^{T}{\bm{\mathbf{{X}}}}$ may be biased if the assumed linear model (2) is misspecified. Therefore in this paper, we propose a subgroup treatment effect test, which is able to test the existence of a beneficial subgroup and does not require specifying the form of treatment effects.

Suppose we have i.i.d. data ${\bm{\mathbf{{O}}}}_{i}=\left\{Y_{i},A_{i},{\bm{\mathbf{{X}}}}_{i}\right\}$, $i=1,2,\cdots$, where $Y_{i}$, $A_{i}$, ${\bm{\mathbf{{X}}}}_{i}$ respectively denote the observed outcome, binary treatment indicator, and user covariates vector. Here, we consider a flexible generalized linear model:

where baseline covariates effect $\mu(\cdot)$ and treatment-covariates interaction effect $\theta(\cdot)$ are completely unspecified functions, and $g(\cdot)$ is a prespecified link function. For example, we use the identity link $g(\mu)=\mu$ for normal responses and the logit link $g(\mu)=\log\{\mu/(1-\mu)\}$ for binary outcomes.

Assuming $Y$ is coded such that larger values indicate a better outcome, we consider the following test of subgroup treatment effects:

where $\mathcal{X}_{0}$ is the beneficial subgroup with $P({\bm{\mathbf{{X}}}}\in\mathcal{X}_{0})>0$. Note that the above subgroup test is very different from the covariate-treatment interaction test considered in (3) and is much more challenging due to several aspects. First, both $\mu(\cdot)$ and $\theta(\cdot)$ are nonparametric and need to be estimated. Second, the considered hypotheses (5) are moment inequalities that are nonstandard. Third, it allows a nonregular setting, i.e. $P(\theta({\bm{\mathbf{{X}}}})=0)>0$, which makes associated inference difficult. Here, we propose a test based on the value difference between the optimal treatment rule and a fixed treatment rule.

Let $Y^{*}(a)$ denote the potential outcome had the subject received treatment $a$ ($a=0,1$), and $V(d)=\mathbb{E}_{\{Y^{*}(a),{\bm{\mathbf{{X}}}}\}}\left\{Y^{*}(d({\bm{\mathbf{{X}}}}))\right\}$ denote a value function for a fixed treatment rule $d({\bm{\mathbf{{X}}}})$ which maps the information in ${\bm{\mathbf{{X}}}}$ to treatment $\{0,1\}$, where the subscript represents the expectation is taken with respect to the joint distribution of $\{Y^{*}(a),{\bm{\mathbf{{X}}}}\}$. Consider the value difference $\Delta=V(d^{\text{opt}})-V(0)$ between the optimal treatment rule $d^{\text{opt}}=1\left\{\theta({\bm{\mathbf{{X}}}})>0\right\}$ and the treatment rule that assigns control to everyone $d=0$, where $1\{\cdot\}$ is an indicator function. If the null hypothesis is true, no one would benefit from the treatment and the optimal treatment rule assigns everyone to control, and therefore the value difference is zero. However, if the alternative hypothesis is true, some people would have higher outcomes being assigned to treatment and thus the value difference is positive. In this way, the testing hypotheses (5) can be equivalently transformed into the following pair:

We make the following standard causal inference assumptions: (i) consistency, which states that the observed outcome is equal to the potential outcome under the actual treatment received, i.e. $Y=Y^{*}(1)1(A=1)+Y^{*}(0)1(A=0)$; (ii) no unmeasured confounders, i.e. $Y^{*}(a)\perp\!\!\!\perp A\,|\,{\bm{\mathbf{{X}}}}$ for $a=0,1$, which means that the potential outcome is independent of treatment given covariates; (iii) positivity, i.e. $P(A=a|{\bm{\mathbf{{X}}}}={\bm{\mathbf{{x}}}})>0$ for $a=0,1$ and all ${\bm{\mathbf{{x}}}}\in\mathcal{X}$ such that $P({\bm{\mathbf{{X}}}}={\bm{\mathbf{{x}}}})>0$. Under these assumptions, it can be shown that

We estimate the value function of a given treatment rule $d$ by the augmented inverse probability weighted (AIPW) estimator [17, 18]:

where $p_{A}({\bm{\mathbf{{X}}}})=A\cdot p({\bm{\mathbf{{X}}}})+(1-A)\cdot(1-p({\bm{\mathbf{{X}}}}))$ and $p({\bm{\mathbf{{X}}}})=P(A=1|{\bm{\mathbf{{X}}}})$ is the propensity score. This estimator is unbiased, i.e. $\mathbb{E}_{(Y,A,{\bm{\mathbf{{X}}}})}\{\hat{V}_{\text{AIPW}}(d)\}=V(d)$. Moreover, the AIPW estimator has double robustness, that is, the estimator remains consistent if either the estimator of $\mathbb{E}(Y|A=d,{\bm{\mathbf{{X}}}})$ or the estimator of the propensity score $p({\bm{\mathbf{{X}}}})$ is consistent, which gives much flexibility. Then the value difference $\Delta$ is unbiasedly estimated by

under our assumed model (4), where $g^{-1}(\cdot)$ is the inverse of the link function. That is,

Since $\mu(\cdot)$, $\theta(\cdot)$, and $p(\cdot)$ are usually unknown, we let data come in batches and estimate them based on previous batches of data.

Like SST, the key idea behind SUBTLE is to construct a statistic that has an (asymptotic) normal distribution with different means under the null hypothesis and alternative hypothesis (6), and then build the test statistics as a mixture of (asymptotic) probability ratios of it. Define $\hat{\Delta}_{k}$ as below:

Note that $\hat{\Delta}_{k}$ is an asymptotic unbiased estimator for $\Delta$ by (7). $R_{k}$ (8) is a multiplier of $\hat{\Delta}_{k}$ and thus has the mean related to the value difference $\Delta$. In section III-C we will show that $R_{k}$ has an asymptotic normal distribution with the same variance but different means under the null and local alternative hypotheses so that our test statistics $\Lambda_{k}^{\pi}$ (9) is a mixture of asymptotic probability ratios of $R_{k}$. As we discussed in Section II, the probability ratio has the same martingale structure as the likelihood ratio, so SUBTLE can be shown to control type I error with the decision rule (1). Algorithm 1 shows our complete testing procedures.

In step (ii) of Algorithm 1, we estimate $\mu(\cdot)$ and $\theta(\cdot)$ by respectively building a random forest on control observations and treatment observations in previous batches. The propensity score $p(\cdot)$ is estimated by computing the proportion of treatment observations ($A=1$) in previous batches. In step (iv) we estimate $\sigma_{k}$ with $\hat{\sigma}_{k}=\sqrt{s_{k}^{2}/m}$, where $s_{k}^{2}$ is the sample variance of $D({\bm{\mathbf{{O}}}}_{i};\hat{\mu}_{k-1},\hat{\theta}_{k-1},\hat{p}_{k-1}),\;\forall{\bm{\mathbf{{O}}}}_{i}\in\overline{\mathcal{C}}_{k-1}$. Since the value difference is always nonnegative, in step (v) we choose a truncated normal $\pi(\Delta)=2/\sqrt{2\pi\tau^{2}}\cdot\exp{\left\{-\Delta^{2}/2\tau^{2}\right\}}\cdot 1\{\Delta>0\}$ as the mixture density. The normality of the mixture density guarantees a closed form for our test statistic:

where $F(\cdot)$ is the cumulative distribution function of a normal distribution with mean $(\sqrt{k}\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1}R_{k})/\{(\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1})^{2}+k\}$ and variance $k\tau^{2}/\{\tau^{2}(\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1})^{2}+k\}$. In theory, the choice of mixture density variance $\tau^{2}$ will not have any effect on the type I error control. Johari et al. [12] proved that an optimal $\tau^{2}$ in terms of stopping time is the prior variance times a correction for truncating, so we suggest estimating $\tau^{2}$ based on historical data. In the last step, we add a failure time $M$ to the decision rule to terminate the test externally and accept the null hypothesis if we ever reach it. If the null hypothesis is rejected, we can employ random forests to estimate $\theta(\cdot)$ based on all the data up to the time that the experiment ends. Then the estimated treatment effect $\hat{\theta}({\bm{\mathbf{{x}}}})$ naturally gives the beneficial subgroup $\mathcal{X}_{0}=\{{\bm{\mathbf{{x}}}}:\hat{\theta}({\bm{\mathbf{{x}}}})>0\}$.

In this section, we will show that our proposed test SUBTLE is able to control type I error at any time, that is, $P_{H_{0}}(\Lambda_{k}^{\pi}>1/\alpha)<\alpha$ for any $k\in\mathbb{N}^{+}$. Theorem III.1 gives the respective asymptotic distributions of $R_{k}$ under the null and local alternative hypotheses, which demonstrates that the test statistics $\Lambda_{k}^{\pi}$ is a mixture of asymptotic probability ratios weighted by $\pi(\cdot)$. Proposition 1 shows that this asymptotic probability ratio has a martingale structure under $\mbox{H}_{0}$ when the sample size is large enough. Combining these two results with the demonstration in Section II-A, we can conclude that the type I error of SUBTLE is always controlled at $\alpha$.

The upcoming theorem relies on the following conditions:

(C1)

The number of batches $k$ diverges to infinity as sample size $n$ diverges to infinity.

(C2)

Lindeberg-like condition: for all $\epsilon>0$

	$\displaystyle\frac{1}{k}\sum_{j=1}^{k}\mathbb{E}$	$\displaystyle\left[\left\{\frac{\bar{D}_{j}(\mathcal{C}_{j};\overline{\mathcal{C}}_{j-1})}{\hat{\sigma}_{j}}\right\}^{2}\right.$
		$\displaystyle\times\left.1\left\{\frac{|\bar{D}_{j}(\mathcal{C}_{j};\overline{\mathcal{C}}_{j-1})|}{\hat{\sigma}_{j}}>\epsilon\sqrt{k}\right\}\bigg{|}\overline{\mathcal{C}}_{j-1}\right]=o_{p}(1)$

(C3)

$\frac{1}{k}\sum_{j=1}^{k}\frac{\sigma_{j}^{2}}{\hat{\sigma}_{j}^{2}}\xrightarrow{P}1$.

(C4)

	$\displaystyle\frac{1}{k}\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1}\left[\mathbb{E}\left\{\bar{D}_{j}(\mathcal{C}_{j};\overline{\mathcal{C}}_{j-1})|\overline{\mathcal{C}}_{j-1}\right\}\right.$
	$\displaystyle\left.-\mathbb{E}\left\{\bar{D}_{j}(\mathcal{C}_{j};\hat{d}^{opt}_{j-1},\mu,\theta,p)|\overline{\mathcal{C}}_{j-1}\right\}\right]=o_{P}(k^{-1/2}),$

where $\bar{D}_{j}(\mathcal{C}_{j};\hat{d}^{opt}_{j-1},\mu,\theta,p)$ is $\bar{D}_{j}(\mathcal{C}_{j};\overline{\mathcal{C}}_{j-1})$ with $\hat{\mu}_{j-1},\hat{\theta}_{j-1},\hat{p}_{j-1}$ replaced by the true value $\mu,\theta,p$, but $\hat{d}^{opt}_{j-1}=1\{\hat{\theta}_{j-1}({\bm{\mathbf{{X}}}})>0\}$ unchanged.

(C5)

$\frac{1}{k}\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1}\left[\mathbb{E}\left\{\bar{D}_{j}(\mathcal{C}_{j};\hat{d}^{opt}_{j-1},\mu,\theta,p)|\overline{\mathcal{C}}_{j-1}\right\}-\Delta\right]=o_{P}(k^{-1/2})$.

Similar conditions are also used in the literature for studying the properties of doubly robust estimators. Their appropriateness was discussed in Section 7 of [19]. Both (C4) and (C5) rely on the convergence rate of the estimator of $\mu,\theta,p$. Wager and Athey [20] showed that under certain constraints on the subsampling rate, random forest predictions converge at the rate $n^{s-1/2}$, where $s$ is chosen to satisfy some conditions. We assume that under this rate, (C4) and (C5) hold.

The proofs of the above results are given in the appendices.

We consider five data generation models in the form of (4) with logistic link $g(\cdot)$. Data are generated in batches with batch size $m=20$ and are randomly assigned to two groups with a fixed propensity score $p({\bm{\mathbf{{X}}}})=0.5$. Each experiment is repeated 1000 times to estimate the type I error and power. For the first four models, we consider

Five covariates

:

	$\displaystyle X_{1}\overset{iid}{\sim}\mbox{Ber}(0.5),\;$	$\displaystyle X_{2}\overset{iid}{\sim}\mbox{U}[-1,1],$
	$\displaystyle X_{3},X_{4},X_{5}$	$\displaystyle\overset{iid}{\sim}\mbox{N}(0,1).$

Two baseline effects

:

	$\displaystyle\mu_{1}({\bm{\mathbf{{X}}}})$	$\displaystyle=-2-X_{1}+X_{3}^{2},$
	$\displaystyle\mu_{2}({\bm{\mathbf{{X}}}})$	$\displaystyle=-1.3+X_{1}+0.5X_{2}-X_{3}^{2}.$

Two treatment-covariates interaction effects

:

	$\displaystyle\theta_{1}({\bm{\mathbf{{X}}}})$	$\displaystyle=c\cdot 1\{X_{1}+2X_{3}>0\},$
	$\displaystyle\theta_{2}({\bm{\mathbf{{X}}}})$	$\displaystyle=c\cdot 1\{X_{2}>0\;\text{or}\;X_{5}<-0.5\}.$

Table I displays which covariates, $\mu({\bm{\mathbf{{X}}}})$, and $\theta({\bm{\mathbf{{X}}}})$ are employed in each model. For model V, we consider the following high-dimensional setting:

where $c$ varies among $\left\{-1,0,0.6,0.8,1\right\}$ indicating the intensity of the value difference. When $c=-1$ and 0, the null hypothesis is true and the type I error is estimated; while when $c=0.6,0.8,1$, the alternative is true and the power is estimated.

Table II shows that the SUBTLE is able to control type I error and achieve competing detection power, especially under high-dimensional setting (Model V); however, SST couldn’t control type I error especially when $c=-1$. This can be explained by two things: (i) the linearity of the model (2) is violated; (ii) SST is testing if there is a difference between treatment and control groups among any subjects, instead of the existence of a beneficial subgroup. Specifically, SST is testing if the least false parameter ${\bm{\mathbf{{\theta}}}}^{*}$, to which the MLE of ${\bm{\mathbf{{\theta}}}}$ under model misspecification converges, is zero or not. We also perform experiments with batch size $m=40$, and the results (Table III) do not have much difference.

To evaluate the robustness of SUBTLE to the mixture density variance $\tau^{2}$, we conduct additional simulations with $\tau^{2}\in\{0.0001,0.001,0.01,0.1,1,10\}$. The data are generated from Model I with $c=0$ or $c=1$, and batch size $m=200$ for computation efficiency. When $c=0$ we estimate the type I error, while when $c=1$ we estimate the power. The results in Table IV show that the type I error is always controlled below significance level 0.05 and the power has considerable robustness to the choice of $\tau^{2}$.

It is common in practice that a large number of covariates are incorporated in the experiment whereas the actual outcome only depends on a few of them. Some covariates do not have any effect on the response, like $X_{4}$ in Model II, III, IV, and we call them noise covariates. In the following simulation, we explore the impact of noise covariates on the detection power of SUBTLE and SST. We choose Model I with $c=0.8$ as the base model, and at each time add three noise covariates which are respectively from normal N$(0,1)$, uniform U$[-1,1]$, and Bernoulli Ber$(0.5)$ distributions. The batch size is set to $m=40$ for computation efficiency. Fig. 1 shows that SST has continuously decreasing powers as the number of noise covariates increases, while the power of SUBTLE is more robust to the noise covariates.

A key feature of sequential testing is that it has an expected smaller sample size than fixed-horizon testing. For comparison, we consider a fixed-horizon version of SUBTLE, which leverages Theorem III.1 and rejects the null hypothesis $\mbox{H}_{0}:\Delta=0$ when $R_{k}>Z_{\alpha}$ for some predetermined $k$, where $Z_{\alpha}$ denotes the $(1-\alpha)$ quantile of standard normal distribution. We assume $\sigma^{-1}=\lim_{k\rightarrow\infty}k^{-1}\sum_{j=1}^{k}\hat{\sigma}_{j}^{-1}$, then the required number of batches $k$ can be calculated as $k=\sigma^{2}(Z_{\alpha}+Z_{1-\mbox{power}})^{2}/\Delta^{2}$, and thus the required sample size is $n=k\times m+l$. The true value difference $\Delta$ can be directly estimated from data generated under the true model with two treatment rules, while $\sigma^{2}$ is estimated by the sample variance of $\hat{\Delta}_{{k}^{\prime}}$ times ${k}^{\prime}$ for some fixed large ${k}^{\prime}$. Here, we choose Model V with $c=1$ and batch size $m=20$. The stopping sample size of our sequential SUBTLE over 1000 replicates is shown in Fig. 2, and the dashed vertical line indicates the required sample size for the fixed-horizon SUBTLE with the same power of 0.997 (seen from Table II) under the same setting. We can find that most of the time our sequential SUBTLE arrives at the decision early than the fixed-horizon version, but occasionally it can take longer. The distribution of the stopping time for sequential SUBTLE is right-skewed, which is in line with the findings in [12] and  [2].

Let $\mathcal{F}_{j}$, $0\leq j\leq k$, denote a filtration generated by observations in first $(j+1)$ batches $\overline{\mathcal{C}}_{j}=\cup_{r=0}^{j}\mathcal{C}_{r}$, and $\bar{D}_{j}(\mathcal{C}_{j};\hat{d}^{opt}_{j-1},\mu,\theta,p)$ denote an AIPW estimator for $\Delta$ with only optimal decision rule estimated by previous batches:

Then

Above (11) follows by condition (C5) and (12) follows by condition (C4). For $j=1,2,\cdots,k$, let

It is obvious that for each $k$, $M_{k,j}$, $1\leq j\leq k$, is a martingale with respect to the filtration $\mathcal{F}_{j}$. In particular, for all $j\geq 1$, $\mathbb{E}(M_{k,j}|\mathcal{F}_{j-1})=0$ and $\sum_{j=1}^{k}\mathbb{E}(M_{k,j}^{2}|\mathcal{F}_{j-1})=k^{-1}\sum_{j=1}^{k}\sigma_{j}^{2}/\hat{\sigma}_{j}^{2}\xrightarrow{p}1$ as $k\xrightarrow{}\infty$ by (C3). The conditional Lindeberg condition holds in (C2), so the martingale central limit theory for triangular arrays gives

Plugging it back into (12), we can get