Variable importance for causal forests: breaking down the heterogeneity of treatment effects

1.1.1 Contributions

Our main contribution is thus the introduction of a variable importance algorithm for causal random forests, following the drop and retrain principle, which is well-established for regression problems [17], [18], [19], [20]. The main idea is to retrain the learning algorithm without a given input variable, and measure the drop of accuracy to get its importance. In particular, such approach ensures that irrelevant variables get a null importance asymptotically. In the context of causal inference, the main obstacle is to retrain the causal forest without a confounding variable, since the unconfoundedness assumption can be violated, leading to inconsistent forest estimates and biased importance values, as explained in Section 2. However, we will see that the local centering of the outcome and treatment assignment leads to consistent estimates, provided that the removed variable is not involved in the treatment effect heterogeneity. Otherwise, to handle a confounder involved in heterogeneity, we introduce a corrective term in the retrained causal forest. Overall, we will show in Section 3, that our proposed variable importance algorithm is consistent, under standard assumptions in the literature about the theoretical analysis of random forests. Next, in Section 4, we run several batches of experiments on simulated, semi-synthetic, and real data to show the good performance of the introduced method compared to the existing competitors. Additionally, we take advantage of the experimental section to illustrate that the extension of our approach to group of variables is straightforward and provides powerful insights in practice. The remaining of this first section is dedicated to the mathematical formalization of the problem.

2.3.1 Identifiability of treatment effect

When a variable X ^(j) is removed from the input variables, the moment Equation (3) does not necessarely hold anymore, since unconfoundedness Assumption (1) may be violated with a reduced set of inputs, even for j ∉ H . However, an important feature of causal forests is the preliminary step of local centering of the observed outcome and treatment assignment, explained above. The following proposition shows that the treatment effect is well identified by the local moment equation of causal forests including only variables in H , provided that the data is centered with all inputs. We recall that m ( X ) = E [ Y ∣ X ] and π ( X ) = E [ W ∣ X ] .

On the other hand, removing an influential and confounding variable j ∈ H to learn a causal forest is more delicate. Indeed, a local moment equation to identify the mean CATE over X ^(j) exists if the treatment effect is uncorrelated to the squared centered treatment assignment.

Athey and Wager [25], Footnote 5, page 42] conduct an empirical analysis using causal forests, and state in a footnote, that local centering “eliminates confounding effects. Thus, we do not need to give the causal forest all features X ^(j) that may be confounders. Rather, we can focus on features that we believe may be treatment modifiers”. However, Propositions 4 and 5 show that this statement must be completed. Indeed, Proposition 4 states that confounders not involved in the heterogeneity of the treatment effect, i.e. confounders that do no belong to H , may be dropped without hurting the identifiability of τ, thanks the local centering step. On the other hand, Proposition 5 shows that this is clearly not the case for confounders involved in heterogeneity, as the treatment effect is not properly identified by the local moment equation of causal forests, even with local centering. To overcome this problem, we introduce a corrective term in the retrained forest.

2.3.2 Corrected causal forests

The additional covariance term in Proposition 5 can be estimated using the original causal forest fit with all inputs. Therefore, we propose the corrected causal forest estimate when removing a confounding variable X ^(j) with j ∈ H . Recall that the weights α′(x ^(−j)) are generated by the causal forest using centered data and dropping variable X ^(j), to define τ M , n ( − j ) ( x ) . We define the corrected causal forest estimate θ M , n ( − j ) ( x ) as

where W α ′ 2 ‾ = ∑ i = 1 n α i ′ ( x ( − j ) ) W ˜ i 2 , W ‾ α ′ = ∑ i = 1 n α i ′ ( x ( − j ) ) W ˜ i , and the mean treatment effect is τ ‾ α ′ = ∑ i = 1 n α i ′ ( x ( − j ) ) τ M , n ( X i ) . More precisely, the corrective term of Equation (5) is the forest estimate of the third term of the equation of Proposition 5 divided by V [ W − π ( X ) ∣ X ( − j ) ] . Consequently, the corrected causal forest θ M , n ( − j ) ( x ) retrained without a confounding variable is an estimate of E [ τ ( X ( H ) ) ∣ X ( − j ) = x ( − j ) ] , which is the targeted quantity, and is consistent, as we will show in Section 3. However, note that the correction term can be small in practice, as demonstrated in the experimental Section 4.

2.3.3 Variable importance estimate

Using D n ′ = X i ′ , Y i ′ , W i ′ i = 1 n an independent copy of D n , we define

