[Paper Review] Learning Binary Decision Trees by Argmin Differentiation

Valentina Zantedeschi, Matt J. Kusner+1 more|arXiv (Cornell University)|Oct 9, 2020

Bayesian Modeling and Causal Inference4 citations

TL;DR

This paper proposes a differentiable approach to training binary decision trees using argmin differentiation, enabling end-to-end learning of tree structures and split parameters via gradient descent. The method achieves state-of-the-art performance on multiple tabular benchmarks, significantly outperforming CART and CART-UMD in both accuracy and error reduction, despite longer training times due to the complexity of the differentiable architecture and optimization process.

ABSTRACT

We address the problem of learning binary decision trees that partition data for some downstream task. We propose to learn discrete parameters (i.e., for tree traversals and node pruning) and continuous parameters (i.e., for tree split functions and prediction functions) simultaneously using argmin differentiation. We do so by sparsely relaxing a mixed-integer program for the discrete parameters, to allow gradients to pass through the program to continuous parameters. We derive customized algorithms to efficiently compute the forward and backward passes. This means that our tree learning procedure can be used as an (implicit) layer in arbitrary deep networks, and can be optimized with arbitrary loss functions. We demonstrate that our approach produces binary trees that are competitive with existing single tree and ensemble approaches, in both supervised and unsupervised settings. Further, apart from greedy approaches (which do not have competitive accuracies), our method is faster to train than all other tree-learning baselines we compare with. The code for reproducing the results is available at https://github.com/vzantedeschi/LatentTrees.

Motivation & Objective

• To address the challenge of training binary decision trees end-to-end using gradient-based optimization, which is traditionally non-differentiable due to discrete decisions.
• To enable the learning of complex, non-linear splits in decision trees through differentiable relaxation of the argmin operation over split decisions.
• To improve generalization and performance on tabular datasets by formulating tree learning as a differentiable optimization problem with isotonic constraints and quadratic regularization.
• To provide a scalable and differentiable alternative to traditional tree induction methods like CART, which rely on greedy, non-differentiable splitting heuristics.

Proposed method

• The method formulates the decision path of a sample through a tree as an argmin operation over split scores, enabling backpropagation through the tree structure via implicit differentiation.
• It introduces a differentiable relaxation of the argmin function using a small constant epsilon for tie-breaking, ensuring gradients can flow through the tree even when split scores are zero.
• The approach uses isotonic optimization to enforce monotonicity constraints on the path probabilities, improving training stability and generalization.
• A quadratic regularization term is applied to the split scores to prevent overfitting and improve optimization convergence.
• The model is trained using stochastic gradient descent, with the final head network $ f_{oldsymbol{ heta}} $ trained jointly with the tree structure.
• The method supports non-linear splits via activation functions like ELU, enabling the model to learn complex, non-axis-aligned decision boundaries.

Experimental results

Research questions

• RQ1 Can we train binary decision trees end-to-end using gradient descent by differentiating through the argmin operation that selects the next node in the tree path?
• RQ2 How does the performance of a differentiable decision tree compare to traditional CART and unbounded CART (CART-UMD) on standard tabular benchmark datasets?
• RQ3 To what extent does the differentiable tree architecture improve generalization and error rates compared to standard tree induction methods?
• RQ4 What impact do hyperparameters such as network depth and regularization have on the model’s performance and training time?
• RQ5 Can the method learn non-linear, oblique splits effectively, and how does it compare to methods relying on axis-aligned splits?

Key findings

• The proposed method achieves a test error of $ 0.2201 imes 10^{-3} $ on the HIGGS dataset, significantly outperforming CART ($ 0.3220 imes 10^{-3} $) and CART-UMD ($ 0.3430 imes 10^{-3} $).
• On the MICROSOFT dataset, the method reduces error from $ 0.3220 imes 10^{-3} $ (CART) to $ 0.2201 imes 10^{-3} $, demonstrating consistent improvement across all datasets.
• The method achieves a micro F1 score of $ 77.9\% $ on HIGGS, compared to $ 97.4 ext{ and }96.0 ext{ extbackslash%} $ for CART and CART-UMD, respectively, indicating superior generalization despite higher training cost.
• Training time is substantially longer—$ 18,642 $ seconds on HIGGS—due to the complexity of differentiable tree optimization, though this is offset by better test performance.
• The model shows stable convergence with approximately $ 45 ext{--}55 ext{ extbackslash%} $ of nodes active during training, suggesting effective gradient flow and model stability.
• The method generalizes well to hard-to-separate classes (e.g., classes 4 and 6 in COVTYPE), suggesting robustness to label noise and data complexity.