Source-linked AI summary
Learning Sparse Decision Trees via Transformer Variational Auto-Encoders
Giacomo Fidone, Alessio Cascione, Riccardo Guidotti
TL;DR
Existing decision-tree learners largely optimize prediction while overlooking structural sparsity, and optimal discrete-tree search is difficult. TREVIS maps trees into a TTVAE latent space and uses a differentiable surrogate for gradient-based optimization. It matches near-optimal predictive performance while producing structurally sparser trees, though its multi-stage optimization is less efficient than greedy methods.
Problem
Existing decision-tree algorithms give limited attention to structural sparsity, while searching the discrete tree space for optimal solutions is computationally difficult.
Method
TREVIS maps decision trees into a TTVAE latent space and uses a differentiable surrogate model for gradient-based optimization of multiple tree properties.
Results
TREVIS achieves predictive performance comparable to near-optimal learning algorithms while improving structural sparsity.
Takeaways & Limitations
TREVIS offers a favorable performance-sparsity trade-off for decision-tree learning by navigating a locally smooth latent space with directions capturing tree properties.
Takeaways & Limitations
TREVIS requires multiple training and optimization stages, making it less efficient than greedy methods.
Abstract
from arXiv · showhide
Decision trees are among the most widely used models in machine learning, largely due to their transparent decision logic, making them well-suited for high-stakes decision-making contexts. However, most existing learning algorithms focus on predictive performance, overlooking the joint optimization of other desirable properties, such as structural sparsity. In this work we propose TREVIS, an approach for learning decision trees with respect to complex objectives, based on the exploration of the latent space of a Tree Transformer Variational Auto-Encoder (TTVAE). By mapping decision trees onto latent representations, TREVIS replaces the discrete search space with a continuous one, enabling gradient-based optimization via a differentiable surrogate model. We experiment with TREVIS for learning decision trees that jointly optimize predictive performance and sparsity. Results show that TREVIS discovers decision trees matching the predictive performance of existing near-optimal algorithms while improving their structural sparsity.
I. INTRODUCTION
Decision-tree learning must balance predictive performance with interpretability and other properties, but discrete search is either greedy and suboptimal or globally optimal yet prohibitively expensive. TREVIS addresses this by exploring a continuous latent space, targeting performance and structural sparsity together.
- Decision trees support transparent, rule-based decisions in high-stakes settings, but many machine-learning models remain difficult to interpret.
- The discrete space of optimal decision trees grows exponentially, making greedy methods typically suboptimal and globally optimal methods computationally prohibitive.
- Structural sparsity, measured through tree depth and numbers of nodes and leaves, is central to decision-tree interpretability alongside predictive performance.
- TREVIS embeds decision trees in a continuous latent space and uses surrogate-model gradients instead of expensive black-box optimization.
- TREVIS targets predictive performance and structural sparsity, while fairness, privacy, and robustness remain future extensions.
- TREVIS learns locally smooth latent directions for tree properties, achieving near-optimal predictive performance while improving structural sparsity.
III. METHODOLOGY
TREVIS learns decision-tree representations with a Tree Transformer Variational Auto-Encoder, then navigates the resulting latent space using gradient ascent on a differentiable surrogate objective.
- TREVIS linearizes decision trees, learns their latent space with a Tree Transformer Variational Auto-Encoder, and optimizes that space with gradient ascent.
A. Tree Linearization
TREVIS converts binary decision trees into invertible depth-first token sequences that preserve split and terminal-node information while adding tree-aware positional structure.
- Tree Linearization: Each internal node is encoded by consecutive feature and threshold tokens, while each terminal node receives the special token ⟨L⟩.
- Tree Linearization: The encoding function π maps trees to token sequences using depth-first pre-order traversal, enabling inverse reconstruction.
- Tree Linearization: Threshold tokens use canonical feature-value intervals rather than unrestricted real-valued thresholds, reducing redundant search representations.
- Tree Linearization: The token vocabulary contains feature names, distinct feature values below each maximum, and transformer control tokens such as ⟨BOS⟩ and ⟨EOS⟩.
- Tree Linearization: Tree positional embeddings encode each token’s root-to-node path, compensating for self-attention’s lack of inherent hierarchical structure.
B. Tree Transformer Variational Auto-Encoder
TTVAE learns a structured latent representation of linearized decision trees with transformer encoder-decoder components. It samples latent vectors, conditions autoregressive decoding, and reconstructs trees from generated token sequences.
- Tree Transformer Variational Auto-Encoder: The encoder and decoder use stacked transformer blocks, with separate source and shifted target sequences for representation learning and autoregressive prediction.
- Tree Transformer Variational Auto-Encoder: The encoder maps the ⟨CLS⟩ representation to latent mean µ and variance σ, then samples z through the reparameterization trick.
- Tree Transformer Variational Auto-Encoder: The latent representation is projected into decoder blocks through layer-specific Multi-Head Cross-Attention projections.
- Tree Transformer Variational Auto-Encoder: TTVAE trains on linearized trees with a weighted β-ELBO that balances reconstruction accuracy against latent-space regularization.
- Tree Transformer Variational Auto-Encoder: At inference, the decoder starts from z and ⟨BOS⟩, generates tokens until ⟨EOS⟩, and maps the resulting sequence back to a decision tree.
C. Latent Space Optimization via Surrogate Model
TREVIS searches for decision trees in a learned latent space using a surrogate model to optimize predictive performance and structural sparsity.
- TREVIS searches its TTVAE latent space for representations of decision trees that optimize a specified objective.
- The objective combines weighted F1-score on training data with a penalty proportional to the tree’s number of leaves.The sparsity hyperparameter λ controls the penalty strength.
- Because the tree objective is nondifferentiable in latent space, TREVIS trains a differentiable surrogate model to approximate it.
- Gradient-based optimization of the surrogate identifies latent regions expected to maximize the objective from randomly initialized candidate solutions.
IV. EXPERIMENTS
The experiments evaluate whether TREVIS can optimize the predictive-performance and sparsity trade-off and compare it with existing decision-tree learning algorithms.
- The study evaluates TREVIS on the predictive-performance and structural-sparsity trade-off of decision trees.
- Experiments compare TREVIS with existing decision-tree learning algorithms.
- The evaluation also analyzes latent-space quality and TTVAE generation through a sensitivity analysis.
A. Experimental Setting
The experimental setting uses diverse benchmark datasets, generated tree collections, configured TTVAE and surrogate models, and tuned candidate-selection procedures.
- Datasets: 18 benchmark datasets cover diverse sample sizes, feature types, and numbers of classes, using stratified 80/20% train-test splits.A stratified 10% of the training data is reserved for validation.
- Datasets: The datasets include original and discretized representations, with categorical features one-hot encoded and all features scaled to [0, 1].
- Tree Datasets: For each representation, four disjoint tree collections support TTVAE training, validation, testing, and early stopping.
- Tree Datasets: Each tree collection contains 20,000 randomly generated trees with depths γ ∈[1, 5], restricting complexity to an interpretable range.
- Model Configurations: The TTVAE uses two encoder and decoder blocks, embedding size k = 120, two attention heads, and a 60-dimensional latent space.
- Optimization: The surrogate is a 128-unit MLP, while TREVIS samples n′ = 50,000 latent candidates and applies 10 gradient-ascent steps to each.
- Optimization: Separate surrogate runs use λ ∈ {0.0, 0.0001, 0.0005, 0.001, 0.005, 0.01}, and validation weighted F1-score selects the final tree.
B. Performance-Sparsity Trade-off
TREVIS is evaluated against greedy and near-optimal learners on predictive performance and structural sparsity. TREVISd offers the most favorable balance, matching strong predictive performance while producing simpler trees.
- TREVIS is compared with greedy and near-optimal decision-tree learners using weighted F1-score and number of leaves.Competitors include CART, FLOW, DL8.5, GOSDTlb, and DL8.5lb, with near-optimal methods run under a one-hour limit.
- TREVISd achieves the second-best average predictive performance, matching GOSDTlb and outperforming DL8.5lb.
- Both TREVIS variants consistently produce simpler trees than competing methods, with the lowest average number of leaves.The reduction is especially evident relative to DL8.5 and DL8.5lb, which produce the most complex solutions.
- TREVISd provides the best trade-off between predictive performance and structural sparsity, occupying the plot’s upper-left region.The upper-left region represents high weighted F1-score and few leaves; the dashed line marks the Pareto frontier.
- TREVISd ranks better than CARTd on predictive performance, remains statistically indistinguishable from the best near-optimal methods, and ranks best for structural sparsity.The rankings are based on Critical Difference diagrams and the Nemenyi test at α = 0.1.
- TREVISc and TREVISd have total runtimes competitive with near-optimal tree-learning approaches, particularly relative to GOSDTlb on larger datasets.
C. Latent Space Analysis
The latent-space analysis tests whether nearby latent points decode to similar trees and whether latent directions correspond to interpretable tree properties. Results support local smoothness and structured directions for varying decision-tree properties.
- Local smoothness: Local smoothness is assessed by perturbing latent points and correlating perturbation magnitude with normalized edit distance between decoded tree sequences.The analysis uses z′ = z + ϵσ, with Gaussian noise and a uniformly sampled perturbation magnitude ϵ.
- Local smoothness: Triplet Accuracy measures whether nearby latent points decode to more similar outputs than farther points.Positive and negative perturbations are constructed so the negative point is farther from the anchor by a factor α > 1.
- Local smoothness: Across 1000 latent points or triplets, positive correlations and positive Triplet Accuracy scores indicate that latent and output-space neighborhoods are largely preserved.The reported correlations have p-values below 0.05, and α is set to 2.0.
- Global coherence: Three Lasso surrogate models predict weighted F1-score, number of leaves, and depth from latent representations, then define gradient-based trajectories from the origin.Each trajectory moves along a normalized surrogate gradient with a factor t sampled from U(−3, 3).
- Global coherence: UMAP projections show smooth directions connecting regions with the lowest and highest values of the analyzed tree properties.The projections cover four datasets and color latent points by decoded-tree properties; structural properties appear smoother and more structured than predictive performance.
D. Sensitivity Analysis
The sensitivity analysis examines how TTVAE architecture, positional encoding, and training-set size affect generation quality and computational cost. Smaller models improve loss, while 20,000 training trees provide a favorable overall trade-off.
- Architecture and Positional encoding: The analysis varies embedding size, transformer depth, attention heads, positional encoding, and training-tree count across original and discretized feature spaces.Evaluation uses β-ELBO components, validity, novelty, and diversity across six datasets.
- Architecture and Positional encoding: Smaller TTVAE models provide better loss in the evaluated configurations.Embedding sizes k ∈ {120, 240}, depths K ∈ {2, 4}, and attention heads H ∈ {2, 4} are tested with |Ttr| = 60,000.
- Training Trees Size: Training time increases almost linearly with |Ttr|, while validation F1-score remains largely stable after the first few thousand trees.The comparison covers both original and discretized features.
- Training Trees Size: Reconstruction loss improves with more training trees but becomes marginal beyond 20,000 trees.The same threshold is used to select the main-experiment training-set size.
- Training Trees Size: Diversity rises rapidly for small training sets and then saturates, motivating |Ttr| = 20,000 as a cost-quality trade-off.The selected size balances training cost, reconstruction quality, predictive performance, and diversity.
V. CONCLUSION
TREVIS learns decision trees by navigating a TTVAE latent space to optimize multiple properties. It matches near-optimal algorithms in predictive performance while improving structural sparsity, but its effectiveness and efficiency have important limitations.
- Contribution: TREVIS navigates a TTVAE latent space and uses a differentiable surrogate for gradient-based joint optimization of decision-tree properties.The framework is applied to the performance-sparsity trade-off.
- Results: TREVIS finds trees with predictive performance comparable to near-optimal algorithms while improving structural sparsity.The authors characterize this as the most favorable performance-sparsity trade-off.
- Future work: Future work targets fairness, privacy, robustness, counterfactual exploration, and alternatives among almost-equally-optimal decision trees.These extensions go beyond the current performance-sparsity trade-off.
- Limitations: Some decision-tree properties may be less organized in the latent space, making useful optimization directions harder to find.This limitation is reported for properties such as predictive performance.
- Limitations: TREVIS is less efficient than greedy methods because it requires TTVAE training, surrogate training, and gradient-based optimization.On the original feature space, it is less effective when considering the full set of possible splits, although prior feature selection and discretization substantially improve it.