# GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent

arXiv: https://arxiv.org/abs/2305.03515 · HTML: https://mlsanity.com/paper/2305.03515/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 39 references that are in the map.
- First posted on arXiv 5 May 2023.

## Builds on (most cited first, 19 of 19)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Distilling a Neural Network Into a Soft Decision Tree](https://mlsanity.com/paper/1711.09784/) (2017)
- [Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data](https://mlsanity.com/paper/1909.06312/) (2019)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [Generalized and Scalable Optimal Sparse Decision Trees](https://mlsanity.com/paper/2006.08690/) (2020)
- [NODE-GAM: Neural Generalized Additive Model for Interpretable Deep Learning](https://mlsanity.com/paper/2106.01613/) (2021)
- [Deep Neural Decision Trees](https://mlsanity.com/paper/1806.06988/) (2018)
- [Adaptive Neural Trees](https://mlsanity.com/paper/1807.06699/) (2018)
- [MurTree: Optimal Classification Trees via Dynamic Programming and Search](https://mlsanity.com/paper/2007.12652/) (2020)
- [PolyLoss: A Polynomial Expansion Perspective of Classification Loss Functions](https://mlsanity.com/paper/2204.12511/) (2022)
- [Quant-BnB: A Scalable Branch-and-Bound Method for Optimal Decision Trees with Continuous Features](https://mlsanity.com/paper/2206.11844/) (2022)
- [Efficient non-greedy optimization of decision trees](https://mlsanity.com/paper/1511.04056/) (2015)
- [Learning Accurate Decision Trees with Bandit Feedback via Quantized Gradient Descent](https://mlsanity.com/paper/2102.07567/) (2021)
- [Sparsity in Optimal Randomized Classification Trees](https://mlsanity.com/paper/2002.09191/) (2020)
- [Learning Binary Decision Trees by Argmin Differentiation](https://mlsanity.com/paper/2010.04627/) (2020)
- [Learning Representations for Axis-Aligned Decision Forests through Input Perturbation](https://mlsanity.com/paper/2007.14761/) (2020)

## Built on by (most cited first, 11 of 11)

- [GRANDE: Gradient-Based Decision Tree Ensembles for Tabular Data](https://mlsanity.com/paper/2309.17130/) (2023)
- [Optimal or Greedy Decision Trees? Revisiting their Objectives, Tuning, and Performance](https://mlsanity.com/paper/2409.12788/) (2024)
- [Decision Tree Induction Through LLMs via Semantically-Aware Evolution](https://mlsanity.com/paper/2503.14217/) (2025)
- [Mitigating Information Loss in Tree-Based Reinforcement Learning via Direct Optimization](https://mlsanity.com/paper/2408.08761/) (2024)
- [Hierarchical Retrieval at Scale: Bridging Transparency and Efficiency](https://mlsanity.com/paper/2502.07971/) (2025)
- [MOSIC: Model-Agnostic Optimal Subgroup Identification with Multi-Constraint for Improved Reliability](https://mlsanity.com/paper/2504.20908/) (2025)
- [LHT: Statistically-Driven Oblique Decision Trees for Interpretable Classification](https://mlsanity.com/paper/2505.04139/) (2025)
- [Decision Trees That Remember: Gradient-Based Learning of Recurrent Decision Trees with Memory](https://mlsanity.com/paper/2502.04052/) (2025)
- [Neurosymbolic Decision Trees](https://mlsanity.com/paper/2503.08762/) (2025)
- [Approximation-Free Differentiable Oblique Decision Trees](https://mlsanity.com/paper/2605.07837/) (2026)
- [Adaptive Multi-Branching for Shallow Decision Tree Induction](https://mlsanity.com/paper/2608.29262/) (2026)

## Abstract (arXiv)

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 method is available under: https://github.com/s-marton/GradTree