Estimating Treatment Effect Heterogeneity in Psychiatry:
A Review and Tutorial with Causal Forests

An important question in psychiatry as well as other scientific disciplines that evaluate effects of interventions on the health and well-being of individuals is the extent to which the treatment effects estimated in typical experimental evaluations differ across individuals (Angus and Chang,, 2021). This is an important question because we have long known that the effects of interventions may vary across different segments of the population (Kravitz et al.,, 2004). Because of these differences, some interventions that are estimated not to be effective in the entire population based on the typical assumption of constant treatment effects are nonetheless effective in some important segments of the population, whereas some interventions found to be effective in the entire population are either not effective or in some cases even harmful in important segments of the population (Kent et al.,, 2016). In addition, when multiple interventions are available, comparative effectiveness can differ across individuals and population segments (Velentgas et al.,, 2013).

Research on heterogeneity of treatment effects (HTE), as this variation in intervention effects has come to be known, has proliferated in recent years in psychiatry (e.g., Feczko et al.,, 2019; Allsopp et al.,, 2019; Kaiser et al.,, 2022; Cohen and DeRubeis,, 2018), and many other disciplines. The basic approach is to search for significant variations in treatment effects across subsamples. This is most easily done by estimating interactions between some hypothesized effect modifiers assessed prior to randomization and random assignment. When the interaction is non-zero, HTE is said to exist with respect to the specified variable—meaning that the treatment effect differs depending on the value of the variable.

Numerous small clinical trials have reported HTE detections of this sort with respect to such fundamental baseline variables as patient age, sex, education, and symptom severity (e.g., Cuijpers et al.,, 2014; Driessen et al.,, 2010). But a question arises in trying to synthesize all this information when numerous significant interactions of this sort emerge, as meta-analyses show that they do (e.g., Maj et al.,, 2020). What is the best way to combine the information about all these specifiers into a single composite? A commonly used method is to apply the potential outcomes framework (Imbens and Rubin,, 2015) in analyzing the results of a moderated multiple regression model that includes multiple interactions. Counterfactual logic is used here to generate two predicted outcome scores from this model for each participant in the study, the first assuming that the participant was assigned to the intervention group and the second assuming that the participant was assigned to the control group (DeRubeis et al.,, 2014). The two within-person predicted outcome scores are then compared at the individual level and used as an estimate of the treatment effect for an individual based on the multivariate specifier profile of the individual. The distribution of these differences scores is then examined to investigate the pattern of significant individual differences in the treatment effect.

A drawback of this approach is that in trying out many different regression models, applications of the procedure that do not account for multiple testing may lead to false positives (van Klaveren et al.,, 2019). This problem can be addressed, though, using data-driven machine learning methods to search for interactions with cross-validation to address the problem of over-fitting (VanderWeele et al.,, 2019; Wager and Athey,, 2018). The use of data-driven methods in this way also addresses the problem that the moderated regression approach only allows for a limited range of typically two-way interactions, whereas HTE can involve much more complex forms of interaction. Machine learning methods, in comparison, allow for these more complex forms of interaction to be discovered using non-parametric methods. A large body of work has made tremendous progress in adapting a variety of machine learning algorithms to estimate HTE. Examples include Athey and Imbens, (2016); Athey et al., (2019); Hahn et al., (2020); Kennedy, (2023); Künzel et al., (2019); Luedtke and van der Laan, (2016); Nie and Wager, (2021); Tian et al., (2014).

To communicate statements about causal effects it is helpful to employ a universal language with well-defined constructs. We employ the potential outcomes framework (Imbens and Rubin,, 2015) where we posit the existence of potential outcomes $Y_{i}(0)~{}\text{and}~{}Y_{i}(1)$ which records the hypothetical clinical outcome $Y_{i}$ for patient $i$ in either treatment state. We typically denote the control state by $W_{i}=0$ and the treatment state by $W_{i}=1$. Naturally, we only observe one potential outcome for any single patient: the outcome corresponding to the treatment arm she was assigned to. Table 1 shows an example of hypothetical potential outcomes along with realized outcomes.

