Neyman meets causal machine learning: Experimental evaluation of individualized treatment rules

3.1 Setup

Suppose that we use an ML algorithm to create a ITR,

Most commonly, researchers first estimate the conditional ATE [see, e.g., 7,17–21]:

They then derive an ITR as the treatment rule that assigns the treatment to everyone who is predicted to have a positive CATE, i.e., f ( x ) = 1 { τ ˆ ( x ) > 0 } , where τ ˆ ( x ) is an estimate of the CATE. Throughout this article, we assume, without loss of generality, that a positive effect implies that the treatment is beneficial. One may also consider a cost associated with the administration of treatment or a budget constraint that limits the proportion of individuals who can receive the treatment.

Our goal is to evaluate the empirical performance of an ITR without assuming that an ITR is indeed optimal. In the aforementioned example, we do not assume that τ ˆ ( x ) is an accurate estimate of the CATE. In fact, we do not make any assumption about how the ITR is constructed and how accurate it is. The ITR may be derived from an application of an ML algorithm or be even based on heuristics. For now, we only assume that the ITR to be evaluated is given. For example, it may be estimated from an external dataset to be used in a downstream decision-making context. In Section 5, we discuss how to use the same experimental data for both learning and evaluating an ITR.

To measure the performance of an ITR, we consider two quantities. The first is the population average value (PAV), which is defined as

This is the standard metric of ITR’s overall performance. The second quantity is the population average prescriptive effect (PAPE) [14,22], which measures the benefit of ITR and is defined as follows:

where p f = Pr ( f ( X i ) = 1 ) represents the proportion of individuals who are treated by the ITR f .

The PAPE compares the performance of ITR against the non-ITR that treats the same proportion of randomly selected individuals. This contrasts with other quantities considered in the literature such as the targeting operator characteristic (TOC) [23], which compares the performance of ITR against the non-individualized rule that treats everyone. Unlike the TOC, the PAPE focuses on the benefit of determining which individuals should be treated while holding the proportion of those who receive the treatment constant.

To gain additional intuition about the PAPE, consider the following alternative but equivalent expression of the same quantity:

This alternative expression shows that the PAPE measures how well the ITR agrees with the true individual treatment effect (ITE). To compare across datasets, we can further normalize the PAPE as the correlation between the ITR and the true ITE, i.e.,

Although this provides a scale-invariant quantity to understand the performance of ITR, it is not identifiable from the data because we cannot identify the variance of ITE, i.e., V ( Y i ( 1 ) − Y i ( 0 ) ) .

The aforementioned equality further implies the following inequality by applying Cauchy-Schwarz twice:

Therefore, PAPE is bounded provided that the second moments of the potential outcomes exist. Thus, given a fixed variance of the potential outcomes, τ f is most likely largest around p f = 0.5 . This is because when p f is around 0.5, the ITR has the greatest room to deviate from the randomized treatment rule.

3.2 Estimation and inference

To estimate the PAV and PAPE under Neyman’s repeated sampling framework, we consider the following “difference-in-means”-type estimators:

where p ˆ f = ∑ i = 1 n f ( X i ) ∕ n is the estimated population proportion of individuals who receive the treatment assignment and Z n = { Y i , T i , X i } i = 1 n represents the experimental data of sample size n . The factor n ∕ ( n − 1 ) in the PAPE estimator represents the loss in one degree-of-freedom due to estimating the population-level quantity p f .

The following theorems, reproduced from our previously published work [14], show that, under Neyman’s repeated sampling framework, these two estimators are unbiased and the finite-sample variances can be derived.

The properties of the PAV estimator shown above follow immediately from Neyman’s classic results by replacing the potential outcome Y i ( t ) with the the potential outcome that incorporates the ITR, i.e., 1 { f ( X i ) = t } Y i ( t ) . The form of the estimator does not mean, however, that we are ignoring observations whose prescribed treatment status differs from the observed one, i.e., f ( X i ) ≠ T i . These observations are still used when estimating the variance. An important implication of this subtle fact is discussed in Section 3.4.

