Double machine learning to estimate the effects of multiple treatments and their interactions

The development of causal inference methods has made significant progress to address confounding issues in observational studies (Hernan and Robins, 2024; Ding, 2024). However, most causal inference methods focus on a binary treatment, such as comparing the effect of a treatment versus a control (Rosenbaum and Rubin, 1983; Gutman and Rubin, 2015). These methods are inadequate for more complex settings involving multiple treatments, especially when treatment interactions are present. Such settings are not only methodologically challenging but also practically important. For example, in an HIV study, researchers seek to evaluate the side effect of multiple treatments and their interactions on participants’ kidney function. To draw meaningful conclusions in this study, it is essential to develop robust causal inference methods that can handle the settings of multiple treatments and their interactions.

Besides the lack of proper methods for multiple treatments, analyzing observational data presents additional challenges. Key issues include high-dimensional confounding covariates (Berisha et al., 2021) and complex confounding relationships—such as non-linear patterns and higher-order interactions—that influence both treatments and outcomes (Huber, Lechner and Wunsch, 2013; Wan and Mitra, 2018). These challenges complicate the estimation process and increase the potential for bias (D’Amour et al., 2021; Johnstone and Titterington, 2009); the complexity increases further with multiple treatments, as the number of variables and interactions grows. Given these challenges, machine learning methods become a natural choice because of their strengths in modeling high-dimensional data, non-linear patterns, and interactions (Jiang, Gradus and Rosellini, 2020). However, direct use of machine learning in causal inference is not a “free-lunch”. As machine learning methods specialize in prediction instead of estimation, they are subject to overfitting and regularization bias (Neyshabur, Tomioka and Srebro, 2014; Mehrabi et al., 2021; Kernbach and Staartjes, 2022), which can lead to biased estimates of causal effects.

To address aforementioned challenges, we propose a class of new methods for multiple treatments based on double machine learning (DML) (Chernozhukov et al., 2018), a framework that uses machine learning methods to estimate causal effects while overcoming the potential bias induced by regularization and overfitting. Specifically, we contribute to the literature in two broadly applicable settings: (1) a DML partial linear model to estimate the effects of multiple concurrent treatments of different types (binary, categorical, and continuous) and their interactions, and (2) a DML interactive model to estimate the average treatment effect (ATE) across categories of multi-valued regimens. For both setting, we define estimands, propose estimators, and develop corresponding algorithms for effect estimation. Notably, our approach for multiple treatments explicitly models treatment interactions, which are important in practice but have rarely been explored in the causal inference literature. We further describe the asymptotic normality of the proposed estimators and thereby derive the analytical variance estimator for inference.

We apply our method to investigate the effect of three different treatments on chronic kidney disease in an adult HIV cohort of 2455 participants in northern Nigeria. This application is motivated by the rising incidence of HIV-associated kidney disease, despite the improved survival rates among people with HIV because of the increased use of antiretroviral therapy (ART) (Wudil et al., 2021). Since people with HIV are typically on multiple treatments, studying the side effects of these treatments on kidney function can improve our understanding of HIV-associated kidney disease.

The goal is to estimate the effects of three treatments for HIV and other health conditions, along with their interactions, on kidney function. Kidney function is assessed using the estimated glomerular filtration rate (eGFR), a measure of the severity of chronic kidney disease The study measured a total of 19 demographic and clinical variables, which may exhibit complex relationships with treatment assignment and kidney function. This application demonstrates how our method addresses challenges arising from multiple treatments, complex confounding structures, and model misspecification.

While most causal inference methods focus on binary treatment, recent years have seen a growth in methods for multi-valued treatments in observational studies. Imbens (2000) and Imai and Van Dyk (2004) extended the work of Rosenbaum and Rubin (1983) by proposing the generalized propensity score for multi-valued treatments. Thereafter, many methods for binary treatment have been adapted for multi-valued treatments using the generalized propensity scores: inverse propensity score weighting (IPW) (Feng et al., 2012; McCaffrey et al., 2013), propensity score adjustment (Linden et al., 2016), and propensity score matching (PSM) (Yang et al., 2016; Lopez and Gutman, 2017). In addition to the methods based on propensity scores, other methods have also been extended for multi-valued treatments, including balancing weights (Li and Li, 2019), Bayes additive regression trees (Hu et al., 2020, 2022), and targeted maximum likelihood estimation (Rose and Normand, 2019).

