GRACE (Gradient-Based Causal Tree Ensembles) - ICML 2026

This repository contains the official Pytorch implementation of Gradient-Based Causal Tree Ensembles: A Backbone Architecture for Heterogeneous Treatment Effects (ICML 2026)

Authors: Yusuke Kano, Jeremy P Voisey, Mihaela van der Schaar

---

Overview

TL;DR: GRACE is gradient‑based causal tree ensembles with multi‑way, oblique, and soft splits that can provide a flexible, robust, and accurate new approach to heterogeneous treatment effect estimation.

Quick Integration

To use GRACE in your model, replace a standard fully connected layer:

import torch.nn as nn
layer = nn.Linear(in_features=32, out_features=32)

with the GRACE layer:

from GRACE_layer import GRACE_layer
layer = GRACE_layer(in_dim=32, out_dim=32, n_trees=100, depth=4, n_split=3, var_dim=0.8, dropout=0.2, relaxed=True)

Key Parameters

• n_trees: Number of trees in the ensemble
• depth: Depth of each tree
• n_split: Number of splits per node
• var_dim: Feature subsampling ratio
• dropout: Dropout rate
• relaxed: Whether to use relaxed routing (GRACE-R)

Environment

• Python 3.x
• Required packages are listed in requirements.txt.

Data Generation

Synthetic datasets and glucose simulation used in the experiments can be generated automatically by running data_generator.py.
This script will automatically create the dataset and store it in the data/ directory.

Training

To run the training process, execute main.py.

Configuration

All hyperparameters and experiment settings can be modified via the configuration file config.yaml.

Citation

If you find this work useful, please cite:

@inproceedings{kano2026grace,
  title={Gradient-Based Causal Tree Ensembles: A Backbone Architecture for Heterogeneous Treatment Effects},
  author={Kano, Yusuke and Voisey, Jeremy P. and van der Schaar, Mihaela},
  booktitle={Proceedings of the International Conference on Machine Learning (ICML)},
  year={2026}
}

License

This project is licensed under the GNU General Public License v2.0 (GPL-2.0).
See the LICENSE file for details.