We can also compare the results of the PAPE estimator with Neyman’s classic results. We observe that the relevant potential outcome is given by Y ˜ f i ( t ) = ( f ( X i ) − p ˆ f ) Y i ( t ) , which directly compares the ITR with the randomized treatment rule. When compared to the PAV estimator, the variance of the PAPE estimator has an additional term, which comes from the fact that Y ˜ f i ( t ) is correlated across observations due to the estimation of the proportion p f .

The correlation can be broadly decomposed into two components: (1) a negative component caused by the negative correlation within the mean-adjusted ITR ( f ( X i ) − p ˆ f ) and ( f ( X j ) − p ˆ f ) , and (2) the remaining component that arises due to the interaction between f ( X i ) and τ i . The negative component generally dominates if the ITR of interest treats roughly 50% of the population ( p f ≈ 0.5 ) and the ATE τ is not too small. This suggests that the mean adjustment could lead to a variance reduction when the ITR treats a roughly half of the population.

On the other hand, for p f ≉ 0.5 and n ≫ 1 , this additional term is positive if and only if the following condition holds:

Thus, this additional term is only likely to be positive under a scenario where p f is away from 0.5 and the magnitude of the PAPE is much greater than that of the ATE (i.e., the ITR is performing well). Equation (5) suggests that this is unlikely unless the variance of the individualized treatment effect is large, implying a large degree of treatment effect heterogeneity.

3.3 Performance comparison among multiple ITRs

While the PAPE compares an ITR with a random treatment assignment rule that treats the same proportion of units, researchers are often interested in comparing the performance of multiple ITRs. In such cases, we recommend estimating the difference in PAV between two ITRs that are subject to the same budget constraint. Imai and Li [14] provided the details of estimation and inference regarding this quantity.

The use of PAPE for comparison of two ITRs is sometimes inappropriate. To see this, note that the difference in PAPE can be written as the covariance between the agreement of two ITRs and the true ITE:

This expression shows that for ITRs with similar treatment proportions, the sign of this difference indicates the relative capability of the ITRs in identifying the optimal individuals to treat. However, if the two ITRs f and g have significantly different treatment proportions, then this comparison is difficult. Figure 1 shows an example, in which ITR f has a higher PAV than ITR g , but f has a negative PAPE, while g has a positive PAPE. In this case, f is not an effective ITR as it performs significantly worse than random treatment, but practitioners might still choose to not use g as it is only able to identify a small percentage of good patients to target.

Figure 1

Illustration of PAPE for two different ITRs f and g . Here, the x axis is the proportion of individuals treated, and y axis is PAV. The PAV of f is higher than PAV of g , but ITR g has a positive PAPE τ g and the ITR f has a negative PAPE τ f .

3.4 Lack of invariance

Unlike Neyman’s ATE estimator, the PAV and PAPE estimators are not invariant to a constant shift of the outcome variable. One might expect that adding a constant δ to Y would shift the PAV estimator by δ and not affect the PAPE estimator at all. Unfortunately, this is not the case. For both of these estimators, a constant shift of the outcome will result in an additional change of the equal magnitude. Let λ ˆ f δ ( Z ) and τ ˆ f δ ( Z ) be the new PAV and PAPE estimators under a constant δ shift, i.e., λ ˆ f δ ( Z ) = λ ˆ f ( Z ) + δ and τ ˆ f δ ( Z ) = τ ˆ f ( Z ) + δ . Then, we have,

Since this term equals zero in expectation, the two estimators remain unbiased. However, this shift affects their variances because the ITR f ( X i ) is only balanced on average due to the randomized treatment assignment. Intuitively, this is because the relevant potential outcomes in the case of both PAV and PAPE are not invariant to a constant shift of the outcome variable. For PAV, the relevant potential outcome is Y f i ( t ) = 1 { f ( X i ) = t } Y i ( t ) , so under a constant shift of δ , only samples with f ( X i ) = t would be shifted by δ , while the remaining samples still have a value of zero. The intuition is similar for PAPE. Therefore, we expect that balancing the observed outcomes Y i close to zero could result in increased efficiency.

We now derive the constant shift δ that minimizes the resulting variances. We can show that a constant shift of δ to potential outcomes creates the following additional variance terms for both estimators:

