Gradient-Based Causal Tree Ensembles: A Backbone Architecture for Heterogeneous Treatment Effects

Estimating Heterogeneous Treatment Effects (HTE) from observational data is essential in fields such as healthcare and policy-making, where randomized experiments are often impractical. While representation learning-based methods have shown promise, recent studies suggest that tree-based approaches may offer superior performance on tabular data, particularly in the presence of uninformative features. We introduce GRAdient-based Causal tree Ensembles (GRACE), a novel tree-based architecture for HTE estimation that incorporates multi-way, oblique, and soft splits, enabling end-to-end training via backpropagation. GRACE can be seamlessly integrated into existing models as a replacement for fully-connected neural network layers. Across diverse benchmarks involving binary and non-binary treatment settings, GRACE consistently surpasses neural network and tree-based baselines, often by a substantial margin. We further analyze GRACE as an extension of fully-connected neural network layers and conduct ablation studies to isolate and quantify the contribution of each architectural component to the improvement in performance. These results position GRACE as a promising new foundation for flexible, robust, and accurate HTE estimation. GRACE code is available at https://github.com/ysk-kano/GRACE

In fields such as healthcare and public policy, it is important to understand how the effects of a treatment or intervention differ from person to person. These individual-level effects are known as heterogeneous treatment effects (HTE). Until now, neural networks, which are a popular type of model in machine learning, have mainly been used to estimate HTE. On the other hand, in standard machine learning tasks, methods based on tree structures are also known to perform well. Based on this, we wondered whether tree-based models might also be useful for estimating HTE. In this study, we propose a new type of model that retains the advantages of tree structures, while also being able to improve step by step like neural networks by adjusting its settings many times. We tested this method in a wide range of situations, from simple cases with two treatment choices to more complex cases with multiple treatments. The results show that our method performs better and more consistently than both traditional tree-based methods and neural networks. In addition, our method can be easily added to existing neural network models by simply replacing one part, which makes it practical and easy to use.