where τ M , n ‾ = ∑ i = 1 n τ M , n X i ′ / n , and I n ( 0 ) is the mean squared difference between the initial forest predictions and the predictions of the corrected forest θ M , n ( 0 ) X i ′ , Θ M ′ , retrained with still all the inputs variables involved but a new randomization Θ M ′ , i.e.,

In fact, I n ( 0 ) partially removes the bias of the first term of I n ( j ) , due to the randomization of the forest training, and vanishes as the sample size increases if the causal forest converges. Notice that the above definition is formalized with D n ′ for the sake of clarity, but that such additional data is usually not available in practice. Instead, out-of-bag causal forest estimates are rather used to define I n ( j ) , as summarized in Algorithm 1 below.

4.1.1 Experiment 1

We consider a first example of simulated data to highlight the good performance of the proposed importance measure. The input is a Gaussian vector of dimension p = 8, defined by X ∼ N ( 0 , Σ ) , where Σ is the identity matrix. The treatment assignment W is simply a Bernoulli random variable of parameter 1/2, and the response Y follows

where ε ∼ N ( 0,0.1 ) . This first experiment gives a baseline in a quite simple case, since input variables are independent, there are no confounding effects, and the ratio V [ τ ( X ( H ) ) ] / V [ Y ] is large, with a value of about 50 %. The targeted theoretical values can be easily computed from the defined distributions and Equation (8), and are reported in the left column of Table 1. Next, we draw a sample of size n = 3000, and the causal forest is fit with all default settings, including the number of trees M = 2000. Table 1 shows the estimated mean importance values over 10 repetitions for our proposed algorithm I n ( j ) , TE-VIM, and the metric from the grf package, along with the standard deviation of the mean importance in brackets. Notice that for each repetition, a new data sample is drawn from the data generating process described above, and then is used to run each variable importance algorithm. As expected, both I n ( j ) and TE-VIM provide accurate estimates of the theoretical importance values I^(j) in this quite simple setting. The fast metric provided by grf-vimp also provides a good approximation of the relative variable importance, even if irrelevant variables get higher values than with I n ( j ) and TE-VIM. Next, Figure 1 displays the importance I n ( j ) with respect to the sample size n. The estimated values get closer to the theoretical quantities, as expected from the convergence results of Section 3. Notice that variables X ⁽³⁾, …, X ⁽⁶⁾ all have a small importance, and therefore overlap on Figure 1. We omit X ⁽⁷⁾ and X ⁽⁸⁾ on the figure, since they are symmetric with X ⁽⁶⁾.Finally, we also take advantage of this first experiment to analyze a case of higher dimension. We thus consider the same settings, but we add 32 independent Gaussian variables of unit variance to get a final input dimension of p = 40. Table 2 shows that the impact of this large number of noisy variables is small for I n ( j ) and TE-VIM, but strong for grf-vimp, with a high decrease of the importance value of X ⁽¹⁾. Such phenomenon is expected, since the forest quite often splits on noisy variables because of the split randomization at each tree node, leading to a lower split frequency of the relevant inputs. Besides, we can also notice that TE-VIM assigns a negative bias to irrelevant variables in this noisy setting, whereas I n ( j ) still provides values very close to 0.

Figure 1:

Importance values I n ( j ) with respect to the sample size n for Experiment 1. Non-null theoretical importance are displayed as solid lines.

4.1.2 Experiment 2

We consider the same settings as in Experiment 1 with the original dimension p = 8, but we introduce three modifications to make the problem more difficult. In particular, we set W ∼ B e r n o u i l l i ( 0.4 + 0.2 1 X ( 1 ) > 0 ) to introduce a confounding factor, the covariance Σ is now the identity matrix except that Cov(X ⁽¹⁾, X ⁽⁵⁾) = 0.9 to get a strong correlation between two variables, and we finally increase the weight of ( X ( 3 ) × X ( 4 ) ) 2 to 1 to reduce the ratio V [ τ ( X ( H ) ) ] / V [ Y ] to about 5 %. Such a quite small ratio is realistic, and makes the treatment effect quite difficult to estimate in practice. The response Y is now defined by

where ε ∼ N ( 0,0.1 ) . We then take a sample size n = 3000, and the causal forest is again fit with default settings and a number of trees M = 2000. Here, both X ⁽¹⁾ and X ⁽²⁾ are involved in heterogeneity, i.e.  H = { 1,2 } , but only X ⁽¹⁾ is also a confounder. Results are averaged over 50 repetitions, and are reported in Table 3. Additionally, the standard deviation of the mean importance for each variable is displayed in brackets, except for negligible values ( < 0.001 ) . The first column of Table 3 is the oracle importance value, precisely estimated using Equation (2), the closed-form of τ given by Equation (9), and a Monte-Carlo method with a large sample drawn from the joint distribution of (Y, W, X), known by construction.

