PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Multi-Branching for Shallow Decision Tree Induction

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

arXiv:2608.29262v1cs.LGcs.AI

TL;DR

Strict depth limits can make binary decision trees under-expressive, motivating a way to increase local partitioning capacity without lengthening decision paths. MBNDT trains a single axis-aligned tree end-to-end with differentiable multi-way splits, adaptive branch masks, and post-hoc pruning, then evaluates it across 21 OpenML binary-classification benchmarks. It achieves the best mean balanced accuracy among depth-constrained single-tree learners, while realizing more leaves than competing single-tree baselines.

  • Problem

    Under strict maximum-depth budgets, conventional binary trees can be under-expressive because each internal node makes only one threshold decision.

  • Method

    MBNDT is a single axis-aligned tree trained end-to-end with ordered differentiable multi-way splits, learnable branch masks, and post-hoc train-path pruning.

  • Results

    Across 21 OpenML binary-classification benchmarks, MBNDT achieves the highest mean balanced accuracy among depth-constrained single-tree learners; ternary routing reaches 0.833 versus 0.820 for binary routing.

  • Takeaways & Limitations

    MBNDT is best suited to settings prioritizing accuracy under short, bounded decision paths over minimal global tree size.

  • Takeaways & Limitations

    MBNDT realizes more leaves than the other single-tree baselines, creating an explicit global-size cost.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Under strict shallow-depth budgets, binary trees can be under-expressive because each node makes only one threshold decision. MBNDT addresses this with jointly optimized multi-way splits, adaptive branch masks, and deterministic pruning while evaluating accuracy, path length, and realized leaves.

  • Motivation: Shallow-depth induction constrains maximum decision depth because predictions are often inspected or implemented as short, sequential decision paths.The paper reports leaf count alongside path length to expose the trade-off between local inspection effort and global tree size.
  • Motivation: Binary trees can be under-expressive at small depths because each internal node makes only one threshold decision, limiting partitioning capacity.The paper characterizes this as an expressivity bottleneck for heterogeneous tabular decision boundaries.
  • Method: MBNDT learns differentiable axis-aligned multi-way interval splits by jointly optimizing selected features, ordered thresholds, and leaf predictions.This replaces node-wise greedy split selection with end-to-end gradient-based optimization.
  • Evaluation: Across the evaluated benchmarks, MBNDT achieves the best average rank and highest mean balanced accuracy among depth-constrained single-tree learners.The evaluation compares greedy, solver-based, and gradient-based baselines under shared depth budgets.
  • Evaluation: MBNDT realizes more leaves than other single-tree baselines, trading minimal global tree size for accuracy under short decision paths.This is the paper’s explicit accuracy–depth–size trade-off.
  • Method: Learnable branch masks adapt each node’s effective arity, while a leaf-budget penalty and post-hoc train-path pruning convert the nominal tree into a compact deterministic inference tree.Each input ultimately follows one root-to-leaf path.

II. RELATED WORK

Prior work covers path complexity, greedy and globally optimized tree induction, multi-way splitting, and differentiable training. MBNDT unifies these ingredients into a single axis-aligned learner with adaptive effective branching.

  • Interpretability: Interpretability research distinguishes executed path complexity from minimal explanations, and this paper focuses on bounding decision steps rather than claiming subset-minimal explanations.Path length is treated as a proxy for inspection effort, while leaf count exposes global size.
  • Greedy induction: Classical learners such as ID3, C4.5, and CART recursively optimize local impurity reduction, which can produce globally suboptimal trees under shallow depth budgets.Their limitation is especially consequential when later corrections are restricted by the depth budget.
  • Optimal trees: Optimal-tree methods formulate routing, split selection, and leaf prediction as discrete global optimization under depth, leaf, or feature budgets.Mixed-integer and dynamic-programming approaches improve global search and scalability in different ways.
  • Multi-way trees: Multi-way numerical splits can improve compactness or expressivity, but higher arity may fragment data and increase overfitting, motivating pruning.Earlier node-wise methods generally do not jointly optimize structure across levels.
  • Differentiable trees: Differentiable tree methods replace discrete split selection with relaxations or surrogate gradients, enabling backpropagation-based optimization.Existing approaches include sparse feature learning, oblique trees, and soft-routing formulations.
  • Positioning: MBNDT’s contribution is to unify ordered multi-way numerical splits, continuous thresholds, adaptive arity, and gradient-based training within one axis-aligned tree.The individual ingredients have been studied separately, but their combination is the paper’s stated contribution.

III. MBNDT: ARCHITECTURE, TRAINING, AND INFERENCE

