TuiML Documentation

NODE: Neural Oblivious Decision Ensembles.

Implements the architecture of Popov et al. (ICLR 2020), Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data. A NODE layer is an ensemble of oblivious decision trees -- trees that use the same split feature and threshold at every node of a given depth, so a tree of depth $D$ is a lookup table with $2^D$ leaves -- made differentiable by replacing hard splits with $\alpha$ -entmax. Layers are stacked DenseNet-style, each seeing the raw features plus every earlier layer's output.

This module also contains a dependency-free implementation of entmax15, the sparse alternative to softmax that gives the trees their near-hard splits and near-hard feature choices.

PyTorch is an optional dependency -- pip install 'tuiml[torch]'. Nothing in this module imports torch until it is called.

Classes

NODEClassifier

class algorithms.tabular_deep.node.NODEClassifier(_NODECore, _DeepTabularClassifierMixin, Classifier)

NODE: differentiable oblivious decision trees, stacked and dense.

Gradient boosting wins on tabular data because axis-aligned splits fit tables; deep learning wins because layers compose. NODE takes both: it replaces the hard split $\mathbb{1}[x_f > b]$ with an entmax relaxation, so a whole ensemble of trees becomes one differentiable layer, and then stacks those layers DenseNet-style so later trees split on earlier trees' outputs -- something a boosted ensemble cannot do.

The trees are oblivious: every node at a given depth shares one split feature and one threshold. A depth- $D$ tree is therefore a lookup table with $2^{D}$ leaves, evaluated for the whole batch with two einsums and no branching.

Overview

Each of n_trees trees picks, per depth level, a split feature via

entmax15 over the input columns -- near one-hot, but differentiable.

The chosen value is compared to a learned threshold and squashed by a

learned temperature, then entmax15 over [go right, go left] gives a soft, often exactly-hard, decision.

The outer product of the tree_depth decisions gives a distribution

over the

2^{D}

leaves; the response table is read with it.

Layers are concatenated to their own input (dense connectivity) and the

responses of all trees in all layers are averaged into class logits.

Theory

For tree $t$ at depth level $d$ , the split score is

h_{td}(x) = \frac{\langle x, \mathrm{entmax}_{1.5}(\theta_{td}) \rangle - b_{td}}{\tau_{td}}

and the leaf-membership weight of leaf $\ell = (c_1, \dots, c_D)$ is the product of the per-level choices

w_{t\ell}(x) = \prod_{d=1}^{D} \mathrm{entmax}_{1.5}\big([h_{td}, -h_{td}]\big)_{c_d} .

The layer output is $\sum_{\ell} w_{t\ell}(x) R_{t\ell}$ with a learned response table $R$ . Because entmax returns exact zeros, most leaves receive weight zero: the relaxation is soft enough to train and sharp enough to behave like a tree.

Parameters

n_layers

int = 1

Densely connected tree-ensemble layers. Real use wants 2-8; the default keeps a fit on toy data well under a second.

n_trees

int = 32

Oblivious trees per layer. Real use wants 128-2048.

tree_depth

int = 4

Depth of each tree, so :math:`2^{\text{tree\_depth}}` leaves. Cost grows exponentially in it; 6 is the practical ceiling.

learning_rate

float = 1e-2

AdamW step size. Higher than the attention models want, because the threshold and response parameters start far from useful values.

weight_decay

float = 0.0

AdamW decoupled weight decay; NODE is usually trained without it.

batch_size

int = 64

Mini-batch size for training and inference.

n_epochs

int = 60

Passes over the training set; real use wants several hundred.

early_stopping

bool = False

Hold out a validation split and restore the best weights.

validation_fraction

float = 0.15

Fraction held out when early_stopping is enabled.

patience

int = 10

Epochs without validation improvement before training stops.

device

{"cpu", "auto", "cuda", "mps"} = "cpu"

Compute device; "auto" trades reproducibility for speed.

random_state

int = 0

Seed for initialisation, shuffling and the validation split.

