6.6 Trees and forests

A tree asks a series of yes/no questions and predicts the average of the cases that end up in the same box. One tree is easy to read but unstable; hundreds of trees, averaged or added up, predict far better.

Code: data and style used in this chapter
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from sklearn import datasets
from sklearn.base import BaseEstimator
from sklearn.ensemble import BaggingRegressor, GradientBoostingRegressor, RandomForestRegressor
from sklearn.inspection import permutation_importance
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import KFold, cross_val_score
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

from statly import style

style.use()
rng = np.random.default_rng(1)

diabetes = datasets.load_diabetes(as_frame=True, scaled=False)
X, y = diabetes.data, diabetes.target.to_numpy()  # 442 patients: 10 measurements, progression
FOLDS = KFold(10, shuffle=True, random_state=1)


def cv_error(model: BaseEstimator, features: np.ndarray) -> float:
    """10-fold cross-validated prediction error (RMSE)."""
    return float(
        -cross_val_score(
            model, features, y, cv=FOLDS, scoring="neg_root_mean_squared_error"
        ).mean()
    )

6.6.1 A tree of questions

Take two of the diabetes measurements, body mass index and s5 (a blood fat). A regression tree looks for the single question that best splits the patients into two groups with different average progression: “is s5 below 4.6?” Then it asks a new question inside each group, and so on. The prediction for a new patient is the average of the training patients in the box where the patient lands.

Code
FEATURES = ["bmi", "s5"]
small_tree = DecisionTreeRegressor(max_depth=2).fit(X[FEATURES].to_numpy(), y)


def node_label(tree: DecisionTreeRegressor, node: int) -> str:
    """The text for one node: its question, or the leaf's average and size."""
    t = tree.tree_
    if t.children_left[node] == -1:
        return f"{t.value[node][0][0]:.0f}\n({t.n_node_samples[node]} patients)"
    return f"{FEATURES[t.feature[node]]} < {t.threshold[node]:.1f}?"


def draw_tree(
    ax: Axes,
    tree: DecisionTreeRegressor,
    node: int = 0,
    x: float = 0.5,
    y: float = 0.9,
    width: float = 0.5,
) -> None:
    """Draw the tree from the root down: questions in grey, leaves in red."""
    t = tree.tree_
    is_leaf = t.children_left[node] == -1
    color = style.CORAL if is_leaf else style.INK
    ax.text(
        x,
        y,
        node_label(tree, node),
        ha="center",
        va="center",
        fontsize=9.5,
        color=color,
        bbox=dict(boxstyle="round", fc="none", ec=color, lw=0.8),
    )
    if is_leaf:
        return
    for child, dx, answer in (
        (t.children_left[node], -width / 2, "yes"),
        (t.children_right[node], width / 2, "no"),
    ):
        ax.annotate(
            "",
            xy=(x + dx, y - 0.32),
            xytext=(x, y - 0.06),
            arrowprops=dict(arrowstyle="-", color=style.INK, lw=0.8),
        )
        ax.text(
            x + dx / 2 + (-0.03 if dx < 0 else 0.03),
            y - 0.18,
            answer,
            fontsize=8.5,
            ha="center",
            color=style.INK,
        )
        draw_tree(ax, tree, child, x + dx, y - 0.38, width / 2)


def plot_boxes(ax: Axes, tree: DecisionTreeRegressor) -> None:
    """The prediction of the tree over the two measurements, with the patients."""
    gx, gy = np.meshgrid(np.linspace(17, 43, 300), np.linspace(3.2, 6.7, 300))
    zones = tree.predict(np.c_[gx.ravel(), gy.ravel()]).reshape(gx.shape)
    ax.pcolormesh(gx, gy, zones, cmap="Blues", alpha=0.45, shading="auto")
    ax.scatter(X.bmi, X.s5, c=y, cmap="Reds", s=10, edgecolors=style.INK, linewidths=0.2)
    ax.set(xlabel="body mass index", ylabel="s5 (log of blood triglycerides)")


