Decision trees

What it is and when you reach for it

A classification tree is a piecewise-constant function built by recursive partitioning: start with all the training data in one node, pick the single question of the form "is feature j ≤ threshold t" that best separates the classes, split the data into the two answers, and recurse on each side until the nodes are pure or some budget runs out. Prediction is a walk from the root to a leaf, answering one question per level, and the leaf's class distribution is the output. Against its neighbors: logistic regression draws one global hyperplane, k-NN memorizes the training set, and kernel methods bend space with a similarity function, while a tree asks a short adaptive sequence of single-feature questions. That buys three things nothing else in the classical toolbox offers together: no feature scaling, no assumed functional form, and a model a domain expert can read line by line. What it costs is smoothness and stability, since a small perturbation of the data can flip an early split and rebuild the entire subtree below it. You reach for a single tree when interpretability is the product, and for ensembles of trees, which fix the instability by averaging or boosting, when accuracy on tabular data is the product. In practice the second case dominates: the single decision tree matters today mostly because it is the weak learner inside random forests and gradient boosting, the models that still win on tabular data.

The math

Recursive partitioning as an objective

The tree defines a partition of feature space into regions R₁, ..., R_M, and in region R_m it predicts the majority class (or the full class-frequency vector when you want probabilities). The natural objective is to choose the partition minimizing the total misclassification, or more usefully a smooth surrogate of it summed over leaves: Σ_m n_m · impurity(R_m), where n_m counts the samples in leaf m. Finding the globally optimal tree under this objective is NP-complete (Hyafil and Rivest proved this in 1976), so every practical algorithm since CART and ID3 has been greedy: optimize one split at a time, locally, and never look back. The whole design of the algorithm falls out of two decisions: what "impurity" means, and how to search for the best single split quickly.

Impurity: Gini and entropy, with numbers

Let p_k be the fraction of samples in a node belonging to class k. The two standard impurity measures are the Gini index, G = 1 − Σ_k p_k², and the entropy, H = −Σ_k p_k log₂ p_k. Both are zero for a pure node, maximal for a uniform mix, and strictly concave, and the concavity is the load-bearing property: it guarantees that any split producing children with different class mixtures strictly decreases the weighted impurity. Gini has a probabilistic reading, the chance that two samples drawn from the node disagree in class, and entropy has an information-theoretic one, the expected number of bits needed to encode a label drawn from the node.

Concretely, take a node with 10 samples, 6 of class A and 4 of class B, so p = (0.6, 0.4). The Gini index is 1 − (0.6² + 0.4²) = 1 − 0.36 − 0.16 = 0.48. The entropy is −(0.6 log₂ 0.6 + 0.4 log₂ 0.4) = 0.6 · 0.737 + 0.4 · 1.322 ≈ 0.971 bits. Now split this node into a left child with 5 samples, all class A, and a right child with 5 samples, 1 A and 4 B. The left child is pure, so both of its impurities are 0. The right child has p = (0.2, 0.8), giving Gini = 1 − (0.04 + 0.64) = 0.32 and entropy = 0.2 · 2.322 + 0.8 · 0.322 ≈ 0.722 bits. The weighted post-split impurity is (5/10) · 0 + (5/10) · 0.32 = 0.16 for Gini and (5/10) · 0.722 ≈ 0.361 for entropy.

Information gain and the split search

The quality of a candidate split is the impurity decrease it buys: gain = impurity(parent) − (n_L/n) · impurity(left) − (n_R/n) · impurity(right). For the example above, the Gini gain is 0.48 − 0.16 = 0.32 and the information gain (the same quantity with entropy) is 0.971 − 0.361 ≈ 0.610 bits, which has a clean reading: this one yes/no question extracts about 0.6 bits of the 0.97 bits of label uncertainty in the node. The split search then just maximizes gain over every feature and every threshold. For a numeric feature with n distinct values only n − 1 thresholds matter, the midpoints between consecutive sorted values, because the partition of the data is constant between them. Exact implementations sort each feature once and sweep class counts across the boundary in O(n) per feature; histogram implementations (LightGBM made this standard) bucket each feature into 32 to 256 quantile bins and only try bin edges, trading a provably negligible loss in split resolution for a large constant-factor win and a much more vectorizable inner loop. The implementations below use the histogram form because it maps perfectly onto batched tensor operations.

