Learning Sparse Decision Trees via Transformer Variational Auto-Encoders

Giacomo Fidone, Alessio Cascione, Riccardo Guidotti

University of Pisa · ISTI-CNR

Published September 1, 2026 · on the feed September 15, 2026 · ▲ 2 · arXiv 2609.01430

Efficiency & inference No code arXiv ↗Hugging Face ↗

The abstract

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.

The Brief

1:50

A short podcast summary of the paper. Written and voiced by AI.

Transcript

Learning Sparse Decision Trees via Transformer Variational Auto-Encoders. By Giacomo Fidone and two others, from the University of Pisa and the ISTI institute of the Italian National Research Council. Published on the first of September. No code repository is linked to it.

Decision trees are among the most widely used models in machine learning, and the reason is largely their transparent decision logic, which makes them well suited to high-stakes decisions. Yet most existing learning algorithms focus on predictive performance. Other desirable properties, structural sparsity among them, are overlooked, because they are not optimized jointly with accuracy. And a tree is a discrete object, so searching over trees for a complex objective means searching a discrete space.

TREVIS learns decision trees with respect to complex objectives by exploring a latent space instead. That space belongs to a Tree Transformer Variational Auto-Encoder, abbreviated TTVAE. Decision trees are mapped onto latent representations, and the discrete search space is replaced with a continuous one. Every point in that continuous space corresponds to a tree, so moving through the space is a way of moving between trees.

Once trees live in a continuous space, gradients become available. TREVIS optimizes through a differentiable surrogate model, so the search can follow gradients rather than stepping from one discrete structure to the next. In the experiments, the objective asks for two things at once: predictive performance and sparsity.

The results show TREVIS discovering decision trees that match the predictive performance of existing near-optimal algorithms, while improving on their structural sparsity. The same accuracy, in other words, from sparser trees.

The move is a change of search space. A problem that was combinatorial becomes continuous by passing through a learned latent representation, and that is what lets a second objective such as sparsity sit alongside accuracy instead of being left out.

Paper Quest · 7 chapters · about 39 min

Play this paper

Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.

Chapter 1 🧩 · The problem

Accurate trees are easy to find. Accurate and small ones are not.

You need a decision tree for a high-stakes setting, so it must be accurate and also small enough for people to read.

The road ahead

  1. 1 · The problem ~6 min You need a decision tree that is both accurate and small enough to read, but the usual learners either chase accuracy greedily or take far too long.
  2. 2 · What came before ~3 min
  3. 3 · The key idea ~8 min Turn trees into points in a continuous space, learn a smooth model of how good each point is, climb its gradient, and decode the best point into a tree.
  4. 4 · Choose your own research ~3 min
  5. 5 · Build the method ~8 min Encode trees into a latent space, fit a surrogate of the objective, climb its gradients, then decode the optimised points back into trees.
  6. 6 · The experiment ~6 min You’ll guess: Compared with existing near-optimal tree learners, how did the trees TREVIS found fare on predictive performance and sparsity?
  7. 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Join free to play the whole quest ›

Free, with Google or an email address. Your stars and progress are kept.

KipQuest: the day’s AI research, as podcasts and quests. All papers Android app Privacy Contact Episodes are written and voiced by AI from each paper. Summaries can contain mistakes: the paper is the source.