Adaptive Multi-Branching for Shallow Decision Tree Induction
Organizations: Department of Information Convergence Engineering, Pusan National University, Busan, Rep. of Korea · School of Computer Science and Engineering, Pusan National University, Busan, Rep. of Korea
Abstract
Decision trees are attractive for tabular prediction tasks because each prediction follows an interpretable sequence of feature-threshold tests. Under a strict maximum-depth budget, however, conventional binary trees can be under-expressive, since each internal node makes only a single threshold decision. We study shallow-depth tree induction, where the goal is to improve accuracy while keeping root-to-leaf paths short. We propose the Multi-Branch Neural Decision Tree with Adaptive Pruning (MBNDT), a single axis-aligned tree trained end-to-end with differentiable multi-way splits. Each internal node learns ordered thresholds over a selected feature and a branch mask that adapts its effective arity, and the trained model is converted to a deterministic single-path tree for inference. Across 21 OpenML binary-classification benchmarks, MBNDT achieves the best average rank and mean balanced accuracy among depth-constrained single-tree baselines; a controlled ablation isolates multi-way splitting as the source of the gain. These gains come with an explicit trade-off: MBNDT realizes more leaves than the other single-tree baselines, making it best suited when accuracy under short, bounded decision paths is prioritized over minimal global tree size.
Figures & tables
| Dataset | Ensemble (ref.) | Single tree | |||
|---|---|---|---|---|---|
| XGBoost | Gradient-based | Solver-based | Conventional | ||
| MBNDT (ours) | GradTree | SPLIT | CART | ||
| Small datasets ( ) | |||||
| hepatitis | 0.745 0.113 | 0.675 0.123 (1) | 0.555 0.098 (4) | 0.673 0.081 (2) | 0.665 0.147 (3) |
| colic | 0.768 0.049 | 0.787 0.055 (1) | 0.772 0.052 (3) | 0.766 0.049 (4) | 0.775 0.054 (2) |
| vote | 0.954 0.019 | 0.955 0.013 (1) | 0.914 0.038 (4) | 0.952 0.018 (2) | 0.939 0.025 (3) |
| Dataset | MBNDT (ours) | GradTree | SPLIT | CART | ||||||||
| Path | Leaves | Time (s) | Path | Leaves | Time (s) | Path | Leaves | Time (s) | Path | Leaves | Time (s) | |
| Small datasets ( ) | ||||||||||||
| hepatitis | 2.40 | 6.8 | 10.73 | 2.80 | 8.0 | 7.55 | 2.94 | 7.0 | 1.07 | 3.55 | 9.4 | |
| colic | 1.66 | 4.0 | 19.81 | 3.00 | 8.8 | 12.75 | 1.92 | 4.6 | 0.70 | 3.36 | 9.2 | |
| vote | 2.29 | 5.6 | 23.27 | 2.40 | 6.8 | 7.17 | 1.10 | 2.4 | 0.74 | 2.85 | 9.0 | |
| breast-w | 2.62 | 18.2 | 20.33 | 3.40 | 11.2 | 7.67 | 2.07 | 5.2 | 0.94 | 3.33 | 10.6 | |
| Comparison | W/T/L | Wilcoxon | Holm | |
|---|---|---|---|---|
| MBNDT vs. GradTree | 21 | 20/0/1 | ||
| MBNDT vs. SPLIT | 21 | 15/2/4 | 0.00628 | 0.01256 |
| MBNDT vs. CART | 21 | 14/2/5 | 0.02633 | 0.02633 |
| Variant | Bal. Acc. | Path | Leaves | Best | |
|---|---|---|---|---|---|
| MBNDT-B2 | 2 | 0.820 | 3.22 | 9.5 | 4 |
| MBNDT-B3 | 3 | 0.833 | 3.12 | 16.0 | 10 |
| MBNDT-B4 | 4 | 0.832 | 3.09 | 23.3 | 7 |
| Masks | PP | Bal. Acc. | Path | Leaves |
|---|---|---|---|---|
| ✗ | ✗ | 0.823 | 3.32 | 60.7 |
| ✓ | ✗ | 0.827 | 3.33 | 56.4 |
| ✗ | ✓ | 0.823 | 3.06 | 22.9 |
| ✓ | ✓ | 0.827 | 3.04 | 20.7 |
| Method / setting | Target | Bal. Acc. | Path | Leaves |
|---|---|---|---|---|
| SPLIT | – | 0.790 | 2.54 | 6.1 |
| CART | – | 0.822 | 3.42 | 10.3 |
| GradTree | – | 0.762 | 3.30 | 11.5 |
| MBNDT, tight budget | 4 | 0.827 | 2.67 | 13.4 |
| MBNDT, moderate budget | 16 | 0.835 | 2.94 | 18.3 |
| MBNDT, large budget | 64 | 0.836 | 3.01 | 20.3 |
| Dataset | OpenML ID | Instances | Attributes |
|---|---|---|---|
| hepatitis | 55 | 155 | 20 |
| colic | 27 | 368 | 23 |
| vote | 56 | 435 | 17 |
| breast-w | 15 | 699 | 10 |
| blood-transfusion-service-center | 1464 | 748 | 5 |
| diabetes | 37 | 768 | 9 |
| Hyperparameter | Values | Scale |
|---|---|---|
| cat. | ||
| cat. | ||
| cat. | ||
| cat. | ||
| log | ||
| for feature, threshold, leaf, mask | log |
| Hyperparameter | Values | Scale |
|---|---|---|
| cat. | ||
| log | ||
| log | ||
| log | ||
| Focal loss | cat. | |
| polyLoss | cat. |
| Hyperparameter | Values | Scale |
|---|---|---|
| Regularization | log | |
| full_depth_budget | fixed | |
| Lookahead, small folds | fixed | |
| Lookahead, other folds | fixed | |
| Binarizer estimators | fixed | |
| Binarizer depth | fixed |
| Hyperparameter | Values | Scale |
|---|---|---|
| criterion | cat. | |
| splitter | cat. | |
| max_depth | cat. | |
| min_samples_split | log | |
| min_samples_leaf | log | |
| max_features | cat. |
| Hyperparameter | Values | Scale |
|---|---|---|
| n_estimators | int | |
| max_depth | cat. | |
| learning_rate | log | |
| min_child_weight | log | |
| gamma | log | |
| reg_alpha | log |
Decision Tree Learning
Decision tree learning builds hierarchical, rule-based models for prediction, balancing accuracy, computational cost, and interpretability.
See all 69 papers