Highly adaptive Lasso for estimation of heterogeneous treatment effects and treatment recommendation

2.1 Setup

Assume we observe an independent and identically distributed sample of observations O i = ( W i , A i , Y i ) ∼ P 0 ∈ ℳ , where P 0 is an unknown distribution that lies in non-parametric model ℳ . W ∈ R p are covariates, A ∈ A = { 0 , 1 , … , K } is a treatment, and Y ∈ R is a real-valued outcome. Furthermore, let V ⊆ W be some subset of covariates with V ∈ R q for q ≤ p . For each a ∈ A , we define a counterfactual data unit X i ( a ) = ( W i , Y i ( a ) ) ∼ P 0 , a corresponding to the data that would have been observed on unit i if, possibly counter to fact, the unit had been assigned treatment a . The conditional average treatment effect (CATE) is defined as E 0 , a [ Y ( a ) ∣ V = v ] − E 0 , a [ Y ( 0 ) ∣ V = v ] , where for a P 0 , a -measurable function f , we write the expectation of f ( O ) as E 0 , a [ f ( O ) ] = ∫ f d P 0 , a . We note that our notation for the CATE explicitly makes level 0 of the treatment a referent level against which all other treatments are compared and we define A 0 ≔ A \ { 0 } .

The CATE quantifies the expected difference in counterfactual outcome under treatment a ∈ A 0 versus treatment 0 for data units with V = v . Also of interest is the OTP, d 0 , ⋅ ( v ) ≔ argmax a ∈ A E 0 , a [ Y ( a ) ∣ V = v ] . Given an estimate ψ ˆ a ( v ) of ψ 0 , a ( v ) , we can construct an estimate of the OTP as

Treatment policies could also be derived from estimators of the CATE by giving level 0 of the treatment unless the additional benefit exceeds a policy-relevant threshold t (in the above, t = 0 ). In our real data analysis mentioned in the following, we consider performance of estimators under various thresholds for defining treatment policies. Estimation of the CATE and OTP requires a set of causal assumptions that allow us to rewrite the parameter without unobservable counterfactuals. One set of such assumptions is as follows.

1. Positivity: P 0 { 0 < P 0 ( A = a ∣ W ) < 1 } = 1 , ∀ a ∈ A ,

2. No unobserved confounding: Y ( a ) ⊥ A ∣ W , ∀ a ∈ A ,

3. Consistency: Y i = ∑ a ∈ A 1 a ( A i ) Y i ( a ) ,

4. No interference: the counterfactual outcome of each individual under each treatment does not depend on treatments assigned to other individuals in the population.

If these assumptions are satisfied, the CATE is identified by a parameter of the distribution of the observed data P 0 . We define μ 0 , a ( w ) ≔ E 0 [ Y ∣ A = a , W = w ] , where E 0 denotes the expectation under P 0 . The CATE is identified by ψ 0 , a ( v ) = E 0 [ μ 0 , a ( W ) − μ 0,0 ( W ) ∣ V = v ] . The object μ 0 , a is referred to as the outcome regression, and we sometimes use the notation μ 0 = ( μ 0 , a : a ∈ A ) to refer collectively to outcome regressions under all treatments. In the following, we also discuss the propensity score for treatment a , π 0 , a ( w ) ≔ P 0 ( A = a ∣ W = w ) .

2.2 Meta-learning with two treatment arms

This CATE estimation problem is commonly presented in the case of two treatment arms, A = { 0 , 1 } , in which case the target parameter for estimation is ψ 0 ( 1 ∣ ⋅ ) , a function only of covariates V . A common approach to estimation of ψ 0 ( 1 ∣ ⋅ ) is the S-learner [4], so-called because it involves a single outcome regression estimator. The estimate μ ˆ of μ 0 can be obtained using any regression technique. For example, many such techniques estimate μ 0 by minimizing an empirical risk criterion based on a loss ( ψ , o ) → L ( μ ) ( o ) . We focus throughout on results for squared-error loss with the understanding that other loss functions could be used. For the purpose of learning μ 0 , squared error loss may be written as L ( μ ) ( o ) ≔ { y − μ a ( v ) } 2 .

