Heterogeneous treatment effect estimation for observational data using model-based forests

TL;DR

This paper introduces a method to estimate how treatment effects vary across individuals in observational data using modified model-based forests.

Contribution

The paper proposes a new approach to reduce confounding in model-based forests for observational data.

Findings

The orthogonalization strategy reduces confounding in simulated studies with various outcome distributions.

The method is applied to assess the heterogeneous effect of Riluzole on Amyotrophic Lateral Sclerosis progression.

The approach works for survival and ordinal outcomes in generalized linear and transformation models.

Abstract

The estimation of heterogeneous treatment effects has attracted considerable interest in many disciplines, most prominently in medicine and economics. Contemporary research has so far primarily focused on continuous and binary responses where heterogeneous treatment effects are traditionally estimated by a linear model, which allows the estimation of constant or heterogeneous effects even under certain model misspecifications. More complex models for survival, count, or ordinal outcomes require stricter assumptions to reliably estimate the treatment effect. Most importantly, the noncollapsibility issue necessitates the joint estimation of treatment and prognostic effects. Model-based forests allow simultaneous estimation of covariate-dependent treatment and prognostic effects, but only for randomized trials. In this paper, we propose modifications to model-based forests to address the confounding issue in observational data. In particular, we evaluate an orthogonalization strategy originally proposed by Robinson (1988, Econometrica) in the context of model-based forests targeting heterogeneous treatment effect estimation in generalized linear models and transformation models. We found that this strategy reduces confounding effects in a simulated study with various outcome distributions. We demonstrate the practical aspects of heterogeneous treatment effect estimation for survival and ordinal outcomes by an assessment of the potentially heterogeneous effect of Riluzole on the progress of Amyotrophic Lateral Sclerosis.

Figures

Model-based forest results for the empirical study (Section 4), Cox means a Cox model applied to the Weibull data. For the Weibull and Cox model, treatment effects τ(x) are estimated as conditional log hazard ratios. Direct comparison of model-based forests without centering (Naive), model-based forests with local centering according to Robinson 34 of Y and W (originally proposed) (Robinson) or only of W ( RobinsonW^ ).

Kaplan-Meier curves of survival probability for both treatment arms.

Distribution of estimated propensities π^(x) (left) and estimated propensities of the centered treatment indicators (right, Robinson’s strategy) as estimated by regression forests for the two treatment groups.

Kernel density estimates of the personalized treatment estimates for the naive model-based forest (Naive) and for the model-based forest with Robinson orthogonalization (Robinson).

Relative frequency distribution plot of the handwriting ability score ( Y6 ) (left) and of changes of the handwriting ability score over 6 months ( Y6−Y0 ) (right) for both treatment arms. Frequencies were calculated relative to the treatment indicator.

Alluvial plot of the progression of the handwriting ability score over 6 months for both treatment arms.

Estimates returned by the regression forest for orthogonalization of the treatment indicator: left for original W as an outcome in the random forest such that it estimates propensity scores π(x,y0) ; right for the centered treatment indicator W−π^(x,y0) as an outcome in the random forest.

Kernel density estimates of the personalized treatment estimates for the naive model-based forest (Naive) vs. the forest with Robinson orthogonalized (Robinson).

Full text

Introduction

Over the past years, there has been emerging interest in methods to estimate heterogeneous treatment effects (HTEs) in various application fields. In healthcare, HTE estimation can be understood as a core principle driving personalized medicine. As opposed to average treatment effects, which assume a constant effect of a treatment on an outcome for the whole population, HTEs account for the heterogeneity in the effect for subgroups or individuals based on their characteristics. Most research on HTE estimation has mainly focused on continuous and binary response variables. These methods have typically built upon Rubin’s potential outcomes framework, a statistical approach to formulating and inferring causal effects in various designs.^1,2^

Traditionally, statistical models were used to estimate the treatment effect, but machine learning methods have been more and more adapted for these tasks over the past decade. Machine learning models rely on weaker assumptions and can automatically learn complex relationships such as higher order interaction effects, resulting in greater predictive performance in a variety of applications. In the case of continuous or binary responses, prominent methods to estimate HTEs are based on random forests,^3????–8^ Bayesian additive regression trees (BART),^9,10^ or neural networks.^11?–13^ Kuenzel et al.^ 14 ^ proposed general frameworks—T-learners, S-learners, U-learners, and X-learners—that base treatment effect estimates on arbitrary machine learning models. Chernozhukov et al.^ 15 ^ coined the term double/debiased machine learning models, which uses machine learning models for nuisance parameter estimation. The approach still relies on parametric models for estimating treatment effects, but Nie and Wager^ 16 ^ derived so-called R-learners that allow for arbitrary (nonparametric or semiparametric) models.

Beyond continuous or binary responses, research on machine learning methods for HTE estimation have primarily focused on (right-censored) survival data. Methods have been proposed based on BARTs,^ 17 ^ random forest-type methods,^18,19^ or deep learning approaches.^12,13^ Theoretically, any machine learning model for survival analysis—such as random survival forests^ 20 ^ or a Cox regression-based deep neural network (deepSurv)^ 21 ^—can estimate HTEs.^ 22 ^ These models can estimate survival or hazard functions in both treatment groups separately; HTEs are then defined as the difference in derived properties of the two functions, e.g. as differences in the median survival time. However, Hu et al.^ 22 ^ found that methods specifically designed for HTE estimation, like the adapted BART,^ 17 ^ produce more reliable estimates.

In general, for a continuous or binary outcome $[eqn]$ conditional on treatment $[eqn]$ and covariates $[eqn]$ , the conditional average treatment effect (CATE) $[eqn]$ can be estimated from the model $[eqn]$ even if the model is misspecified, e.g. when the prognostic effect $[eqn]$ cannot be fully estimated due to missing covariate information. Beyond mean regression, stricter assumptions are necessary both for randomized and for observational studies to estimate HTEs. For example, under a true Cox model with survivor function $[eqn]$ with log-cumulative baseline hazard $[eqn]$ at time $[eqn]$ and log-hazard ratio $[eqn]$ , the prognostic effect $[eqn]$ must be specified correctly, even in a randomized trial. Estimated marginal log-hazard ratios $[eqn]$ – i.e. when the model is fitted under the constraint $[eqn]$ – are shrunken towards zero if this constraint is unrealistic.^ 23 ^ Naturally, this problem carries over to heterogeneous log-hazard ratios $[eqn]$ .

