Learning Tree-Based Models with Gradient Descent

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.

💡 Research Summary

The dissertation “Learning Tree‑Based Models with Gradient Descent” tackles two long‑standing obstacles in decision‑tree research: the combinatorial, non‑differentiable nature of hard, axis‑aligned splits and the incompatibility of classic tree‑learning algorithms with modern gradient‑based machine‑learning pipelines. The author introduces a dense representation of a decision tree in which every internal node and leaf is stored as a fixed‑size tensor. Each split is parameterized by a feature index and a threshold θ. During the forward pass the model uses a hard step function to route samples exactly as a conventional CART tree would, preserving interpretability and crisp decision rules. In the backward pass a straight‑through (ST) estimator replaces the step with a smooth surrogate (e.g., sigmoid), allowing gradients to flow through the routing decisions. Consequently, the whole tree becomes a differentiable computation graph that can be trained end‑to‑end with stochastic gradient descent or Adam, jointly optimizing all split thresholds and leaf predictions.

The core contribution is threefold. First, the method eliminates the greedy, node‑by‑node split selection of CART, enabling a global optimization of the entire tree structure. Second, because the tree is now a differentiable module, it can be plugged into any deep‑learning architecture—multimodal pipelines that combine CNNs and Transformers, or reinforcement‑learning agents that require a policy network. Third, the author extends the single‑tree formulation to ensembles (named GRANDE). Each tree in the ensemble receives an instance‑wise weight produced by a small auxiliary network; L1 regularization and dropout enforce sparsity, yielding a controllable trade‑off between performance and interpretability.

Extensive experiments validate the approach. On small tabular benchmarks (Titanic, Iris, Wine) the gradient‑based trees achieve 3–5 % higher accuracy than CART while maintaining comparable depth and leaf count, and training is accelerated by up to 2× on GPUs. On larger, heterogeneous tabular datasets (Adult, Credit, Higgs) the GRANDE ensembles match or surpass state‑of‑the‑art gradient‑boosted trees (LightGBM, CatBoost) in AUC, yet use roughly 30 % fewer parameters and provide clear, rule‑based explanations. In multimodal settings, a tree‑based feature selector learns jointly with visual and textual encoders, reducing the total parameter budget by 28 % without sacrificing accuracy. For reinforcement learning, the SYMPOL (Symbolic On‑Policy RL) algorithm directly optimizes a hard decision‑tree policy via policy‑gradient methods. Experiments on Pendulum and MiniGrid show that SYMPOL compresses policies without information loss, achieves training stability comparable to PPO and A2C, and offers fast inference (O(tree depth) operations).

Theoretical analysis demonstrates that the dense, differentiable formulation expands the search space from the exponential subset explored by greedy algorithms to a continuous manifold amenable to gradient descent. The author also compares several smooth approximations (sigmoid, tanh, swish) and quantifies their impact on routing fidelity and gradient magnitude. Limitations are acknowledged: very deep trees (>20 levels) can suffer from gradient attenuation through the ST estimator; hard routing remains sensitive to noisy features, potentially leading to over‑fitting; and the current implementation is CPU‑centric, with GPU‑specific memory optimizations left for future work.

Future directions include hierarchical parameter sharing to reduce depth‑related gradient issues, sparse routing mechanisms for memory‑efficient tensors, and automated hyper‑parameter search to select the optimal smooth surrogate for a given dataset. The dissertation concludes that by marrying the interpretability of hard, axis‑aligned decision trees with the flexibility of gradient‑based optimization, it delivers a universal, high‑performing, and readily integrable tree learning framework suitable for high‑stakes domains (healthcare, finance) as well as cutting‑edge AI research.