# The Tree Ensemble Layer: Differentiability meets Conditional Computation

arXiv: https://arxiv.org/abs/2002.07772 · HTML: https://mlsanity.com/paper/2002.07772/

## Key facts

- Cited by 57 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 18 Feb 2020.

## Built on by (most cited first, 30 of 57)

- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [Dynamic Neural Networks: A Survey](https://mlsanity.com/paper/2102.04906/) (2021)
- [On Embeddings for Numerical Features in Tabular Deep Learning](https://mlsanity.com/paper/2203.05556/) (2022)
- [How to DP-fy ML: A Practical Guide to Machine Learning with Differential Privacy](https://mlsanity.com/paper/2303.00654/) (2023)
- [LIFT: Language-Interfaced Fine-Tuning for Non-Language Machine Learning Tasks](https://mlsanity.com/paper/2206.06565/) (2022)
- [DSelect-k: Differentiable Selection in the Mixture of Experts with Applications to Multi-Task Learning](https://mlsanity.com/paper/2106.03760/) (2021)
- [TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023](https://mlsanity.com/paper/2307.14338/) (2023)
- [Transfer Learning with Deep Tabular Models](https://mlsanity.com/paper/2206.15306/) (2022)
- [Boost then Convolve: Gradient Boosting Meets Graph Neural Networks](https://mlsanity.com/paper/2101.08543/) (2021)
- [Revisiting Pretraining Objectives for Tabular Deep Learning](https://mlsanity.com/paper/2207.03208/) (2022)
- [T2G-Former: Organizing Tabular Features into Relation Graphs Promotes Heterogeneous Feature Interaction](https://mlsanity.com/paper/2211.16887/) (2022)
- [Mixture of In-Context Prompters for Tabular PFNs](https://mlsanity.com/paper/2405.16156/) (2024)
- [A Performance-Driven Benchmark for Feature Selection in Tabular Deep Learning](https://mlsanity.com/paper/2311.05877/) (2023)
- [Learning Accurate Decision Trees with Bandit Feedback via Quantized Gradient Descent](https://mlsanity.com/paper/2102.07567/) (2021)
- [Computation-efficient Deep Learning for Computer Vision: A Survey](https://mlsanity.com/paper/2308.13998/) (2023)
- [MIXRTs: Toward Interpretable Multi-Agent Reinforcement Learning via Mixing Recurrent Soft Decision Trees](https://mlsanity.com/paper/2209.07225/) (2022)
- [GRANDE: Gradient-Based Decision Tree Ensembles for Tabular Data](https://mlsanity.com/paper/2309.17130/) (2023)
- [Multimodal Clinical Trial Outcome Prediction with Large Language Models](https://mlsanity.com/paper/2402.06512/) (2024)
- [GCondNet: A Novel Method for Improving Neural Networks on Small High-Dimensional Tabular Data](https://mlsanity.com/paper/2211.06302/) (2022)
- [Weight Predictor Network with Feature Selection for Small Sample Tabular Biomedical Data](https://mlsanity.com/paper/2211.15616/) (2022)
- [Vanilla Gradient Descent for Oblique Decision Trees](https://mlsanity.com/paper/2408.09135/) (2024)
- [TabSeq: A Framework for Deep Learning on Tabular Data via Sequential Ordering](https://mlsanity.com/paper/2410.13203/) (2024)
- [A Neural Tangent Kernel Perspective of Infinite Tree Ensembles](https://mlsanity.com/paper/2109.04983/) (2021)
- [Homological Convolutional Neural Networks](https://mlsanity.com/paper/2308.13816/) (2023)
- [Flexible Modeling and Multitask Learning using Differentiable Tree Ensembles](https://mlsanity.com/paper/2205.09717/) (2022)
- [Yet Another Representation of Binary Decision Trees: A Mathematical Demonstration](https://mlsanity.com/paper/2101.07077/) (2021)
- [Tree in Tree: from Decision Trees to Decision Graphs](https://mlsanity.com/paper/2110.00392/) (2021)
- [Analyzing Tree Architectures in Ensembles via Neural Tangent Kernel](https://mlsanity.com/paper/2205.12904/) (2022)
- [COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search](https://mlsanity.com/paper/2306.02824/) (2023)

## Abstract (arXiv)

Neural networks and tree ensembles are state-of-the-art learners, each with its unique statistical and computational advantages. We aim to combine these advantages by introducing a new layer for neural networks, composed of an ensemble of differentiable decision trees (a.k.a. soft trees). While differentiable trees demonstrate promising results in the literature, they are typically slow in training and inference as they do not support conditional computation. We mitigate this issue by introducing a new sparse activation function for sample routing, and implement true conditional computation by developing specialized forward and backward propagation algorithms that exploit sparsity. Our efficient algorithms pave the way for jointly training over deep and wide tree ensembles using first-order methods (e.g., SGD). Experiments on 23 classification datasets indicate over 10x speed-ups compared to the differentiable trees used in the literature and over 20x reduction in the number of parameters compared to gradient boosted trees, while maintaining competitive performance. Moreover, experiments on CIFAR, MNIST, and Fashion MNIST indicate that replacing dense layers in CNNs with our tree layer reduces the test loss by 7-53% and the number of parameters by 8x. We provide an open-source TensorFlow implementation with a Keras API.