Decision Tree Embedding by Leaf-Means
Cencheng Shen, Yuexiao Dong, Carey E. Priebe
TL;DR
Decision Tree Embedding (DTE) addresses the variance and interpretability trade-offs of standard trees and forests by turning leaf-region means into a data-driven anchor-based embedding. A single tree (or a small ensemble) yields an embedding Z that, when combined with a linear classifier like LDA, achieves competitive accuracy with substantially reduced training time compared to random forests and shallow neural networks. The authors prove population-level guarantees: (i) conditional label distributions are preserved up to an epsilon bound under Bayes-homogeneous partitions, and (ii) a population-classification error bound for the induced embedding. Empirically, DTE demonstrates strong performance on synthetic and real datasets, offering a scalable, interpretable alternative to traditional tree ensembles with favorable accuracy-time trade-offs.
Abstract
Decision trees and random forest remain highly competitive for classification on medium-sized, standard datasets due to their robustness, minimal preprocessing requirements, and interpretability. However, a single tree suffers from high estimation variance, while large ensembles reduce this variance at the cost of substantial computational overhead and diminished interpretability. In this paper, we propose Decision Tree Embedding (DTE), a fast and effective method that leverages the leaf partitions of a trained classification tree to construct an interpretable feature representation. By using the sample means within each leaf region as anchor points, DTE maps inputs into an embedding space defined by the tree's partition structure, effectively circumventing the high variance inherent in decision-tree splitting rules. We further introduce an ensemble extension based on additional bootstrap trees, and pair the resulting embedding with linear discriminant analysis for classification. We establish several population-level theoretical properties of DTE, including its preservation of conditional density under mild conditions and a characterization of the resulting classification error. Empirical studies on synthetic and real datasets demonstrate that DTE strikes a strong balance between accuracy and computational efficiency, outperforming or matching random forest and shallow neural networks while requiring only a fraction of their training time in most cases. Overall, the proposed DTE method can be viewed either as a scalable decision tree classifier that improves upon standard split rules, or as a neural network model whose weights are learned from tree-derived anchor points, achieving an intriguing integration of both paradigms.