cs.LGMar 11, 2026

Learning Tree-Based Models with Gradient Descent

Authors: Sascha Marton

Organizations: UNIVERSITY OF MANNHEIM

Abstract

Tree-based models are widely recognized for their interpretability and have proven effective in various application domains, particularly in high-stakes domains. However, learning decision trees (DTs) poses a significant challenge due to their combinatorial complexity and discrete, non-differentiable nature. As a result, traditional methods such as CART, which rely on greedy search procedures, remain the most widely used approaches. These methods make locally optimal decisions at each node, constraining the search space and often leading to suboptimal tree structures. Additionally, their demand for custom training methods precludes a seamless integration into modern machine learning (ML) approaches. In this thesis, we propose a novel method for learning hard, axis-aligned DTs through gradient descent. Our approach utilizes backpropagation with a straight-through operator on a dense DT representation, enabling the joint optimization of all tree parameters, thereby addressing the two primary limitations of traditional DT algorithms. First, gradient-based training is not constrained by the sequential selection of locally optimal splits but, instead, jointly optimizes all tree parameters. Second, by leveraging gradient descent for optimization, our approach seamlessly integrates into existing ML approaches e.g., for multimodal and reinforcement learning tasks, which inherently rely on gradient descent. These advancements allow us to achieve state-of-the-art results across multiple domains, including interpretable DTs rees for small tabular datasets, advanced models for complex tabular data, multimodal learning, and interpretable reinforcement learning without information loss. By bridging the gap between DTs and gradient-based optimization, our method significantly enhances the performance and applicability of tree-based models across various ML domains.

Figures & tables

Appendix figures & tables73 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. Approximation-Free Differentiable Oblique Decision Trees

    May 8, 2026Subrat Prasad Panda, Blaise Genest, Arvind EaswaranInterpretable MLTabular ML

  2. Learning Sparse Decision Trees via Transformer Variational Auto-Encoders

    Sep 1, 2026Giacomo Fidone, Alessio Cascione, Riccardo GuidottiVariational AutoencodersStructured Sparsity

  3. Multistage Defer Trees for Hybrid Interpretability: If at First You Can't Succeed, Tree Again

    Jun 30, 2026Zakk Heile, Hayden McTavish, Margo Seltzer +1Ensemble LearningLearning to Defer