Proof is in Appendix B. The proposition implies that if we wish to minimize variance across a range of f ( X i ) , then when n 1 = n 0 = n ∕ 2 , the optimal value of δ approximately balances the two potential outcomes around zero after a constant shift, i.e.,

4.1 Setup

For the ex-ante evaluation of the PAPE, we assume a simple random of n units from the same target population, P . Consider a completely randomized experiment, in which a total of n f units are randomly assigned to an ITR f , while the remaining units n r = n − n f are assigned to the random treatment rule with the probability of treatment assignment equal to n r 1 ∕ n r . Let F i be an indicator variable, which is equal to 1 if unit i is assigned to the ITR f and is equal to 0 otherwise. Under the random treatment rule, the number of units that are randomly assigned to the treatment condition is n r 1 , while n r 0 = n r − n r 1 units are assigned to the control condition. As before, we use T i to represent the treatment indicator. We formally state these assumptions.

Using this experimental data, we wish to estimate the PAPE defined in equation (4). For simplicity, we have the number of treated units under the random treatment rule to equal that under the ITR condition, i.e., p ˆ f = n r 1 ∕ n r , where p ˆ f = ∑ i = 1 n f ( X i ) ∕ n , so that the evaluation estimator would not need to be further adjusted. Fortunately, in practice, this can be easily accomplished so long as the covariates are available prior to the randomization of treatment assignment among the group assigned to the random treatment rule. Finally, the so-called Neyman allocation implies that if the variances of Y i ( 1 ) and Y i ( 0 ) differ significantly from one another, one can gain additional statistical efficiency by allocating more units to the treatment condition whose potential outcome has a greater variance. This optimal design, however, does require the availability of external data to estimate these variances. For the sake of simplicity, we do not consider such optimal designs here.

We consider the following estimator of the PAPE for the ex-ante experimental evaluation that accounts for a potential difference in the proportion of treated units between the ITR and the random treatment rule by appropriately weighting the latter:

The ex-ante evaluation differs from the ex-post evaluation in two ways. First, the ex-ante estimator requires two separate random assignments ( T i and F i ), while the ex-post estimator only involves one. Intuitively, an additional layer of randomization increases variance. Second, the ex-ante evaluation requires a separate group that follows an ITR, whereas all individuals under the ex-post evaluation are simply randomly assigned either to the treatment or control group. As a result, under the ex-post evaluation, we use the samples identically, which could further reduce the variance. Together, we expect the ex-ante evaluation to be less efficient than the ex-post evaluation as the full sample is not utilized for every part of the estimation. In the following, we use Neyman’s repeated sampling framework to confirm this intuition under a set of simplifying assumptions.

4.2 Comparison of the two experimental designs

Before comparing two modes of evaluation, we derive the bias and variance of the ex-ante evaluation estimators under the Neyman’s repeated sampling framework. In the current case, the uncertainty comes from three types of randomness: (1) the random assignment to the individualized or random treatment rule, (2) the randomized treatment assignment under the random assignment rule, and (3) the simple random sampling of units from the target population. The next theorem shows that this estimator is unbiased and the variance is identifiable. Proof is given in Appendix A.

Given these results, we examine the relative statistical efficiency of the ex-post and ex-ante experimental evaluations. To facilitate the comparison, we assume n 1 = n 0 = n f = n r = n ∕ 2 . In words, the ex-post evaluation sets the treatment assignment probability to 1/2, and the ex-ante evaluation also sets the probability of being assigned to the ITR to 1/2. In the same fashion, we also assume n r 1 = n r 0 = n ∕ 4 , implying that the ex-ante evaluation sets the treatment assignment probability under the random treatment rule to 1/2 as well. Although our result below may not be applicable beyond this simplified setting, we believe that this equal allocation setting is a common choice in practice and therefore is worthy of investigation.

Under this simplified setting, the difference in the variance of the PAPE estimator between the ex-ante and ex-post evaluations is given by

The details of the derivation are given in Appendix C. Suppose now that the ITR correctly assigns individuals on average, i.e., E ( Y i ( t ) ∣ f ( X i ) = t ) ≥ E ( Y i ( t ) ∣ f ( X i ) = 1 − t ) for t = 0 , 1 . Under this assumption, the last two terms in the square bracket are positive, i.e.,

