Decision trees that learn complex shapes more accurately and optimally
Literati: Towards Anytime Optimal Shape Generalized Trees via AO*
Machine Learning
Summary
Decision trees help computers make decisions by splitting data along certain rules, but the usual methods can create overly complicated trees that aren’t the best possible. The authors created a new method called Literati that finds the best possible decision tree using more flexible rules for splitting data. This approach uses a special search technique to explore all options efficiently and guarantees an optimal tree structure. They tested Literati on many real datasets and found it makes better predictions than current top methods.
What this means in practice
- •For data scientists: Create more accurate and simpler decision trees by optimally selecting complex splitting rules improving model performance on tabular datasets.
- •For machine learning engineers: Implement an anytime algorithm that guarantees optimal decision trees with enhanced splits for interpretable model deployment in production.
Authors
Nakul Upadhya, Eldan Cohen
Abstract
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.