Regardless of the learning approach adopted to generate μ ˆ , the S-learner CATE estimate can be obtained post hoc. If V = W , then the estimate is ψ ˆ S , 1 ( w ) ≔ μ ˆ 1 ( w ) − μ ˆ 0 ( w ) . If instead, V ⊂ W , then a second stage of learning is required. In particular, we define a pseudo-outcome D μ ˆ ( O i ) ≔ μ ˆ 1 ( W i ) − μ ˆ 0 ( W i ) and regress this pseudo-outcome on covariates V to obtain the CATE estimate ψ ˆ S , 1 ( v ) ≔ E ˆ [ D μ ˆ ( O ) ∣ V = v ] . Again, the regression learning technique is interchangeable. For example, we could minimize empirical risk based on squared error loss L S ( ψ ) ( o ) = { D μ ( o ) − ψ ( 1 ∣ v ) } 2 .

An alternative strategy sometimes adopted is to split the data and regress Y on W separately for each treatment group to obtain estimates μ ˆ 0 ( w ) and μ ˆ 1 ( w ) of μ 0 , 1 ( w ) and μ 0,0 ( w ) , respectively. This method is referred to as the T-learner [4]. Again, if V = W , we can estimate the CATE by subtracting the two regression estimates: ψ ˆ T , 1 ( w ) = μ ˆ 1 ( w ) − μ ˆ 0 ( w ) . If V ⊂ W , this difference constitutes a pseudo-outcome, which we further regress V to obtain our CATE estimate. The T-learner typically has lower performance due to the lack of information sharing across treatment arms when estimating the outcome regression [6]. Thus, it is not considered further here.

Van der Laan [10] described a doubly robust (DR)-learner, which involves estimation not only of the outcome regression, but also the propensity score. The two nuisance parameter estimates μ ˆ or π ˆ are both used to construct the pseudo-outcome:

D μ ˆ , π ˆ ( O ) is then regressed on V to obtain the CATE estimate ψ ˆ DR ( 1 ∣ v ) ≔ E ˆ [ D μ , π ( O ) ∣ V = v ] . We can obtain this estimate, for example, via empirical risk minimization based on a loss L DR ( ψ ) ( o ) . We focus on squared-error loss, L DR ( ψ ) ( o ) = { D μ , π ( o ) − ψ 1 ( v ) } 2 . This estimator is referred to as doubly robust because E 0 [ D μ , π ( O ) ∣ V = v ] = ψ 0 , 1 ( v ) if either μ = μ 0 or π = π 0 . Thus, we generally expect that ψ ˆ DR will be consistent for ψ 0 if either μ ˆ is consistent for μ 0 or π ˆ is consistent for π 0 .

2.3 Meta-learning with multiple treatment arms

These methods readily extend to the case when ∣ A ∣ > 2 . We describe analogs of each of the meta-learners when there are multiple treatment arms. First, we must choose a reference treatment against which we will compare all others. This may be a control group or it may just be one arm chosen arbitrarily. Throughout this article, we will let A = 0 be the reference treatment.

In the case when V = W , the S-learner procedure is largely unchanged. We regress Y on A and W to obtain μ ˆ a ( w ) , which can then be used to construct ψ ˆ S , a ( w ) ≔ μ ˆ a ( w ) − μ ˆ 0 ( w ) for each a ∈ A 0 .

In the case when V ⊂ W , the S-learner requires the estimation of μ 0 ( a , w ) and second-stage regression as mentioned earlier. For a given a , we could learn ψ ˆ S , a ( v ) by empirical risk minimization, based on the loss L S , μ a ( ψ ) ( o ) , such as mean squared-error, L S , μ a ( ψ ) ( o ) = { D μ a ( o ) − ψ ( v ) } 2 .

Alternatively, we could learn each ψ ˆ a ( v ) simultaneously by considering a summed loss function L S , μ ( ψ ) ≔ ∑ a ∈ A 0 ω ( a ) L S μ a ( ψ ) . Here, ω ( a ) is any positive weight. We index the loss by μ to enable distinction between the regression using the estimated nuisance parameters and the regression using the true nuisance parameters. To minimize this summed loss, we create ∣ A ∣ − 1 pseudo-outcomes for each individual, with the form of the pseudo-outcome for the i th individual associated with ψ a ( v ) calculated as

