GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent
Sascha Marton, Stefan Lüdtke, Christian Bartelt, Heiner Stuckenschmidt
Abstract
Decision Trees (DTs) are commonly used for many machine learning tasks due to their high degree of interpretability. However, learning a DT from data is a difficult optimization problem, as it is non-convex and non-differentiable. Therefore, common approaches learn DTs using a greedy growth algorithm that minimizes the impurity locally at each internal node. Unfortunately, this greedy procedure can lead to inaccurate trees. In this paper, we present a novel approach for learning hard, axis-aligned DTs with gradient descent. The proposed method uses backpropagation with a straight-through operator on a dense DT representation, to jointly optimize all tree parameters. Our approach outperforms existing methods on binary classification benchmarks and achieves competitive results for multi-class tasks. The implementation is available under: https://github.com/s-marton/GradTree
BibTeX
@article{Marton_Lüdtke_Bartelt_Stuckenschmidt_2024, title={GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent}, volume={38}, url={https://ojs.aaai.org/index.php/AAAI/article/view/29345}, DOI={10.1609/aaai.v38i13.29345}, abstractNote={Decision Trees (DTs) are commonly used for many machine learning tasks due to their high degree of interpretability. However, learning a DT from data is a difficult optimization problem, as it is non-convex and non-differentiable. Therefore, common approaches learn DTs using a greedy growth algorithm that minimizes the impurity locally at each internal node. Unfortunately, this greedy procedure can lead to inaccurate trees. In this paper, we present a novel approach for learning hard, axis-aligned DTs with gradient descent. The proposed method uses backpropagation with a straight-through operator on a dense DT representation, to jointly optimize all tree parameters. Our approach outperforms existing methods on binary classification benchmarks and achieves competitive results for multi-class tasks. The implementation is available under: https://github.com/s-marton/GradTree}, number={13}, journal={Proceedings of the AAAI Conference on Artificial Intelligence}, author={Marton, Sascha and Lüdtke, Stefan and Bartelt, Christian and Stuckenschmidt, Heiner}, year={2024}, month={Mar.}, pages={14323-14331} }