Learning a Decision Tree Algorithm with Transformers (2402.03774v2)

Abstract: Decision trees are renowned for their ability to achieve high predictive performance while remaining interpretable, especially on tabular data. Traditionally, they are constructed through recursive algorithms, where they partition the data at every node in a tree. However, identifying a good partition is challenging, as decision trees optimized for local segments may not yield global generalization. To address this, we introduce MetaTree, a transformer-based model trained via meta-learning to directly produce strong decision trees. Specifically, we fit both greedy decision trees and globally optimized decision trees on a large number of datasets, and train MetaTree to produce only the trees that achieve strong generalization performance. This training enables MetaTree to emulate these algorithms and intelligently adapt its strategy according to the context, thereby achieving superior generalization performance.

• The paper presents MetaTree, a transformer model that learns to generate decision trees by integrating greedy and global optimization strategies.
• The methodology leverages an attention mechanism that captures row and column dependencies with learnable positional biases for refined split decisions.
• Results demonstrate that MetaTree achieves lower variance and superior generalization over traditional decision tree algorithms on real-world datasets.

Introduction

Decision trees stand as one of the quintessential methods in the machine learning landscape, praised for their interpretability and efficacy, particularly in the domain of tabular data. Traditional construction of decision trees relies on recursive algorithms, with splits determined to optimize local performance, which may not always align with global generalization goals. The crux of enhancing decision trees has thus revolved around finding the most appropriate partitions, a nontrivial task given the NP-hard nature of the problem.

Methodology

In the pursuit of building decision trees with a more global perspective, a novel approach titled MetaTree has been proposed. MetaTree is essentially a transformer-based model trained on a diverse assembly of datasets, assimilating knowledge from both greedy decision trees and globally optimized decision trees, specifically learning to generate trees that show strong generalization.

The transformative aspect of MetaTree lies in its adaptable architecture, granting it the intelligence to not only imitate traditional algorithms but to also transition between greedy and global optimization strategies based on contextual dataset nuances. Key to MetaTree's strategy is a unique attention mechanism that accounts for both row and column dependencies within tabular data along with a thoughtful injection of learnable positional biases.

Results

MetaTree's performance is quantitatively assessed on numerous real-world datasets outside its training scope, demonstrating impressive generalization across two tree depths. Results indicate that MetaTree consistently surpasses traditional algorithms by intelligently switching between emulation and strategy adaption. Notably, MetaTree's performance flourishment is also observed within the field of LLM-generated data, speaking to the robustness and versatility of the method.

Analysis and Discussion

Analytical probing into MetaTree reveals a sophisticated internal decision-making process whereby the model dynamically refines its splits across layers. This behavior is reminiscent of the refinement seen in LLM token predictions, showcasing potential for early exit strategies which could significantly bolster efficiency.

A bias-variance decomposition paints MetaTree as a model with substantially lower variance compared to benchmarks, implying greater stability against training data fluctuations.

While MetaTree showcases unprecedented potential, constraints remain tied to the maximum sequence lengths manageable by Transformers. Nonetheless, the paper is a testament to the expanding capabilities of generative AI, with MetaTree providing a glimpse into a future where deep learning not only forecasts outcomes but also cultivates stronger, context-aware machine learning models.