PaperlayerView exact revision on arXiv

Source-linked AI summary

The Tree Ensemble Layer: Differentiability meets Conditional Computation

Hussein Hazimeh, Natalia Ponomareva, Petros Mol, Zhenyu Tan, Rahul Mazumder

arXiv:2002.07772v2cs.LGcs.CVstat.ML

TL;DR

Differentiable trees combine neural and tree-model advantages but are slow because they lack conditional computation. The paper introduces TEL with sparse routing and specialized propagation, achieving competitive performance with substantially more efficient and compact models. Its scope includes joint first-order optimization of deep and wide ensembles, with feature sparsity left for future work.

  • Problem

    Differentiable trees offer promising neural-network integration but lack conditional computation, causing computation to involve every tree node and hindering efficient training of large ensembles.

    Source: §1 · p. 1, §2 · p. 3, §3 · p. 4

  • Method

    TEL is an additive ensemble of differentiable decision trees using smooth-step sparse routing and specialized forward and backward passes for end-to-end first-order optimization.

    Source: §1 · p. 1, §3 · p. 4, §5 · p. 9

  • Results

    TEL achieves competitive predictive performance while providing over 10x speed-ups over differentiable trees and substantially more compact models than gradient boosted trees.

    Source: §1 · p. 1, §5 · p. 9

  • Takeaways & Limitations

    Sparse conditional computation makes jointly optimized deep and wide tree ensembles practical within neural networks while reducing computational and model-size costs.

    Source: §1 · pp. 1–2, §5 · p. 9

    Further speedups from exploiting feature sparsity remain an open direction for TEL.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Tree ensembles offer robust, interpretable, easy-to-tune learning with conditional computation, while neural networks provide representation learning but lack natural conditional computation. TEL combines differentiable tree ensembles with neural networks and sparse routing to support efficient joint optimization.

1. Introduction

  • TEL is an additive layer of differentiable decision trees that can be inserted into neural networks and trained end-to-end with gradient-based methods.It combines neural-network representation learning with tree-ensemble structure.Source: §1 · p. 1
  • TEL introduces sparse differentiable routing and specialized forward and backward algorithms to realize conditional computation during both training and inference.The approach routes each sample through only a small part of the tree while exploiting sparsity computationally.Source: §1 · pp. 1–2
  • TEL achieves over 10x speed-ups versus current differentiable trees without sacrificing predictive performance across 23 real datasets.The introduction presents this as the main empirical efficiency result.Source: §1 · p. 1
  • Jointly updating all trees with first-order methods addresses the computational difficulty of optimizing deep and wide ensembles simultaneously.Earlier differentiable-tree training complexity grew exponentially with depth, while gradient boosting typically updated trees stage-wise.Source: §1 · pp. 1–2
  • The contribution includes an activation function for routing through small tree regions, optimal-complexity sparse propagation algorithms, broad evaluations, and an open-source TensorFlow/Keras implementation.The experiments compare TEL with differentiable trees, gradient boosted decision trees, and dense CNN layers.Source: §1 · p. 2
  • Differentiable trees usually process every node because soft routing assigns positive probability to every branch, making computation expensive.Prior work therefore used inference heuristics or different routing models, creating discrepancies between training and inference.Source: §1 · p. 2

2. The Tree Ensemble Layer

The Tree Ensemble Layer is an additive ensemble of differentiable decision trees that preserves gradient-based training while introducing sparse, conditional routing. Its smooth-step activation produces exact routing zeros and ones, enabling efficient computation over reachable tree regions.

The Tree Ensemble Layer

  • TEL is an additive model of differentiable decision trees that can be inserted into neural networks and trained with gradient-based optimization.The layer sums the outputs of multiple trees to produce a vector of raw predictions.Source: §2 · p. 3

Soft Routing

  • Soft trees use trainable hyperplane splits and probabilistic left/right routing, while their deterministic prediction is the expected value of leaf outputs.The routing probability at an internal node is computed from an activation applied to the inner product between the input and the node weight vector.Source: §2 · p. 3