This results in a modified dataset with n ( ∣ A ∣ − 1 ) rows that can then be used with any common regression software to learn a regression of the stacked pseudo-outcome values onto covariates plus one-hot-encoded a values. This approach may provide some benefit when some levels of treatment are rarely observed.

To construct the DR-learner with multiple treatments, we similarly create ∣ A ∣ − 1 pseudo-outcomes for each individual. In this case, the pseudo-outcome for the i th individual associated with ψ 0 ( a ∣ v ) is

Note that if an individual i has A i = a ≠ 0 , (3) reduces to

If instead A i = 0 , then (3) reduces to

Finally, if individual i has A i ∉ { 0 , a } , (3) reduces to

This pseudo-outcome form is doubly robust as established in the following.

The proof of Theorem 2.1 can be found in Section S1.1 of Supplementary Material. As with the S-learner, the CATE estimate is obtained by regressing the stacked pseudo-outcomes D μ ˆ , π ˆ a ( O ) against the covariates V and one-hot-encoded a values.

2.4 Causal trees and forests

Athey and Imbens [7] developed causal trees for estimating the CATE in a way that can be represented as a decision tree. The method involves recursively partitioning the feature space into regions that have constant treatment effect. They employ what they call “honest estimation” of the CATE by using separate sets of data to partition the feature space and estimate the treatment effects within the resulting partitions. Partitioning is done in a greedy way to minimize an estimate of the expected mean square error of the treatment effect. The optimal depth of the tree is selected via cross-validation, and the tree is then pruned to the selected depth. Finally, treatment effects are estimated within each region of the partition using a held-out-dataset. An extension of this method, causal forests, applies bootstrap aggregation to causal trees in order to obtain smoother estimates of the CATE [8]. We compare the performance of both causal trees and causal forests to our method in Section 5.2.

3.1 Background

Our method centers on the use of HAL for regression in the final stage of the chosen meta-learning algorithm. We briefly introduce HAL and its properties in the context of a general supervised learning task. Later, we extend these properties to estimation of the CATE.

HAL is a general approach for the estimation of functional parameters using regularized empirical risk minimization. The theory of HAL is built around two assumptions about the underlying functional parameter: (i) that it is right-continuous with left-hand limits and (ii) that it has a finite variation norm. The assumption of smoothness defined by a function’s variation norm has been long studied in the statistical learning literature [11].

Benkeser and van der Laan [12] and van der Laan [13] discussed an approach for finding the minimum loss-based estimator (MLE) in the class of functions with variation norm smaller than a finite constant. The MLE is computed based on the fact that any cadlag function with finite variation norm can be arbitrarily well approximated by a tensor product of indicator basis functions. If learning, for example, the mean of a continuous-valued random variable Z conditional on a univariate continuous-valued variable C , then a HAL estimate is created by first mapping C into n basis functions of the form 1 [ C i , ∞ ) ( C ) for i = 1 , … , n . Thus, the functional form of the model can be expressed as E [ Z ∣ C ] = β 0 + β 1 , C 1 [ C 1 , ∞ ) ( C ) + β 2 , C 1 [ C 2 , ∞ ) ( C ) + … + 1 [ C n , ∞ ) ( C ) for some ( β 0 , β 1 , C , … , β n , C ) ∈ R n + 1 . In this way, the HAL model parameterizes a step function that steps at each observed value of C . To fit the estimator, a relevant empirical risk criterion (e.g., mean-squared error) is optimized over model coefficients while penalizing the L 1 norm of the ( n + 1 ) -dimensional coefficient vector [14]. The coefficient for the penalty is selected using cross-validation. As the dimension of covariates increases, HAL is extended by including tensor product basis functions. For example, in the bivariate case C = ( C 1 , C 2 ) the functional form of the model is extended to E [ Z ∣ C 1 ] = β 0 + β 11 , C 1 [ C 11 , ∞ ) ( C 1 ) + ⋯ + 1 [ C 1 n , ∞ ) ( C 1 ) + β 21 , C 1 [ C 21 , ∞ ) ( C 2 ) + ⋯ + 1 [ C 2 n , ∞ ) ( C 2 ) + β 1 , C 1 [ C 11 , ∞ ) , [ C 21 , ∞ ) ( C 1 , C 2 ) + ⋯ + 1 [ C 1 n , ∞ ) , [ C 2 n , ∞ ) ( C 1 , C 2 ) , where C j i is the observed value of C j for the i -th observation in the data. In this way, the HAL model parameterizes a histogram regression that jumps at observed values of ( C 1 , C 2 ) . The idea generalizes to arbitrary covariate dimensions. At covariate dimension p , the HAL model includes at most n ( 2 p − 1 ) basis functions included (see Benkeser and van der Laan [12] for additional examples and details).

