Identify subgroups.

[linyingyang]

Problem: I think there should be some ways to identify subgroups suffer heterogeneous treatment effects. e.g tree structure can help us to tell that. It seems that SHAP targets only on 1-order (or maybe 2-order interaction), but what if we want more information, say higer orders?

[vasilismsr]

I agree interpretability is a question that is on our plate but have yet to implement some basic ways.

SHAP is a good first attempt at interpretability. It also offers feature importances, which can hint at which variables create heterogeneity and in which direction.

If your final model is a linear model you can simply inspect the coefficients in front of each variable.

Another way to sort of identify which subgroups suffer more is to post-process the final model and train a highly regularized classifier that predicts the sign of the heterogeneous treatment effect output by our CATE model (weighting also every sample by the magnitude of the effect). You can train a single tree for this and then print the structure of the tree. You could also then print the mean CATE for the samples contained in each leaf of the tree to get an understanding of the magnitude of the CATE of each leaf of the tree.

This solution is similar to the final step of Algorithm 1 proposed in the causal boosting paper here: https://arxiv.org/pdf/1707.00102.pdf

[vasilismsr]

Here is for instance one example post-processing with a decision tree. Consider the following DGP:

import numpy as np
n = 10000
d = 100
d_x = 10
X = np.random.binomial(1, .5, size=(n, d))
T = np.random.binomial(1, .5, size=(n,))
y = (X[:, 0] + .5)*T + X[:, 0] + (0*X[:,0] + 1)*np.random.normal(0, 1, size=(n,))

We can estimate the heterogeneous treatment effect with the linear double ml estimator as follows:

from econml.dml import DMLCateEstimator, LinearDMLCateEstimator
from sklearn.linear_model import LinearRegression, LassoCV, LogisticRegressionCV
from sklearn.dummy import DummyClassifier
est = LinearDMLCateEstimator(model_y = LassoCV(cv=3),
                             model_t = LogisticRegressionCV(cv=3, solver='lbfgs'),
                             n_splits=6,
                             linear_first_stages=False,
                             discrete_treatment=True)
est.fit(y, T, X[:, :d_x], X[:, d_x:])

Then we can post-process the model to identify important subgroups by simply fittting a single regression tree trained on the outcome of the fitted model. We can also export the tree to graphviz for visualization:

import sklearn
from sklearn.tree import DecisionTreeRegressor
reg = DecisionTreeRegressor(max_depth=3).fit(X, est.effect(X[:, :d_x]))
sklearn.tree.export_graphviz(reg, 'test.dot')

Then in a shell script you can visualize the tree (if you don't have graphviz, do sudo apt install graphviz):

dot -Tpng test.dot -o test.png

This will produce the following image:

Where you can see which variables are chosen on each split and you can also see the mean treatment effect at each node (the "value" number). Thus you can see that most of the heterogeneity is produced by splitting on variable X[0] and every subsequent splits don't create much difference.