Most of the aforementioned literature does not truly address the settings where multiple treatments are assigned simultaneously. Instead, they focus on “multi-valued treatments,” where each subject receives a single level of treatment that is modeled as a categorical variable $A_{i}=k$ for $k\in\{1,…,K\}$. Therefore, the generalized propensity score is estimated for each category, and the focus is typically on the pairwise contrast, such as ATE between treatment categories. However, in many fields— such as clinical (Wudil et al., 2021), economic (Frölich, 2004), and environmental studies (Zhu et al., 2024) — subjects receive multiple treatments at the same time. For example, with two treatments $a_{1}$ and $a_{2}$, the causal relationship may take the form:

In this setting, the methods for “multi-valued treatment” show substantial limitations:

• •

  It is not straightforward to estimate $\theta_{3}$, the interaction effect between $a_{1}$ and $a_{2}$ on the outcome.

• •

  It is particularly challenging when $a_{1}$ or $a_{2}$ is continuous. Research on continuous treatments is still being developed (Zhu, Coffman and Ghosh, 2015; Kallus and Zhou, 2018; Brown et al., 2021), where a key difficulty is that $P(A=a)=0$ for a specific $a$ due to its continuity. Moreover, propensity score methods may lead to worse covariate balance for continuous treatments (Brown et al., 2021).

In this article, we address these limitations by presenting a DML partial linear model (Section 3) that accommodates (1) multiple treatments assigned simultaneously, (2) treatment interactions, and (3) treatments of varying types (binary, categorical, and continuous). To our knowledge, this is the first work to leverage machine learning methods for the robust estimation of causal effects involving both multiple treatments and their interactions.

Related work on such settings remains limited, and most of them overlook the presence of treatment interactions. Siddique et al. (2019) studied effect estimation for multiple concurrent treatments but still treated the combination of treatments as a multi-valued categorical variable. Wang and Blei (2019, 2021) proposed the de-confounder, focusing on causal identification for multiple treatments with unobserved confounders. From a computer science perspective, Wang et al. (2024) reviewed complex treatment settings, including neural network-based methods for multiple treatments (Qian, Curth and van der Schaar, 2021; Zou et al., 2020; Tanimoto et al., 2021).

Additionally, we describe the DML interactive model for the settings of multi-valued regimens, as detailed in Section 4. While this setting aligns with the classic causal inference literature of multi-valued treatments, we adopt the term “multi-valued regimen” to better reflect its characteristics. See detailed description for this setting in Section 2.

Section 2 describes two settings: (1) mixed-type multiple treatments and (2) multi-valued regimens, along with notation and assumptions. Section 3 presents the DML partial linear model for estimating the effects of multiple treatments and their interactions. Section 4 presents DML interactive model for estimating the pairwise ATE of multiple regimens. Section 5 describes the Neyman orthogonality score function for DML in the context of multiple treatments and multi-valued regimens and derives the analytical variance of our proposed estimator. Section 6 presents simulation studies to demonstrate the performance of DML in different settings. Section 7 applies the DML partial linear model to study the effect of ART level and ART duration on HIV-associated kidney disease. Finally, a discussion section concludes this article.

We outline two settings in this article, as illustrated in Figure 1. In the first setting of multiple treatments, subject $i$ receives $D$ treatments simultaneously, denoted as a vector of treatments

Here, $A_{id}$ denotes an individual treatment of different types: binary, categorical, or continuous for $d\in\{1,...,D\}$. In addition, we explicitly model the vector of $T$ treatment interactions

where $A^{\circ}_{it}$ represents the $t$-th interaction term, formed as the product of two or more individual treatments, e.g., $A^{\circ}_{i1}=A_{1}A_{2}$, $A^{\circ}_{i2}=A_{1}A_{2}A_{3}$, etc. The maximum number of interactions is $2^{D}-D-1$, but since not all treatments necessarily interact with each other, the actual number of interactions satisfies: $0\leq T\leq 2^{D}-D-1$.