Due to the construction of the basis functions, the L 1 -norm of the model coefficients equals the variation norm of the estimate. Benkeser and van der Laan [12] showed that HAL estimates converge in terms of regret at a rate of O p ( n − [ 1 ⁄ 4 + 1 ⁄ { 8 p + 1 } ] ) , where p is the dimension of the feature vector. The authors also demonstrated competitive performance of HAL with other commonly used machine learning frameworks across a variety of empirical experiments. This bound was later refined to include sieves that allow the variation norm of the function to grow with sample size while achieving a rate of O p ( a n n − 1 ⁄ 3 log ( n ) ( 2 p − 1 ) ⁄ 3 ) , where a n is a factor that controls the rate of growth of the variation norm [15].

3.2 Highly adaptive regression trees

Previous work has shown that regression fits generated by the HAL algorithm have an identical representation as trees [9] referred to as HART. This result is not entirely surprising given HAL’s representation as a histogram regression. Any histogram regression could be represented as a tree, where each leaf in the tree would correspond to a single “bar” in the histogram (i.e., a region of the covariate space over which the model returns a constant prediction). The challenge in representing general histogram regressions as trees is that the partitioning of the covariate space implied by an arbitrary histogram regression is not necessarily recursive. This is in contrast to the typical approach for building decision trees in regression problems, which builds trees by explicitly recursively partitioning the covariate space until a certain goodness-of-fit criterion is reached. For general histogram regressions, achieving a tree representation is more challenging, given that there may be many non-rectangular “bars” in the histogram regression that must be represented by multiple leaves in the tree. Nizam and Benkeser [9] solved this problem specifically for HAL by introducing an algorithm that maps a fitted HAL model into a decision tree representation. The algorithm involves generating a set of candidate values implied by the non-zero coefficients in a HAL fit and using those values to recursively partition the feature space. The algorithm is structured such that in each sub-region of the feature space, the set of candidate values is appropriately pruned to minimize the amount of redundancy in the final tree representation, i.e., the algorithm yields the fewest trees leaves possible for representing non-rectangular regions in the covariate space.

These tree representations make HAL an ideal candidate for the estimation of CATEs and, by extension, treatment policies.

4.1 Theoretical guarantees for S and DR learners

Let μ ˆ and π ˆ be given estimates of μ 0 and π 0 . Furthermore, let R ˆ p * ≔ O p ( n − [ 1 ⁄ 4 + 1 ⁄ { 8 p + 1 } ] ) , which is the regret rate of the HAL estimator when estimating a function of bounded variation. As discussed in the introduction, we could equivalently set R ˆ p * ≔ O p ( a n n − 1 ⁄ 3 log ( n ) ( 2 p − 1 ) ⁄ 3 ) as in Bibaut and van der Laan [15]. Finally, for general P 0 -measurable function f of O , we define the L m ( P 0 ) norm of f to be ∣ ∣ f ∣ ∣ m , P 0 ≔ ∫ f m d P 0 1 ⁄ m .

A proof of Theorem 4.1 can be found in Section S1.2. Recall that when V = W , the S-learner only involves the estimation of the outcome regression μ 0 ( a , v ) . Intuitively then, when HAL is used to estimate μ 0 ( a , v ) , the S-learner CATE estimate converges in terms of regret to the true CATE at a rate of R n , p * . When V ⊂ W , the S-learner requires the estimation of μ 0 as a nuisance parameter. We consider the rate at which the S-learner CATE estimate constructed with the true outcome regression converges in terms of regret to the true CATE. We see that when HAL is used for the second-stage regression, this rate depends both on the regret rate of the outcome regression estimator and that of the HAL pseudo-outcome regression estimator.

We have the following result for the DR learner.