Attributes

classes_

np.ndarray of shape (n_classes,)

Sorted class labels seen during fit.

network_

torch.nn.Module

The fitted network, rebuilt lazily after unpickling.

feature_mean_, feature_scale_

np.ndarray

Feature standardisation statistics; the thresholds are initialised on the standardised scale, so this is not optional here.

loss_curve_

np.ndarray of shape (n_iter_,)

Mean training loss per epoch.

n_iter_

int

Epochs actually run.

n_features_in_

int

Number of features seen during fit.

Notes

Requires PyTorch. Install with pip install 'tuiml[torch]'. The class constructs and introspects without torch; fit raises ImportError naming the install command.

Complexity. Per epoch, $O(n L T (m D + 2^{D} k))$ for $L$ layers, $T$ trees, depth $D$ and $k$ outputs -- exponential in the depth, linear in everything else. Memory is $O(b T 2^{D})$ for the leaf weights.

When to use. NODE is the deep model to try when the problem looks like one gradient boosting would win: axis-aligned structure, thresholds, moderate feature counts. It keeps that inductive bias while remaining differentiable, so it can be trained jointly with other neural components -- which is the reason to prefer it over an actual boosted ensemble.

References

Popov2020

Popov, S., Morozov, S., & Babenko, A. (2020). Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data. International Conference on Learning Representations (ICLR). :doi:`10.48550/arXiv.1909.06312`

Peters2019

Peters, B., Niculae, V., & Martins, A. F. T. (2019). Sparse Sequence-to-Sequence Models. Proceedings of ACL 2019, 1504-1519. :doi:`10.18653/v1/P19-1146`

Methods

get_capabilities (cls) -> List[str]

Return supported capabilities.

get_complexity (cls) -> str

Return complexity analysis.

get_references (cls) -> List[str]

Return academic citations.

NODERegressor

class algorithms.tabular_deep.node.NODERegressor(_NODECore, _DeepTabularRegressorMixin, Regressor)

NODE for regression: differentiable oblivious tree ensembles.

The regression counterpart of NODEClassifier. The architecture is unchanged -- entmax-relaxed oblivious trees, densely stacked -- and only the objective changes: a single response dimension trained with mean squared error against a standardised target.

Overview

Standardise features and target.
Each tree picks split features with entmax15, compares them to

learned thresholds, and reads a leaf response.

Layers concatenate their outputs to their input, so later trees can

split on earlier trees' responses.

Average every tree's response and undo the target standardisation.

Theory

The prediction is an average over all $LT$ trees of soft leaf lookups,

\hat{y}(x) = \frac{1}{LT} \sum_{t} \sum_{\ell} w_{t\ell}(x) \, R_{t\ell}, \quad w_{t\ell}(x) = \prod_{d=1}^{D} \mathrm{entmax}_{1.5}\big([h_{td}, -h_{td}]\big)_{c_d}

trained by minimising $\|\hat{y} - \tilde{y}\|^2$ on the standardised target. Because $w$ is piecewise-smooth rather than piecewise-constant, the fitted surface is continuous -- unlike a tree ensemble's staircase, which is often the practical difference on smooth targets.

Parameters

n_layers

int = 1

Densely connected tree-ensemble layers; real use wants 2-8.

n_trees

int = 32

Oblivious trees per layer; real use wants 128-2048.

tree_depth

int = 4

Depth of each tree, so :math:`2^{\text{tree\_depth}}` leaves.

learning_rate

float = 1e-2

AdamW step size.

weight_decay

float = 0.0

AdamW decoupled weight decay.

batch_size

int = 64

Mini-batch size for training and inference.

n_epochs

int = 60

Passes over the training set; real use wants several hundred.

early_stopping

bool = False

Hold out a validation split and restore the best weights.

validation_fraction

float = 0.15

Fraction held out when early_stopping is enabled.

patience

int = 10

Epochs without validation improvement before training stops.

device

{"cpu", "auto", "cuda", "mps"} = "cpu"

