Distilling heterogeneous treatment effects: Stable subgroup estimation in causal inference

TL;DR

This paper introduces causal distillation trees (CDT), a method combining machine learning and decision trees to identify interpretable subgroups with heterogeneous treatment effects, validated on HIV treatment data.

Contribution

The paper proposes CDT, a novel approach that distills black-box treatment effect estimates into interpretable subgroups with theoretical guarantees and stability diagnostics.

Findings

CDT outperforms existing methods in stability and clinical relevance.

Theoretical guarantees ensure consistency of subgroup estimation.

Application to HIV trial demonstrates practical utility.

Abstract

Recent methodological developments have introduced new black-box approaches to better estimate heterogeneous treatment effects; however, these methods fall short of providing interpretable characterizations of the underlying individuals who may be most at risk or benefit most from receiving the treatment, thereby limiting their practical utility. In this work, we introduce \textit{causal distillation trees} (CDT) to estimate interpretable subgroups. CDT allows researchers to fit any machine learning model to estimate the heterogeneous treatment effect, and then leverages a simple, second-stage tree-based model to "distill" the estimated treatment effect into meaningful subgroups. As a result, CDT inherits the improvements in predictive performance from black-box machine learning models while preserving the interpretability of a simple decision tree. We derive theoretical guarantees for the consistency of the estimated subgroups using CDT, and introduce stability-driven diagnostics for researchers to evaluate the quality of the estimated subgroups. We illustrate our proposed method on a randomized controlled trial of antiretroviral treatment for HIV from the AIDS Clinical Trials Group Study 175 and show that CDT out-performs state-of-the-art approaches in constructing stable, clinically relevant subgroups.