Activation Functions

  • Unlike logistic routing, the smooth-step activation outputs exact zeros and ones, allowing differentiable trees to perform true conditional computation.The logistic function assigns every sample positive probability of reaching every node, making computation exponential in tree depth; smooth-step routing avoids this behavior.Source: §2 · pp. 3–4
  • The smooth-step function is a continuously differentiable cubic interpolation between hard-routing regions, constrained to reach 0 and 1 with zero endpoint derivatives.Its parameter γ controls the width of the transition interval and therefore the balance between hard routing and gradient availability.Source: §2 · p. 4

Conditional Computation

  • For a depth-10 tree, small enough γ values cause the average number of reachable leaves per sample to converge rapidly to 1 during training.The reported experiment used the diabetes dataset and measured reachable leaves on a logarithmic y-axis.Source: §2 · p. 4, Figure 2 · p. 4

Efficient Propagation

  • The smooth-step’s sparse activations and gradients are used to design efficient forward and backward propagation algorithms, including a backward pass whose complexity can be independent of tree depth.Higher-order polynomial variants can also use the proposed algorithms.Source: §2 · p. 4

3. Conditional Computation

The paper develops sparse forward and backward algorithms for differentiable tree ensembles, enabling conditional computation and first-order joint optimization. Their costs depend on the reachable portion of the tree rather than its full depth, with backward propagation often cheaper than forward propagation.

3. Conditional Computation

  • The algorithms optimize TEL with first-order methods by exploiting sparsity in the smooth-step activation and its gradient.The method targets gradient computation, whose conventional time and memory costs can grow exponentially with tree depth.Source: §3 · p. 4
  • The forward pass assumes each sample reaches relatively few leaves, enforced by choosing the smooth-step parameter γ sufficiently small.With U reachable leaves and N reachable internal nodes, memory complexity is O(d) for inference and O(d + U) for training.Source: §3 · p. 5, §3.1 · p. 5

3.1. Conditional Forward Pass

  • O(Np + Uk) conditional-forward time replaces O(2^dp + 2^dk) dense evaluation when the sample reaches U sub-exponential leaves.The traversal prunes descendants whenever an edge has zero probability, visiting each reachable node once.Source: §3.1 · p. 5

3.2. Conditional Backward Pass

  • Only fractional ancestors of reachable leaves and the reachable leaves themselves are needed to compute gradients, while gradients for unreachable leaves are zero.The backward pass builds and traverses a fractional tree containing these relevant nodes, with O(U) memory and constant memory when U is bounded.Source: §3.2 · pp. 5–6
  • O(Up + Uk) worst-case time and O(k) best-case time make the conditional backward pass optimal in the number of non-zero gradient entries.Its worst-case complexity is generally lower than the conditional forward pass and can be significantly lower than dense backward propagation when reachable leaves are sparse.Source: §3.2 · p. 6
  • The fractional-tree backward pass can be cheaper than forward propagation because it traverses a simplified graph containing fewer nodes than the reachable forward tree.This contrasts with standard neural-network backpropagation, where forward and backward passes traverse the same computational graph.Source: §3.2 · p. 6

4. Experiments

Experiments evaluate TEL as a standalone learner and CNN layer across classification datasets, comparing its predictive performance, computational cost, and compactness with soft trees, GBDT, and dense layers. TEL preserves competitive predictive performance while enabling faster training and smaller CNN models.

4.1. Soft Trees: Smooth-step vs. Logistic Activation

  • Smooth-step and logistic activations have comparable predictive performance across 23 PMLB datasets, with each outperforming the other on 7 datasets and ties on the remainder.Smooth-step wins on 7 datasets, including 5 statistically significant wins; logistic wins on 7, including 4 statistically significant wins.Source: §4.1 · p. 7, Table 2 · p. 7
  • Over 10x speed-up at depth 10 is achieved by smooth-step routing compared with logistic activation, whose training time increases steeply after depth 6.The comparison uses a fixed ensemble size of 10 and γ = 1 for smooth-step, representing its worst-case training time within the tuning range.Source: §4.1 · p. 7, Figure 3 · p. 8