The assumption that our column of realized outcomes is consistent with the columns of potential outcomes, $Y_{i}=Y_{i}(W_{i})$, is called the Stable Unit Treatment Value Assumption (SUTVA). In some complicated social network settings, this assumption may not be viable. Consider, for example, if patient A is given behavioral therapy and patient B is not. If patient A interacts with patient B through some social functioning, like a support group, then A’s newly acquired behavioral knowledge may influence B. Here, however, we rely on SUTVA and thus rule out scenarios like these by assuming no interference or spillover between patients.

The type of causal effects we are interested in quantifying are average differences in potential outcomes. Given some hypothetical target population, the average treatment effect (ATE) is

$$\tau=\mathbb{E}\left[Y_{i}(1)-Y_{i}(0)\right],$$

where $\mathbb{E}\left[\cdot\right]$ denotes an average taken over a population. This would quantify, following the setup in Table 1, for example, the increase in healthy prevalence when assigning every patient the intervention, as opposed to withholding the intervention from everyone.

Even though for any single patient $i$ we never observe the difference $Y_{i}(1)-Y_{i}(0)$, we still have enough information to identify the average difference. If we collect a representative sample from our population and then randomly assign an intervention to some units in the sample while withholding it from others, then the simple difference in means (i.e., the prevalence among treated patients minus the mean prevalence among control patients) provides an unbiased estimate of the average treatment effect. The fact that the treatment assignment is randomized means that our hypothetical potential outcomes and treatment status are decoupled, which justifies this simple calculation.

The availability of such a randomized treatment assignment is typically constrained to randomized controlled trials. In many important settings, due to ethical and resource considerations, we don’t have access to a randomized intervention. A dataset where the intervention assignment is not randomized is often referred to as an observational dataset. As an example, via electronic health records, we could assemble a dataset that records an intervention, such as psychiatric hospitalization, together with an outcome, such as subsequent suicide, in a sample of patients presenting at an emergency department with suicidality. The decision whether to hospitalize such patients is clearly not randomized but will rather depend on attending physician and hospital protocols, patient severity and perceived risk of suicidal behavior, available outpatient alternatives to hospitalization (e.g., day treatment programs), and availability of hospital beds. Some of these determinants of hospitalization are also likely to be determinants of subsequent suicidal behavior, which means that a simple comparison of the rates of suicidal behavior in, say, six months after emergency department presentation cannot be interpreted as providing evidence of the effects of hospitalization on these behaviors. At the same time, constructing a randomized controlled trial in a setting of this sort poses ethical challenges.

We can, however, under suitable circumstances, make causal inferences from such data if we are willing to make certain assumptions about access to information about the determinants of treatment assignment that are also determinants of subsequent suicidal behavior in the absence of treatment assignment. The key assumption in this regard is known as the unconfoundedness assumption (Rosenbaum and Rubin,, 1983), which posits that if we account for a set of observed patient characteristics $X_{i}$ known prior to intervention, then we can think of the treatment assignment as if it were randomized,

$$\left(Y_{i}(1),~{}Y_{i}(0)\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}W_{i}\right)\mid X_{i}~{}~{}\text{(unconfoundedness)}.$$

An important object under this setup is the propensity score,

$$e(x)=\mathbb{E}\left[W_{i}\mid X_{i}=x\right],$$

which is the probability that a patient with confounder variables equal to $x$ is assigned the treatment. To estimate causal effects, a necessary condition is that the association between baseline predictors and intervention assignment is not so strong that no individuals have a non-zero probability of assignment to any of the interventions under study. This is the overlap assumption, which requires that the propensity score $e(x)$ is strictly above 0 and below 1 for every unit $x$. For example, if 100% of the suicidal patients presenting at an emergency department with severe thought disturbance were hospitalized, it would be impossible to use observational methods to evaluate the effect of hospitalization in preventing subsequent suicidal behaviors among such patients. It is only in the subset of patients whose probability of hospitalization is less than 100% and greater than 0% that data-driven estimates of treatment effects can be made. Assessing this requirement empirically translates into ensuring that estimated propensity scores are not close to either zero or one, recognizing that estimation of treatment effects is impossible in the subset of patients whose scores are close to zero or one.

