Research map / Optimization / Optimizers
GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes GradTree, learning hard axis-aligned decision trees by joint gradient descent on a dense tree representation with straight-through entmax+hardmax and Heaviside-logistic split operators.
Problem
Learning decision trees is non-convex and non-differentiable, so greedy CART-style growth yields locally optimal, sometimes inaccurate trees.
Approach
Represent a fully-grown axis-aligned tree densely (one-hot feature selection matrix plus per-feature thresholds and leaf class distributions) and train all parameters jointly by backpropagation, using straight-through operators (entmax followed by hardmax, and rounded logistic) to keep hard splits in the forward pass.
Key results
- GradTree achieves the best mean reciprocal rank (0.758) and mean relative difference (0.008) on 20 binary classification datasets, outperforming CART (MRR 0.556).
- On multi-class datasets GradTree is second-best (MRR 0.619) behind CART (MRR 0.774), with CART best on high-dimensional multi-class tasks.
- ST entmax ablation beats ST Gumbel softmax (binary macro F1 0.771 vs 0.569) and temperature annealing (0.759) with optimized parameters.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use GradTree instead of CART | Binary classification benchmarks · Binary classification benchmarks | MRR: 0.556 ± .293 -> 0.758 ± .306 | Helped |
| Use GradTree instead of CART | Multi-class classification benchmarks · Multi-class classification benchmarks | MRR: 0.774 ± .274 -> 0.619 ± .303 | Hurt |
| Replace ST Entmax with ST Gumbel Softmax (ablation) | GradTree · Binary classification benchmarks, optimized parameters | macro F1: 0.771 ± .036 -> 0.569 ± .094 | Hurt |
| Replace ST Entmax with temperature annealing (ablation) | GradTree · Binary classification benchmarks, optimized parameters | macro F1: 0.771 ± .036 -> 0.759 ± .044 | Hurt |
| Replace ST Entmax with ST Gumbel Softmax (ablation) | GradTree · Multi-class benchmarks, optimized parameters | macro F1: 0.699 ± .064 -> 0.297 ± .093 | Hurt |
| Replace ST Entmax with temperature annealing (ablation) | GradTree · Multi-class benchmarks, optimized parameters | macro F1: 0.699 ± .064 -> 0.601 ± .067 | Hurt |
| Replace ST Entmax with ST Gumbel Softmax (ablation) | GradTree · Binary classification benchmarks, default parameters | macro F1: 0.764 ± .039 -> 0.560 ± .097 | Hurt |
| Replace ST Entmax with temperature annealing (ablation) | GradTree · Binary classification benchmarks, default parameters | macro F1: 0.764 ± .039 -> 0.757 ± .042 | Hurt |
| Replace ST Entmax with ST Gumbel Softmax (ablation) | GradTree · Multi-class benchmarks, default parameters | macro F1: 0.638 ± .076 -> 0.272 ± .092 | Hurt |
| Replace ST Entmax with temperature annealing (ablation) | GradTree · Multi-class benchmarks, default parameters | macro F1: 0.638 ± .076 -> 0.602 ± .073 | Hurt |
Surprising findings
- With default hyperparameters, CART attains the highest multi-class MRR (0.929) and GradTree falls to second (0.470), reversing its binary-task advantage.
- GradTree's train-test gap on binary tasks (0.051) is small, but its train performance is lower than CART's, so its advantage comes from reduced overfitting rather than better fit.
Limitations
- Dense representation makes optimization harder as the number of features and classes grows, leading to weaker multi-class and high-dimensional performance.
- Only fully-grown trees are trained, with post-hoc pruning rather than learned pruning during training.
- DNDT was limited to datasets with at most 12 features, so comparisons are incomplete on high-dimensional data.
Experimental setup
Datasets: Blood Transfusion Banknote Authentication Titanic Raisins Rice Echocardiogram Wisconsin Breast Cancer Loan House Heart Failure Heart Disease Adult Bank Marketing
Benchmarks: Blood Transfusion Banknote Authentication Titanic Adult Iris Car Splice UCI tabular datasets (binary and multi-class)
Metrics: macro F1 mean reciprocal rank mean relative difference train-test difference tree size runtime
Baselines: CART DNDT GeneticTree DL8.5 Gumbel softmax ST temperature annealing
Builds on: DNDT straight-through estimator entmax DL8.5 CART
Closest papers by meaning
- Branches: Efficiently Seeking Optimal Sparse Decision Trees via AO* (ICML 2025 · similarity 0.88)
- Learning Small Decision Trees with Few Outliers: A Parameterized Perspective (AAAI 2024 · similarity 0.88)
- Symbolic Regression Enhanced Decision Trees for Classification Tasks (AAAI 2024 · similarity 0.87)
- Near-Optimal Decision Trees in a SPLIT Second (ICML 2025 · similarity 0.86)
- Witty: An Efficient Solver for Computing Minimum-Size Decision Trees (AAAI 2025 · similarity 0.86)
- TREE-G: Decision Trees Contesting Graph Neural Networks (AAAI 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.