In the second setting of multi-valued regimens, each subject receives one regimen from $D$ regimen categories. Let $R_{i}=d$ be a categorical indicator of observed treatment status for $d\in\{1,2,..,D\}$; that is, $R_{i}=d$ if subject $i$ received the $d$-th regimen. This setting is often referred to as multi-valued treatments in causal inference literature. However, we refer to it as “multi-valued regimen”, since it is common in practice that subjects receive not only a single level of treatment but rather a regimen, which consists of a systemic plan or a combination of treatments.

Let $Y_{i}$ denotes the observed continuous outcome for subject $i$. $\bm{X}_{i}$ denotes a vector of $p$-dimensional pretreatment confounding covariates. We identify and estimate the causal effects based on the potential outcome framework (Rosenbaum and Rubin, 1983). For multiple concurrent treatments, each subject $i$ has a set of potential outcomes, denoted as $Y_{i}^{(a_{1},a_{2},…,a_{D})}$, which represents the outcome that would have been observed had subject $i$ received the treatment combination $(A_{i1}=a_{1},A_{i2}=a_{2},…,A_{iD}=a_{D})$. Similarly, for multi-valued regimens, subject $i$ has $D$ potential outcomes, where $Y_{i}^{(r=d)}$ denotes the potential outcome under the $d$-th regimen for $d\in\{1,...,D\}$.

Under the potential outcome framework, assumptions are required to identify causal effects. For multiple concurrent treatments, we introduce the following assumptions:

• •

  Strong unconfoundedness: $Y_{i}^{(a_{1},a_{2},...,a_{D})}\perp\!\!\!\!\perp(A_{i1},A_{i2},...,A_{iD})\rvert\bm{X}_{i}$, for $(a_{1},a_{2},...,a_{D})\in\bm{\mathcal{A}}$, where $\bm{\mathcal{A}}$ denotes the space for all possible treatment combinations.

• •

  Positivity: $0<P((A_{i1},A_{i2},...,A_{iD})\in\bm{\mathcal{A}}\rvert\bm{X}_{i})$.

Strong unconfoundedness asserts that, given confounders $\bm{X}_{i}$, potential outcomes under all possible treatment combinations are independent of the assignment of multiple treatments. Positivity asserts that, given $\bm{X}_{i}$, there is a non-zero probability of observing every possible combination of treatments within $\bm{\mathcal{A}}$. Unlike settings with a binary treatment, these assumptions hold jointly across all treatments. It enables identification of causal effects for any combination of treatments and higher-order interactions (e.g., the interaction effect of three treatments simultaneously). Wan and Mitra (2018) also considered similar assumptions.

For multi-valued regimens, since subjects only receive a single category of regimens, the assumptions are different:

• •

  Weak unconfoundedness: $Y_{i}^{(r=d)}\perp\!\!\!\!\perp\mathbbm{1}_{\{R_{i}=d\}}\rvert\bm{X}_{i}$, for $d\in\{1,…,D\}$.

• •

  Positivity: $0<P(R_{i}=d\rvert\bm{X}_{i})<1$ for $d\in\{1,…,D\}$.

Weak unconfoundedness asserts that, given confounders $\bm{X}_{i}$, the potential outcome under the $d$-th regimen is independent of whether subject $i$ received that regimen, as indicated by $\mathbbm{1}_{\{R_{i}=d\}}$. Unlike strong unconfoundedness, weak unconfoundedness only requires the independence to be ‘local’. That is, independence between the potential outcome and the observed regimen holds within each specific regimen category (Imbens, 2000). Similarly, the positivity assumption for multi-valued regimens only requires the marginal probability of observing each category of the regimen is non-zero.

Theorem 1 below shows the identification for settings of multiple treatments:

Appendix Section 1.2 proves Theorem 1, which establishes that the potential outcome under all possible treatments can be identified from observational data. This result enables the identification of the effects of treatment interactions under the partial linear model in Section 3. For example, consider a outcome model with binary treatments $a_{1}$ and $a_{2}$:

