Empowering Decision Trees via Shape Function Branching

AI summary

Overview

• Research area: Interpretable machine learning, specifically decision tree induction for tabular data. The paper sits at the intersection of greedy tree induction (CART-style), generalized additive models (GAMs) with shape functions, and interpretability-oriented model design.
• Technical level: Advanced. The paper contains formal definitions, two expressiveness theorems, a bi-level optimization formulation, complexity analysis, and a large experimental study. Readers need some familiarity with decision trees, impurity measures, and oblique splits.
• Scope in one sentence: The paper introduces a new family of interpretable decision trees whose internal nodes split on learnable, visualizable shape functions rather than simple thresholds, plus an induction algorithm for learning them and empirical comparisons to axis-aligned and bivariate tree baselines.

What This Paper Is About

Standard decision trees split on a single feature with a simple threshold (for example, "send left if feature d is less than or equal to theta"). Because real feature-target relationships are often non-linear, trees must reuse the same feature at many levels, producing deep, node-heavy trees that are harder for people to understand and visualize. This paper replaces the threshold at each internal node with a learnable shape function over one feature, so a single split can carve out complex non-linear regions, and it provides an algorithm (ShapeCART) for learning such trees from data, including bivariate and multi-way variants.

Key Contributions

1. Shape Generalized Trees (SGTs): A new class of decision trees where each internal node routes samples using an axis-aligned shape function applied to a single selected feature, yielding rich non-linear partitioning within one split while remaining visualizable node by node.
2. ShapeCART: An efficient top-down induction algorithm for building SGTs, structured after the CART framework, that learns each node's shape function through an internal binning tree plus a bin-to-branch assignment solved by coordinate descent.
3. Two framework extensions: S²GT, which allows bivariate shape functions (at most two features per node), and SGT_K, which allows K-way branching; the algorithm is correspondingly extended into Shape²CART and ShapeCART_K, with a bivariate candidate-filtering heuristic to avoid constructing all O(D²) pairwise shape functions.
4. Expressiveness guarantees and empirical evaluation: Two theorems showing SGTs are at least as expressive as binary axis-aligned linear trees with the same number of decision nodes and can be strictly more expressive, plus experiments on 26 real-world classification datasets showing higher accuracy with more compact trees.

Main Findings

• Compactness on a toy problem: On a two-dimensional synthetic "Plus Sign" dataset, a traditional axis-aligned linear tree requires six splits with a maximum depth of four, while the SGT represents the same boundaries with only two nodes and a maximum depth of two.
• Expressiveness (Theorem 1 and Theorem 2): Every function representable by a binary axis-aligned linear tree can be represented by an SGT with the same number of decision nodes. Conversely, for every B in the natural numbers, there exists a function for which a binary axis-aligned tree needs at least B additional decision nodes compared to an SGT. The second proof constructs a family of one-dimensional labeling functions with periodic boundaries.
• Information gain bound (Lemma 1): The information gain from a shape function split at a node is greater than or equal to the information gain from an axis-aligned threshold split at that node in CART. Because of this, ShapeCART inherits the terminal-node information gain bound established for CART by Klusowski and Tian.
• Axis-aligned accuracy gains: Averaged over the datasets tested, SGT-C (best) reaches 86.3 percent test accuracy versus 85.2 for CART, 85.1 for SERDT, 85.1 for HSTree, 85.4 for AxTAO, 86.2 for DPDT, and 80.3 for SPLIT. The paper reports that SGT-C consistently outperforms all baseline TDIDT methods at every evaluated depth.
• Multi-way branching helps: SGT_3-C (best) reaches 87.8 percent, above all the axis-aligned baselines and above SGT-C, which the authors attribute to the benefit of higher branching factors.
• TAO refinement helps further: The TAO-refined axis-aligned models score highest among axis-aligned approaches: SGT-T at 86.8 and SGT_3-T at 88.0.
• Bivariate approaches are stronger still: S²GT-C (best) reaches 91.4 percent versus 89.6 for BiCART; S²GT_3-C reaches 91.7; with TAO refinement, S²GT-T reaches 91.6, S²GT_3-T reaches 91.9, and BiTAO reaches 90.0.
• Ternary trees do not cost more explanation size: The authors note that despite having more nodes, ternary trees have the same size of local explanation as binary trees of the same depth.
• Runtime heuristic is consequential: A naive bivariate approach requires constructing O(D²) shape functions. The proposed delta heuristic retains only the P most promising pairs, needing O((N + CK²)D²) operations for set intersections plus construction of P shape functions, rather than O(D²(NC log N + K²LC)).
• Complexity of a single shape function: Constructing a shape function takes O(NC log N + K²LC) time, where C is the number of classes, per the paper's analysis.

Methodology in Plain English

The authors keep the familiar top-down, recursive way of growing a tree (the TDIDT strategy used by CART) but change what happens inside each node.

At a node, the algorithm treats the choice of split as a two-level problem: for each candidate feature (or pair of features), find the best shape function, then pick whichever candidate produces the lowest weighted impurity across the resulting child branches. Impurity is measured with standard superadditive measures such as Gini or entropy.

