# Learning Sparse Decision Trees via Transformer Variational Auto-Encoders

arXiv: https://arxiv.org/abs/2609.01430 · HTML: https://mlsanity.com/paper/2609.01430/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 61 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [A Survey on Bias and Fairness in Machine Learning](https://mlsanity.com/paper/1908.09635/) (2019)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [Grammar Variational Autoencoder](https://mlsanity.com/paper/1703.01925/) (2017)
- [Scalable and accurate deep learning for electronic health records](https://mlsanity.com/paper/1801.07860/) (2018)
- [Understanding How Dimension Reduction Tools Work: An Empirical Approach to Deciphering t-SNE, UMAP, TriMAP, and PaCMAP for Data Visualization](https://mlsanity.com/paper/2012.04456/) (2020)
- [Variational Autoencoder for Deep Learning of Images, Labels and Captions](https://mlsanity.com/paper/1609.08976/) (2016)
- [Benchmarking and Survey of Explanation Methods for Black Box Models](https://mlsanity.com/paper/2102.13076/) (2021)
- [Generalized and Scalable Optimal Sparse Decision Trees](https://mlsanity.com/paper/2006.08690/) (2020)
- [D-VAE: A Variational Autoencoder for Directed Acyclic Graphs](https://mlsanity.com/paper/1904.11088/) (2019)
- [Exploring the Whole Rashomon Set of Sparse Decision Trees](https://mlsanity.com/paper/2209.08040/) (2022)
- [Directed Acyclic Graph Neural Networks](https://mlsanity.com/paper/2101.07965/) (2021)
- [Tree Transformer: Integrating Tree Structures into Self-Attention](https://mlsanity.com/paper/1909.06639/) (2019)
- [MurTree: Optimal Classification Trees via Dynamic Programming and Search](https://mlsanity.com/paper/2007.12652/) (2020)
- [Fast Sparse Decision Tree Optimization via Reference Ensembles](https://mlsanity.com/paper/2112.00798/) (2021)
- [Strong Optimal Classification Trees](https://mlsanity.com/paper/2103.15965/) (2021)
- [Transformers over Directed Acyclic Graphs](https://mlsanity.com/paper/2210.13148/) (2022)
- [PACE: A Parallelizable Computation Encoder for Directed Acyclic Graphs](https://mlsanity.com/paper/2203.10304/) (2022)
- [Global Counterfactual Directions](https://mlsanity.com/paper/2404.12488/) (2024)
- [AST-Transformer: Encoding Abstract Syntax Trees Efficiently for Code Summarization](https://mlsanity.com/paper/2112.01184/) (2021)
- [Dependency Transformer Grammars: Integrating Dependency Structures into Transformer Language Models](https://mlsanity.com/paper/2407.17406/) (2024)
- [Optimal Decision Trees for the Algorithm Selection Problem: Integer Programming Based Approaches](https://mlsanity.com/paper/1907.02211/) (2019)
- [Optimizing Training Trajectories in Variational Autoencoders via Latent Bayesian Optimization Approach](https://mlsanity.com/paper/2207.00128/) (2022)
- [DisCo-DSO: Coupling Discrete and Continuous Optimization for Efficient Generative Design in Hybrid Spaces](https://mlsanity.com/paper/2412.11051/) (2024)
- [ODTlearn: A Package for Learning Optimal Decision Trees for Prediction and Prescription](https://mlsanity.com/paper/2307.15691/) (2023)
- [Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation](https://mlsanity.com/paper/2507.04003/) (2025)

## Abstract (arXiv)

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.