Learning Sparse Decision Trees via Transformer Variational Auto-Encoders

arXiv:2609.01430 · cs.LG, cs.AI · Submitted 2026-09-01 · Read on arXiv

cs.LG, cs.AI

Submitted: 2026-09-01

Updated: 2026-09-01

Comments: Accepted for publication at the 2026 IEEE International Conference on Data Mining (ICDM 2026)

Code: https://github.com/aia-uclouvain/pydl8.5

License: http://creativecommons.org/licenses/by/4.0/

The gist: 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.

Terminology

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.

Sources

Related papers