While the overlap assumption is testable, the direct conditional independence assumption under unconfoundedness is not, as it involves statements of independence between random variables. This is usually referred to as an identifying assumption and is fundamentally untestable. Instead, the evaluation of identifying assumptions requires domain knowledge and subject-matter expertise to assess how credible they are in a particular application. In this tutorial, we are treating this first step as done or known, i.e., the underlying causal graph (Pearl,, 1995) is agreed upon. In the context of the example of psychiatric hospitalization as a predictor of subsequent suicidal behavior among suicidal emergency department patients, an example where this assumption is invoked is Ross et al., (2023).

To estimate average treatment effects in observational settings under unconfoundedness, we must account for how the different types of patients respond to and are assigned the treatment before making comparisons between average outcome scores. We can do this by estimating the propensity score. It is also possible to perform the confounding adjustment by estimating conditional outcome functions or via a doubly robust combination of estimates of both the propensity and the conditional outcome function. This is the approach grf implements in the function average_treatment_effect; see Ding, (2023), Hernan and Robins, (2023), and Wager, (2024) for recent textbook treatments of these methods.

Causal forests (Athey et al.,, 2019) build on Breiman’s random forest algorithm (Breiman,, 2001), a popular and empirically successful algorithm for prediction that allows for a mix of real- and discrete-valued predictors. Random forests have a strong empirical track record of doing well out-of-the-box in many applications, often with minimal tweaking of tuning parameters. An appealing aspect of random forests is that they essentially act as an algorithmic way of dividing data into subgroups and then making predictions based on which subgroup a patient belongs to. A core component of the random forest algorithm is to scan over each patient characteristic, then decide on a cut point that serves as a “good” candidate for a subgroup partition. The algorithm then aggregates these partitions into neighborhoods where patients with specific characteristics in terms of predictor variables $X_{i}$ have similar values on outcome $Y_{i}$.

At first glance, the algorithmic blueprint of random forests appears like a promising candidate for our task of discovering subgroups with different treatment effects, as we are looking for ways to partition a sample of participants in an intervention experiment according to observable characteristics. The primary difference between causal forests and Breiman’s random forests is that the objective of the original random forests is to partition the sample to optimize discrimination of individual differences in scores on an outcome, whereas our objective is to predict differential treatment response.

The challenge in targeting treatment heterogeneity that does not arise in the basic prediction setting is that scores on an outcome variable are observed for each participant, whereas individual-level treatment effects are not observed. As we saw previously, though, treatment effect estimation with non-experimental data can be carried out by adjusting for baseline characteristics. We do this using a two-step process. In the first step, causal forests account for uneven treatment assignments via the propensity score $e(x)$, as well as baseline effects by estimating the conditional mean (marginalizing over treatment)

$$m(x)=\mathbb{E}\left[Y_{i}\mid X_{i}=x\right].$$

In this first step, baseline effects are essentially stripped away via something known as an orthogonal construction (Chernozhukov et al.,, 2018; Robinson,, 1988). This construction works on the outcomes as well as on treatment indicators to generate residualized outcome $\widetilde{Y}_{i}=Y_{i}-m(X_{i})$ and treatment indicators $\widetilde{W}_{i}=W_{i}-e(X_{i})$ that have been adjusted for confounding effects. In the second step, $\widetilde{Y}_{i}$ and $\widetilde{W}_{i}$ are used to form treatment effects estimates, and a random forest is used to find predictor variable partitions where these treatment effects estimates exhibit heterogeneity. Figure 1 presents a stylized illustration of the underlying logic for this phase of the algorithm. For a candidate predictor variable, such as age, we seek a single axis-aligned split separating units into two groups with different average treatment effect estimates. Here, units with ages above 25 appear to benefit more, and a partition that immediately increases heterogeneity, as measured by the vertical distance between estimates in Figure 1(b), would be at this point. This exercise is repeated recursively for the left and right samples, considering all possible predictor variables (e.g., blood pressure) to partition the covariates $X_{i}$ into regions where treatment effects are expected to differ.

In the case where the analysis is applied to an experiment, where the probability of assignment to each intervention arm is known by construction, these probabilities can simply be supplied in the first step. However, in a non-experimental situation in which we are attempting to approximate estimates of treatment effects by adjusting for nonrandom exposure respective to observed baseline covariates, the probabilities of treatment assignment are estimated in the first step with separate random forests.