MBNDT jointly trains a predefined axis-aligned tree whose nodes select features, thresholds, branch masks, and leaf logits. Training uses differentiable feature selection and binary cross-entropy with leaf-budget regularization.

  • Architecture: MBNDT jointly optimizes a predefined tree with branching factor B and depth D, using node-wise feature parameters, threshold parameters, branch masks, and leaf logits.Each axis-aligned node routes according to the interval containing the selected feature value.
  • Training: The model outputs a scalar logit for binary classification, whose sigmoid is interpreted as the estimated positive-class probability.The training objective combines binary cross-entropy with a leaf-budget regularization term.
  • Feature selection: α-entmax converts node feature logits into a sparse differentiable feature-selection distribution that can assign exact zeros.This supports sparse candidate features while remaining differentiable.
  • Feature selection: A straight-through estimator uses a hard one-hot feature choice in the forward pass while propagating gradients through the differentiable selection during backpropagation.The construction preserves axis-aligned routing at inference-facing computation.

B. Multi-branch routing via ordered thresholds

Each MBNDT node creates ordered thresholds through positive learned gaps, then uses differentiable cumulative-logistic binning to assign soft probabilities across multiple intervals.

  • Ordered thresholds: Each internal node learns B−1 ordered thresholds, partitioning the selected feature into B ordered regions instead of the two regions of a binary node.The thresholds are parameterized as cumulative sums of softplus-transformed gap logits.
  • Ordered thresholds: Softplus-transformed gap logits create strictly positive gaps, ensuring thresholds remain strictly increasing without per-iteration sorting.The transformation remains smooth for gradient-based optimization.
  • Soft routing: The selected feature score is routed with a differentiable soft-binning scheme based on cumulative logistic probabilities.A temperature controls how sharply the approximation approaches interval membership.
  • Soft routing: Adjacent differences of cumulative terms produce the intermediate branch probabilities, with boundary branches defined by the first cumulative term and its complement.These probabilities form a valid distribution over the B branches.

C. Adaptive pruning via branch masks

Branch masks let each MBNDT node suppress or emphasize outgoing branches, adapting effective arity while preserving normalized routing and changing hard-routing selections.

  • Adaptive branch arity: MBNDT treats B as an upper bound on local branching and uses input-independent masks to suppress unnecessary branches during training.This avoids requiring every internal node to use all B outgoing branches.
  • Masked routing: Sigmoid-transformed masks multiply unmasked routing probabilities, after which the results are renormalized across branches.The normalized masked probabilities preserve the partition-of-unity constraint.
  • Masked routing: Hard routing selects the highest masked probability while a straight-through estimator passes gradients through the masked soft probabilities.The forward pass is discrete, but backpropagation remains differentiable.
  • Adaptive pruning: In the toy B=4 example, masks reweight interval probabilities and can change which branches are selected under hard routing without moving ordered thresholds.The figure compares raw probabilities with masked, renormalized probabilities.
  • Adaptive pruning: If a branch mask approaches zero, that branch is effectively never selected and its subtree becomes unreachable at inference.Masks therefore induce adaptive pruning and can suppress weak branches across the input space.

D. Inference

Inference propagates leaf logits upward through the selected child at each node, producing one deterministic root-to-leaf prediction path.

  • Leaf-to-root reduction: Each leaf stores a scalar logit, while internal nodes receive recursively computed subtree values rather than independent output logits.Subtree values are initialized at leaves and reduced upward.
  • Leaf-to-root reduction: Because hard routing is one-hot at inference, exactly one child contributes at each internal node and the selected leaf logit reaches the root.The recursive computation therefore follows a single active path.
  • Prediction: For binary classification, the root output is converted to a predicted probability with σ(fθ(x)).

IV. STRUCTURE SPARSIFICATION IN MBNDT

MBNDT sparsifies structure through branch masks, differentiable leaf-budget control, and post-hoc pruning, converting a nominal tree into a smaller effective structure.

  • Sparsification mechanisms: MBNDT supports structure sparsification through learned branch masks, leaf-budget regularization, and post-hoc train-path pruning.These mechanisms operate during training and after training.
  • Effective sparsification: Branch masks define the effective tree by suppressing weak branches and retaining only hard-routable leaves and their reachable subtrees.The effective tree is the subtree induced by all routable leaves.
  • Leaf-budget regularization: The leaf-budget penalty uses a differentiable surrogate because the exact number of hard-routable leaves is discrete and non-differentiable.The surrogate is constructed from masked soft routing probabilities.
  • Leaf-budget regularization: Effective arity is estimated from batch-averaged masked routing probabilities, increasing as branch usage becomes more distributed across up to B branches.It is close to 1 when routing concentrates on one branch and approaches B when usage is distributed.
  • Leaf-budget regularization: The soft effective leaf count is used to penalize violations of a target leaf budget while avoiding penalties for trees already below budget.The positive-part operation ensures sparse trees are not penalized.
  • Training and construction: Training initializes a full B-ary tree, optimizes masked routing with a budget loss, then identifies used branches and prunes unused ones.The algorithm redirects removed branches to the nearest surviving branch.