fig, axes = style.subplots(2, height=3.6, width=10.5)
draw_tree(axes[0], small_tree)
axes[0].set_axis_off()
axes[0].set(xlim=(0, 1), ylim=(0, 1))
plot_boxes(axes[1], small_tree)
fig.tight_layout()
plt.show()
Figure 6.6.1: A tree with two levels of questions on body mass index and s5. Left: the tree; each leaf shows the average progression and the number of patients. Right: the same tree as boxes, shaded by the prediction, with the patients on top.

How does the tree choose a question? For every measurement and every possible threshold, split the cases into two boxes, R1R_1 and R2R_2, and compute the residual sum of squares around each box’s average. Keep the split that makes

∑i∈R1(yi−y‾R1)2+∑i∈R2(yi−y‾R2)2 \sum_{i \in R_1} (y_i - \bar y_{R_1})^2 + \sum_{i \in R_2} (y_i - \bar y_{R_2})^2

smallest. Then repeat inside each box. This is recursive binary splitting: greedy, because it picks the best split now without looking ahead.

For a class outcome, a classification tree predicts the most common class in each box, and needs a different measure of how mixed a box is. With p̂k\hat p_k the share of class k in a box:

Gini=∑kp̂k(1−p̂k),entropy=−∑kp̂klogp̂k \text{Gini} = \sum_k \hat p_k (1 - \hat p_k), \qquad \text{entropy} = -\sum_k \hat p_k \log \hat p_k

Both are 0 for a pure box and largest for an even mix. The split chosen is the one that lowers the (size-weighted) impurity of the two new boxes the most. The plain error rate is a poorer guide for growing, because it often does not change when a split makes the boxes purer.

Code
share = np.linspace(0.001, 0.999, 300)
cancer = datasets.load_breast_cancer(as_frame=True)
malignant = (cancer.target == 0).astype(int).to_numpy()
CANCER_FEATURES = ["worst radius", "worst concave points"]
class_tree = DecisionTreeClassifier(max_depth=2, random_state=0).fit(
    cancer.data[CANCER_FEATURES].to_numpy(), malignant
)


def plot_impurities(ax: Axes) -> None:
    """Gini, entropy and error rate against the share of one class."""
    gini = 2 * share * (1 - share)
    entropy = -(share * np.log(share) + (1 - share) * np.log(1 - share))
    ax.plot(share, gini, color=style.BLUE, label="Gini")
    ax.plot(
        share, entropy / entropy.max() * gini.max(), color=style.CORAL, label="entropy (scaled)"
    )
    ax.plot(share, np.minimum(share, 1 - share), color=style.GREY, ls="--", label="error rate")
    ax.xaxis.set_major_formatter(style.PERCENT)
    ax.legend(fontsize=8)
    ax.set(title="Impurity", xlabel="share of class 1 in the box", ylabel="impurity")


def plot_class_boxes(ax: Axes) -> None:
    """The boxes of the classification tree, with share malignant and Gini in each."""
    gx, gy = np.meshgrid(np.linspace(7, 37, 300), np.linspace(0, 0.3, 300))
    leaves = class_tree.apply(np.c_[gx.ravel(), gy.ravel()]).reshape(gx.shape)
    ax.scatter(
        cancer.data[CANCER_FEATURES[0]],
        cancer.data[CANCER_FEATURES[1]],
        c=np.where(malignant, style.CORAL, style.GREEN),
        s=6,
        alpha=0.6,
    )
    ax.contour(
        gx, gy, leaves, levels=np.unique(leaves)[:-1] + 0.5, colors=style.INK, linewidths=1
    )
    t = class_tree.tree_
    for leaf in np.unique(leaves):
        ys, xs = np.where(leaves == leaf)
        p = t.value[leaf][0][1] / t.value[leaf][0].sum()
        ax.text(
            gx[0, xs].mean(),
            gy[ys, 0].mean(),
            f"{p:.0%}\nGini {2 * p * (1 - p):.2f}",
            ha="center",
            va="center",
            fontsize=8,
            bbox=dict(boxstyle="round", fc="white", ec="none", alpha=0.7),
        )
    ax.set(title="A classification tree", xlabel="worst radius", ylabel="worst concave points")