In practice the choice between Gini and entropy almost never changes the tree much: both are concave, they rank candidate splits nearly identically, and empirical studies find they disagree on the chosen split only a few percent of the time. CART defaults to Gini partly because it avoids computing logarithms in the innermost loop.

Why greedy, and what greed costs

Greedy splitting is not an approximation you would remove given more compute; it is the only thing that makes the problem tractable, since the space of trees grows super-exponentially with depth. The cost is real, though: greed is myopic, and the classic failure is XOR-structured data. If y = x₁ XOR x₂, then no single split on x₁ or x₂ alone changes the class mix at all, so every candidate has zero gain and a strictly greedy learner is blind to a pattern that a depth-2 tree represents exactly. Real implementations still find such structure in practice because ties get broken and subsequent levels recover, but the lesson stands: a greedy tree can miss interactions whose individual components carry no marginal signal, and this is one of the quiet reasons ensembles over many randomized trees work better than any single greedy tree.

Overfitting, pruning, and depth

An unrestricted tree drives training error to zero by growing one leaf per training sample, which is pure memorization: the model has high variance and its test error is far worse than its training error. Two families of remedies exist. Pre-pruning stops growth early with budgets: max_depth, min_samples_leaf, min_samples_split, or a minimum impurity decrease per split. Post-pruning grows the full tree and then removes subtrees that do not pay for themselves; CART's cost-complexity pruning minimizes error(T) + α · |leaves(T)|, sweeping α from 0 upward to produce a nested sequence of subtrees and picking the best by cross-validation (this is ccp_alpha in scikit-learn). For a single tree meant to be read by humans, pruning matters a great deal. For trees used inside ensembles the calculus flips: random forests intentionally grow deep, low-bias, high-variance trees and let averaging remove the variance, while boosting intentionally grows shallow, high-bias trees (depth 3 to 8) and lets the additive stagewise procedure remove the bias.

The axis-aligned inductive bias

Every split is a question about one feature at a time, so the decision boundary is a union of axis-aligned boxes. That is the tree's inductive bias, and it cuts both ways. When the true boundary aligns with single features, thresholds on income, age, a lab measurement, a count, trees are extraordinarily efficient and invariant to any monotone transformation of each feature, which is why feature scaling is unnecessary. When the true boundary is an oblique line, say x₁ + x₂ > 1, a tree can only approximate it with a staircase, needing many splits for a shape that logistic regression captures with two weights:

oblique truth: x1 + x2 > 1          tree approximation: a staircase
x2                                  x2
 |\      B                           |____
 | \                                 | B  |__
 |  \                                |       |__
 | A \                               |  A       |____
 |____\_____ x1                      |_______________ x1

This bias also explains a practical folk rule: trees do not need one-hot encoding to be avoided or interactions to be hand-built, but they do benefit from features expressed in the coordinates where the signal is axis-aligned, ratios instead of raw pairs, for example. And because leaves predict constants, trees cannot extrapolate: outside the range of the training data the prediction is frozen at the nearest leaf's value.

Implementation, twice

Here is the honest framing before the code. Decision trees are a poor fit for PyTorch and JAX, and that mismatch is the pedagogy. The split objective is piecewise constant in the threshold, so there is no gradient to descend; the model's structure is data-dependent control flow, which is exactly what jax.jit cannot trace and what autograd has nothing to say about. What a tensor framework genuinely offers a tree is the split search: scoring every (feature, threshold) pair in one batched computation, a histogram of class counts over candidate thresholds built with a single einsum, instead of two nested Python loops. So both implementations below draw the same line that production libraries draw in C++: a vectorized, jit-able split search over fixed-shape tensors, and a plain Python recursion for tree growth, which is inherently sequential and data-dependent. These implementations exist to teach that boundary, not to compete with the real libraries.

import torch
import torch.nn.functional as Fn

class Node:
    __slots__ = ("feature", "threshold", "left", "right", "probs")
    def __init__(self, probs):
        self.feature = None
        self.threshold = None
        self.left = None
        self.right = None
        self.probs = probs