It is possible to modify this random forest framework to accommodate many other relevant statistical quantities and settings – thus the name “generalized” random forests–grf. For example, Athey et al., (2019) use this general framework to construct an instrumental forest that estimates HTE in settings with an instrumental variable, Cui et al., (2023) use this framework to construct a causal survival forest for estimating HTE with right-censored survival outcomes, and Friedberg et al., (2020) constructs a local linear forest. While these random forest algorithms are designed to directly target various causally relevant targets, we emphasize that one still needs to exercise caution when interpreting flexible non-parametric point estimates of these quantities. To give some intuition, machine learning models serve as great complements to traditional parametric statistical models as they allow us to be model-agnostic. This agnostic property, however, comes at the cost of introducing higher estimation uncertainty. So, our estimates of individual-level treatment effects come with considerable uncertainty, and in assessing the quality of these estimates, it is practical to move to a less granular summary measure. In the applied Section 4, we cover different approaches to assess if the individual-level estimates capture meaningful heterogeneity via calibration exercises. On a general note, the perspective we are taking is that machine learning tools can serve as effective algorithmic devices that target CATEs, which an analyst can then validate on held-out data and transparently convey potentially interesting scientific findings regarding treatment heterogeneity.

Finally, it is worth pointing out that the underlying identifying causal assumptions used in causal forests are the same as in classical parametric models, such as in the moderated regression approach described earlier. However, a main difference is that whereas classical parametric models (e.g., linear and logistic regressions) specify particular predictors and functional forms of their associations with the outcomes prior to estimation, causal machine learning approaches like causal forest use flexible data-driven forest-based constructions (grf also allow for more advanced use-cases by allowing for user-specified estimates of the first-stage inputs $e(x)$ and $m(x)$, which can be derived from separate machine learning models, such as gradient boosting or neural networks, or through ensemble methods that combine multiple models (van der Laan et al.,, 2007)).

We illustrate the approach by focusing on a recent application of grf to estimate differential resilience to the effects of combat stress in leading to post-traumatic stress disorder (PTSD) among US Army soldiers (Kessler et al.,, 2024). PTSD is the signature mental disorder of war (Paulson and Krippner,, 2007). The roughly 7% of the US population made up of military personnel and veterans are estimated to account for nearly 20% of all cases of PTSD in the US, with an annual societal cost estimated at more than $230B (DeAngelis,, 2023). However great variation is thought to exist in the effects of combat stress on PTSD among service members (Karstoft et al.,, 2015; Schultebraucks et al.,, 2021). We illustrate the value of causal forests by investigating these effects in a secondary analysis of a sample of US Army soldiers in three combat brigades who participated in a self-reported survey shortly before and then again after returning from a combat deployment in Afghanistan in 2012-2014 as part of the Army Study to Assess Risk and Resilience in Servicemembers (Army STARRS, Ursano et al., (2014)). Exposure to combat stress was assessed in the follow-up surveys administered to this sample after they returned from deployment. As expected, high combat stress exposure was found to be associated significantly with elevated prevalence of current PTSD as assessed in the follow-up surveys (Kessler et al.,, 2024).

The effects of exposure to high combat stress can be approximated by defining a dichotomous measure of the extent of combat stress exposure in the sample who had combat arms occupations (e.g., infantry, artillery). The extent to which information about individual differences in both baseline survey reports (including reports about history of and recent symptoms of PTSD) and baseline administrative data predicted this variation in exposure to high combat stress can then be investigated. The net association of high combat stress with the subsequent occurrence of meeting diagnostic criteria for PTSD in the follow-up survey can then be examined to estimate ATE as well as to inspect the distribution of estimated CATE. More details on the dataset (Papini et al.,, 2023) and a more extensive substantive analysis of the data (Kessler et al.,, 2024) are presented elsewhere.