Compute device; "auto" trades reproducibility for speed.

random_state

int = 0

Seed for initialisation, shuffling and the validation split.

Attributes

network_

torch.nn.Module

The fitted network, rebuilt lazily after unpickling.

target_mean_, target_scale_

float

Target standardisation statistics.

feature_mean_, feature_scale_

np.ndarray

Feature standardisation statistics.

loss_curve_

np.ndarray of shape (n_iter_,)

Mean training loss per epoch.

n_iter_

int

Epochs actually run.

n_features_in_

int

Number of features seen during fit.

Notes

Requires PyTorch. Install with pip install 'tuiml[torch]'.

Complexity. $O(n L T (m D + 2^{D}))$ per epoch, exponential in tree_depth and linear in everything else.

When to use. Reach for NODE when the target is a smooth function of threshold-like structure: it keeps the axis-aligned bias of a boosted ensemble but produces a continuous surface, and it can be trained jointly with other neural components.

References

Popov2020

Methods

get_capabilities (cls) -> List[str]

Return supported capabilities.

get_complexity (cls) -> str

Return complexity analysis.

get_references (cls) -> List[str]

Return academic citations.

Functions

Func

entmax15

Line 31

entmax15(inputs, dim: int=...)

Project logits onto the simplex with 1.5-entmax: a sparse softmax.

\alpha

-entmax interpolates between softmax (

\alpha = 1

, always dense) and sparsemax (

\alpha = 2

, aggressively sparse). At

\alpha = 1.5

the solution has a closed form found by sorting, and -- unlike softmax -- it assigns exactly zero to low-scoring coordinates, which is what lets a NODE tree commit to one feature and one side of a split while staying differentiable.

Parameters

inputs

torch.Tensor

Logits. Any shape; the projection is applied along dim.

dim

int = -1

Axis to project over.

Returns

probabilities

torch.Tensor

Same shape as inputs. Non-negative and summing to one along dim, with exact zeros outside the support.

Notes

Solves

\mathrm{entmax}_{1.5}(z) = \big[(\alpha - 1) z - \tau \mathbf{1}\big]_{+}^{1/(\alpha - 1)} \Big|_{\alpha = 1.5} = \big[z/2 - \tau\big]_{+}^{2}

where the threshold $\tau$ is chosen so the result sums to one. The exact algorithm sorts the scores and walks the candidate support sizes, costing $O(k \log k)$ for $k$ coordinates.

The backward pass uses the closed form rather than autograd through the sort. That is not an optimisation: at the edge of the support the threshold search evaluates $\sqrt{0}$ , whose derivative is infinite, so differentiating the search itself produces NaN weights a few hundred steps into training. The exact Jacobian-vector product is

\nabla_z \mathcal{L} = s \odot \left(g - \frac{\langle g, s \rangle} {\langle s, \mathbf{1} \rangle} \mathbf{1}\right), \quad s = \sqrt{p},

for upstream gradient $g$ , which is finite everywhere.

Requires PyTorch: pip install 'tuiml[torch]'.

References

Peters2019

Peters, B., Niculae, V., & Martins, A. F. T. (2019). Sparse Sequence-to-Sequence Models. Proceedings of ACL 2019, 1504-1519. :doi:`10.18653/v1/P19-1146`

python

>>> from tuiml.algorithms.tabular_deep.node import entmax15
>>> from tuiml.utils.torch_backend import has_torch
>>> if has_torch():
...     import torch
...     p = entmax15(torch.tensor([[1.0, 2.0, 9.0]]))
...     print(bool(torch.allclose(p.sum(-1), torch.ones(1))), bool(p[0, 0] == 0))
... else:
...     print(True, True)
True True

node.py

Classes

NODEClassifier

Overview

Theory

Parameters

Attributes

Notes

References

See Also

Methods

NODERegressor

Overview

Theory

Parameters

Attributes

Notes

References

See Also

Methods

Functions

entmax15

Parameters

Returns

Notes

References