Adaptive Multi-Branching for Shallow Decision Tree Induction

Hanul Park; Jeonghoon Choi; Juseong Kim; Sanghun Sel; Giltae Song

cs.LGAug 29, 2026

Adaptive Multi-Branching for Shallow Decision Tree Induction

Authors: Hanul Park, Jeonghoon Choi, Juseong Kim, Sanghun Sel, Giltae Song

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

Fig. 1: Overview of MBNDT with prespecified maximum branching factor $B=3$ . Left : a nominal depth-2 ternary tree with mask-suppressed branches. Right : node-level routing, where a selected feature is partitioned by ordered thresholds into soft interval probabilities, reweighted by branch masks, and hard-routed to a single child.

Fig. 2: Illustration of the branch-mask effect in a toy $B{=}4$ routing node. Top : raw interval probabilities; Bottom : probabilities after applying branch masks and renormalizing. Dashed lines show the original ordered-threshold locations; masks reweight branch probabilities and can change which branches are selected under hard routing.

Small datasets ( $n<10^{3}$ )
Dataset	Ensemble (ref.)	Single tree
	XGBoost	Gradient-based		Solver-based	Conventional
		MBNDT (ours)	GradTree	SPLIT	CART
hepatitis	0.745 $\pm$ 0.113	0.675 $\pm$ 0.123 (1)	0.555 $\pm$ 0.098 (4)	0.673 $\pm$ 0.081 (2)	0.665 $\pm$ 0.147 (3)
colic	0.768 $\pm$ 0.049	0.787 $\pm$ 0.055 (1)	0.772 $\pm$ 0.052 (3)	0.766 $\pm$ 0.049 (4)	0.775 $\pm$ 0.054 (2)
vote	0.954 $\pm$ 0.019	0.955 $\pm$ 0.013 (1)	0.914 $\pm$ 0.038 (4)	0.952 $\pm$ 0.018 (2)	0.939 $\pm$ 0.025 (3)

TABLE I: Per-dataset test balanced accuracy on 21 binary-classification datasets. Values are reported as mean

\pm

standard deviation over five dataset splits. All single-tree methods use a maximum-depth budget of

D\leq 4

; XGBoost is included as a depth-limited ensemble reference. Ties are assigned the same competition rank and counted as wins when tied for first.

Dataset	MBNDT (ours)			GradTree			SPLIT			CART
	Path	Leaves	Time (s)	Path	Leaves	Time (s)	Path	Leaves	Time (s)	Path	Leaves	Time (s)
Small datasets ( $n<10^{3}$ )
hepatitis	2.40	6.8	10.73	2.80	8.0	7.55	2.94	7.0	1.07	3.55	9.4	$<0.01$
colic	1.66	4.0	19.81	3.00	8.8	12.75	1.92	4.6	0.70	3.36	9.2	$<0.01$
vote	2.29	5.6	23.27	2.40	6.8	7.17	1.10	2.4	0.74	2.85	9.0	$<0.01$
breast-w	2.62	18.2	20.33	3.40	11.2	7.67	2.07	5.2	0.94	3.33	10.6	$<0.01$

TABLE II: Structural complexity and training time for depth-constrained single-tree methods. Path length, realized leaves, and time are averaged over five outer splits, with size-stratified means.

Comparison	$N$	W/T/L	Wilcoxon $p$	Holm $p$
MBNDT vs. GradTree	21	20/0/1	$2.86\times 10^{-6}$	$8.58\times 10^{-6}$
MBNDT vs. SPLIT	21	15/2/4	0.00628	0.01256
MBNDT vs. CART	21	14/2/5	0.02633	0.02633

TABLE III: Paired Wilcoxon tests for test balanced accuracy between MBNDT and single-tree baselines with maximum depth

D\leq 4

Variant	$B$	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

TABLE IV: Branching-factor ablation for MBNDT with fixed depth

D=4

and leaf budget

K=16

. Best counts the number of datasets on which each variant achieves the highest balanced accuracy.

Masks	PP	Bal. Acc. $\uparrow$	Path $\downarrow$	Leaves $\downarrow$
✗	✗	0.823	3.32	60.7
✓	✗	0.827	3.33	56.4
✗	✓	0.823	3.06	22.9
✓	✓	0.827	3.04	20.7

TABLE V: Effect of branch masks and post-hoc train-path pruning (PP) in a controlled setting without a leaf-budget penalty.

Method / setting	Target $K$	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

TABLE VI: Leaf-budget ablation for MBNDT with single-tree reference points. MBNDT rows vary the target leaf budget

K

; the HPO-selected row reports the mean selected

K

across dataset splits. Reference methods are shown under the same depth constraint.

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

TABLE VII: Dataset characteristics for the 21 OpenML binary-classification benchmarks.