where $\varepsilon_{i}\sim N(0,1)$. Following Theorem 1, the interaction effect $\theta_{3}$ can be expressed as a combination of average potential outcomes: $\theta_{3}=E[Y^{(a_{1}=1,a_{2}=1)}]-E[Y^{(a_{1}=0,a_{2}=1)}]-E[Y^{(a_{1}=1,a_{2}=0)}]+E[Y^{(a_{1}=0,a_{2}=0)}].$ Similarly, the estimands for main effects $\theta_{1}$ and $\theta_{2}$ can also be expressed as a combination of average potential outcomes. Together, these coefficients allow us to interpret both individual and interaction effects of multiple treatments.

For multi-valued regimens, relevant identification results have been previously studied under the assumptions of weak unconfoundedness and positivity, including Theorem 1 in Imbens (2000) and Lemma 3 in Yang et al. (2016). Though the generalized propensity score is estimated using the full sample of $\bm{X}_{i}$, the identification can be reached for a specific category of regimen. For example, the potential outcome under b-th regimen is identified as

In this setting, researchers are usually interested in pairwise contrasts between regimens, for example, the ATE between the potential outcomes of the $b$-th regimen and $c$-th regimen, with $b,c$ $\in\{1,...,D\}$. The estimand for such contrast, denoted as $ATE_{bc}$, can be expressed as:

We omit the subscript $i$ in the subsequent sections for simplicity.

The Neyman orthogonality ensures the estimation of the effects of multiple treatments/regimens — our primary parameters of interest — remains robust to the bias in the estimation of nuisance parameters (e.g., the treatment model $m(\bm{X})$ and the outcome model $g(\bm{X})$). To formalize this, we consider the method of moments estimator for our parameter of interest $\bm{\Theta}_{0}$:

where $\psi$ is a vector-valued score function that characterizes the moment conditions; $\bm{W}$ denotes the sample space: $(\bm{X},Y,\bm{A})$ for multiple treatments and $(\bm{X},Y,R)$ for multiple regimens; $\bm{\Theta}_{0}$ denotes the true value of parameters of interest, e.g., for multiple treatments and interactions, $\bm{\Theta}_{0}^{\top}=[\bm{\theta}^{\top},\bm{\theta}^{\circ\top}]$; $\bm{\eta}_{0}$ denotes the true values of the nuisance parameters;

The Neyman orthogonality condition requires the score function $\psi$ satisfies:

where $\bm{\delta}=\bm{\eta}-\bm{\eta}_{0}$ represents deviations of the nuisance parameters from their true values. This condition ensures that bias in estimating $\bm{\eta}_{0}$ does not affect the moment condition evaluated at $\bm{\Theta}_{0}$, thereby preserving the unbiased estimation for $\bm{\Theta}_{0}$. Furthermore, under conditions in Chernozhukov et al. (2022) the Neyman orthogonality score function in many settings presents double robustness. That is, if at least one component of $\bm{\eta}$ is consistently estimated, the estimation for $\bm{\Theta}_{0}$ remains unbiased.

The Neyman orthogonality score function for both DML partial linear model and interactive model is linear in $\bm{\Theta}$, and thereby can be expressed as

where $\psi_{a}(\bm{W};\bm{\eta})$ is the component linear in $\theta$ and $\psi_{b}(\bm{W};\bm{\eta})$ is independent of $\theta$. Based on this linear representation, we can also estimate $\bm{\Theta}$ by solving the moment condition:

Furthermore, such linearity representation can facilitate the estimation of the variance, which will be presented in the following subsections.

In this setting, we propose the DML partial linear model, where the score function is formulated based on the “partialling-out” approach. The Lemma 1 shows that such score function satisfies the Neyman orthogonality condition in (9).

In this lemma, $\bm{\Theta}_{0}^{\top}=[\bm{\theta}^{\top},\bm{\theta}^{\circ\top}]$ and the nuisance parameters $\bm{\eta}=(\bm{m},\bm{m}^{\circ},l)$, i.e., $\bm{m}(\bm{X})=E[\bm{A}|\bm{X}]$, $\bm{m}^{\circ}(\bm{X})=E[\bm{A}^{\circ}|\bm{X}]$, and $l(\bm{X})=E[Y|\bm{X}]$. The Appendix Section 1.2 proves the lemma 1.

Based on this score function, we can also estimate the parameter of interest following equation (10) and (11):