fig, axes = style.subplots(2, height=3.3, width=9.5)
plot_impurities(axes[0])
plot_class_boxes(axes[1])
fig.tight_layout()
plt.show()
Figure 6.6.3: Impurity of a box with two classes, by the share of the first class. Gini (blue) and entropy (red, scaled to the same top) are both 0 for a pure box and highest for a 50:50 mix; the error rate (grey) is a sharp tent. Right: a classification tree of depth 2 for breast tumours, with the share of malignant tumours and the Gini impurity in each box.

Trees are easy to explain: the left panel of Figure 2.2.6.1 is the whole model, readable by anyone. They need no scaling, handle groups and numbers alike, and find interactions on their own: the question asked in one branch can differ from the one asked in the other.

6.6.2 How deep?

Every extra level of questions halves the boxes. A deep enough tree puts every patient in a box of its own and fits the training data perfectly, and new data badly:

Code
DEPTHS = np.arange(1, 13)
train = [
    np.sqrt(
        (
            (DecisionTreeRegressor(max_depth=d, random_state=0).fit(X, y).predict(X) - y) ** 2
        ).mean()
    )
    for d in DEPTHS
]
tested = [cv_error(DecisionTreeRegressor(max_depth=d, random_state=0), X) for d in DEPTHS]


def plot_depth(ax: Axes) -> None:
    """Training and cross-validated error against the depth of the tree."""
    ax.plot(DEPTHS, train, color=style.GREY, label="training")
    ax.plot(DEPTHS, tested, color=style.CORAL, label="cross-validated")
    ax.legend(fontsize=8)
    ax.set(xlabel="levels of questions (depth)", ylabel="error (RMSE)", xticks=DEPTHS)


fig, ax = plt.subplots(figsize=(8, 3.2))
plot_depth(ax)
plt.show()
Figure 6.6.4: Training and cross-validated error of trees on all ten measurements, by the number of levels of questions. The cross-validated error is lowest for shallow trees and rises steadily after.

The usual remedy is to grow a big tree and then prune it: cut back branches whose questions do not pay for themselves. Cost-complexity pruning charges α per leaf: it looks for the subtree T that minimises

∑leaves m∑i∈Rm(yi−y‾Rm)2+α|T| \sum_{\text{leaves } m} \; \sum_{i \in R_m} (y_i - \bar y_{R_m})^2 + \alpha\, |T|

where |T| is the number of leaves. α = 0 keeps the whole tree; larger α prunes more. As with λ for the lasso, α is chosen by cross-validation:

Code
path = DecisionTreeRegressor(random_state=0, min_samples_leaf=5).cost_complexity_pruning_path(
    X, y
)
alphas = np.unique(
    np.r_[path.ccp_alphas[::4], path.ccp_alphas[-8:-1]]
)  # include the smallest trees
pruned = [
    DecisionTreeRegressor(random_state=0, min_samples_leaf=5, ccp_alpha=a).fit(X, y)
    for a in alphas
]
leaves = [tree.get_n_leaves() for tree in pruned]
train_rmse = [np.sqrt(np.mean((y - tree.predict(X)) ** 2)) for tree in pruned]
cv_rmse = [
    cv_error(DecisionTreeRegressor(random_state=0, min_samples_leaf=5, ccp_alpha=a), X)
    for a in alphas
]


def plot_pruning(ax: Axes) -> None:
    """Training and cross-validated error along the pruning path."""
    ax.plot(leaves, train_rmse, "o-", ms=3, color=style.GREY, label="training")
    ax.plot(leaves, cv_rmse, "o-", ms=3, color=style.CORAL, label="cross-validated")
    best = int(np.argmin(cv_rmse))
    ax.plot(leaves[best], cv_rmse[best], "o", ms=9, color=style.CORAL)
    ax.annotate(
        f"best: {leaves[best]} leaves",
        (leaves[best], cv_rmse[best]),
        xytext=(10, 10),
        textcoords="offset points",
        fontsize=9,
        color=style.CORAL,
    )
    ax.set_xscale("log")
    ax.legend(fontsize=8)
    ax.set(xlabel="leaves in the pruned tree (log scale)", ylabel="error (RMSE)")