Learning a shape function happens in two stages. First, the algorithm builds an internal decision tree on the relevant feature(s) and the target to sort samples into a bounded number of bins — CART for a single feature, BiCART for a pair of features. Each bin stores its empirical class distribution and its sample count. Second, the algorithm assigns each bin to an output branch. This assignment is a discrete optimization solved by coordinate descent: repeatedly revisit one bin, try each possible branch assignment, keep the one that lowers the impurity, shuffle the update order each pass, and repeat for R iterations. Because coordinate descent is sensitive to initialization, the authors try two starting points and keep the better one — Weighted K-Means clustering of the bins' class distributions, or the left/right position of each bin relative to the root of the internal tree.

Multi-way branching is handled by running the same machinery for each branching factor k from 2 to K, adding a penalty of lambda times (k minus 2) to the weighted impurity, and selecting the k with the lowest penalized impurity. To keep bivariate trees tractable, the authors first build all univariate shape functions, then score each pair of features by how much the Cartesian product of their branch sets reduces impurity relative to the best univariate option, and only fit bivariate shape functions for the top P pairs. Bivariate splits also carry a tunable penalty gamma to discourage them unless they help substantially.

Finally, because greedy top-down growth can be globally suboptimal, the trees are refined with Tree Alternating Optimization (TAO), which the authors extend from binary to K-way branching. Refitting re-learns each internal node's shape function using the predictions of the subtrees below it and prunes nodes that do not meaningfully reduce error.

Why This Matters

Impact on research. The paper opens a bridge between two previously separate lines of interpretable modeling: decision trees and GAM-style shape functions. It shows formally that swapping thresholds for shape functions strictly increases what a fixed number of decision nodes can represent, and it contributes an induction algorithm that fits inside the well-understood CART/TDIDT pipeline, including the theoretical information gain bound. It also extends TAO, a post-processing method originally limited to binary trees, to multi-way branching.

Real-world applications. The paper motivates interpretability through high-stakes domains where decision trees are already favored:

• Healthcare, where the paper cites tree use and where the evaluated datasets include eye-movements and eye-state.
• Finance, where model decisions often need to be explained to regulators, auditors, or customers.
• Manufacturing, where process decisions on tabular sensor data need to be traceable.
• Any tabular workflow where a practitioner must visualize the learned feature-target relationship — the paper's shape functions can be plotted directly per node, and bivariate shape functions can be visualized as heatmaps.

Industry relevance. Interpretable tabular models are widely deployed, and tree size directly affects how much explanation a stakeholder must read: the paper notes that tree comprehensibility is highly sensitive to decision path depth and leaf count, that deeper trees generate larger and less interpretable local explanations, and that large trees are hard to visualize without complex interactive tooling. Producing comparable accuracy with smaller trees, and with a per-node visual summary, is directly useful for model review, documentation, and debugging. The bivariate candidate-filtering heuristic also matters practically, since it targets the main computational bottleneck of pairwise methods.

Future Directions

• Scaling and cost of bivariate splits: The authors leave the choice of P to the user's computational preferences, trading runtime for performance; how to set it automatically, and whether better pair-selection heuristics exist, remains open.
• Beyond K = 3: The authors state the model naturally extends to any K but deliberately limit exploration of SGT_K to K = 3 to preserve sparsity and interpretability. Where the accuracy-interpretability trade-off actually turns is unresolved.
• Shape function learners beyond CART and BiCART: The paper notes the binning step can use a variety of approaches and reports results with an alternative tree induction algorithm (DPDT) in an appendix, suggesting further alternatives are worth exploring.
• Regression and other settings: The paper presents the methods in the classification setting where labels take values in a finite class set, with a regression variant described in an appendix; the truncated content does not report those results, so their behavior is not established here.

Target Audience

Researchers and practitioners in interpretable machine learning, especially those working on decision tree induction, tabular data modeling, and GAM-style additive models. It is also relevant to engineers and analysts who need compact, visualizable models in regulated or high-stakes settings, and to readers interested in the theory of tree expressiveness. Some parts — the bi-level optimization, the theorems, and the complexity analysis — assume a graduate-level background in machine learning.

Authors’ abstract

Decision trees are prized for their interpretability and strong performance on tabular data. Yet, their reliance on simple axis-aligned linear splits often forces deep, complex structures to capture non-linear feature effects, undermining human comprehension of the constructed tree. To address this limitation, we propose a novel generalization of a decision tree, the Shape Generalized Tree (SGT), in which each internal node applies a learnable axis-aligned shape function to a single feature, enabling rich, non-linear partitioning in one split. As users can easily visualize each node's shape function, SGTs are inherently interpretable and provide intuitive, visual explanations of the model's decision mechanisms. To learn SGTs from data, we propose ShapeCART, an efficient induction algorithm for SGTs. We further extend the SGT framework to bivariate shape functions (S$^2$GT) and multi-way trees (SGT$_K$), and present Shape$^2$CART and ShapeCART$_K$, extensions to ShapeCART for learning S$^2$GTs and SGT$_K$s, respectively. Experiments on various datasets show that SGTs achieve superior performance with reduced model size compared to traditional axis-aligned linear trees.