Theorem 4.2 gives a similar result for the DR-learner as with the S. learner, but where the rate now depends on the regret rates of the estimates of π 0 , a ( w ) and μ 0 , a ( w ) . A proof can be found in Section S1.3. The theorem indicates that the DR-learner may enjoy advantages over the S-learner. Consider a scenario in which the CATEs and propensity score are easier to learn than the outcome regression. When using the S-learner, the overall regret rate may be slowed by poor estimation of the outcome regression. However, the rate of the DR-learner involves the product of the regret rates of the propensity score and the outcome regression. Thus, the ability to quickly learn the propensity score may mitigate the impact of poor outcome regression estimation, thereby achieving an overall rate closer or equal to R ˆ q * . This scenario is plausible in several settings including when the data are generated from randomized experiments.

Moreover, the pseudo-outcome regression may be easier to learn than the outcome regression itself simply because the contrast function μ 0 , a − μ 0,0 may be simpler or smoother function than μ 0 itself. In the case of HAL, the most relevant measure of smoothness is the variation norm of the respective functions. In the case when V ⊂ W , we may also expect that ψ 0 will have lower variation norm than μ 0 , a − μ 0,0 , as it conditions on a subset of variables, thereby averaging over covariates across which there may be substantial variation in μ 0 .

4.2 Tree representations

When HAL is used as the final-stage regression estimator for either the S- or DR-learning procedures, the resulting CATE and OTP estimates can be represented using trees. Both of those procedures rely on extensions of the original HART algorithm [9]. In the case when there are two treatments, we can simply apply the original HART algorithm to build a CATE representation and apply a threshold of 0 to the terminal nodes in that tree in order to obtain an OTP estimate representation. In the case when there are more than two treatments, we must account for the treatment indicator variables in the tree. The HART algorithm allows the user to select the order that variables appear in the tree. Typically, we aim to find an ordering that minimizes the size of the tree. However, in this case, it is convenient to enforce that the treatment indicator variables appear first in the tree so that the effects for each treatment arm relative to the reference treatment are represented by distinct regions of the tree. Full details for representing the CATE estimate as a tree are shown in Algorithm 1.

When representing the OTP in the case of more than two treatments, we can enforce that the treatment indicator variables appear last in the tree, just before the terminal nodes in the regions where they are necessary to split on. We can then replace the treatment indicator variable nodes and the terminal nodes with terminal nodes displaying the treatment that has the highest expected counterfactual outcome. Full details for representing the OTP estimate can be found in Algorithm 2.

5.1 Motivating study

Antibiotics for Children with severe Diarrhea (ABCD) was a randomized clinical trial designed to evaluate the efficacy of azithromycin, an antibiotic commonly used against diarrheal pathogens, on mortality and linear growth in dehydrated or undernourished children presenting with acute non-bloody diarrhea [16]. While the trial did not show a survival benefit, there was a small benefit in linear growth. Further analysis identified benefits of azithromycin among children with likely bacterial etiologies of diarrhea and suggested a broader range of factors beyond diarrhea etiology could help determine treatment benefit [17]. Interpretable OTPs could help incorporate these additional factors into treatment guidance, presenting an ideal use case for our an interpretable method for determining treatment.

We applied both the S- and DR-learner versions of HATT to a dataset simulated to closely follow the true ABCD data for linear growth. We also compared the performance of S- and DR-learner to that of causal trees. For all estimators, we used fivefold cross-validation to select any required tuning parameters. Finally, we used repeated training/test sets to evaluate the methods on the ABCD trial data to compare the generalization performance of treatment rules on test data. Finally, we generate a treatment recommendation based on HATT.

Additional simulations are included in the supplemental material that examine the asymptotic properties of the estimator and compare to a wider variety of competitor approaches.

5.2 Simulations

5.3 Real data

5.3.2 Results

We found that across a range of thresholds, treatment rules based on HAL tended to lead to better estimates of expected outcomes than those generated by causal tree, although the differences were relatively modest. The HAL-based learners tended to be more liberal in terms of treatment recommendation. Overall, the results provide some evidence of modestly improved performance of HAL-based estimators over causal trees.