fig, ax = plt.subplots(figsize=(8, 3.2))
plot_pruning(ax)
plt.show()
Figure 6.6.5: Cost-complexity pruning of a large tree on all ten measurements. As the price per leaf α grows, the tree is pruned to fewer leaves (to the left) and the training error rises (grey), while the cross-validated error (red) falls, until the tree becomes too small. The best tree has only a handful of leaves.

Even pruned, a single tree is rarely the best predictor. Its bigger problem is that it is unstable: a few different patients near a threshold change the first question, and everything below it.

6.6.3 Many trees: bagging and random forests

Instability is variance, and variance falls when you average. Bagging grows many deep trees, each on a bootstrap resample of the data (Chapter 3.3), and averages their predictions. A random forest adds one twist: at every question, each tree may only choose among a random few of the measurements. The trees then differ more from each other, and their average gets better.

Code
bmi = X[["bmi"]]
GRID = np.linspace(bmi.bmi.min(), bmi.bmi.max(), 400).reshape(-1, 1)


def plot_single_trees(ax: Axes, count: int = 5) -> None:
    """A few trees, each fitted to its own bootstrap resample."""
    ax.scatter(bmi, y, s=6, color=style.GREY, alpha=0.4)
    for i, color in zip(range(count), style.PALETTE):
        pick = rng.integers(0, len(y), len(y))
        tree = DecisionTreeRegressor(min_samples_leaf=15).fit(
            bmi.iloc[pick].to_numpy(), y[pick]
        )
        ax.plot(GRID, tree.predict(GRID), color=color, lw=1)
    ax.set(title="Five trees", xlabel="body mass index", ylabel="disease progression")


def plot_bagged(ax: Axes) -> None:
    """The average of many such trees."""
    ax.scatter(bmi, y, s=6, color=style.GREY, alpha=0.4)
    forest = BaggingRegressor(
        DecisionTreeRegressor(min_samples_leaf=15), n_estimators=500, random_state=0
    ).fit(bmi.to_numpy(), y)
    ax.plot(GRID, forest.predict(GRID), color=style.CORAL, lw=2)
    ax.set(title="Their average, over 500 trees", xlabel="body mass index")


fig, axes = style.subplots(2, height=3.3, width=9.5, sharey=True)
plot_single_trees(axes[0])
plot_bagged(axes[1])
fig.tight_layout()
plt.show()
Figure 6.6.6: Predicted progression against body mass index alone. Left: five deep trees, each grown on a different bootstrap resample; each one is jagged and they disagree. Right: the average of 500 such trees is much smoother and steadier.

Each tree in a bagged ensemble leaves out about a third of the cases (Chapter 3.3). Those out-of-bag cases are unseen data for that tree, so predicting every case only from the trees that did not see it gives a test-error estimate for free, without cross-validation:

Code
TREE_COUNTS = [10, 20, 40, 80, 150, 300, 500]


def oob_curve(max_features: int | float) -> list[float]:
    """Out-of-bag RMSE for forests of increasing size."""
    errors = []
    for n in TREE_COUNTS:
        forest = RandomForestRegressor(
            n_estimators=n,
            max_features=max_features,
            oob_score=True,
            random_state=0,
            min_samples_leaf=5,
        ).fit(X, y)
        errors.append(np.sqrt(np.mean((y - forest.oob_prediction_) ** 2)))
    return errors


fig, ax = plt.subplots(figsize=(8, 3.1))
ax.plot(TREE_COUNTS, oob_curve(1.0), "o-", ms=4, color=style.GREY, label="bagging (10 of 10)")
ax.plot(
    TREE_COUNTS, oob_curve(3), "o-", ms=4, color=style.GREEN, label="random forest (3 of 10)"
)
ax.legend(fontsize=8)
ax.set(xlabel="number of trees", ylabel="out-of-bag error (RMSE)")
plt.show()
Figure 6.6.7: Out-of-bag error of bagging (all ten measurements available at every split) and of a random forest (3 of 10 at every split), as trees are added. Both settle after about a hundred trees and do not overfit as more are added; the forest’s decorrelated trees average out to a lower error.