Hyperparameter	Values	Scale
$B$	$\{3,4\}$	cat.
$D$	$\{1,2,3,4\}$	cat.
$K\mid B=3$	$D_{1}:\{2,3\}$ $D_{2}:\{3,6,9\}$ $D_{3}:\{6,12,18,24\}$ $D_{4}:\{12,24,36,48\}$	cat.
$K\mid B=4$	$D_{1}:\{2,3,4\}$ $D_{2}:\{4,8,12,16\}$ $D_{3}:\{8,12,16,24,32\}$ $D_{4}:\{16,24,32,48,64\}$	cat.
$\tau_{\mathrm{cdf}}$	$[0.08,1.2]$	log
$\eta_{*}$	$[10^{-3},10^{-1}]$ for feature, threshold, leaf, mask	log

TABLE VIII: MBNDT hyperparameter search space.

Hyperparameter	Values	Scale
$D$	$\{1,2,3,4\}$	cat.
$\eta_{\mathrm{index}}$	$[10^{-3},10^{-1}]$	log
$\eta_{\mathrm{value}}$	$[10^{-3},10^{-1}]$	log
$\eta_{\mathrm{leaf}}$	$[10^{-3},10^{-1}]$	log
Focal loss	$\{0,1\}$	cat.
polyLoss	$\{0,1\}$	cat.

TABLE IX: GradTree hyperparameter search space.

Hyperparameter	Values	Scale
Regularization	$[10^{-6},10^{-1}]$	log
full_depth_budget	$4$	fixed
Lookahead, small folds	$3$	fixed
Lookahead, other folds	$2$	fixed
Binarizer estimators	$50$	fixed
Binarizer depth	$1$	fixed

TABLE X: SPLIT hyperparameter search space.

Hyperparameter	Values	Scale
criterion	$\{\mathrm{gini},\mathrm{entropy},\mathrm{log\_loss}\}$	cat.
splitter	$\{\mathrm{best},\mathrm{random}\}$	cat.
max_depth	$\{1,2,3,4\}$	cat.
min_samples_split	$[10^{-3},0.20]$	log
min_samples_leaf	$[10^{-3},0.20]$	log
max_features	$\{\mathrm{None},\mathrm{sqrt},\mathrm{log2}\}$	cat.

TABLE XI: CART hyperparameter search space.

Hyperparameter	Values	Scale
n_estimators	$[50,500]$	int
max_depth	$\{1,2,3,4\}$	cat.
learning_rate	$[10^{-2},3\times 10^{-1}]$	log
min_child_weight	$[10^{-2},20]$	log
gamma	$[10^{-8},5]$	log
reg_alpha	$[10^{-8},10]$	log

TABLE XII: XGBoost hyperparameter search space.

Decision Tree Learning

Decision tree learning builds hierarchical, rule-based models for prediction, balancing accuracy, computational cost, and interpretability.

See all 69 papers

Literati: Towards Anytime Optimal Shape Generalized Trees via AO*
Sep 8, 2026Nakul Upadhya, Eldan CohenTree Search Decision Tree Learning
Department of Mechanical and Industrial Engineering University of Toronto, Toronto, Canada
MotherTree: Meta-learning on synthetic data improves decision tree training
Oct 7, 2026Ziyuan Wang, Fredrik D. JohanssonDecision Tree Learning Meta-Learning
Chalmers University of Technology University of Gothenburg
A Moving-Horizon Approximate Branch-and-Reduce Method for Deep Classification Trees
Sep 18, 2026Chenxuanyin Zou, Jiayang Ren, Qiangqiang Mao +3Decision Tree Learning Branch and Bound
Department of Chemical and Biological Engineering University of British Columbia · Department of Electrical and Computer Engineering University of British Columbia · Department of Computer Science University of British Columbia

Sep 8, 2026J/K move · Enter open · S save

Literati: Towards Anytime Optimal Shape Generalized Trees via AO*

Nakul Upadhya, Eldan Cohen

Department of Mechanical and Industrial Engineering University of Toronto, Toronto, Canada

Tree Search Decision Tree Learning

Open paper AskPDF arXiv

Decision trees are prized for their interpretability and strong performance on tabular data, but popular greedy top-down induction algorithms can yield suboptimal and unnecessarily complex structures. Optimal decision tree methods address this through global optimization, yet remain restricted to axis-aligned threshold splits, which limit the expressivity of each node and often force deep, complex trees to capture non-linear feature effects. Shape Generalized Trees (SGTs) generalize threshold splits to learnable univariate shape functions, improving expressivity and enabling more compact trees. However, existing SGT induction algorithms are greedy and offer no optimality guarantees. In this work, we introduce Literati, the first algorithm for optimal SGT induction. We propose a novel AND/OR graph formulation of the problem that jointly optimizes tree structure and shape function complexity. To solve this AND/OR graph, we develop an AO*-based algorithm with two enhancements that improve anytime performance while preserving optimality: a secondary heuristic for OR-node selection and a round-robin policy for AND-node exploration. Across 24 real-world datasets, Literati achieves higher training and test accuracy than state-of-the-art tree approaches.

Abstract

Decision Tree Learning

Explore similar work

Literati: Towards Anytime Optimal Shape Generalized Trees via AO*

MotherTree: Meta-learning on synthetic data improves decision tree training

A Moving-Horizon Approximate Branch-and-Reduce Method for Deep Classification Trees