We applied the HATT algorithm to the HAL DR-learner for the lowest non-trivial threshold (0.025) for the diarrhea outcome to assess the scientific feasibility of the resultant tree. The tree is displayed in Figure 1. The tree recommended treatment for all children with baseline LAZ < − 1.7 , indicating that the children with greatest levels of malnourishment may stand to benefit the most from antibiotic therapy. Among the remaining children, with higher LAZ, treatment was recommended in children under 6.6 months of age and children who had Shigella detected in their stool. Overall, the rule would have recommended that 88% of children in the ABCD trial receive antibiotics (Tables 2 and 3).

Figure 1

Treatment recommendation tree for the ABCD data using the HAL DR-learner.

Table 2

Estimated expected day 90 LAZ under treatment rule accounting for baseline LAZ, Shigella, and age

	Threshold	Expected day 90 LAZ	Proportion treated
HAL (DR-learner)	0.000	− 1.602 ( − 1.642 , − 1.573 )	1.000 (1.000, 1.000)
	0.025	− 1.605   ( − 1.646 , − 1.580 )	1.000 (1.000, 1.000)
	0.050	− 1.641   ( − 1.677 , − 1.599 )	0.596 (0.000, 1.000)
	0.075	− 1.669   ( − 1.700 , − 1.639 )	0.000 (0.000, 0.029)
	0.100	− 1.671   ( − 1.700 , − 1.638 )	0.000 (0.000, 0.000)
HAL (S-learner)	0.000	− 1.602 ( − 1.642 , − 1.577 )	0.987 (0.933, 1.000)
	0.025	− 1.626   ( − 1.659 , − 1.584 )	0.800 (0.731, 0.930)
	0.050	− 1.643   ( − 1.670 , − 1.604 )	0.511 (0.353, 0.635)
	0.075	− 1.657   ( − 1.686 , − 1.630 )	0.127 (0.023, 0.274)
	0.100	− 1.667   ( − 1.697 , − 1.637 )	0.023 (0.000, 0.087)
Causal tree	0.000	− 1.610 ( − 1.644 , − 1.579 )	0.680 (0.643, 0.723)
	0.025	− 1.614   ( − 1.655 , − 1.575 )	0.658 (0.619, 0.706)
	0.050	− 1.613   ( − 1.647 , − 1.581 )	0.635 (0.576, 0.676)
	0.075	− 1.622   ( − 1.652 , − 1.583 )	0.587 (0.487, 0.658)
	0.100	− 1.633   ( − 1.662 , − 1.591 )	0.524 (0.462, 0.617)

Median (IQR) over 500 splits of data into 70% training and 30% test.

Table 3

Estimated proportion of children with diarrhea at day 3 post-enrollment under treatment rule accounting for baseline LAZ, Shigella, and age

	Threshold	Proportion with diarrhea at day 3 post-enrollment	Proportion treated
HAL (DR-learner)	0.000	0.129 (0.120, 0.141)	0.973 (0.929, 1.000)
	0.025	0.137 (0.125, 0.148)	0.710 (0.603, 0.906)
	0.050	0.148 (0.137, 0.161)	0.453 (0.382, 0.509)
	0.075	0.173 (0.162, 0.185)	0.102 (0.071, 0.147)
	0.100	0.183 (0.172, 0.195)	0.036 (0.000, 0.065)
HAL (S-learner)	0.000	0.129 (0.121, 0.139)	0.995 (0.992, 0.997)
	0.025	0.136 (0.126, 0.148)	0.900 (0.793, 0.955)
	0.050	0.170 (0.154, 0.184)	0.167 (0.068, 0.235)
	0.075	0.185 (0.176, 0.197)	0.000 (0.000, 0.016)
	0.100	0.185 (0.177, 0.199)	0.000 (0.000, 0.000)
Causal tree	0.000	0.133 (0.123, 0.145)	0.722 (0.663, 0.753)
	0.025	0.135 (0.125, 0.147)	0.617 (0.542, 0.702)
	0.050	0.139 (0.131, 0.153)	0.523 (0.452, 0.569)
	0.075	0.142 (0.133, 0.156)	0.396 (0.327, 0.480)
	0.100	0.153 (0.139, 0.161)	0.296 (0.193, 0.353)

Median (IQR) over 500 splits of data into 70% training and 30% test.