6.6.4 Boosting

Boosting also uses many trees, but grows them one after another. Each new tree is small, and it is fitted to what the trees so far got wrong: the residuals. Its prediction is shrunk by a small learning rate and added on. The model creeps towards the data in small steps.

Code
def plot_boosted(ax: Axes, trees: int) -> None:
    """The boosted fit after a given number of trees."""
    ax.scatter(bmi, y, s=6, color=style.GREY, alpha=0.4)
    model = GradientBoostingRegressor(
        n_estimators=trees, max_depth=2, learning_rate=0.1, random_state=0
    ).fit(bmi.to_numpy(), y)
    ax.plot(GRID, model.predict(GRID), color=style.CORAL, lw=2)
    ax.set(title=f"After {trees} tree{'s' if trees > 1 else ''}", xlabel="body mass index")


fig, axes = style.subplots(3, height=3, width=10.5, sharey=True)
for ax, trees in zip(axes, (1, 20, 500)):
    plot_boosted(ax, trees)
axes[0].set_ylabel("disease progression")
fig.tight_layout()
plt.show()
Figure 6.6.8: Boosting with small trees of two levels and a learning rate of 0.1, predicting progression from body mass index. After one tree the fit is a crude step; after 20 it has the shape; after 500 it starts to chase noise. The number of trees is a knob, chosen by cross-validation.

In steps, with a learning rate λ and B trees:

  1. Start with f̂(x)=0\hat f(x) = 0 and residuals ri=yir_i = y_i.
  2. For b = 1, …, B: fit a small tree f̂b\hat f^b to the residuals; update f̂(x)←f̂(x)+λf̂b(x)\hat f(x) \leftarrow \hat f(x) + \lambda \hat f^b(x) and ri←ri−λf̂b(xi)r_i \leftarrow r_i - \lambda \hat f^b(x_i).
  3. The model is f̂(x)=∑bλf̂b(x)\hat f(x) = \sum_b \lambda \hat f^b(x).

For other outcomes (yes/no, counts), gradient boosting replaces “the residuals” with the gradient of the loss: the direction in which each prediction should move to fit better. Bayesian additive regression trees (BART) are a related method that also sums many small trees, but updates them by random perturbations and averages over many sums, which gives intervals as well.

Forests are hard to overfit by adding trees; boosting is not, so its number of trees, tree size and learning rate all need tuning. Tuned well, boosted trees are among the most accurate methods for tables of data, and they win many prediction competitions.

6.6.5 Which works best here?

Code
COMPARE = [
    ("one tree (depth 3)", DecisionTreeRegressor(max_depth=3, random_state=0)),
    ("bagging", BaggingRegressor(DecisionTreeRegressor(), n_estimators=300, random_state=0)),
    (
        "random forest",
        RandomForestRegressor(
            n_estimators=300, max_features=3, min_samples_leaf=5, random_state=0
        ),
    ),
    (
        "boosting",
        GradientBoostingRegressor(
            n_estimators=150, max_depth=2, learning_rate=0.05, random_state=0
        ),
    ),
    ("linear regression", LinearRegression()),
]
scores = [(name, cv_error(model, X)) for name, model in COMPARE]


def plot_compare(ax: Axes) -> None:
    """One bar per method: its cross-validated error."""
    names, values = zip(*scores)
    best = int(np.argmin(values))
    colors = [style.BLUE if i == best else style.GREY for i in range(len(values))]
    ax.barh(names, values, color=colors, alpha=0.7)
    for i, value in enumerate(values):
        ax.text(value + 0.3, i, f"{value:.1f}", va="center", fontsize=9)
    ax.invert_yaxis()
    ax.set(xlabel="cross-validated error (RMSE)", xlim=(50, max(values) + 4))


