Search⌘ K
AI Features
Home
Courses
Data Science Bootcamp

Tree Depth, Pruning, and Overfitting

Explore the impact of tree depth and pruning on decision trees to manage overfitting. Understand how deeper trees can memorize training data leading to poor generalization, and how pre-pruning and cost-complexity pruning techniques help create simpler, more reliable models. Learn to evaluate models using training and validation scores, cross-validation, and generalization gaps to select optimal tree settings and improve prediction stability for new data.

We'll cover the following...

A credit-default tree can score 0.99 on the customers used for training and only 0.78 on unseen customers. The training score looks excellent, but the gap shows that the tree learned details that do not repeat for new applicants.

A too-deep tree keeps splitting until small groups of training records receive highly specific rules. Those rules can memorize random variation rather than credit-default patterns that generalize. As with the earlier regression example, training performance is not the goal. We want a tree complex enough to retain predictive signal and simple enough to make stable predictions for new customers.

Use the generalization gap as the diagnostic

Before we restrict a tree, we need evidence that its complexity is the problem. The credit-default data below contains application and account features with a binary default label. Both models train on the same records and face the same held-out test set. The only difference is tree depth.

Python
from sklearn.datasets import make_classification
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
def main() -> None:
    X, y = make_classification(
        n_samples=500,
        n_features=6,
        n_informative=4,
        n_redundant=1,
        n_repeated=0,
        n_classes=2,
        weights=None,
        flip_y=0.03,
        class_sep=1.0,
        random_state=42,
    )
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=0.3,
        random_state=42,
        stratify=y,
    )
    models = [
        (
            "Unconstrained (max_depth=None)",
            DecisionTreeClassifier(random_state=42),
        ),
        (
            "Shallow (max_depth=3)",
            DecisionTreeClassifier(max_depth=3, random_state=42),
        ),
    ]
    rows = []
    for name, model in models:
        model.fit(X_train, y_train)
        train_acc = accuracy_score(y_train, model.predict(X_train))
        test_acc = accuracy_score(y_test, model.predict(X_test))
        gap = train_acc - test_acc
        rows.append((name, train_acc, test_acc, gap))
    name_w = max(len(r[0]) for r in rows)
    header = (
        f"{'Model':<{name_w}}  "
        f"{'Train acc':>9}  "
        f"{'Test acc':>8}  "
        f"{'Gap (train-test)':>16}"
    )
    print(header)
    print("-" * len(header))
    for name, train_acc, test_acc, gap in rows:
        print(f"{name:<{name_w}}  {train_acc:9.3f}  {test_acc:8.3f}  {gap:16.3f}")
if __name__ == "__main__":
    main()

  • Lines 8–19: The dataset is generated with a fixed random_state=42 so the train/test accuracy and resulting generalization gap are stable run-to-run, making the depth comparison reproducible.

  • Lines 21–27: stratify=y preserves the class balance across train and test splits, so differences in test accuracy are less likely to be an artifact of a skewed split.

  • Lines 29–38: The two classifiers differ only by max_depth (None vs. 3), ...