The third equality in the equation above uses that $\bm{A}-\bm{m}(\bm{X})=\bm{\varepsilon}_{\bm{A}}$ and $Y-l(\bm{X})=\bm{\varepsilon}_{Y|X}$. Note that such equation leads to an estimator of the same format as the one derived using the residual regression in Section 3.

Proposition 1 introduces the the asymptotic normality of this estimator for $[\bm{\theta}^{\top},\bm{\theta}^{\circ\top}]$.

As a consequence of Proposition 1, after computing the estimated score function $\hat{\psi}(\bm{W};\hat{\bm{\Theta}},\hat{\bm{\eta}})$ and $\hat{\bm{J}}$, we can estimate the variance using formula (13):

In this setting, we proposed the DML interactive model to estimate the ATE between two regimen categories. The Lemma 2 below shows the score function for this model satisfies the Neyman orthogonality condition.

In this lemma, the $\theta$ denotes parameter of interest – ATE between the $b$-th regimen and the $c$-th regimen, $b,c\in{1,...,D}$. The nuisance parameters $\bm{\eta}=(m,g)$ are estimated for both $b$-th and $c$-th regimens. Appendix Section 1.4 also proves the Lemma 2.

Based on this Neyman orthogonal score, Proposition 2 introduce the the asymptotic normality of the proposed estimator for ATE between two categories of regimen.

As a consequence of Proposition 2, after computing the estimated score function $\hat{\psi}(\bm{W};\hat{\theta},\hat{\bm{\eta}})$, we can estimate the variance using formula (14):

This subsection evaluates the performance of the DML partial linear model on a simulated dataset. We generate covariates $\bm{X}$ of 10 dimensions. Specifically, covariates $X_{1}$ to $X_{5}$ are drawn independently from $N(0,1)$, and $X_{6}$ to $X_{10}$ are drawn independently from a bernoulli distribution with different probabilities of success. We generate two treatments:

• •

  a binary treatment generated as $A_{1}\sim\text{Bernoulli}(\pi_{1})$, where $\text{logit}(\pi_{1})=m_{1}(\bm{X})$.

• •

  a continuous treatment $A_{2}$ generated by $A_{2}=m_{2}(\bm{X})+\varepsilon_{A_{2}}$, where $\varepsilon_{A_{2}}\sim N(0,1)$.

The functions $m_{1}(\bm{X})$ and $m_{2}(\bm{X})$ incorporate a mixture of linear, nonlinear, and interaction effects of the covariates, as detailed in Appendix Section 2.1.

The outcome $Y$ is generated by a model of these two treatments and their interaction, as well as other covariates:

where $\varepsilon$ is a noise term of normal distribution $N(0,1)$. $g(\bm{X})$ is a function of $\bm{X}$ that includes various types of effects (linear, non-linear, interaction, etc.), as defined in Appendix Section C. The true effects for treatments are set as

We use five different methods to estimate the treatment models and the outcome model: an ordinary regression method, a $\ell_{1}$-penalized method (Lasso), two tree-based methods (random forest and boosting trees) and a neural network method. To evaluate whether these methods can capture the complex pattern of confounding covariates, the input covariate matrix $\bm{X}$ is used in its original form without transformations. The details of parameter choices and model specification are provided in Appendix Section 2.1.

We use Algorithm 1 to estimate $\theta_{1}$, $\theta_{2}$ and $\theta_{3}$, where we consider both 2-fold and 5-fold cross-fitting in our algorithm. Since the estimates can vary depending on the random sample splits in cross-fitting, we make 50 different splits on this dataset and thereby run 50 estimations. We use the median estimate and median standard error across 50 splits as the final estimates. Table 1 summarizes the biases and rMSE of the simulation results. Based on these results:

• •

  Regarding the bias, tree-based models and the neural network show lower bias than ordinary regression and Lasso, suggesting their strength in capturing nonlinear relationships in the data generation process. The biases for $\theta_{3}$ are slightly higher than the biases for $\theta_{1}$ and $\theta_{2}$, possibly because the treatment model for interaction term is implicit and more complex than the model of the individual treatment.

• •

  Regarding the rMSE, ordinary regression and Lasso show higher rMSE for $\theta_{1}$, $\theta_{2}$, and $\theta_{3}$ compared to other methods. For all methods, the rMSE for $\theta_{1}$ is relatively higher, suggesting that estimating the binary treatments shows more variability.