class CARTClassifier:
    """CART with a histogram split search done as batched tensor ops.

    Candidate thresholds are per-feature quantiles computed once on
    the full data and reused at every node (as LightGBM does). The
    split search scores every (feature, threshold) pair at once; the
    recursion stays in Python because tree growth is data-dependent
    control flow that autograd and vectorization cannot help with.
    """

    def __init__(self, max_depth=6, min_samples_leaf=2, n_bins=32):
        self.max_depth = max_depth
        self.min_samples_leaf = min_samples_leaf
        self.n_bins = n_bins

    def fit(self, X, y):
        self.n_classes = int(y.max()) + 1
        Y = Fn.one_hot(y, self.n_classes).float()
        qs = torch.linspace(0, 1, self.n_bins + 1)[1:-1]
        self.thresholds = torch.quantile(X, qs, dim=0).T   # (d, T)
        self.root = self._grow(X, Y, depth=0)
        return self

    def _best_split(self, X, Y):
        n = X.shape[0]
        # left[i, j, t]: does sample i go left under threshold t of feature j
        left = (X.unsqueeze(-1) <= self.thresholds.unsqueeze(0)).float()
        left_counts = torch.einsum("ndt,nk->dtk", left, Y)   # (d, T, K)
        total = Y.sum(0)                                      # (K,)
        right_counts = total - left_counts
        n_left = left_counts.sum(-1)                          # (d, T)
        n_right = n - n_left
        # clamp only guards 0/0 on empty sides; those are masked out below
        gini_l = 1 - (left_counts / n_left.clamp(min=1).unsqueeze(-1)).pow(2).sum(-1)
        gini_r = 1 - (right_counts / n_right.clamp(min=1).unsqueeze(-1)).pow(2).sum(-1)
        score = (n_left * gini_l + n_right * gini_r) / n      # weighted child impurity
        valid = (n_left >= self.min_samples_leaf) & (n_right >= self.min_samples_leaf)
        score = torch.where(valid, score, torch.tensor(float("inf")))
        j, t = divmod(int(score.argmin()), score.shape[1])
        parent = 1 - (total / n).pow(2).sum()
        if not torch.isfinite(score[j, t]) or parent - score[j, t] <= 1e-12:
            return None                                        # no impurity decrease
        return j, float(self.thresholds[j, t])

    def _grow(self, X, Y, depth):
        node = Node(probs=Y.sum(0) / max(len(Y), 1))
        if (depth >= self.max_depth
                or len(Y) < 2 * self.min_samples_leaf
                or node.probs.max() == 1.0):
            return node
        split = self._best_split(X, Y)
        if split is None:
            return node
        node.feature, node.threshold = split
        m = X[:, node.feature] <= node.threshold
        node.left = self._grow(X[m], Y[m], depth + 1)
        node.right = self._grow(X[~m], Y[~m], depth + 1)
        return node

    def predict_proba(self, X):
        out = torch.empty(len(X), self.n_classes)
        stack = [(self.root, torch.arange(len(X)))]
        while stack:                       # route index sets, not single rows
            node, idx = stack.pop()
            if node.feature is None:
                out[idx] = node.probs
                continue
            m = X[idx, node.feature] <= node.threshold
            stack.append((node.left, idx[m]))
            stack.append((node.right, idx[~m]))
        return out

    def predict(self, X):
        return self.predict_proba(X).argmax(dim=-1)

import jax
import jax.numpy as jnp
from functools import partial

# Tree growth is data-dependent control flow, which is exactly what
# jit cannot trace, so the recursion runs in plain Python on concrete
# arrays. The split search IS a fixed-shape computation, so it jits
# cleanly. Drawing that line is the whole lesson of writing a tree
# in an autodiff framework.

@partial(jax.jit, static_argnames="min_leaf")
def best_split(X, Y, thresholds, min_leaf=2):
    """Score every (feature, threshold) pair at once via a class histogram."""
    n = X.shape[0]
    left = (X[:, :, None] <= thresholds[None]).astype(Y.dtype)  # (n, d, T)
    left_counts = jnp.einsum("ndt,nk->dtk", left, Y)            # (d, T, K)
    total = Y.sum(0)
    right_counts = total[None, None] - left_counts
    n_left = left_counts.sum(-1)                                 # (d, T)
    n_right = n - n_left
    gini_l = 1 - jnp.sum((left_counts / jnp.clip(n_left, 1)[..., None]) ** 2, -1)
    gini_r = 1 - jnp.sum((right_counts / jnp.clip(n_right, 1)[..., None]) ** 2, -1)
    score = (n_left * gini_l + n_right * gini_r) / n
    valid = (n_left >= min_leaf) & (n_right >= min_leaf)
    score = jnp.where(valid, score, jnp.inf)
    j, t = jnp.unravel_index(jnp.argmin(score), score.shape)
    parent = 1 - jnp.sum((total / n) ** 2)
    gain = parent - score[j, t]                # -inf gain if nothing valid
    return j, thresholds[j, t], gain

