Configure DecisionTreeRegressor "max_depth" Parameter

The max_depth parameter in scikit-learn’s DecisionTreeRegressor limits the maximum depth of the decision tree, which can prevent overfitting.

Decision trees learn by recursively splitting the data based on feature values until a stopping criterion is met. The max_depth parameter sets the maximum number of splits allowed from the root to a leaf node.

By default, max_depth is set to None, which allows the tree to grow until all leaves contain samples from a single class or all splits result in leaves with the minimum number of samples (controlled by min_samples_split).

In practice, max_depth values between 1 and 10 are commonly used depending on the complexity of the dataset.

from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import mean_squared_error

# Generate synthetic dataset
X, y = make_regression(n_samples=200, n_features=1, noise=20, random_state=42)

# Split into train and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Train with different max_depth values
max_depth_values = [None, 2, 5, 10]
mse_scores = []

for depth in max_depth_values:
    dt = DecisionTreeRegressor(max_depth=depth, random_state=42)
    dt.fit(X_train, y_train)
    y_pred = dt.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    mse_scores.append(mse)
    print(f"max_depth={depth}, MSE: {mse:.2f}")

Running this example outputs:

max_depth=None, MSE: 913.73
max_depth=2, MSE: 1865.39
max_depth=5, MSE: 750.58
max_depth=10, MSE: 863.12

The key steps in this example are:

Generate a synthetic regression dataset with a clear underlying relationship
Split the data into train and test sets
Train DecisionTreeRegressor models with different max_depth values
Evaluate the mean squared error of each model on the test set

Some tips and heuristics for setting max_depth:

Start with the default (None) and lower max_depth until performance on a validation set decreases
Use lower max_depth values for small datasets to avoid overfitting
Increase max_depth for more complex datasets to capture intricate patterns

Issues to consider:

Setting max_depth too low can lead to underfitting
Setting max_depth too high can cause overfitting, especially for small datasets
Deeper trees require more time and memory to train and store

See Also