4.2. TEL vs. Gradient Boosted Decision Trees

  • TEL is competitive with GBDT across 23 PMLB datasets, outperforming it on 9 datasets and losing on 8, with ties on the remaining 6.Five TEL wins and seven GBDT wins are statistically significant; ties typically occur on easy tasks with AUC near 1.Source: §4.2 · p. 7, Table 3 · p. 8
  • More than 20x fewer parameters are required on pima when TEL reaches comparable AUC with 5 trees while GBDT requires more than 100.Across heart-c, pima, and spambase, TEL peaks at a smaller number of trees and is less sensitive to the tree count.Source: Figure 4 · p. 8, §4.2 · p. 9
  • TEL’s compactness and reduced sensitivity to tree count are attributed to joint optimization, which can produce more expressive ensembles than GBDT’s stage-wise optimization.The comparison matches GBDT and TEL approximately in parameters per tree.Source: Figure 4 · p. 8, §4.2 · p. 9

4.3. TEL vs. Dense Layers in CNNs

  • Replacing CNN dense layers with TEL yields lower test loss on CIFAR-10, MNIST, and Fashion MNIST while maintaining comparable classification accuracy.The experiments compare CNN-Dense with CNN-TEL after tuning convolutional, dense, optimization, and TEL parameters over five random initializations.Source: §4.3 · p. 9, Table 4 · p. 9
  • Test-loss improvements on CIFAR and Fashion MNIST are 28% and 53%, respectively, and are statistically significant.Because the metric is cross-entropy loss, the authors interpret this as higher average scores on correct predictions.Source: §4.3 · p. 9
  • CNN-TEL uses approximately 8x fewer parameters than CNN-Dense, while representation layers allow TEL to outperform GBDT on image datasets such as MNIST and CIFAR-10.The CNN experiments use convolutional representation layers before the TEL component.Source: §4.3 · p. 9

5. Conclusion and Future Work

TEL combines differentiable tree ensembles with neural networks while supporting conditional computation through sparse routing and specialized propagation. The authors report competitive predictive performance and compact models, with future work focused on exploiting feature sparsity.

  • TEL combines additive differentiable trees with neural networks, supports end-to-end first-order training, and routes each sample through only a small part of the tree.Conditional computation uses smooth-step routing and specialized forward and backward passes.Source: §5 · p. 9
  • TEL achieves competitive performance against gradient boosted trees and dense layers while producing significantly more compact models.The layer also significantly outperforms GBDT on multiple image classification datasets when leveraging convolutional layers.Source: §5 · p. 9
  • Future work includes adding feature-sparsity mechanisms to further speed up TEL computation.Source: §5 · p. 9

A. Notation

The appendix provides a notation table for the paper.

B. Appendix for Section 2

The appendix describes the setup for a smooth-step routing experiment using a single depth-10 tree on the diabetes dataset.

B. Appendix for Section 2

  • Figure 2 trains a depth-10 tree with smooth-step activation using Adam to optimize cross-entropy loss.The base learning rate is 0.1 and the batch size is 256.Source: B. Appendix for Section 2 · p. 11
  • The figure’s y-axis reports the average number of reachable leaves per sample, averaged over each batch.Source: B. Appendix for Section 2 · p. 11

C.1. Example of Conditional Forward and Backward Passes

The appendix illustrates conditional forward and backward passes on a depth-4 regression tree, showing sparse traversal and a smaller backward-pass tree.

C.1. Example of Conditional Forward and Backward Passes

  • Only 3 of the 16 leaves in the depth-4 regression tree are reachable during the conditional forward pass.The reachable leaves store scalar regression values used for the tree output.Source: §C.1 · p. 11
  • The backward-pass tree is substantially smaller because hard-routing internal nodes are removed.Source: §C.1 · p. 11

C.2. Memory Complexity of the Conditional Forward Pass

Conditional computation keeps inference memory linear in tree depth and training memory linear in depth plus reachable leaves, while the backward pass scales with the active fractional subtree.