Consequently, HTE estimation in more complex models requires the simultaneous estimation of both the prognostic part $[eqn]$ and the predictive HTE $[eqn]$ . Model-based forests allow simultaneous estimation of $[eqn]$ and $[eqn]$ ^24????–29^ in the same forest model and have been demonstrated to perform at least on par with the best competitor in several independent studies evaluating HTE estimation in randomized trials.^30??–33^ In a nutshell, model-based forests combine the parametric modeling framework with random forests to estimate individual treatment effects.^ 25 ^ By using generalized linear models and transformation models, model-based forests can be adapted for survival data,^24?–26^ ordinal data,^ 27 ^ or clustered data.^ 28 ^

In observational studies, the treatment group assignment is not under control of the researcher and confounding effects could bias the estimation of HTEs. Robinson’s^ 34 ^ orthogonalization strategy has shown to be instrumental for achieving robustness to confounding effects for HTE estimation in observational studies.^5,35,36^ In this work, we propose and evaluate novel model-based forest variants based on the orthogonalization strategy that are suitable for HTE estimation in the observational setting. Compared to previous work, our work focuses on the utilization of model-based forests for HTE estimation and covers a richer class of models—(generalized) linear and transformation models. Simulation and application studies demonstrate the efficacy of our method for estimating HTEs under confounding biases yielding a versatile model framework applicable to a wide range of application cases. We contribute the first holistic approach to HTE estimation for observational studies which is applicable to many clinically relevant outcomes, including survival times, ordinal scores, binary events, or counts.

We review key components of model-based forests for HTE estimation in randomized trials in Section 2. In Section 3, we start introducing the orthogonalization approach by Robinson.^ 34 ^ We motivate previous developments using linear models^ 36 ^ and leverage adaptations to more complex models discussed by Gao and Hastie^ 35 ^ to define novel model-based forest variants suitable for HTE in the observational setting. These variants’ performances are empirically assessed in a simulation study with a range of outcome distributions in Section 4. Finally, in Section 5 presenting a re-analysis of the patient-specific effect of Riluzole in patients with amyotrophic lateral sclerosis (ALS), practical aspects of model estimation and interpretation are discussed.

Review of model-based forests for randomized trials

We are interested in estimating HTEs based on i.i.d. observations $[eqn]$ with $[eqn]$ , where $[eqn]$ , $[eqn]$ and $[eqn]$ are realizations of the outcome $[eqn]$ , covariates $[eqn]$ , and control vs. treatment indicator $[eqn]$ . For the remainder of this manuscript (with the exception of equation (4)), we will omit the subscript $[eqn]$ . $[eqn]$ and $[eqn]$ denote the potential outcomes under the two treatment conditions $[eqn]$ . Throughout this paper, we assume that $[eqn]$ includes all relevant variables to explain heterogeneity both in the treatment effect and the outcome $[eqn]$ , and that the base model underlying model-based forests is correctly specified.

We review model-based forests for HTE estimation based on randomized trials as introduced by Seibold et al.^ 25 ^ and Korepanova et al.^ 26 ^ Within this section, we only consider settings where the treatment assignment is randomized and, therefore, follows a binomial model $[eqn]$ with constant propensities $[eqn]$ . We omit discussion of the abstract framework underlying model-based forests and instead discuss the important linear, generalized linear,^ 25 ^ and transformation models^ 26 ^ in detail.

Linear model

2.1.

For a continuous outcome $[eqn]$ with symmetric error distribution, a model-based forest might be defined based on the model

$$[eqn]$$

where the residuals are given by the error term $[eqn]$ with $[eqn]$ and standard deviation $[eqn]$ .^ 36 ^ We are mainly interested in estimating $[eqn]$ , the treatment effect that depends on predictive variables in $[eqn]$ . With model-based forests, however, we also obtain an estimated value for the prognostic effect $[eqn]$ , which depends on prognostic variables in $[eqn]$ . A variable might be predictive and prognostic at the same time. We refer to these situations as “overlays.”

Because we assume in this section that $[eqn]$ applies, $[eqn]$ holds. Consequently, $[eqn]$ can be interpreted as a CATE

$$[eqn]$$

on the absolute scale. To estimate $[eqn]$ the $[eqn]$ loss

$$[eqn]$$

is minimized w.r.t. $[eqn]$ and $[eqn]$ using an ensemble of trees. Inspired by recursive partitioning techniques,^37,38^ split variable and split point selection are separated. Specifically in every node, $[eqn]$ and $[eqn]$ in the constant model $[eqn]$ are estimated using the observations in the node. The split variable is then the variable that has the lowest $[eqn]$ -value for the bivariate permutation tests for the $[eqn]$ -hypothesis that $[eqn]$ and $[eqn]$ are constant and independent of any split variable. The cut-point is the point of the chosen split variable at which the score functions

$$[eqn]$$

in the two resultant subgroups differ the most; details are available in Appendix 2 of Seibold et al.^ 25 ^

Once $[eqn]$ trees were fitted to subsamples of the training data, predictions for the treatment effect for a new observation $[eqn]$ are obtained via local maximum likelihood aggregation.^5,29,39–41^ First, for the $[eqn]$ th training sample, the frequency $[eqn]$ with which it falls in the same leaf as $[eqn]$ over all $[eqn]$ trees is measured. The obtained weighting vector $[eqn]$ is used as an input for minimizing

$$[eqn]$$

where $[eqn]$ denotes the loss for the $[eqn]$ th sample. Model-based forests easily allow adaptions if HTEs for an outcome variable $[eqn]$ that is not well represented by equation (1) should be estimated. In this case, model-based forests can build on generalized linear models or transformation models in the recursive partitioning framework.^ 38 ^ As detailed in the following sections, the loss function $[eqn]$ in equation (3) changes from the squared error to the negative (partial) log-likelihood of some appropriate model.

Generalized linear models

2.2.

When the conditional outcome distribution is better described through a generalized linear model

$$[eqn]$$

with parameter $[eqn]$ depending on the additive function $[eqn]$ , the conditional mean

$$[eqn]$$

is linear on the scale of a link function $[eqn]$ . Thus, the interpretation of $[eqn]$ as CATE (2) generally no longer holds. Instead, the predictive effect is understood as the difference in natural parameters (DINA)^ 35 ^

$$[eqn]$$

