Differentiable Decision Tree via "ReLU+Argmin" Reformulation

RADDT ("ReLU+Argmin"-based Differentiable Decision Tree) is designed to learn oblique decision trees by optimizing the entire tree structure via gradient-based optimization. The algorithm introduces a novel unconstrained reformulation for decision tree training, leveraging "ReLU+Argmin" followed by softmin approximation to overcome the inherent non-differentiability of decision trees. It supports both regression and classification tasks. For detailed information on the algorithm, please refer to our Spotlight 🏴 paper published at NeurIPS 2025, available at https://neurips.cc/virtual/2025/poster/119074.

Features in this version:

This repository provides the source code for the RADDT algorithm, including both single-GPU/CPU and distributed multi-GPU versions.

We are currently working on wrapping up the code into a more organized and user-friendly interface like pip package. Stay tuned for updates!

If you require some specific features not included in the current release, please feel free to contact us for access to more additional codes :)

Requirements

• PyTorch 2.0.1
• Python 3.9.6
  • scikit-learn 1.0.2
  • numpy 1.21.6
  • pandas 1.3.5
  • h5py 3.8.0

Script Description

• src folder contains the scripts of "ReLU+Argmin"-based Differentiable Decision Tree Optimization for oblique trees, termed as RADDT.

  • ancestorTF_File subfolder contains the deterministic tree path routing (simply tree path of sample assignment) in h5 file format. These files can be generated by the function of treePathCalculation in treeFunc.py script.
  • treeFunc.py includes utility functions.
  • dataset.py is to load the dataset.
  • warmStart.py generates the warm-start initialization based on CART method.
  • modifiedScheduler.py is a learning rate scheduler with initial linear warmup. We acknowledge the original contributor for this type of scheduler as posted in PyTorch issues #80308.
  • RADDT.py includes the main functions.

• test folder contains the script of running these algorithms.

  • test_RADDT.py is to test the RADDT method.

• data folder contains the datasets. These publicly-available datasets can be obtained from the UCI Machine Learning repository and OpenML. Each dataset is shuffled and split into training, validation and testing sets, and saved in the *.csv format.

• sh_narval_MultiGPU folder contains job submission script for running experiments on the "Narval" cluster in Computer Canada. This version utilizes the Distributed Data Parallel (DDP) strategy using multi-GPUs (e.g. eight GPUs in the example).

We release the source code in two versions, located in separate folders: singleGPUorCPUVersion and distributedMultiGPUVersion. The scripts are nearly identical, but the distributedMultiGPUVersion is adapted for multi-GPU training using the Distributed Data Parallel (DDP) strategy. Scripts modified for this purpose are distinguished by a "_DDP" suffix.

Usage Example

The examples for single GPU/CPU computing can be implemented via:

# test the RADDT method 
python .\test\test_RADDT.py 3 3 1 1 2 3000 "cuda" 10 5

For distributed multi-GPU computing, please refer to the job submission script in the sh_narval_MultiGPU folder for an example of running on the "Narval" cluster in Computer Canada.

Others

If you encounter any errors or notice unexpected tree performance, please don't hesitate to contact us.

License

This repository is published under the terms of the GNU General Public License v3.0 .