C.2. Memory Complexity of the Conditional Forward Pass

  • O(d) inference memory is sufficient because depth-first traversal discards each node after it is traversed.Training additionally stores reachable-leaf probabilities and internal-node activations, giving O(d+U) worst-case memory.Source: §C.2 · p. 11
  • O(d+U) training memory stores one probability per reachable leaf and inner products for the U−1 fractional ancestors.These stored quantities support efficient backward propagation.Source: §C.2 · p. 11

C.3. Time Complexity of the Conditional Backward Pass

  • The conditional backward pass operates on a fractional tree in which most internal splitting nodes from the forward reachable subtree are eliminated.Source: Figure C.1 · p. 12
  • O(Up+Uk) is the worst-case time and gradient sparsity complexity of the conditional backward pass.It follows from traversing U−1 internal nodes and producing at most p+p|F|+Uk nonzero gradient entries.Source: §C.3 · p. 13
  • O(k) is the best-case backward complexity when exactly one leaf is reachable and the fractional tree has one node.Source: §C.3 · p. 13

C.4. Proof of Theorem 1

The proof derives gradients with respect to inputs, leaf outputs, and split weights by applying the chain rule and exploiting reachability and fractional-node sparsity.

Gradient of Loss w.r.t. x

  • The input-gradient derivation begins with the chain rule and expands the tree output over leaves and ancestor routing probabilities.Source: §C.4 · p. 13
  • Unreachable leaves contribute zero, so the gradient summation restricts to reachable leaves R.Source: §C.4 · p. 13
  • Only fractional routing nodes contribute input derivatives, restricting ancestor terms to A(l) ∩ F.Non-fractional routing probabilities have zero input derivative.Source: §C.4 · p. 13

Gradient of Loss w.r.t. w_i

  • For split weights outside the fractional-node set, the relevant summation term is zero, further enforcing conditional sparsity in the gradients.Source: §C.4 · p. 13

Gradients of Loss w.r.t. O and w_i

  • Gradients with respect to leaf outputs and split weights are obtained by applying the chain rule, using backpropagated terms and simplified routing derivatives.Source: §C.4 · pp. 13–14

D. Appendix for Section 4

The appendix evaluates the models on 26 datasets, using PMLB classification benchmarks alongside CIFAR-10, MNIST, and Fashion MNIST under specified splitting and normalization protocols.

Datasets

  • 23 classification datasets come from the Penn Machine Learning Benchmarks repository, with no additional preprocessing and random 70%/30% train/test splits.Source: D. Appendix for Section 4 · p. 14
  • CIFAR-10, MNIST, and Fashion MNIST retain their original training/testing splits and have pixel values normalized to [0, 1].MNIST and Fashion MNIST use 60K/10K splits, while CIFAR-10 uses 50K/10K.Source: D. Appendix for Section 4 · p. 14

Computing Setup

  • Experiments run on CentOS 7 hardware with model- and dataset-specific CPU and memory limits, using TensorFlow, XGBoost, Scikit-learn, Hyperopt, and custom compiled propagation passes.Source: D. Appendix for Section 4 · p. 14

D.1. Tuning Parameters and Architectures:

The appendix specifies hyperparameter search spaces and CNN architectures for the predictive, comparison, and CNN replacement experiments.

Tuning Parameters

  • TEL tuning searches tree depth [2, 8], tree count [1, 100], L2 regularization, and learning rates from 10^-1 through 10^-5.Source: §D.1 · pp. 14–15
  • A cited tree-model search sampled tree depth uniformly over [2, 20] and number of trees uniformly over [1, 500].Source: §D.1 · p. 15

CNN Architectures

  • CNN-Dense consists of two convolutional layers, dropout, tunable ReLU dense layers, and a softmax output layer.Source: §D.1 · p. 15
  • CNN-TEL keeps the convolutional and dropout components but uses a tunable TEL dense layer before softmax.Source: §D.1 · p. 15

CNN Hyperparameters

  • The CNN-TEL dense layer searches 16, 32, or 64 units, while dropout ranges from 0.1 to 0.5.Source: §D.1 · p. 15

Related papers

Loading 2002.07772v2…