Hyperparameter Tuning with Optuna

Why Grid Search Is Not Enough

Grid search evaluates every combination of hyperparameters. With 4 values for learning rate, 4 for max depth, 3 for subsample, and 3 for regularization, that is 144 combinations. Each requires cross-validation. Most of those 144 trials explore regions of the space that are clearly suboptimal.

Random search is better. It samples randomly and tends to find good configurations faster because it covers more of the important dimensions. But it still wastes trials on bad regions.

Bayesian optimization, which Optuna implements, builds a model of the objective function and uses it to choose the next set of hyperparameters to try. It focuses on promising regions and avoids ones that are likely to perform poorly.

Installing Optuna

pip install optuna

Your First Optuna Study

The core concept in Optuna is the study, which manages a collection of trials. Each trial evaluates one set of hyperparameters.

import optuna
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier

data = load_breast_cancer()
X, y = data.data, data.target

def objective(trial):
    # Define the search space
    n_estimators = trial.suggest_int("n_estimators", 50, 500)
    max_depth = trial.suggest_int("max_depth", 2, 10)
    min_samples_split = trial.suggest_int("min_samples_split", 2, 20)
    min_samples_leaf = trial.suggest_int("min_samples_leaf", 1, 10)

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf,
        random_state=42,
    )

    scores = cross_val_score(model, X, y, cv=5, scoring="f1")
    return scores.mean()

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)

print(f"Best F1: {study.best_value:.4f}")
print(f"Best params: {study.best_params}")

Optuna maximizes or minimizes the return value of the objective function. Set direction="maximize" for metrics like accuracy and F1. Set direction="minimize" for metrics like MSE or log loss.

Search Space Types

Optuna supports several parameter types that map to common hyperparameter patterns.

def objective(trial):
    # Integer parameters
    n_estimators = trial.suggest_int("n_estimators", 50, 500, step=50)

    # Float parameters
    learning_rate = trial.suggest_float("learning_rate", 0.001, 0.3, log=True)

    # Float with uniform distribution
    subsample = trial.suggest_float("subsample", 0.5, 1.0)

    # Categorical parameters
    booster = trial.suggest_categorical("booster", ["gbtree", "dart"])

    # Conditional parameters (only sample if condition is met)
    if booster == "dart":
        dropout_rate = trial.suggest_float("dropout_rate", 0.0, 0.5)
    else:
        dropout_rate = 0.0

    # ... build and evaluate model ...

The log=True flag for floats samples on a logarithmic scale. This is essential for parameters like learning rate where the difference between 0.001 and 0.01 matters more than the difference between 0.2 and 0.3.

Conditional parameters are a major advantage over grid search. Optuna only samples dropout_rate when the booster is “dart”, keeping the search space efficient.

Tuning XGBoost with Optuna

import xgboost as xgb
from sklearn.model_selection import StratifiedKFold
import numpy as np

def xgboost_objective(trial):
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 100, 1000),
        "max_depth": trial.suggest_int("max_depth", 3, 8),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
        "subsample": trial.suggest_float("subsample", 0.6, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.6, 1.0),
        "reg_alpha": trial.suggest_float("reg_alpha", 1e-8, 10.0, log=True),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-8, 10.0, log=True),
        "min_child_weight": trial.suggest_int("min_child_weight", 1, 10),
        "gamma": trial.suggest_float("gamma", 0.0, 5.0),
        "random_state": 42,
        "eval_metric": "logloss",
    }

    model = xgb.XGBClassifier(**params)

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    scores = cross_val_score(model, X, y, cv=skf, scoring="f1")

    return scores.mean()

study = optuna.create_study(direction="maximize")
study.optimize(xgboost_objective, n_trials=100, show_progress_bar=True)

print(f"Best F1: {study.best_value:.4f}")
print(f"Best params: {study.best_params}")

Pruning: Stop Bad Trials Early

Pruning is one of Optuna’s most powerful features. It stops a trial early if intermediate results suggest it will not beat the current best. This can cut total compute time by 50% or more.