The results displayed in Table 3 show that both our algorithm and TE-VIM provide the accurate variable ranking, where X ⁽²⁾ is the most important variable, and X ⁽¹⁾ the second most important one. However, the standard deviations of the mean importance values over 50 repetitions are higher for TE-VIM, which is induced by a high instability across repetitions. This can be a limitation in practice with real data, as importance values are computed only once. We also observe that variables with a theoretical null importance get a quite strong negative bias. On the other hand, the importance measure from the grf package underestimates the importance of variable X ⁽²⁾, and identifies X ⁽³⁾, X ⁽⁴⁾, and X ⁽⁵⁾ as slightly more important than X ⁽²⁾, although these three variables are not involved in the treatment heterogeneity by construction. In particular, X ⁽⁵⁾ is not involved at all in the response Y, but is strongly correlated to the influential input X ⁽¹⁾. Because of this dependence, X ⁽⁵⁾ is frequently used in the causal forests splits, leading to this quite high importance given by the grf package. On the other hand, I n ( j ) gives an importance close to 0 for X ⁽⁵⁾. This result is expected, since the removal of X ⁽⁵⁾ does not lead to any loss of information regarding the treatment heterogeneity, by definition. An additional interesting phenomenon is the non-negligible importance for variables X ⁽³⁾ and X ⁽⁴⁾ given by all procedures. In fact, the interaction term in the baseline function μ, which takes the form of a squared product, is rather difficult to estimate by regression forests. Then, the local centering of Y is only partial, and X ⁽³⁾ and X ⁽⁴⁾ still have impact on the variance of treatment estimates. Besides, notice that the corrective term of Equation (5) is negligible in this experiment, and that using the original causal forest retrained with one variable removed, gives the same result as in Table 3 for I n ( j ) , up to the displayed digits. Finally, Figure 2 also displays the importance values I n ( j ) with respect to the sample size n. Results are consistent with the convergence results of the previous section. In particular, for a large sample of n = 10000, the relative errors of I n ( j ) happen to be really small, whereas irrelevant variables get high importance values for n = 500.

Figure 2:

Importance values I n ( j ) with respect to the sample size n for Experiment 2. Non-null theoretical importance are displayed as solid lines.

4.1.3 Experiment 3

This third experiment has the same setting than Experiment 2, except that variable X ⁽¹⁾ is only a confounder and is not involved in the treatment effect heterogeneity anymore. Now, the response writes

The results are provided in Table 4. Clearly, I n ( j ) outperforms the competitors. Indeed, X ⁽²⁾ is well-identified by I n ( j ) as responsible for most of the heterogeneity of the treatment effect, whereas TE-VIM is strongly biased, with some values exceeding 1 because of the variance of estimates involved in the ratio of TE-VIM. The importance procedure of the grf package outputs quite close values for X ⁽²⁾, X ⁽⁴⁾, and X ⁽³⁾. As expected, the importance of these last two variables is relatively larger than in Experiment 1, since the ratio V [ τ ( X ( H ) ) ] / V [ Y ] drops to 1 % in this case. As in the previous experiments, Figure 3 displays the importance values I n ( j ) with respect to the sample size n, and shows that the errors decrease as n increases, following the theory.

Figure 3:

Importance values I n ( j ) with respect to the sample size n for Experiment 3. Non-null theoretical importance are displayed as solid lines.

4.1.4 Experiment 4

The goal of this fourth simulated experiment is to highlight a case where the corrective term in the retrained causal forest has a strong influence, as opposed to Experiments 1, 2, and 3. We consider p = 5 inputs uniformly distributed over [0,1], except X ⁽¹⁾ defined as X ⁽¹⁾ = U ³, where U ∼ U ( 0,1 ) . The treatment assignment W is a Bernoulli variable defined from π(X) = X ⁽¹⁾, and the response is given by

where ε ∼ N ( 0,0.1 ) . We still use n = 3000 and M = 2000 trees in the causal forests. Next, we compute our importance measure I n ( j ) for all inputs, as well as its counterpart I n ( j ) , where the corrective term is removed, and with 10 repetitions for uncertainties. Results are reported in Table 5, and clearly show the high bias of the importance of X ⁽¹⁾ when the corrective term in the retrained forest is removed. Indeed, we get I n ( 1 ) = 1.57 , whereas the target quantity is I⁽¹⁾ = 1, since X ⁽¹⁾ is the only variable involved in the treatment effect heterogeneity and X ⁽¹⁾ is independent of the other inputs. With the correction, we recover an importance value of 0.98 for X ⁽¹⁾ as expected. Notice that the asymptotic bias exhibited in Theorem 5 takes values 0.72 for this case, which explains the empirical results. Importantly, this bias takes small values in practice in most cases. Here, we take the treatment effect as τ ( X ( H ) ) = 10 π ( X ) ( 1 − π ( X ) ) to maximize the covariance term involved in the bias of Theorem 5.