• •

  Results between 2-fold and 5-fold cross-fitting are similar.

Consider the following data generation process of covariates $\bm{X}$ of 10 dimensions and regimens $R$ of three categories. Covariates $X_{1}$ to $X_{5}$ follow a normal distribution of $N(0,1)$ and $X_{6}$ to $X_{10}$ follow a binomial distribution with different probabilities of success. $R=1,2,$ or $3$, is generated by a multinomial/softmax regression model:

with the probability

where $l_{1}(\bm{X})=1$ and $l_{d}(\bm{X})=\bm{X}^{L}\beta^{L}_{d}+\bm{X}^{NL}\beta^{NL}_{d}$ for $d=2,3$. See detailed model specification for $l_{d}(\bm{X})$ in the appendix. In this setting, $R=1$ can be treated as a control regimen due to $l_{1}(\bm{X})=1$ and the relationships between regimens are

$Y$ is a continuous outcome generated by the outcome model $g(R,\bm{X})$ that allows for the interaction between $R$ and $\bm{X}$. Hence the effects of regimens are heterogeneous:

where $\mathbbm{1}_{R=3}$ interacts with $X_{9}$. This simulation is to estimate causal estimand $ATE_{bc}$, as described in Section 4. Here, the true values of these estimands are

We compare the performance of DML interactive model with two other previously proposed methods based on the generalized propensity score: IPW and PSM. In DML, the outcome and treatment models are both fitted using the boosting trees. We use the Python package DoubleML to implement DML interactive model (Bach et al., 2022). For IPW and PSM, the generalized propensity score model is estimated using both Generalized Linear Models (GLMs) and the boosting trees, respectively. We use the R package PSweight (Zhou et al., 2022) to implement them.

The performance of all methods is measured by relative bias. All methods use the original covariate matrix $\bm{X}$ as the input. 200 datasets are generated with $N=1000$. The Figure 3 shows the simulation results.

The simulation results show that DML interactive model using boosting trees shows the lowest relative bias across all ATE comparisons, and the 5-fold cross-fitting presents a slightly less variability of relative bias than the 2-fold cross-fitting. In contrast, traditional IPW and PSM methods show larger relative biases and greater variability, even when the propensity score is estimated using boosting trees. These findings suggest that DML interactive model account features double robustness by estimating both treatment models and outcome models, which provides more robust estimates for ATE across multiple regimens than traditional methods.

We use the DML partial linear model in Algorithm 1 to estimate the effects of $A_{\text{ART}}$, $A_{\text{TDF}}$, $A_{\text{HTN}}$, and the interaction term $A^{\circ}_{\text{ART*TDF}}$ , on the outcome $Y_{\text{eGFR}}$. An interesting aspect of this application is performing a residual regression that includes $A_{\text{ART}}$—a three-class categorical variable and $A^{\circ}_{\text{ART*TDF}}$—a six-class categorical variable. As mentioned in Section 3, for $A_{\text{ART}}$, we fit a multi-class classification model to estimate the conditional probability of each class of ART treatment. We then use dummy variable encoding on $A_{\text{ART}}$ to create a set of binary indicators: $A_{\text{ART: NNRTI}}$ (reference level), $A_{\text{ART: bPI}}$, and $A_{\text{ART: DTG}}$. The residual is computed as

where $\hat{p}[A_{\text{ART:}d}|\bm{X}]$ is the estimated probability of ART type $d$ given covariates $\bm{X}$, for $d\in\{\text{NNRTI, bPI, DTG}\}$. Such approach also applies to the treatment interaction $A^{\circ}_{\text{ART*TDF}}$, which is a six-level categorical variable, as detailed in Appendix.

Importantly, we must ensure that the coefficients in the residual regression can be interpreted as effects relative to the reference levels. Since the NNRTI serves as a reference level for the ART treatment, the residual $\tilde{A}_{\text{ART: NNRTI}}$ is omitted from the residual regression model. The same approach applies to the interaction term: we omit the terms involving either contain the level of NNRTI or the level of “no TDF use”. Therefore, the final residual regression is specified as