fig, ax = plt.subplots(figsize=(8, 2.8))
plot_compare(ax)
plt.show()
Figure 6.6.9: Cross-validated error on the diabetes data, all ten measurements. The ensembles beat the single tree by far, but plain linear regression beats them all here.

That last bar is a lesson. With 442 patients and effects that are close to straight lines, a flexible method has nothing extra to find, and pays for its flexibility in variance. Ensembles of trees shine with more data and with interactions and curves that a linear model would miss. Always include a simple model in the comparison.

6.6.6 Which measurements matter?

A forest of 300 trees cannot be read like one tree. Permutation importance asks the question another way: shuffle one measurement, breaking its link with the outcome, and see how much worse the predictions get. The more they suffer, the more the model relied on it.

Code
forest = RandomForestRegressor(
    n_estimators=300, max_features=3, min_samples_leaf=5, random_state=0
).fit(X, y)
importance = permutation_importance(
    forest, X, y, n_repeats=20, random_state=0, scoring="neg_root_mean_squared_error"
)
order = np.argsort(importance.importances_mean)


def plot_importance(ax: Axes) -> None:
    """One bar per measurement, longest at the top."""
    ax.barh(
        X.columns[order],
        importance.importances_mean[order],
        xerr=importance.importances_std[order],
        color=style.GREEN,
        alpha=0.7,
        ecolor=style.INK,
    )
    ax.set(xlabel="increase in error when shuffled (RMSE)")


fig, ax = plt.subplots(figsize=(7, 3.4))
plot_importance(ax)
plt.show()
Figure 6.6.10: Permutation importance in the random forest: how much the prediction error grows when each measurement is shuffled. Body mass index and s5 carry most of the information.

Importance says what the model uses, not what causes what: two correlated measurements share their importance, and a measurement can be important because it stands in for another.

Importance says how much a measurement matters, not how. A partial dependence plot shows the shape: set every patient’s body mass index to 20, predict, and average; then 21, and so on. The curve is the forest’s average prediction as one measurement varies with everything else as it is.

Code
def plot_dependence(ax: Axes, name: str, label: str) -> None:
    """The forest's average prediction as one measurement varies."""
    grid = np.linspace(X[name].quantile(0.02), X[name].quantile(0.98), 40)
    averages = []
    for value in grid:
        changed = X.copy()
        changed[name] = value
        averages.append(forest.predict(changed).mean())
    ax.plot(grid, averages, color=style.GREEN, lw=2)
    ax.plot(X[name], np.full(len(X), min(averages) - 5), "|", ms=6, color=style.INK, alpha=0.3)
    ax.set(xlabel=label)


fig, axes = style.subplots(3, height=2.9, width=10.5, sharey=True)
for ax, (name, label) in zip(
    axes, (("bmi", "body mass index"), ("s5", "s5"), ("bp", "blood pressure"))
):
    plot_dependence(ax, name, label)
axes[0].set_ylabel("average predicted progression")
fig.tight_layout()
plt.show()
Figure 6.6.11: Partial dependence of predicted disease progression on body mass index, s5 and blood pressure in the random forest. Progression rises with all three, steeply over the middle of each range and flattening at the ends, where the trees have few cases to split on. The ticks show where the patients are.

Partial dependence averages over interactions; when the effect of one measurement depends strongly on another, plots of single predictions (individual conditional expectation curves) or two-measurement partial dependence show more.

6.6.7 In one picture

  • A tree splits the data with yes/no questions and predicts the average in each box: easy to read, finds interactions, but unstable.
  • Bagging averages trees grown on bootstrap resamples; a random forest also limits the measurements each question may use.
  • Splits are chosen greedily: smallest RSS for numbers, lowest Gini or entropy for classes; cost-complexity pruning cuts back, with α chosen by cross-validation.
  • Out-of-bag cases give a free test error for bagging and forests.
  • Boosting adds small trees one after another, each fixing the errors of the ones before.
  • Ensembles predict much better than one tree, but a simple model can still win on small or linear problems.
  • Permutation importance shows which measurements a model relies on; partial dependence shows how.

Next: support vector machines, which look for the widest gap between classes.