4.3.1 Welfare data

For a first experiment with real data, we use the “Welfare” dataset from a GSS survey, introduced in Green and Kern [28] and available at https://github.com/gsbDBI/ExperimentData. The goal of this survey is to analyze the impact of question wording about the support of Americans to the government welfare spending. Respondents are randomly assigned one of two possible questions, with the same introduction and response options, but using the phrasing “welfare” or “assistance to the poor”. In fact, this slight wording difference has a quite strong impact on the survey answers, and defines the treatment. The output of interest indicates if respondents have answered that “too much” is spent. Our objective is to identify the main characteristics of individuals that have an impact on the heterogeneity of the treatment effect. The considered dataset is of size n = 13198 with p = 31 input variables, and basic data preparation steps were used to drop rows with missing values. Notice that imputing missing values may improve estimates. We leave this topic for future work, as handling missing values for variable importance is of high practical interest.

Table 8 displays the top 10 most important variables for Welfare data using our algorithm I _n and also the importance from the grf package. The ranking provided by the two algorithms are close, but I _n has a clear meaning as the variance proportion of the treatment effect lost when a given variable is removed, whereas grf-vimp can only be used as a relative importance between covariates, without an intrinsic meaning.

Notice that the sum of the importance of all input variables, i.e.  ∑ j I n ( j ) , adds to 0.45, which is far from 1. Indeed, when inputs are independent, we have ∑ _j I^(j) ≥ 1. Such a low value is explained by the correlation within input variables. We run a simple hierarchical clustering of the input variables in 10 groups based on correlation, to enforce a small correlation between these groups. More precisely, the hierarchical clustering uses the dissimilarity matrix defined as 1 − C, with C the correlation matrix of the inputs, and the cutting threshold is set to get 10 groups of variables. Then, we run the group variable importance I n ( J ) for each group of variables J ⊂ {1, …, p}. The results are displayed in the following Table 9, and are quite straightforward to read. Indeed, half of the treatment heterogeneity is explained by political orientations of individuals, almost a quarter of the heterogeneity is given by variables mostly related to education and degrees. Then, several groups have a small impact, especially a group about income and working status, and a second one about family information.

4.3.2 NHEFS health data

For the second case study, we use the NHEFS real data about body weight gain following a smoking cessation, extensively described in the causal inference book of Hernan and Robins [29]. As highlighted in the introduction of Chapter 12, these data help to answer the question “what is the average causal effect of smoking cessation on body weight gain?”. According to the authors, the unconfoundedness assumption holds. Here, we go a step further to analyze the heterogeneity of this causal effect with respect to health and personal data of individuals who have stopped smoking, using causal forests and our variable importance algorithm. The data record the weight of individuals, first measured in 1971, and then in 1982. The treatment assignment W indicates whether people have stopped smoking during this period, and the observed output Y is the weight difference between 1971 and 1982. We take the dataset of size n = 1566 used in Hernan and Robins [29, Chapter 12]. Notice that 63 rows with the output missing were removed, introducing a small bias, as discussed by the authors. They include 9 variables in their analysis, sufficient for unconfoundedness. To better estimate heterogeneity, we also include all variables of the original dataset, that do not contain missing values and are not related to the response, and obtain p = 41 input variables. As already mentioned, handling missing values is out of scope of this article, and is left for future work. We run our variable importance algorithm and the grf importance, using M = 4000 trees.

The results are displayed in Table 10. Clearly, the original weight of individuals in 1971 has a strong causal effect on weight gain following smoking cessation, with half of the treatment effect variance lost when this variable is removed. The intensity and duration of smoking, as well as personal characteristics, such as height and age are also involved in treatment heterogeneity, according to both algorithms. Notice that grf-vimp underestimates the importance of wt71 with respect to other variables. Next, we group together variables that are highly correlated, to compute group variable importance. We use the same clustering procedure as the Welfare case, except that we increase the number of groups to 30, since most variables have a very weak dependence with the others. Sex, height, and birth control are highly correlated with the weight in 1971, and this group explains two third of the treatment effect heterogeneity. In fact, age and smoke years also have a quite strong impact with a quarter of heterogeneity explained. Also notice that sex and birth control belong to the most influential group, but have a low importance I n ( j ) , which means that they do not contain unique information regarding the treatment effect. This shows how single and group variable importance provide complementary information (Table 11).