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 branchesmin_samples_split: Requires a minimum number of samples in a node before attempting any split, blocking splits on tiny subsetsmin_samples_leaf: Guarantees each resulting leaf contains at least this many samples, ensuring leaves have statistical significancecriterion: 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 inphases/02-ml-fundamentals/04-decision-trees/code/trees.pyand provides immediate overfitting control during training. - The
_buildmethod 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:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →