Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data

In plain words

On tables of mixed numeric columns, the usual winners are gradient-boosted tree packages such as CatBoost and XGBoost. Neural networks have not reliably beaten them. This paper proposes NODE, a neural network built from "oblivious" decision trees: every level of a tree asks the same question, and the answers index a small lookup table. Here the choices are made soft, so each question is a weighted blend of features and each threshold is a smooth switch. That makes the whole stack trainable end-to-end with gradient descent. The layers are stacked the way DenseNet stacks convolutions, so later trees can use the outputs of earlier ones. The authors report that it beats the tree packages on most of their tests.

TL;DR

NODE is a deep, differentiable version of an ensemble of oblivious decision trees, the kind of tree that CatBoost uses. It uses entmax with alpha=1.5 to select features and to make soft splits, and it is trained with QHAdam. Layers are densely connected, and the prediction averages every tree in every layer §3.2. Author-reported: it beats default CatBoost and XGBoost on all six datasets and tuned CatBoost and XGBoost on four of six, with Yahoo and Microsoft going to tuned XGBoost §4.1 tab. 1 tab. 2. Training and inference are slower than the GBDT baselines in the authors' timing table tab. 4. Consider it when you want a differentiable tree model that can sit inside a larger network §5.

Problem

Deep networks dominate vision, language and speech, but on heterogeneous tabular data the best results usually come from shallow models such as gradient-boosted decision trees. The authors say no deep method had been shown to beat the leading GBDT packages consistently §1. Differentiable trees exist, but the choice functions they use, REINFORCE and Gumbel-softmax, are noted to need long training §3.1. Multi-layer tree stacks such as mGBDT and Deep Forest are not trained end-to-end §2. The paper targets a model that is end-to-end differentiable and competitive with tuned GBDTs §1 §2.

Method

Differentiable oblivious tree. A tree of depth d uses the same feature and threshold at every node of a given depth, so it is a table of 2^d responses §2. The paper makes its choices soft. Feature selection uses entmax over a learnable matrix F eq. 2. Each split is a two-class entmax of a scaled distance from a learnable threshold b with a learnable scale tau eq. 3. The leaf weights are products of the per-level choice probabilities eq. 3, and the output is the weighted sum of learned leaf responses R eq. 4. The relaxation equals a classic oblivious tree when every choice is one-hot §3.1. For multi-class problems each tree outputs a vector with one entry per class §3.1.

NODE layer. A layer is m such trees of equal depth, all reading the same input vector §3.1 fig. 1. Its output is the concatenation of the m tree outputs §3.1.

Architecture. The model is k NODE layers in the DenseNet pattern. Each layer's input is the original features plus all earlier layer outputs, and the prediction is the average of every tree in every layer §3.2 fig. 2. Only the first |C| coordinates of each tree output are averaged for classification, and the first one for regression §3.2.

Training. Features are mapped to a normal distribution with a quantile transform §3.3. Before training, the feature selection matrix is drawn uniformly from (0, 1), thresholds are placed at feature values from the first batch, scales are chosen so that every first-batch point lies in the linear region of the soft threshold, and responses are drawn from N(0, 1) §3.3. Training is mini-batch optimisation of F, b, tau and R jointly, with cross-entropy for classification and MSE for regression §3.3. The optimiser is Quasi-Hyperbolic Adam with the parameters recommended by its authors §3.3. Parameters are averaged over c=5 consecutive checkpoints, and the stopping point is chosen on a validation set §3.3.

Inference. After training, the sparse feature selectors can be precomputed and stored in a sparse format §3.3.

# NODE forward pass: dense stack of oblivious differentiable tree layers
# Input:  features (B, n); k layers of m trees each, depth d
# Output: predictions (B, out_dim)

def node_forward(features, layers):
    # Dense stacking: each layer sees the input and all earlier outputs [§3.2]
    stacked = features
    all_trees = []
    for layer in layers:
        responses = node_layer(stacked, layer)         # (B, m, l)  [eq. 4]
        all_trees.append(responses)
        stacked = concat(stacked, flatten(responses))  # width grows by m*l
    # Average over every tree of every layer, first out_dim coords [§3.2]
    trees = concat(all_trees, dim=1)                   # (B, M, l), M = k*m
    return mean(trees[:, :, :out_dim], dim=1)         # (B, out_dim)

# One NODE layer: m oblivious trees of depth d  [§3.1] [fig. 1]
# Parameters: F (n, m, d), thresholds b (m, d), scales tau (m, d),
#             responses R (m, l, 2**d)