C. Post-hoc train-path pruning

Post-hoc train-path pruning produces a train-supported subtree, while evaluation compares MBNDT with single-tree baselines under a shared shallow-depth protocol.

  • Post-hoc train-path pruning: Post-hoc pruning retains branches observed at least once during hard routing on the training set.This yields a train-path-supported subtree.
  • Post-hoc train-path pruning: Inputs selecting removed branches are redirected to the nearest surviving branch in the original ordering, while original leaf logits remain unchanged.
  • Baselines: The comparable baselines are CART, SPLIT, and GradTree, each producing one axis-aligned tree with one root-to-leaf path per input.XGBoost is reported only as a black-box reference rather than a comparable single-tree model.
  • Experimental setting: Table I reports mean ± standard deviation test balanced accuracy over five splits for 21 binary-classification datasets under the shared depth budget.Ties receive the same competition rank and count as wins when tied for first.
  • Experimental setting: All single-tree methods are evaluated with maximum depth D ≤4, while MBNDT tunes maximum branching factor B ∈{3, 4}.The protocol uses nested stratified evaluation with shared splits and budgets.
  • Experimental setting: The evaluation uses five 80/20 outer train-test splits, with each outer test set held out until final evaluation.

B. Depth-constrained predictive performance

Under a shared shallow-depth evaluation, MBNDT delivers the strongest predictive performance among single-tree learners while making the accuracy–complexity trade-off explicit. Ablations attribute gains to multi-way branching and show that pruning and masks improve size or accuracy without erasing performance.

  • Predictive performance: MBNDT achieves the best average rank (1.48), 15 of 21 dataset wins, and highest mean balanced accuracy (0.832) among depth-constrained single-tree learners.Mean balanced accuracy is 0.822 for CART, 0.790 for SPLIT, and 0.762 for GradTree.
  • Structural trade-offs: Mean decision-path length is 2.95 for MBNDT, while realized leaves reach 20.6, making its accuracy advantage a larger-tree trade-off.SPLIT has the smallest trees but the lowest accuracy; MBNDT’s refit time is higher than CART and GradTree but lower than SPLIT on average.
  • Branching ablation: Increasing branching from B=2 to B=3 raises balanced accuracy from 0.820 to 0.833 and shortens mean paths from 3.22 to 3.12, while increasing leaves from 9.5 to 16.0.B3 outperforms B2 on 16 of 21 datasets, supporting greater local partitioning capacity rather than longer paths as the source of improvement.
  • Sparsification ablation: Post-hoc pruning reduces unmasked realized leaves from 60.7 to 22.9 and mean paths from 3.32 to 3.06 without statistically changing balanced accuracy (0.823).Masks provide a separate small accuracy benefit, and the masked pruned model reaches 0.827 balanced accuracy with 20.7 leaves versus 0.823 and 22.9 leaves without masks.
  • Leaf-budget ablation: Increasing the leaf budget from K=4 to K=64 raises mean balanced accuracy from 0.827 to 0.836 while lengthening paths from 2.67 to 3.01 and increasing leaves from 13.4 to 20.3.K=16 is close to K=64, with 0.835 accuracy, path length 2.94, and 18.3 leaves; the budget does not make MBNDT smaller than reference trees.

VI. CONCLUSION

The paper concludes that MBNDT improves accuracy under bounded decision depth by increasing local partitioning capacity through differentiable multi-way splits. This advantage has a global-size cost and remains evaluated only for binary classification.

  • Across 21 OpenML binary-classification benchmarks, MBNDT achieves the highest mean balanced accuracy among depth-constrained single-tree learners.
  • A branching-factor ablation shows that multi-way splits outperform a binary variant under the same depth and budget setting.
  • MBNDT realizes more leaves than other single-tree baselines, trading minimal global tree size for accuracy under bounded per-instance decision depth.
  • The evaluation is limited to binary classification, with multiclass classification and regression identified as future extensions.

APPENDIX

The appendix documents refitting and leaf-budget procedures and provides dataset and hyperparameter search-space tables for the evaluated methods.

  • Final refits use five random restarts, selecting checkpoints by validation loss before early stopping on validation balanced accuracy.The reported final-refit limits are 40/8 and 500/25 epochs/patience for the two stages.
  • The selected leaf budget K is handled with a log-scale violation based on soft leaf count Lsoft.The minibatch objective adds a penalty involving the violation term.
  • The appendix includes dataset characteristics for 21 OpenML binary-classification benchmarks.
  • Separate tables specify hyperparameter search spaces for MBNDT, GradTree, SLP, CART, and XGBoost.
Loading 2608.29262v1…