def grow(X, Y, thresholds, depth, max_depth=6, min_leaf=2):
    probs = Y.sum(0) / max(X.shape[0], 1)
    leaf = {"probs": probs}
    if depth >= max_depth or X.shape[0] < 2 * min_leaf or probs.max() == 1.0:
        return leaf
    j, thr, gain = best_split(X, Y, thresholds, min_leaf=min_leaf)
    if not bool(gain > 1e-12):               # forces a sync; fine outside jit
        return leaf
    m = X[:, j] <= thr
    return {
        "feature": int(j), "threshold": float(thr), "probs": probs,
        "left":  grow(X[m], Y[m], thresholds, depth + 1, max_depth, min_leaf),
        "right": grow(X[~m], Y[~m], thresholds, depth + 1, max_depth, min_leaf),
    }

def fit(X, y, n_classes, max_depth=6, min_leaf=2, n_bins=32):
    Y = jax.nn.one_hot(y, n_classes)
    qs = jnp.linspace(0, 1, n_bins + 1)[1:-1]
    thresholds = jnp.quantile(X, qs, axis=0).T                   # (d, T)
    return grow(X, Y, thresholds, depth=0,
                max_depth=max_depth, min_leaf=min_leaf)

def predict_proba(tree, X):
    """Fully vectorized over rows: every leaf is evaluated for every row
    and jnp.where selects the right one. O(leaves * n), but branch-free."""
    if "feature" not in tree:
        return jnp.broadcast_to(tree["probs"], (X.shape[0],) + tree["probs"].shape)
    m = (X[:, tree["feature"]] <= tree["threshold"])[:, None]
    return jnp.where(m, predict_proba(tree["left"], X),
                        predict_proba(tree["right"], X))

def predict(tree, X):
    return predict_proba(tree, X).argmax(-1)

Both versions share one design decision worth noticing: candidate thresholds are global per-feature quantiles computed once, not recomputed per node. That is the histogram trick, and it is what turns the split search into a fixed-shape tensor program.

Using it on a real shape of problem

A realistic smoke test: 4,000 samples, 20 features, 3 classes, with class structure on a handful of informative features and the rest noise, which is roughly the shape of a small tabular problem.

import torch

torch.manual_seed(0)
n, d, k = 4000, 20, 3
centers = torch.randn(k, d) * 2.0
y = torch.randint(0, k, (n,))
X = centers[y] + torch.randn(n, d)          # blobs + noise dims
X_tr, y_tr, X_te, y_te = X[:3000], y[:3000], X[3000:], y[3000:]

for depth in (2, 4, 6, 10):
    tree = CARTClassifier(max_depth=depth).fit(X_tr, y_tr)
    acc_tr = (tree.predict(X_tr) == y_tr).float().mean()
    acc_te = (tree.predict(X_te) == y_te).float().mean()
    print(f"depth {depth:2d}  train {acc_tr:.3f}  test {acc_te:.3f}")

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
n, d, k = 4000, 20, 3
k1, k2, k3 = jax.random.split(key, 3)
centers = jax.random.normal(k1, (k, d)) * 2.0
y = jax.random.randint(k2, (n,), 0, k)
X = centers[y] + jax.random.normal(k3, (n, d))
X_tr, y_tr, X_te, y_te = X[:3000], y[:3000], X[3000:], y[3000:]

for depth in (2, 4, 6, 10):
    tree = fit(X_tr, y_tr, n_classes=k, max_depth=depth)
    acc_tr = (predict(tree, X_tr) == y_tr).mean()
    acc_te = (predict(tree, X_te) == y_te).mean()
    print(f"depth {depth:2d}  train {acc_tr:.3f}  test {acc_te:.3f}")

Expect the classic pattern, with exact numbers varying by machine and seed: training accuracy climbs monotonically with depth and approaches 1.0 by depth 10, while test accuracy rises, peaks around a moderate depth (typically 4 to 6 on data like this), and then flattens or degrades as the deeper splits start fitting noise. If you plot both curves against depth you are looking at the bias-variance tradeoff in its cleanest form; that picture is the single most useful diagnostic a tree gives you.

Applications