for t = 0 , 1 . Hence, the only term that is possibly negative in equation (9) is the third term in the square bracket. For simplicity, further assume that we shift the outcomes to minimize variance of the ex-post estimator and achieved E ( Y i ( 1 ) + Y i ( 0 ) ∣ f ( X i ) = 1 ) = E ( Y i ( 1 ) + Y i ( 0 ) ∣ f ( X i ) = 0 ) = 0 (see equation (1)). This guarantees that the optimal choice of δ is zero, and hence, no adjustment in variance is necessary. Under this assumption, we can bound equation (9) from below as follows (see Appendix D for details):

The result implies that under a set of simplifying assumptions, the ex-post evaluation is more efficient than the ex-ante evaluation. We note, however, that this conclusion may not hold if the ex-ante and ex-post setups have sample allocation different from the setting considered here.

5.1 Setup

Consider a generic ML algorithm, which we define as a deterministic function mapping the space of training data of finite size, denoted by Z , to the space of all possible scoring rules S ,

Typically, the scoring rule of interest is the estimated CATE such that the largest value indicates the highest treatment prioritization. Alternatively, the scoring rule may be based on the estimated baseline risk, i.e., E ( Y i ( 0 ) ∣ X i = x ) . We do not, however, assume that the ML algorithm used to generate the scoring rule accurately estimates either the CATE or baseline risk. Indeed, we essentially impose no assumption on how the scoring rule is created. Once the scoring rule is estimated by an ML algorithm, the ITR is given by

where the notation makes it explicit that the ITR depends on the specific training data Z n ∈ Z of sample size n .

Next, consider the following standard cross-fitting procedure. First, we randomly split the experimental data of size n into K subsamples of equal size m = n ∕ K , where, for notational simplicity, we assume n is a multiple of K . Then, for each k = 1 , 2 , … , K , we use the k th subsample as an evaluation dataset Z m ( k ) = { X i ( k ) , T i ( k ) , Y i ( k ) } i = 1 m , while the remaining ( K − 1 ) subsamples are used as the training dataset Z n − m ( − k ) = { X i ( − k ) , T i ( − k ) , Y i ( − k ) } i = 1 n − m . Without loss of generality, we assume that the number of treated (control) units is identical across K folds and denote it using m 1 ( m 0 = m − m 1 ).

Then, for each fold k , we estimate an ITR by applying the ML algorithm F to the training data Z n − m ( − k ) , which we denote by f ˆ ( − k ) = f ˆ Z n − m ( − k ) . We then evaluate the performance of the ML algorithm F by computing an evaluation metric of interest based on the test data Z m ( k ) . Repeating this process K times for each k and averaging the results gives a cross-fitting estimator of the evaluation metric. Here, we focus on the cross-fitting PAV estimator:

where λ ˆ f ( ⋅ ) is defined in equation (6). We now discuss the estimand, for which this cross-fitting estimator is unbiased.

5.2 Evaluation metrics under cross-fitting

To extend Neyman’s repeated sampling framework to cross-fitting with K ≥ 2 folds, we begin by noting that the ITR in this setting varies as a function of training data. Thus, we consider the performance measure that averages over the random sampling of training data as well as the randomization of treatment assignment and random sampling of units. In other words, we evaluate the average performance of ITR that is generated by the application of ML algorithm F across K different (but overlapping) training datasets. This contrasts with the performance evaluation metric of a fixed ITR discussed in earlier sections.

For the PAV under cross-fitting, we consider an average ITR over across training data of size n − m :

which represents the proportion of times the estimated ITR would assign the treatment to a unit with a specific value of covariates. The notation makes explicit the dependence on the size of training data n − m as well as the ML algorithm F .

Under Neyman’s repeated sampling framework, one can view each estimated ITR as another random sampling from a population of ITRs based on the ML algorithm F with training dataset of size n − m . Thus, the PAV under cross-fitting can be defined as

For the PAPE, we consider the cross-fitting version of the proportion treated by ITR p f as follows:

Then, the PAPE under cross-fitting can be defined as