from optuna.pruners import MedianPruner

def objective_with_pruning(trial):
    params = {
        "n_estimators": 500,
        "max_depth": trial.suggest_int("max_depth", 3, 8),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
        "subsample": trial.suggest_float("subsample", 0.6, 1.0),
        "random_state": 42,
    }

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

    fold_scores = []
    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]

        model = xgb.XGBClassifier(**params)
        model.fit(X_train, y_train)
        score = model.score(X_val, y_val)
        fold_scores.append(score)

        # Report intermediate result and check if we should prune
        trial.report(np.mean(fold_scores), fold)
        if trial.should_prune():
            raise optuna.TrialPruned()

    return np.mean(fold_scores)

study = optuna.create_study(
    direction="maximize",
    pruner=MedianPruner(n_startup_trials=10, n_warmup_steps=2),
)
study.optimize(objective_with_pruning, n_trials=100)

The MedianPruner prunes a trial if its intermediate value is worse than the median of previous trials at the same step. The n_startup_trials parameter lets the first 10 trials run to completion so the pruner has a baseline.

Analyzing Results

Optuna provides built-in visualization tools.

# Parameter importance
importance = optuna.importance.get_param_importances(study)
for param, imp in importance.items():
    print(f"  {param}: {imp:.4f}")

# History of trials
fig = optuna.visualization.plot_optimization_history(study)
fig.show()

# Parameter relationships
fig = optuna.visualization.plot_parallel_coordinate(study)
fig.show()

# Hyperparameter importance
fig = optuna.visualization.plot_param_importances(study)
fig.show()

The parameter importance plot tells you which hyperparameters matter most. If learning_rate has 60% importance and gamma has 2%, you know where to focus future tuning efforts.

Saving and Resuming Studies

Optuna can persist studies to a database so you can resume interrupted experiments.

# Save to SQLite
study = optuna.create_study(
    study_name="xgboost_tuning",
    storage="sqlite:///optuna_study.db",
    direction="maximize",
    load_if_exists=True,
)

# Run 50 trials now
study.optimize(objective, n_trials=50)

# Later, resume with 50 more trials
study = optuna.load_study(
    study_name="xgboost_tuning",
    storage="sqlite:///optuna_study.db",
)
study.optimize(objective, n_trials=50)
print(f"Total trials: {len(study.trials)}")

Multi-Objective Optimization

Sometimes you need to balance multiple objectives, like maximizing accuracy while minimizing inference time.

import time

def multi_objective(trial):
    n_estimators = trial.suggest_int("n_estimators", 50, 500)
    max_depth = trial.suggest_int("max_depth", 2, 10)

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        random_state=42,
    )

    scores = cross_val_score(model, X, y, cv=3, scoring="f1")
    f1 = scores.mean()

    # Measure inference time
    model.fit(X, y)
    start = time.time()
    model.predict(X[:100])
    inference_ms = (time.time() - start) * 1000

    return f1, inference_ms

study = optuna.create_study(
    directions=["maximize", "minimize"],
)
study.optimize(multi_objective, n_trials=50)

# Get Pareto-optimal trials
for trial in study.best_trials:
    print(f"F1={trial.values[0]:.4f}, Inference={trial.values[1]:.2f}ms")

Training the Final Model

After tuning, train the final model on the full training set with the best parameters.

best_params = study.best_params
best_params["random_state"] = 42
best_params["eval_metric"] = "logloss"

final_model = xgb.XGBClassifier(**best_params)
final_model.fit(X_train, y_train)

test_score = final_model.score(X_test, y_test)
print(f"Test accuracy with tuned params: {test_score:.4f}")

Summary

Optuna implements Bayesian optimization with a tree-structured Parzen estimator that focuses trials on promising regions of the hyperparameter space. Pruning stops bad trials early, often cutting total compute in half. Conditional search spaces let you define parameters that depend on other parameter values. Persist studies to a database to resume long-running experiments. After tuning, always train a final model on the full training set and evaluate on a held-out test set.