In order to use methods for CATE estimation to capture heterogeneity in resilience, we need to somewhat reconceptualize the notion of an “intervention”: The intervention here, i.e., member combat deployment, is not a therapeutic intervention as is often considered in medical settings, but is nonetheless one that can be manipulated in meaningful ways. Only a fraction of all military personnel in a unit (often in the range 40-60%) are assigned to deploy, with the remainder held in reserve (referred to by the military as the “rear detachment”). Commanders take a wide range of factors into consideration in deciding which soldiers to deploy and which to hold in reserve. Estimated mental fitness based on supervisor assessment is one of these considerations. A CATE estimate based on a comprehensive assessment of administrative data and an optimized algorithm based on causal forests could be a useful addition to such assessments either to help commanders decide which members of their units to deploy and which to keep in the rear detachment or, in cases where it is necessary to deploy a non-resilient service member based on other considerations, to help target special resilience-enhancing resources (e.g., special pre-deployment resilience training programs (Thompson and Dobbins,, 2018)).

Considering deployment with few traumatic events (“low combat stress”) as a treatment arm ($W_{i}=1$) and deployment with “high combat stress” as a control arm ($W_{i}=0$), we can construct a causally motivated empirical measure of resilience using the language of potential outcomes. Table 2 shows an example of potential outcome configurations, or principal strata, describing outcomes a unit would experience in both treatment states. These principal strata are not observable (because, like the ITE, they depend on both potential outcomes); however, they are still useful for interpretation (Frangakis and Rubin,, 2002). The type of soldier who is healthy regardless of combat stress exposure can be deemed to have high resilience, whereas the type of soldier who develops PTSD under high combat stress but is healthy under low combat stress can be deemed to have low resilience. Finally, the type of soldier that develops PTSD regardless of combat stress exposure can be termed high risk. We assume that there are no soldiers who would be healthy under high stress and sick under low stress; this assumption is mathematically related to the widely used “no defiers” assumptions in clinical trials with non-compliance (Angrist et al.,, 1996).

We record realized healthy outcomes by 1 and PTSD by 0,

$$Y_{i}=\begin{cases}1&\text{if Healthy}\\ 0&\text{if PTSD},\end{cases}$$

and so our average treatment effect measures the fraction of healthy soldiers when no one is deployed minus the counterfactual fraction of healthy soldiers when all soldiers are deployed, which we expect to be larger than zero.

$\displaystyle\begin{split}\tau&=\mathbb{E}\left[Y_{i}(1)-Y_{i}(0)\right]\\ &=\mathbb{P}\left[Y_{i}(\text{Low combat stress})=\text{Healthy}\right]-\mathbb{P}\left[Y_{i}(\text{High combat stress})=\text{Healthy}\right]\\ &=\mathbb{P}\left[\text{Low resilience}\right].\end{split}$

The CATE then captures heterogeneity in resilience:

$$\tau(x)=\mathbb{P}\left[\text{Low resilience}\mid X_{i}=x\right].$$

This type of taxonomy gives us two approaches to measuring resilience: a causal measure based on CATEs (low vs. high resilience) and a predictive measure based on risk (high vs. low risk). How these two approaches differ in practice is an empirical question we address in Section 4.6.

We have surveyed how causal forests can be used to generate both principled estimates of ATE in observational studies and to detect treatment heterogeneity in experimental or observational studies. For more details on grf, including additional functionality such as missing values support, tree-based policy learning, loss-to-follow-up corrections, and more, we refer to vignettes and journal references on the grf website at grf-labs.github.io/grf.

We have also outlined a general framework for analyzing and comparing output from causal machine learning algorithms on a level playing field with tools such as TOC and Qini curves. As noted in the introduction, there are a number of alternative machine learning approaches for CATE estimation that are popular in practice (e.g., Athey and Imbens,, 2016; Hahn et al.,, 2020; Kennedy,, 2023; Künzel et al.,, 2019; Luedtke and van der Laan,, 2016; Nie and Wager,, 2021; Tian et al.,, 2014). Customizable Python libraries that implement several of these methods include causalML (Chen et al.,, 2020) and econML (Battocchi et al.,, 2019); and the software toolkit for machine learning-based causal inference is still rapidly growing. In practice, it can sometimes be a good idea to try multiple different machine learning-based approaches to CATE estimations; the TOC, Qini curves, and related tools can then be used to compare and benchmark their performance.

We are grateful to the STARRS-LS collaborators for the use of their survey data and to Ron Kessler for providing invaluable feedback and advice.