def node_layer(inputs, layer):
    # Soft feature choice per tree and depth: entmax over the n features [eq. 2]
    selectors = entmax15(layer.feature_logits, dim=0)  # (n, m, d)
    # Soft-selected feature value per sample, tree and depth [eq. 2]
    values = einsum("bi,imd->bmd", inputs, selectors)  # (B, m, d)
    # Soft split: right = sigma_alpha((f - b) / tau), left = 1 - right [eq. 3]
    right = entmoid15((values - layer.thresholds) / layer.scales)
    left = 1 - right
    # Choice weight of each leaf: product over depth of left/right [eq. 3]
    leaf_weights = leaf_products(left, right)          # (B, m, 2**d)
    # Output: leaf responses weighted by the choice weights [eq. 4]
    return einsum("bmk,mlk->bml", leaf_weights, layer.responses)  # (B, m, l)

# Data-aware init and training loop  [§3.3]
# Code for the same settings: notebooks/year_node_8layers.ipynb#cell-3

def train(model, train_x, train_y, valid_x, valid_y, batch_size):
    # Set thresholds and scales from the first batch, before training [§3.3]
    model.initialize(train_x[:first_batch_size])
    # Learning rate 1e-3 [app. A.2.3]; nus and betas per the code
    optimizer = QuasiHyperbolicAdam(
        model.parameters(), lr=1e-3,
        nus=(0.7, 1.0), betas=(0.95, 0.998))
    batches = minibatches(train_x, train_y, batch_size)
    for step, (batch_x, batch_y) in enumerate(batches):
        # MSE for regression, cross-entropy for classification [§3.3]
        optimizer.zero_grad()
        loss = loss_function(model(batch_x), batch_y)
        loss.backward()
        optimizer.step()
        # Periodically average the last 5 checkpoints, then validate
        # [§3.3] [code: notebooks/year_node_8layers.ipynb#cell-6]
        if step % 100 == 0:
            averaged = average_last_checkpoints(n=5)
            score = validation_score(averaged, valid_x, valid_y)
            keep_if_best(averaged, score)
        # Stop once the validation score has stopped improving [§3.3]
        if no_improvement_for(patience_steps):
            break
    return best_checkpoint

Hyperparameters

Section numbers are from the LaTeX source. Appendix numbers follow the \appendix structure in the source.

Results

Author-reported. Setup: six public tabular datasets, mean and standard deviation over ten seeds, with CatBoost and XGBoost as the main baselines §4.1. Metrics are classification error (Epsilon, Higgs, Click) and MSE (YearPrediction, Microsoft, Yahoo) §4.1. In tables 1 and 2, bold marks the best value in each row.

Default hyperparameters tab. 1

Tuned hyperparameters tab. 2

The paper reports mGBDT's YearPrediction value without a standard deviation tab. 2. DeepForest is only available for classification §4.1.

Ablation: choice function and depth tab. 3. Values are MSE for YearPrediction and error for Epsilon.

The paper's text says entmax is best across all experiments §4.2. In this table sparsemax beats entmax on YearPrediction at 2 layers (76.81 vs 77.05).

Runtime tab. 4. Models with 1024 trees of depth 6 on YearPrediction, averaged over five runs. The paper states NODE inference is on par with heavily optimised GBDT libraries.

Dataset sizes tab. 5: Epsilon 400K/100K train/test, 2000 features; YearPrediction 463K/51.6K, 90; Higgs 10.5M/500K, 28; Microsoft 723K/241K, 136; Yahoo 544K/165K, 699; Click 800K/200K, 11.

Feature importance fig. 3. On Higgs, the paper reports that the first layer's features have the highest permutation importance, that importance falls with depth, and that trees in deeper layers contribute more to the final response.

Paper vs code

