How Decision Trees Handle Overfitting Through Pruning Strategies in AI Engineering From Scratch

Decision trees prevent overfitting through pre-pruning hyperparameters that constrain growth during training and post-pruning techniques that remove unnecessary branches from fully grown trees.

Decision trees naturally tend to overfit by splitting recursively until leaves contain single samples, capturing training noise rather than generalizable patterns. The AI Engineering From Scratch repository implements robust pre-pruning controls directly in the DecisionTree class while providing detailed guidance on post-pruning strategies for learners to extend the base implementation. Understanding these complementary approaches, as implemented in phases/02-ml-fundamentals/04-decision-trees/code/trees.py, helps you balance model complexity with predictive accuracy.

Pre-Pruning: Early Stopping Controls

The most efficient way to control overfitting is preventing excessive growth before it occurs. The DecisionTree class in phases/02-ml-fundamentals/04-decision-trees/code/trees.py accepts several hyperparameters that terminate recursion early in the _build method (lines 108-122).

Critical Pre-Pruning Parameters

The implementation checks four key constraints at the start of each recursive split:

  • max_depth: Halts splitting when the tree reaches the specified depth limit, preventing deep, overly specific branches
  • min_samples_split: Requires a minimum number of samples in a node before attempting any split, blocking splits on tiny subsets
  • min_samples_leaf: Guarantees each resulting leaf contains at least this many samples, ensuring leaves have statistical significance
  • criterion: Selects the impurity measure (Gini or entropy) that guides split quality, indirectly affecting tree complexity

According to the source code, these parameters are validated at lines 75-80 during initialization and enforced at lines 108-120 within the tree construction logic. When any stopping condition is met, the _build method returns a leaf node immediately rather than recursing deeper.

Enforcing Constraints During Training

The following example demonstrates how to instantiate a tree with pre-pruning constraints:

from phases.02_ml_fundamentals.04_decision_trees.code.trees import DecisionTree, generate_classification_data, train_test_split, accuracy

# Generate toy data

X, y = generate_classification_data()
X_train, y_train, X_test, y_test = train_test_split(X, y)

# Initialize with pre-pruning parameters

tree = DecisionTree(
    max_depth=5,          # Stop after 5 levels

    min_samples_leaf=5,   # Each leaf must contain ≥5 samples

    criterion="gini"
)

# Fit and evaluate

tree.fit(X_train, y_train)
preds = tree.predict(X_test)
print("Test accuracy (pre-pruned):", accuracy(y_test, preds))

Setting max_depth and min_samples_leaf simultaneously creates a regularization cascade that limits both vertical depth and horizontal granularity, effectively constraining the model's capacity to memorize noise.

Post-Pruning: Trimming Fully Grown Trees

While the base implementation focuses on pre-pruning, the curriculum's documentation in phases/02-ml-fundamentals/04-decision-trees/docs/en.md describes post-pruning techniques that first grow a complete tree, then selectively remove branches that don't improve validation performance.

Cost-Complexity Pruning

Cost-complexity pruning (also known as weakest link pruning) adds a penalty term proportional to the number of leaves. The algorithm evaluates subtrees and removes those where the increase in training error is outweighed by the reduction in complexity. This approach requires solving for an optimal trade-off parameter that balances tree size against impurity.

Reduced-Error Pruning

Reduced-error pruning uses a validation set to make pruning decisions. The algorithm evaluates each subtree bottom-up, replacing it with a leaf if the validation accuracy does not decrease. This method is intuitive to implement and often yields better generalization than pre-pruning alone because it allows the tree to first capture the richest structure before removing non-contributing components.

Implementing Manual Post-Pruning

Since the repository provides conceptual guidance rather than automated post-pruning, you can implement reduced-error pruning as an extension:

def reduced_error_prune(node, X_val, y_val):
    """Recursively prune subtrees that do not improve validation accuracy."""
    if node["leaf"]:
        return node
    
    # Create leaf version of current node

    leaf_node = {"leaf": True, "value": majority_vote(y_val)}
    
    # Evaluate both versions

    before = _evaluate_tree(node, X_val, y_val)
    after = _evaluate_tree(leaf_node, X_val, y_val)
    
    if after >= before:  # Pruning does not hurt performance

        return leaf_node
    
    # Otherwise keep children and recurse

    node["left"] = reduced_error_prune(node["left"], X_val, y_val)
    node["right"] = reduced_error_prune(node["right"], X_val, y_val)
    return node

def _evaluate_tree(root, X, y_true):
    """Compute accuracy of a tree stored as nested dictionaries."""
    def predict_one(x, n):
        if n["leaf"]:
            return n["value"]
        return predict_one(x, n["left"] if x[n["feature"]] <= n["threshold"] else n["right"])
    
    preds = [predict_one(x, root) for x in X]
    return sum(p == t for p, t in zip(preds, y_true)) / len(y_true)

# Usage: Train full tree, then prune

full_tree = DecisionTree()  # No pre-pruning

full_tree.fit(X_train, y_train)

# Create validation split

X_tr, y_tr, X_val, y_val = train_test_split(X_train, y_train, test_ratio=0.2)

# Apply post-pruning

pruned_root = reduced_error_prune(full_tree.tree, X_val, y_val)
full_tree.tree = pruned_root

Why Pruning Works: The Bias-Variance Trade-off

Pre-pruning limits model complexity upfront, reducing variance at the cost of potentially increased bias. By stopping early, you prevent the tree from fitting noise in the training data. Post-pruning, conversely, allows the algorithm to first capture the richest possible structure, then selectively removes components that do not contribute to predictive performance. This often yields a superior bias-variance trade-off because decisions about which branches to remove are informed by actual validation performance rather than arbitrary depth limits.

Summary

  • Pre-pruning (max_depth, min_samples_leaf, min_samples_split) is implemented directly in phases/02-ml-fundamentals/04-decision-trees/code/trees.py and provides immediate overfitting control during training.
  • The _build method checks stopping conditions at lines 108-120, returning leaf nodes when constraints are violated.
  • Post-pruning strategies (cost-complexity and reduced-error) are conceptualized in the lesson documentation but require manual implementation to trim fully grown trees.
  • Combining both approaches—using pre-pruning for efficiency and post-pruning for refinement—typically yields the best generalization performance.

Frequently Asked Questions

What is the difference between pre-pruning and post-pruning?

Pre-pruning (early stopping) halts tree growth during training by enforcing constraints like max_depth or min_samples_split before splits occur. Post-pruning grows a complete tree first, then removes branches that do not improve validation performance. Pre-pruning is computationally cheaper but may stop too early, while post-pruning allows the tree to discover complex patterns before simplifying.

How does min_samples_leaf prevent overfitting?

The min_samples_leaf parameter guarantees that each terminal node contains at least a specified number of training samples. This prevents the creation of leaves that represent outliers or noise, ensuring that each prediction is based on a statistically significant subset of data. In the repository's implementation, this check occurs alongside other stopping conditions in the _build method.

Is post-pruning implemented in the DecisionTree class?

No, the base DecisionTree class in trees.py does not include automated post-pruning algorithms. The curriculum describes cost-complexity and reduced-error pruning as conceptual extensions in the documentation (phases/02-ml-fundamentals/04-decision-trees/docs/en.md), encouraging learners to implement these techniques manually using the tree structure returned by the fit method.

When should I use cost-complexity pruning versus reduced-error pruning?

Use cost-complexity pruning when you need a systematic way to generate a sequence of subtrees of decreasing complexity and select the optimal size via cross-validation, particularly with large datasets. Use reduced-error pruning when you have a dedicated validation set and want an intuitive, bottom-up approach that removes branches only when they demonstrably hurt validation accuracy. Reduced-error pruning is often easier to implement for educational purposes, while cost-complexity pruning is more common in production libraries like scikit-learn.

Have a question about this repo?

These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:

Share the following with your agent to get started:
curl -s "https://instagit.com/install.md"

Works with
Claude Codex Cursor VS Code OpenClaw Any MCP Client

Maintain an open-source project? Get it listed too →