In contrast to the linear model case, HTEs $[eqn]$ are now defined on relative scales, such as odds ratios in binary logistic regression models or multiplicative mean effects in a Poisson or Gaussian model with a log-link. The negative log-likelihood contribution of some observation $[eqn]$ is

$$[eqn]$$

with $[eqn]$ as the conditional density of an exponential family distribution

$$[eqn]$$

Model-based trees and forests^38,24,25^ jointly estimate the prognostic effect $[eqn]$ and the predictive effect $[eqn]$ . The procedure simultaneously minimizes the negative log-likelihood with respect to $[eqn]$ and $[eqn]$ . In each node of the model-based forest, $[eqn]$ , $[eqn]$ , and potentially $[eqn]$ are estimated by minimizing

$$[eqn]$$

and regressing the bivariate gradient

$$[eqn]$$

on $[eqn]$ . This means that one is not explicitly looking for changes in the scale parameter $[eqn]$ , but this could be implemented by looking at the three-variate gradient

$$[eqn]$$

for example, in a heteroscedastic normal linear model

$$[eqn]$$

After the tree fitting phase, an HTE is estimated with equation (4) with $[eqn]$ of equation (7) as the corresponding loss function.

Thus, model-based forests can be directly applied to estimate HTEs on relative scales for binary outcomes (binary logistic or probit regression, for example), counts (Poisson or quasi-Poisson regression), or continuous outcomes where a multiplicative effect is of interest (normal model with log-link).

Transformation models

2.3.

More complex responses like ordered categorical or time-to-event outcomes are not covered by generalized linear models but can be analyzed using transformation models; corresponding model-based forests for survival analysis have been introduced by Korepanova et al.^ 26 ^ For some at least ordered outcome $[eqn]$ , we write the conditional distribution function as

$$[eqn]$$