Code: official repository Qwicen/node (linked from the paper's introduction §1), at commit 3bae6a8a63f0205683270b6d566d9cfa659403e4, MIT licence. The notebooks correspond to the paper's YearPrediction and Epsilon settings.

• Default choice function. The paper uses entmax with alpha = 1.5 §4.1. The ODST class defaults to sparsemax and sparsemoid odst.py. The notebooks pass entmax15 and entmoid15 explicitly year_node_8layers.ipynb epsilon_node_multigpu.ipynb. Implement entmax, not the class default.
• Learning rate. The paper states 1e-3 for all runs app. A.2.3. The notebooks do not pass lr. The QHAdam default is in the external qhoptim package and is not in this repository, so it is not verified.
• Target scaling (YearPrediction). The notebook standardises the target with training mean and standard deviation [code: notebooks/year_node_8layers.ipynb#cell-1]. The paper does not mention this. The notebook reports test MSE in raw units by multiplying by std squared year_node_8layers.ipynb.
• Output dimension (classification). The paper says each tree outputs one value per class §3.1. The Epsilon notebook uses num_classes + 1 channels per tree and averages only the first num_classes epsilon_node_multigpu.ipynb. The extra channel feeds later layers.
• Early stopping and batch size. The paper sets the batch to the largest size that fits GPU memory app. A.2.3. The notebooks use fixed batches of 1024 and 512, and patience of 5000 or 10000 steps year_node_8layers.ipynb epsilon_node_multigpu.ipynb.
• Checkpoint averaging. The paper says c = 5 consecutive checkpoints §3.3. The code saves one every 100 steps and averages the last five trainer.py year_node_8layers.ipynb.
• Threshold initialisation. The paper says thresholds take random first-batch feature values §3.3. The code draws a random percentile of the batch, with a Beta(beta, beta) distribution and beta = 1 by default odst.py.
• Quantile noise. The code adds Gaussian noise with relative std 1e-3 when fitting the quantile transform, and uses random state 1337 data.py year_node_8layers.ipynb. The paper does not mention noise.
• Feature selection. The code uses entmax over features for each tree and depth odst.py. This matches eq. 2.
• Forward pass. Each tree output is flattened and concatenated to the running feature tensor, which is the DenseNet pattern [code: lib/arch.py#L21-L37]. The prediction averages the first channel over all trees year_node_8layers.ipynb.
• Single-run logs. The notebook logs show one run each: Epsilon test error 0.10372 with 2 layers epsilon_node_multigpu.ipynb; YearPrediction test MSE 76.18 in raw units for 8 layers year_node_8layers.ipynb; 77.51 for the shallow notebook year_node_shallow.ipynb. They are not means over ten seeds and are not a reproduction of tables 1 and 2.
• Tuned configurations. The notebooks show one configuration per dataset, and the paper does not give the tuned optimum for each dataset. The notebooks are not evidence of the paper's chosen settings.

Pitfalls

• Wrong default. A library-default implementation uses sparsemax, not entmax, and so does not reproduce the paper's results. The paper's own ablation reports entmax as the best choice function tab. 3.
• Quantile transform and normalisation are not in the paper. The YearPrediction target is standardised in the notebook, and the transform fits on noisy features. Both change optimisation. Implement them, or report that you did not.
• Dense widths. Layer inputs grow by m*l per layer. The printed Epsilon model shows input widths of 2000 and 5072, that is 2000 + 1024 * 3 epsilon_node_multigpu.ipynb. Input width, and so per-layer cost, grows with depth of the stack.
• Memory. Every sample materialises a 2**d leaf-weight vector per tree odst.py. The README warns that the implementation is memory-inefficient and may need a lot of GPU memory README.md.
• Data-aware init needs data. Initialisation runs on the first batch and warns below 1000 samples odst.py. The threshold and scale parameters are NaN until that first call, so they must be set before training.
• Speed. Training and inference are slower than the GBDT baselines in the paper's timing table tab. 4. The paper describes NODE inference as on par with optimised GBDT libraries.
• Tuning budget. The paper gives 50 TPE steps for CatBoost, XGBoost and FCNN app. A.2. For NODE it gives only the grid app. A.2.3, so the tuning budgets may not be equal.
• Yahoo and Microsoft. NODE does not win there in the tuned regime tab. 2. The authors speculate that oblivious trees are a poor inductive bias for Yahoo and leave non-oblivious trees for future work §4.1.
• Target is pointwise. Ranking datasets are treated as regression, the pointwise approach to learning-to-rank §4.1 app. A.1.

Lineage

• Builds on: oblivious decision trees, as used in CatBoost §1 §2; entmax, arXiv:1905.05702, as cited in the code nn_utils.py; DenseNet-style concatenation §3.2; checkpoint averaging, arXiv:1803.05407 §3.3; Quasi-Hyperbolic Adam §3.3.
• Compared against: CatBoost, XGBoost and FCNN §4.1; mGBDT and Deep Forest, which the authors could only partly run §4.1; Gumbel-softmax, arXiv:1611.01144, and sparsemax as choice functions §4.2.
• Superseded by: not known from content in this wiki.