Soft Decision Tree classifier: explainable and extendable PyTorch implementation

TL;DR

This paper presents a PyTorch implementation of Soft Decision Trees and an extension with short-term memory, demonstrating comparable performance to XGBoost and better results than traditional methods on various datasets.

Contribution

The authors introduce a PyTorch-based implementation of Soft Decision Trees and an extension with short-term memory, highlighting their explainability and competitive performance.

Findings

SDT and SM-SDT achieve similar AUC to XGBoost.

Methods outperform Random Forest, Logistic Regression, and Decision Tree.

Clinical dataset results show comparable performance across methods.

Abstract

We implemented a Soft Decision Tree (SDT) and a Short-term Memory Soft Decision Tree (SM-SDT) using PyTorch. The methods were extensively tested on simulated and clinical datasets. The SDT was visualized to demonstrate the potential for its explainability. SDT, SM-SDT, and XGBoost demonstrated similar area under the curve (AUC) values. These methods were better than Random Forest, Logistic Regression, and Decision Tree. The results on clinical datasets suggest that, aside from a decision tree, all tested classification methods yield comparable results. The code and datasets are available online on GitHub: https://github.com/KI-Research-Institute/Soft-Decision-Tree