The transformation function $[eqn]$ is monotone non-decreasing and the inverse link function $[eqn]$ governs the interpretability of $[eqn]$ as log-odds ratios ( $[eqn]$ ), log-hazard ratios ( $[eqn]$ ), log-reverse time hazard ratios ( $[eqn]$ ), or shift effects ( $[eqn]$ , the cumulative distribution function of the standard normal). The shift term $[eqn]$ differs between the two treatment groups $[eqn]$ . The distribution functions of the potential outcomes are $[eqn]$ for $[eqn]$ and $[eqn]$ for $[eqn]$ . The negative log-likelihood of a discrete or interval-censored observation ( $[eqn]$ (where $[eqn]$ is the lower interval bound, $[eqn]$ is the upper) is

$$[eqn]$$

For a continuous datum $[eqn]$ , we obtain

$$[eqn]$$

details are given in Hothorn et al.^ 42 ^ Transformation forests apply the model-based recursive partitioning principle and estimate $[eqn]$ in each node along with the transformation function $[eqn]$ (a “nuisance” parameter) by minimizing $[eqn]$ .^ 29 ^ Because $[eqn]$ contains an intercept term, the parameter $[eqn]$ is not identified. We thus estimate the model under the constraint $[eqn]$ . Variable and cut-points are selected using the bivariate gradient

$$[eqn]$$

This model family includes proportional odds logistic regression (for ordered categorical, count or continuous outcomes), Box-Cox type models, Cox proportional hazards model, Weibull proportional hazards models for discrete and continuous outcomes, reverse time proportional hazards models relying on Lehmann alternatives, and many more.^ 42 ^ Forests for ordinal outcomes were evaluated by Buri and Hothorn,^ 27 ^ and a general approach to “transformation forests” is described in Hothorn and Zeileis.^ 29 ^

Application of the ideas underlying model-based forests allows HTEs to be estimated for such outcomes under all types of random censoring and truncation.^ 26 ^ For example, for Weibull distributed outcomes under right censoring, $[eqn]$ is chosen for the conditional distribution function in equation (8).^ 42 ^

In this case, we define $[eqn]$ as the event time, $[eqn]$ as the censoring time and $[eqn]$ as the observed time. For identification of $[eqn]$ under potential censoring, the following assumption must hold^ 18 ^: Assumption 1(Ignorable censoring)Censoring time $[eqn]$ is independent of survival time $[eqn]$ conditional on treatment indicator $[eqn]$ and covariates $[eqn]$

$$[eqn]$$

An important special case represents the Cox proportional hazards model, where the profile likelihood over the baseline hazard function defines the partial log-likelihood $[eqn]$ with $[eqn]$ . The scores with respect to the constant $[eqn]$ are known as martingale residuals. Model-based forests for such models, and extensions to time-varying prognostic and predictive effects, are discussed in Korepanova et al.^ 26 ^

Noncollapsibility

2.4.

As mentioned in the introduction, one problem with the Cox model is that misspecifications of prognostic effects $[eqn]$ lead to biased estimates such that the estimated hazard ratios cannot be interpreted causally. This issue arises from the noncollapsiblity of the Cox model, the notion of which is characterized by the fact that in these models, the mean of the conditional effect estimates defined over covariates $[eqn]$ does not coincide with the marginal effect over $[eqn]$ . Because the noncollapsiblity of the Cox model arises from its nonadditivity of the hazard function, models such as a Weibull model (parameterized in form of an accelerated failure time model) do not suffer from this issue because they satisfy the additivity condition. Consequently, misspecifications of prognostic effects do not affect treatment effect estimates.^ 23 ^

The noncollapsibility issue is not limited to the Cox model but also affects members of the exponential family without identity or linear link functions. Without adjustments, effect estimates can only be interpreted causally if there is no treatment effect ( $[eqn]$ ) or there are no prognostic covariates.^ 43 ^

If this is not the case, specific methods are needed; ignoring the estimation of $[eqn]$ at all and only focusing on $[eqn]$ does not solve the problem. Conditioning on available prognostic variables is a common solution and is already applied by model-based forests because they estimate both the prognostic effect $[eqn]$ and $[eqn]$ . The ensemble of trees used to estimate these effects provides a high degree of flexibility and might therefore retain some of the potential complexity in the underlying $[eqn]$ to mitigate misspecification. Whether conditioning resolves the noncollapsibility issue depends heavily on the assumption that all prognostic variables are known which is often not the case in the real world.^ 23 ^

For members of the exponential family and the Cox model, Gao and Hastie^ 35 ^ derived a method to account for noncollapsibility in the context of observational data with confounding effects. While we consider the noncollapsibility issue beyond the scope of this work, we briefly review the work of Gao and Hastie and discuss its applicability to model-based forests in Section A of the Supplemental Material.

Model-based forests for observational studies

In the previous section, we described model-based forests in the randomized setting under the assumption that $[eqn]$ . In observational studies, in which the treatment group assignment is not under the control of the researcher, the propensity score (and therefore, the probability of being in the treatment group) often depends on covariates $[eqn]$

$$[eqn]$$

In this case, confounding effects could bias the estimation of treatment effects $[eqn]$ , and stricter assumptions are necessary in order to interpret $[eqn]$ causally.^ 44 ^ Assumption 2(Ignorability/Unconfoundedness)The treatment assignment is independent of the potential outcomes conditional on covariates $[eqn]$

$$[eqn]$$

Assumption 3(Positivity)The propensity score $[eqn]$ must be bounded away from 0 and 1

$$[eqn]$$

Assumption 2 could be violated by an unmeasured confounder, while Assumption 3 could be violated if all observations in a certain group (defined via $[eqn]$ ) are in the treatment group. Effects and potential solutions to these violations are well-studied in the literature.^45??–48^

Dandl et al.^ 36 ^ showed for mean regression models that model-based forests are not robust to confounding effects and need further adaptions to estimate causal effects in case of observational data. One strategy for dealing with confounding effects is the orthogonalization strategy originally introduced by Robinson,^ 34 ^ which has received considerable attention in recent years.^15,5,16^ The reformulation of the linear model

$$[eqn]$$

to

$$[eqn]$$

given the conditional mean function

$$[eqn]$$

motivates this approach.^5,36^

Overall, the orthogonalization strategy consists of two steps: First, nuisance parameters $[eqn]$ and $[eqn]$ are estimated. Originally, Robinson^ 34 ^ used kernel estimators, but any machine learning method could be employed.^15,16^ Regressing $[eqn]$ on $[eqn]$ then yields unbiased estimates for $[eqn]$ . Subtracting $[eqn]$ and $[eqn]$ from $[eqn]$ and $[eqn]$ , respectively, partially eliminates the association between $[eqn]$ and $[eqn]$ and between $[eqn]$ and $[eqn]$ . The orthogonalization strategy has the distinct advantage over other methods against confounding—such as inverse propensity weighting and matching—that it is stable for extreme propensity scores and forgoes stratification.^ 35 ^

Robinson^ 34 ^ and Chernozhukov et al.^ 15 ^ use parametric models to estimate treatment effects based on residualized $[eqn]$ and $[eqn]$ , but these models could be replaced by non-parametric or local parametric models^16,49^—such as model-based forests. For mean regression, Dandl et al.^ 36 ^ adapted the orthogonalization strategy to model-based forests. Their approach closely follows causal forests, which were the first to combine the orthogonalization strategy with tree-based estimators for $[eqn]$ .

Gao and Hastie^ 35 ^ proposed extensions of Robinson’s strategy to members of the exponential family and the Cox model, where DINA (6) is of interest. Gao and Hastie^ 35 ^ assume $[eqn]$ and use parametric models to estimate $[eqn]$ , but they conclude that non-parametric or local parametric models could be applied instead. We review model-based forests in combination with linear models for observational data in the next section and summarize the idea by Gao and Hastie^ 35 ^ in Section 3.2. On this basis, we assess how the orthogonalization strategy could be employed in model-based forests beyond mean regression with generalized linear models and transformation models as base models.

Review of Dandl et al.

36

3.1.

As noted above, Athey et al.^ 5 ^ were the first to combine the orthogonalization strategy of Robinson with tree-based estimators to estimate $[eqn]$ . First, the marginal model $[eqn]$ and propensity score $[eqn]$ are estimated by regression forests. Afterwards, causal forests estimate individual treatment effects $[eqn]$ in the model

$$[eqn]$$

using the “locally centered” outcomes $[eqn]$ and treatment indicators $[eqn]$ .

Equation (13) shows that causal forests and model-based forests share common foundations for mean regression. The main difference is that the splitting scheme of model-based forests allows splitting according to heterogeneity in both treatment and prognostic effects, whereas causal forests only split with respect to heterogeneity in treatment effects (in equation (11), $[eqn]$ cancels out).

Dandl et al.^ 36 ^ identified which elements of both approaches lead to improved performance in randomized trials and observational studies. They defined and evaluated the following blended versions of model-based forests and causal forests:

• (1)* $[eqn]$ *, which uses model-based forests to estimate the model

$$[eqn]$$

i.e. after centering the treatment indicator $[eqn]$ and the outcome $[eqn]$ . Both parameters $[eqn]$ and $[eqn]$ are estimated simultaneously.

• (2)* $[eqn]$ *, which uses model-based forests to estimate the model

$$[eqn]$$

i.e. after only centering the treatment indicator $[eqn]$ but not outcome $[eqn]$ . Both $[eqn]$ and $[eqn]$ are estimated.

• (3)cfmob, a method which uses model-based forests to estimate the model

$$[eqn]$$

i.e. after only centering the treatment indicator $[eqn]$ and splitting only according to $[eqn]$ . That is, only the parameters $[eqn]$ are estimated in this variant.

Their blended approaches competed with the original implementations of (uncentered) model-based forests and causal forests in an extensive simulation study. In case of confounding, the authors identified local centering of treatment indicator $[eqn]$ and simultaneous estimation of both predictive and prognostic effects of the treatment indication (* $[eqn]$ ) as the key driver for good performance. Additionally, centering $[eqn]$ ( $[eqn]$ *) is recommended, because it further improved performances in some cases. Splitting only according to $[eqn]$ but not $[eqn]$ (cfmob) resulted in lower predictive performance.

Review of Gao and Hastie (2022)

3.2.

Robinson^ 34 ^ derived the orthogonalization strategy only for semi-parametric additive models with $[eqn]$ . Gao and Hastie^ 35 ^ extended the idea to a broader class of distributions including the exponential family and Cox model.

Local centering of the treatment indicator works analogously to mean regression. First, propensity scores $[eqn]$ are estimated. The effects of the covariates $[eqn]$ on the treatment assignment are then regressed out by subtracting $[eqn]$ from $[eqn]$ .

Orthogonalization of $[eqn]$ is not straightforward due to the link function that relates the linear predictor $[eqn]$ in equation (5) to the outcome $[eqn]$ . To understand how Gao and Hastie derived $[eqn]$ to center $[eqn]$ , we consider equation (10) as a model of the exponential family with identity link function $[eqn]$ . Now we can rewrite equation (12) to

$$[eqn]$$

Accordingly, Gao and Hastie derived $[eqn]$ for all other distributions of the exponential family by

$$[eqn]$$

We can regard the estimated $[eqn]$ as an offset in the linear predictor

$$[eqn]$$

Note that equation (14) states that (only) for the Gaussian distribution we can directly estimate $[eqn]$ without estimating $[eqn]$ and $[eqn]$ . We can also derive $[eqn]$ for transformation models based on the definition of $[eqn]$ and $[eqn]$ in equation (8). As mentioned in Section 2.4, compared to the difference in conditional means, the DINA additionally suffers from the noncollapsibility issue.^ 50 ^ Gao and Hastie^ 35 ^ also extend the Robinson strategy to tackle not only the confounding but also the noncollapsibility issue for members of the exponential family (without a linear or log-link function, otherwise confounding is not an issue) and the Cox model. While it is beyond the scope of this work to address the noncollapsibility issue, Section A of the Supplemental Material briefly summarizes and discusses the applicability of Gao and Hastie’s approach to model-based forests.

Novel model-based forests for observational data

3.3.

As stated above, our main goal is to assess how the orthogonalization strategy proposed for continuous outcomes could be extended to models beyond mean regression, specifically generalized linear models and transformation models. Based on Dandl et al.^ 36 ^ and Gao and Hastie^ 35 ^ we propose two different versions of model-based forests, which should be more robust against confounding. Following Dandl et al.,^ 36 ^ we formulate research questions for these versions, which we aim to answer empirically in Section 4. An overview of all proposed versions is given in Table 1.

Table 1.: Overview of proposed model-based forest versions.

The first version of model-based forests directly applies Robinson’s orthogonalization strategy: First, we estimate propensities $[eqn]$ as well as $[eqn]$ and $[eqn]$ to derive $[eqn]$ . Then, we update the linear predictor of equation (5) by centering $[eqn]$ by $[eqn]$ and by adding the offset $[eqn]$ . For generalized linear models, we obtain

$$[eqn]$$

and for the conditional distribution function of equation (8) in case of transformation models

$$[eqn]$$

Based on the updated models, both prognostic and predictive effects $[eqn]$ and $[eqn]$ are simultaneously estimated by model-based forests.

In the simulation study and practical example in Sections 4 and 5, we use regression forests to estimate $[eqn]$ and gradient boosting machines (with tailored loss functions) to estimate $[eqn]$ and $[eqn]$ . In the following, we denote this version of model-based forests as Robinson forests in recognition of Robinson^ 34 ^ while model-based forests without centering $[eqn]$ and without offset $[eqn]$ are called Naive forests.

Similar to Dandl et al.–who saw an improvement in performance when only centering $[eqn]$ (compared to the naive model-based forests)—we define an approach called $[eqn]$ forests that applies model-based forests to models with linear predictors

$$[eqn]$$

Empirical evaluation

In a simulation study, we aim to compare the different model-based forest versions (Table 1) by answering the following research questions:

• RQ 1To what extent does centering $[eqn]$ by $[eqn]$ and including $[eqn]$ as an offset (Robinson) affect the performance of model-based forests in the presence of confounding?
• RQ 2Do centered treatment indicator model-based forests ( $[eqn]$ ) perform better than uncentered model-based forests in the presence of confounding?
• RQ 3Are model-based forests with centered treatment indicators ( $[eqn]$ ) relevantly outperformed by model-based forests with $[eqn]$ as an additional offset (Robinson) in the presence of confounding?

The study includes different outcome types, different predictive and prognostic effects, and a varying number of observations and covariates. Model-based forests were fitted with the model4you R add-on package.^ 51 ^ Similar to Dandl et al.,^ 36 ^ we base our study settings on the four setups (A, B, C and D) of Nie and Wager.^ 16 ^ In addition, in Section B of the Supplemental Material, we show the results for simulation settings first proposed by Wager and Athey^ 49 ^ and later reused by Athey et al.^ 5 ^

Data generating process

4.1.

Given $[eqn]$ , for Setup A, we sampled $[eqn]$ . For all other setups, we used $[eqn]$ . The treatment indicator was binomially distributed with $[eqn]$ . The propensity function $[eqn]$ differed for the four considered setups:

$$[eqn]$$

$[eqn]$ in Setup B implies a randomized study. The treatment effect function $[eqn]$ and the prognostic effect function $[eqn]$ also differed between the setups

$$[eqn]$$ $$[eqn]$$

Setup A has extensive confounding that must be eliminated before estimating an easily predictable treatment effect function $[eqn]$ . Setup B needs no confounding adjustment for reliable estimation of $[eqn]$ . Although Setup C contains strong confounding, the propensity score function is easier to estimate than the prognostic effect, while the treatment effect is constant. In Setup D, the treatment and control arms are unrelated, and therefore, learning the conditional expected outcomes of both arms jointly is not beneficial.^16,36^

To assess robustness of the methods to missing prognostic covariates, we additionally created Setup A’ from Setup A by removing covariate $[eqn]$ from the training data. Therefore, the DGP of Setup A and Setup A’ are identical, the only difference being that the training data did not contain $[eqn]$ although $[eqn]$ affects the prognostic effect.

We studied four different simulation models

$$[eqn]$$

Model (15a) is a normal linear regression model, model (15b) is a binary logistic regression model, model (15c) is a 4-nomial model with log-odds function $[eqn]$ with threshold parameters $[eqn]$ for $[eqn]$ , and model (15d) is a Weibull model with log-cumulative hazard function $[eqn]$ . We added $[eqn]$ random right-censoring to the Weibull-generated data. Additionally, we applied a Cox proportional hazards model to the Weibull data to determine if the performance of model-based forests degrades when the forests do not take the true underlying model as their base model.

Due to $[eqn]$ in all scenarios, half of the (negative) predictive effect $[eqn]$ was added to the prognostic effect. We refer to the implied scenario—where one variable which is both prognostic (impact in $[eqn]$ ) and predictive (impact in $[eqn]$ ) exists—as overlay. Apart from Setup C in which the treatment effect is constant and independent of any covariate, overlay was present for all scenarios.

Like Dandl et al.,^ 36 ^ we compared all study settings and outcome types for a varying number of samples $[eqn]$ and dimensions $[eqn]$ . All model-based forests were grown with the same hyperparameter options specified in Section 7. We used random forests as implemented in the grf package to estimate $[eqn]$ for centering $[eqn]$ .^ 52 ^ To estimate $[eqn]$ and $[eqn]$ to derive $[eqn]$ , we relied on different tree-based estimators depending on the outcome type. For normally distributed outcomes (models (15a)), we used grf regression forests.^ 52 ^ For all other outcomes, we relied on gradient boosting machines (with adapted loss functions) as implemented in mboost and gbm.^53,54^ The employed distribution varied depending on the outcome type.

In accordance with Dandl et al.,^ 36 ^ we evaluated the models with respect to the mean squared error $[eqn]$ , bias $[eqn]$ and standard error $[eqn]$ on a test sample of size $[eqn]$ . The results for the mean squared error are shown in Figure 1 and were statistically analyzed by means of a normal linear mixed model with a log-link. The model explained the estimated mean squared error for $[eqn]$ by a four-way interaction of the data generating process, sample size $[eqn]$ , dimension $[eqn]$ , and random forest variant. We estimated the mean squared error ratios between different model-based forest versions according to the research questions stated in Section 4. The corresponding tables are given in Tables 2? to 4. In the Supplemental Material, the results for the bias and standard error are presented in Figures S. 2 and S. 3, as well as the statistical analysis of the relative efficiency in Tables S. 4 to S. 6.

Model-based forest results for the empirical study (Section 4), Cox means a Cox model applied to the Weibull data. For the Weibull and Cox model, treatment effects τ(x) are estimated as conditional log hazard ratios. Direct comparison of model-based forests without centering (Naive), model-based forests with local centering according to Robinson 34 of Y and W (originally proposed) (Robinson) or only of W ( RobinsonW^ ).

Table 2.: Results of RQ 1 for the experimental setups in Section 4.

Table 3.: Results of RQ 2 for the experimental setups in Section 4.

Table 4.: Results of RQ 3 for the experimental setups in Section 4.

Results

4.2.

The results for the normal distribution coincide with the results obtained by Dandl et al.^ 36 ^ summarized in Section 3.1. To some degree, they also hold for the other distributions. The boxplots are not directly comparable between different data generating processes because of different signal-to-noise ratios. In general, a more informative outcome (binary < ordered < right-censored < exact normal), more data (higher $[eqn]$ ), and less noise (lower $[eqn]$ ) leads to better results. Using a Cox model compared to a Weibull model (last two rows of Figure 1) did not lead to a major decrease in performance, although knowledge of the true functional form of the transformation function did not enter the Cox modeling process.

For Setup A, model-based forests without centering (Naive) were unable to cope with complex confounding, but solely centering of the treatment indicator ( $[eqn]$ ) was valuable. Additionally adding $[eqn]$ as an offset (Robinson) did not further improve the results for the normal, binomial, and Weibull distributions, but an improvement was observed for the multinomial distribution. Suppression of $[eqn]$ in the training dataset (Setup A’) had little to no effect on the ranking of the methods. Slight deterioration in the performance of all methods was only observed for the binomial data with the logistic regression model as a base model. This is expected since this model is noncollapsible (see Section 2.4 for details).

For Setup B, the Robinson forests performed slightly better in disentangling the more complicated prognostic and predictive effects compared to Naive and $[eqn]$ model-based forests. An exception is the binomial model: $[eqn]$ forests performed similarly to Robinson forests.

In Setup C, over all distributions, uncentered model-based forests (Naive) failed to overcome the strong confounding effect and therefore did not provide accurate estimates for the treatment effect. The performance was fundamentally improved by centering the treatment indicator ( $[eqn]$ ) and was further improved by additionally adding $[eqn]$ as an offset (Robinson).

In Setup D – with unrelated treatment and control arms – all methods had a higher mean squared error than in the other setups, as jointly modeling the expected conditional outcomes for both arms has no benefit. Apart from the normal distributions, Robinson forests were inferior to the $[eqn]$ and Naive model-based forests.

The empirical evidence of our simulation study can be summarized as follows: If confounding was present, model-based forests performed better when centering $[eqn]$ by $[eqn]$ ( $[eqn]$ ) compared to not centering $[eqn]$ (Naive). Adding $[eqn]$ as an offset (Robinson) further improved the performance – especially in cases with very strong confounding.

Effect of Riluzole on progression of ALS

ALS is a progressive nervous system disease causing loss of muscle control. The status of the disease as well as the rate of progression is commonly evaluated by the ALS functional rating scale (ALSFRS).^55,56^ Here, physical abilities such as speaking, handwriting, and walking are assessed and rated on a scale from 0 (inability) to 4 (normal ability). In 1995, the FDA approved the first drug to manage and slow progression of ALS, named Riluzole. The largest database for study results on the effect of Riluzole offers the Pooled Resource Open-Access Clinical Trials (PROACT) database – initiated by the non-profit organization Prize4Life (http://www.prize4life.org). The data comes from different randomized and observational studies not disclosed in the data. Thus, the assumption of random treatment assignment is quite hard to justify in an analysis. Patient characteristics and treatment group sizes might vary greatly between the centers, which affect both the probability of receiving treatment as well as the outcome. To account for these potential confounding effects, we compared the treatment effects estimated by the naive model-based forests to the ones estimated with local centering by Robinson. As in Section 4, we use random forests to estimate the propensity scores to center $[eqn]$ and gradient boosting machines (with adapted loss functions) to estimate the values of the linear predictors $[eqn]$ and $[eqn]$ to center $[eqn]$ . Model-based forests, random forests, and gradient boosting machines rely on the hyperparameter values stated in Section 7. As for Seibold et al.^ 25 ^ and Korepanova et al.,^ 26 ^ 16 phase II and phase III randomized trials and one observational study from the PROACT database serve as a training dataset. We analyze the effect of Riluzole with respect to two outcome variables: Survival time and the handwriting ability score approximately 6 months after treatment—an item of the ALSFRS. We omitted observations with missing outcome values. As splitting variables, Seibold et al.^ 25 ^ used demographic, medical history, and family history data, which were informative in the sense that not more than half of their values were missing.

Survival Time

5.1.

The dataset for the survival time contains 3306 observations and 18 covariates. Of the 3306 observations, 2199 received Riluzole. Because very few patients had event times that exceed those of the others by a factor of two, we artificially censored five observations with (censoring or event) times of more than 750 days. The Kaplan-Meier estimates of survival probabilities for both treatment arms of the preprocessed dataset are shown in Figure 2. Overall, the estimated survival curves are very close to each other, and the treated group has only a slight survival advantage compared to the untreated group. As a base model, we use a Cox proportional hazards model. We compared treatment effects from two approaches: The naive uncentered model-based forests (Naive) and the model-based forest with Robinson’s orthogonalization (Robinson).

Kaplan-Meier curves of survival probability for both treatment arms.

Personalized models

5.1.1.

For the naive model-based forests, the underlying Cox proportional hazards base model for the survival outcome $[eqn]$ was, on the hazard scale,

$$[eqn]$$

Because $[eqn]$ contains an intercept term, $[eqn]$ is not identified (and was constraint to $[eqn]$ ). The treatment effect $[eqn]$ is the log-hazard ratio of the treated versus untreated patients and our aim is to replace a constant marginal effect $[eqn]$ with a heterogeneous (and thus conditional) log-hazard ratio $[eqn]$ and, simultaneously, to estimate prognostic effects $[eqn]$ .

For Robinson’s strategy, we first centered the treatment indicator $[eqn]$ by estimating the propensity scores $[eqn]$ using a regression forest. Figure 3 compares the distributions of estimated propensity scores (left) and of the estimated centered treatment $[eqn]$ (Robinson’s strategy, right), both obtained from regression forests. We can already see a decent overlap of propensity scores in the two treatment arms without centering, but the overlap increases if the strategy by Robinson was applied.

Distribution of estimated propensities π^(x) (left) and estimated propensities of the centered treatment indicators (right, Robinson’s strategy) as estimated by regression forests for the two treatment groups.

In addition to centering $[eqn]$ , Robinson’s strategy requires the estimation of $[eqn]$ to use as an offset (see Section 3). As in Section 4, we used gradient boosting machines (with the negative log partial likelihood of the Cox proportional hazards model as a loss) to estimate the natural parameters $[eqn]$ and $[eqn]$ for the control and treatment group, respectively.^ 57 ^ The offset $[eqn]$ for each observation is equal to the sum of natural parameter estimates weighted by $[eqn]$ (see equation (14)). The final base model for model-based forests using Robinson’s orthogonalization is

$$[eqn]$$

Model-based forests

5.1.2.

The corresponding base models serve as an input for the model-based forests to estimate personalized effects of Riluzole. Figure 4 compares the kernel density estimates of $[eqn]$ for each forest version (Naive and Robinson). The naive approach reveals that on average the treatment reduced the hazard compared to no treatment, whereas the model-based forest with centering according to Robinson obtained weaker effects of Riluzole with more mass centered around 0.

Kernel density estimates of the personalized treatment estimates for the naive model-based forest (Naive) and for the model-based forest with Robinson orthogonalization (Robinson).

A meta-analysis of previous studies by Andrews et al.,^ 58 ^ also yielded a mixed picture: only eight of the 15 studies meeting their inclusion criteria showed a statistically significant increase of median survival time due to Riluzole.

Over all strategies, for both approaches there were some patients for which Riluzole was estimated to increase the hazard. The dependency plots in Figures S. 6 and S. 7 in the Supplemental Material provide indications of the characteristics of the group of harmed individuals. For example, both the naive and centering approach agree that for patients with atrophy or fasciculation, Riluzole intake would increase the hazard. The estimated effects differed most between the uncentered forest (Naive) and the orthogonalized forest (Robinson) for the covariate sex (Figure S. 6 (c)), the covariate of whether patients swallow, and for the covariate specifying whether cases in the same generation exist (Figure S. 7 (f) and (i)).

For the variables time onset treatment, age, height and weakness the dependency plots (Figure S. 6 (a), (d), (e) and Figure S. 7 (g)) of the Naive forest agree with the ones of Seibold et al.^ 25 ^: For middle-aged people with a longer time between disease onset and start of treatment, lower height, and no weakness, the treatment appears to be more beneficial. By considering confounding effects due to orthogonalization (Robinson forests), these effects diminished. For Korepanova et al.^ 26 ^ the effect of Riluzole was also rather weak and showed low heterogeneity across covariates.

Handwriting ability score

5.2.

The dataset for the handwriting ability score—an ordinal outcome with five categories—contains 2538 observations and 58 covariates. Besides the covariate age, all covariates had missing values (but less than 50 % of the values were missing per variable enforced by the preprocessing step stated at the beginning of this section). Of the 2538 observations, 1754 received Riluzole, and 784 did not. Figure 5 displays the frequency of the ability scores for both treatment groups. Most of the patients have an ability score of 3 or 4 (normal ability); only a few have ability scores less than 2. Note that the plot shows the conditional proportions given the treatment indicator. We chose a proportional odds logistic regression model as a base model for the model-based forests—once without further adaptions (Naive), and once parameterized with centered $[eqn]$ and with an offset (Robinson).

Relative frequency distribution plot of the handwriting ability score ( Y6 ) (left) and of changes of the handwriting ability score over 6 months ( Y6−Y0 ) (right) for both treatment arms. Frequencies were calculated relative to the treatment indicator.

In addition to the handwriting ability score after 6 months, the ability score values at treatment start are also available. In the following, we denote $[eqn]$ as the handwriting score after 6 months and $[eqn]$ as the handwriting score at the beginning of the treatment period. To account for the ability level at treatment start, $[eqn]$ served as an additional splitting variable for both model-based forests (Naive and Robinson) and was included in $[eqn]$ . The alluvial plot in Figure 6 breaks down the change in each ability class over 6 months. Overall, for most patients, the handwriting ability remained constant over the 6 months or worsened slightly. Rarely, patients experienced a progression to both extremes (0 to 4, or 4 to 0). These results hold regardless of whether patients received Riluzole or not.

Alluvial plot of the progression of the handwriting ability score over 6 months for both treatment arms.

Personalized models

5.2.1.

The proportional odds logistic regression model for the naive model-based forests is defined as^59,60^

$$[eqn]$$

with $[eqn]$ as the ordinal ability score classes. The parameters $[eqn]$ are increasing thresholds, depending on covariates $[eqn]$ and the initial score $[eqn]$ . Due to the proportional odds assumption, the treatment effect $[eqn]$ is the same for all scores $[eqn]$ . Negative $[eqn]$ indicate a negative effect of Riluzole, as treated patients are expected to have a higher odds of low writing ability scores compared to untreated patients.

We first applied regression forests to estimate propensity scores $[eqn]$ and a gradient boosting machine (with adapted loss functions for the proportional odds model) to estimate the natural parameters $[eqn]$ and $[eqn]$ . The personalized model for the model-based forest with Robinson orthogonalization was specified as

$$[eqn]$$

with $[eqn]$ as defined in equation (14).

Figure 7 compares the estimated treatment indicators with $[eqn]$ as the outcome in the random forest without centering (left), with $[eqn]$ as the outcome in the random forest (right). Before centering, there is a lack of overlap of the propensity scores; the distribution of $[eqn]$ for the control group is bimodal, and the distribution for the treatment group is heavily left-skewed. After centering, the distributions of the estimated $[eqn]$ for the treatment groups move closer together and have a similar unimodal shape. However, there is still a lack of overlap of the groups, which indicates that important covariates to explain the remaining heterogeneity in the two treatment groups seem to be missing.

Estimates returned by the regression forest for orthogonalization of the treatment indicator: left for original W as an outcome in the random forest such that it estimates propensity scores π(x,y0) ; right for the centered treatment indicator W−π^(x,y0) as an outcome in the random forest.

Model-based forests

5.2.2.

The proportional odds logistic regression models served as a base model for the (Naive and Robinson) model-based forests to derive personalized treatment effects. Figure 8 displays the kernel density estimates of $[eqn]$ for each forest version (Naive and Robinson).

Kernel density estimates of the personalized treatment estimates for the naive model-based forest (Naive) vs. the forest with Robinson orthogonalized (Robinson).

Both random forests estimate on average a negative effect of Riluzole. Naive model-based forest estimated on average a log-odds of $[eqn]$ , which indicates that treated patients have a 0.08 points higher log-odds for low writing scores than untreated patients. The distribution of $[eqn]$ for the model-based forest relying on the Robinson orthogonalization is slightly shifted to the left ( $[eqn]$ ). For a larger subgroup of patients, the naive approach estimates a negative effect of Riluzole ( $[eqn]$ ), meaning that patients receiving treatment with Riluzole have higher odds of low writing scores than untreated patients. According to the dependency plots (Figures S. 8 to S. 13 in the Supplemental Material), this subgroup could be identified as having low initial ability scores (left side of Figure S. 8 (a)). For all other splitting variables, the distributions of estimated treatment effects are very similar.

Discussion and outlook

HTE estimation is a challenging problem, especially for observational studies and even more when the outcome cannot be modeled by a linear model. Covering diverse outcomes is not straightforward due to the noncollapsibility issue. Measures like the hazard ratios are shrunken towards zero and are incomparable between individuals if the full model including prognostic effects is not estimated. Model-based forests allow for the full model estimation given randomized trial data. In this work, we investigated several versions of model-based forests for the estimation of potentially complex HTEs $[eqn]$ based on observational data with various outcome types based on the orthogonalization strategy by Robinson.^ 34 ^ These investigations suggest the following workflow for model-based forests: (1) Estimate propensities $[eqn]$ using some machine learning procedure (binary random forests are a good default), (2) center the treatment indicator $[eqn]$ for each observation, (3) set up an appropriate model for the outcome conditioning on the centered treatment and – if possible – add an offset for centering $[eqn]$ , (4) use model-based forests to estimate predictive and prognostic effects $[eqn]$ and $[eqn]$ simultaneously. Notably, $[eqn]$ is the CATE only in specific models, especially a linear or log-linear model. We demonstrate these steps by estimating the individual effects of Riluzole for ALS patients using survival times and ordinal ability scores as outcomes.

Many conceptual and technical issues remain to be addressed, for example the question how model-based forests perform for survival data for which the censoring procedure is not randomized but depends on $[eqn]$ , or how ( $[eqn]$ -fold) cross-fitting influences the performance, where only one part of the data is used to estimate nuisance parameters and the other part to estimate $[eqn]$ .^ 15 ^ We leave investigations to these questions to future research.

Last but not least, we want to emphasize that all approaches for estimating HTEs – including those presented in this work – rely on strong and typically untestable assumptions. For example, for models beyond mean regression, $[eqn]$ cannot be expected to be robust against missing covariates or other violations of model assumptions due to noncollapsibility. Consequently, results from these approaches in practical applications should be evaluated with the utmost caution, reservation, and humility.

Computational details

For all computations, we used R version 4.1.1,^ 61 ^ with the following add-on packages: model4you,^ 62 ^ trtf,^ 63 ^ partykit,^ 64 ^ grf,^ 52 ^ mboost,^ 53 ^ and gbm.^ 54 ^

Model-based forests were always grown with M $[eqn]$ trees (model4you::pmforest default) with a minimum node size of node $[eqn]$ , number of chosen variables per split mtry $[eqn]$ , and subsampling. These settings were also used by Dandl et al.^ 36 ^ Transformation forests implemented in the trtf package fitted the Weibull transformation forests of Section 4.^63,29^

Propensity scores $[eqn]$ were estimated with grf (honest) regression forests with 125 trees, a minimum node size of 5, and subsampling. Natural parameters $[eqn]$ and $[eqn]$ and probability of not being censored were estimated with gradient boosting machines implemented in the mboost or gbm packages. The used maximum tree depth was 2 (default of mboost::blackboost), and a loss function that differed depending on the outcome type was also employed.^53,54^

Ratios and confidence intervals presented in Table 2 were calculated using generalized linear mixed models of the glmmTMB package.^ 65 ^ Post hoc inference relied on the multcomp package.^ 66 ^

All study settings are available in a dedicated R package called htesim.^ 67 ^ It is published on Github: https://github.com/dandls/htesim.

Supplemental Material

sj-pdf-1-smm-10.1177_09622802231224628 - Supplemental material for Heterogeneous treatment effect estimation for observational data using model-based forestsSupplemental material, sj-pdf-1-smm-10.1177_09622802231224628 for Heterogeneous treatment effect estimation for observational data using model-based forests by Susanne Dandl, Andreas Bender and Torsten Hothorn in Statistical Methods in Medical Research