Three roles, in increasing order of real-world weight. First, interpretable rules. A pruned tree of depth 3 or 4 is one of the few models that can be printed in a report and audited clause by clause, which is why shallow trees and tree-derived rule lists keep appearing where decisions must be explained to a regulator or a clinician: credit decisioning notices, triage protocols, and clinical decision rules (the emergency-medicine Ottawa ankle rules are literally a small decision tree). Second, feature-importance triage. Fitting a quick forest of trees and reading impurity-based or permutation importances is a standard first pass over a new tabular dataset with hundreds of columns; it is cheap, needs no scaling or encoding ceremony, and reliably surfaces the dozen features worth real attention, with the caveat about cardinality bias covered in the traps below. Third, and this is the role that dominates, the tree as weak learner: random forests average hundreds of deep decorrelated trees, and gradient boosting adds thousands of shallow ones, and together those two ensemble recipes are most of applied machine learning on tabular data, from fraud scores and credit risk to search ranking and ad click prediction. Nearly every quality gain in that world over the past two decades, XGBoost, LightGBM, CatBoost, came from better ways to grow and combine exactly the object built on this page. The boosting side of that story is developed fully on the gradient boosting page.

Against the real libraries

The reference implementation to measure against is scikit-learn's DecisionTreeClassifier, a Cython implementation of CART. Over the code above it adds the exact split search (sorted features with an O(n) sweep over all midpoints rather than quantile bins), support for sample weights and multi-output targets, both Gini and entropy criteria, cost-complexity pruning via ccp_alpha, sparse input support, and years of edge-case hardening around ties, missing-value handling (native since version 1.3), and degenerate splits. The same repository supplies the ensemble layer a single tree is usually just a part of: RandomForestClassifier and ExtraTreesClassifier for bagging-style variance reduction, and HistGradientBoostingClassifier for boosting. Beyond scikit-learn sit the boosting engines that actually get deployed, each with its own highly engineered tree grower: XGBoost (second-order boosting, sparsity-aware splits, GPU training), LightGBM (the histogram method this page borrows, plus leaf-wise growth), and CatBoost (ordered boosting and native categorical handling). Those are compared properly on the gradient boosting page; the point here is that all of them are, at the bottom, a very fast loop around the histogram split search you just read.

When is the from-scratch version enough? For teaching, for reading, and for embedding a tiny interpretable rule set into a system where a dependency is unwelcome; a depth-3 tree can be transcribed into a dozen lines of if-statements. For anything where accuracy matters, use the library: the Cython exact search is orders of magnitude faster than the Python recursion here, and the ensembles matter more than any single-tree refinement. Verification is straightforward and worth actually doing: fit both on the same data with the same max_depth and min_samples_leaf, and compare test-set predictions. Because the implementation above bins thresholds at 32 quantiles while scikit-learn searches every midpoint, expect agreement on the large majority of test points and near-identical accuracy (within a percentage point or two on data like the blobs above) rather than an exact match; raise n_bins and watch the two converge. Also compare the root split: with enough bins, both implementations should choose the same feature and nearly the same threshold, which is a sharp, single-number check that the gain computation is right.

Traps and misconceptions

"Deeper is better." Depth is capacity, and an unconstrained tree will memorize the training set. Test accuracy almost always peaks at moderate depth for a single tree. The correct depth also depends on the tree's job: deep for forests, shallow for boosting, pruned for human reading.

"Gini versus entropy is an important choice." It almost never is; they rank splits nearly identically and disagree on a small fraction of splits. Spend the tuning budget on depth, leaf-size minimums, and pruning instead.

"Feature importances tell you what matters." Impurity-based importances are biased toward high-cardinality and continuous features, which offer more candidate thresholds and thus more chances to look good by luck, and they are computed on training data, so a feature the tree overfit to can rank highly. Prefer permutation importance on held-out data when the ranking will drive decisions, and treat correlated features as sharing credit.

"Trees generalize like other models outside the training range." They cannot extrapolate at all. Every leaf predicts a constant, so beyond the range of the training data the prediction is frozen. For regression on trending signals (prices, growth curves) this is a structural failure no ensemble of trees fixes.

"A good single tree is a stable model." Trees are high-variance: resampling the training data can change an early split and rebuild the whole subtree below it, which is precisely why bagging and random feature selection, i.e. random forests, work as well as they do. If you need stability and use one tree, you are relying on luck; if you need one tree for interpretability, prune hard and validate that the top splits survive resampling.