Optuna: Complete Guide — Hyperparameter Optimization

Summary

Optuna is an open-source hyperparameter optimization framework developed by Preferred Networks, a Japanese company specializing in artificial intelligence. Since its initial release in 2019, Optuna has established itself as one of the most popular and powerful tools for automating the search for the best hyperparameters in machine learning models.

Unlike classical approaches such as grid search or random search, Optuna uses a sophisticated Bayesian sampler called TPE (Tree-structured Parzen Estimator) that learns from each previous trial to guide subsequent proposals toward the most promising regions of the search space. This approach makes the optimization process considerably more efficient, especially when the hyperparameter space is vast and evaluating each configuration is computationally expensive.

Optuna’s distinctive feature is its define-by-run API, which allows the search space to be defined dynamically inside the objective function itself. This approach contrasts with static frameworks where the search space must be defined in advance, offering far greater flexibility for complex models with conditional hyperparameters.

Mathematical Principle

TPE (Tree-structured Parzen Estimator)

The algorithmic core of Optuna is based on the Tree-structured Parzen Estimator, a Bayesian optimization method that models the conditional probability p(x|y) — that is, the probability of a hyperparameter x given an observed performance y.

Instead of directly modeling p(y|x) as a classical Bayesian method would with Gaussian processes, TPE takes a dual approach by building two distinct density distributions:

• l(x) (lower): the density distribution of hyperparameters that led to the best performances, those above a quantile threshold y* (typically the top 10 to 25% of trials).
• g(x): the density distribution of hyperparameters that led to other performances, those below the threshold y*.

The acquisition ratio is then expressed as the Expected Improvement (EI), which is proportional to the ratio of these two densities:

$$EI(x) \propto \frac{l(x)}{g(x)}$$

Intuitively, this ratio measures how “typical” a hyperparameter x is of good performances compared to mediocre ones. The higher this ratio, the more likely x is to lead to good performance. The algorithm therefore maximizes this acquisition ratio to choose the next configurations to evaluate.

Each density is modeled as a mixture of Gaussian distributions (for continuous hyperparameters) or categorical distributions (for discrete hyperparameters), with kernels centered on the values observed during previous trials. This modeling allows TPE to capture complex and multimodal structures in the hyperparameter space.

Pruning: Median Pruner

One of Optuna’s major assets is its automatic pruning system. The Median Pruner is the default implementation: at each intermediate step of a trial, it compares the current performance with the median of intermediate performances from previous trials at the same stage. If the current trial’s performance is below this median, the trial is stopped prematurely, freeing up computational resources to explore more promising configurations.

Formally, a trial t is pruned at intermediate step s if:

$$v(t, s) < \text{median}({v(t’, s) \mid t’ \in \text{previous trials}})$$

where v(t, s) represents the intermediate value of trial t at step s. This mechanism can considerably reduce total optimization time, often by 30 to 50%, by quickly abandoning configurations that prove mediocre.

Sampling: Define-by-Run

Optuna’s define-by-run API allows the search space to be defined dynamically, inside the objective function itself. This means that hyperparameters can be conditioned on the values of other hyperparameters, creating a structurally dependent search tree. For example, if the kernel hyperparameter of an SVM takes the value 'rbf', then the gamma hyperparameter becomes relevant; if kernel takes the value 'linear', then gamma no longer makes sense. This expressiveness is natural with Optuna and represents a significant advantage over statically defined frameworks.

Intuition

Imagine Optuna as an expert sommelier who must find the best wines in a vast cellar containing thousands of bottles. Instead of tasting every bottle to the last drop (which a naive grid search would do), our sommelier develops a much smarter strategy:

From the first sips, they identify bottles that show signs of mediocrity — a closed nose, unbalanced acidity, a lack of length on the palate. They abandon them immediately (pruning), saving their time and palate. Then, they analyze the common characteristics of promising bottles — perhaps wines from the Rhône Valley with 18 months of aging and a dominant Syrah grape variety systematically achieve better scores. They then focus their tastings on regions that share these attributes (TPE sampling), while occasionally exploring unexpected leads to miss nothing.

After each tasting, the sommelier refines their understanding of what makes a great wine. They learn from each experience, correct their intuitions, and gradually become more precise in their selections. This is exactly what Optuna does with each trial: it builds an increasingly fine understanding of the hyperparameter space, guiding the search toward optimal configurations with remarkable efficiency.

Python Implementation

Integration with scikit-learn: Random Forest

Here is a complete example of optimizing a Random Forest with Optuna, including pruning with intermediate reporting and result visualization:

import optuna
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import StratifiedKFold

# Load data
X, y = load_breast_cancer(return_X_y=True)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

def objective(trial):
    # Dynamic hyperparameter definition
    n_estimators = trial.suggest_int('n_estimators', 50, 500, step=50)
    max_depth = trial.suggest_int('max_depth', 3, 30)
    min_samples_split = trial.suggest_int('min_samples_split', 2, 20)
    min_samples_leaf = trial.suggest_int('min_samples_leaf', 1, 10)
    max_features = trial.suggest_categorical('max_features', ['sqrt', 'log2', None])
    bootstrap = trial.suggest_categorical('bootstrap', [True, False])

    if bootstrap:
        max_samples = trial.suggest_float('max_samples', 0.5, 1.0)
    else:
        max_samples = None

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf,
        max_features=max_features,
        bootstrap=bootstrap,
        max_samples=max_samples,
        random_state=42,
        n_jobs=-1
    )

    # Cross-validation with intermediate reporting for pruning
    scores = []
    for fold, (train_idx, val_idx) in enumerate(cv.split(X, y)):
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]

        model.fit(X_train, y_train)
        score = model.score(X_val, y_val)
        scores.append(score)

        # Intermediate reporting for pruning
        trial.report(score, step=fold)
        if trial.should_prune():
            raise optuna.TrialPruned()

    return float(np.mean(scores))

# Create study with pruner and sampler
study = optuna.create_study(
    direction='maximize',
    sampler=optuna.samplers.TPESampler(seed=42),
    pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=2)
)

# Launch optimization
study.optimize(objective, n_trials=100, timeout=600, show_progress_bar=True)

# Display results
print(f"Best score: {study.best_value:.4f}")
print(f"Best hyperparameters: {study.best_params}")

Visualization with optuna.visualization

Optuna offers a complete built-in visualization ecosystem for analyzing study results:

import optuna.visualization as vis

# Plot optimization history
fig_history = vis.plot_optimization_history(study)
fig_history.show()

# Contour map of main hyperparameters
fig_contour = vis.plot_contour(
    study,
    params=['n_estimators', 'max_depth', 'min_samples_split']
)
fig_contour.show()

# Hyperparameter importance
fig_importance = vis.plot_param_importances(study)
fig_importance.show()

# Pareto front (for multi-objective optimization)
# fig_pareto = vis.plot_pareto_front(multi_objective_study)
# fig_pareto.show()

# Slice plot
fig_slice = vis.plot_slice(study)
fig_slice.show()

# Intermediate values evolution plot
fig_inter = vis.plot_intermediate_values(study)
fig_inter.show()

# Rank plot (EDF - Empirical Distribution Function)
fig_edf = vis.plot_edf(study)
fig_edf.show()

Integration with PyTorch

Optuna integrates seamlessly with PyTorch through its callback and intermediate reporting system:

import torch
import torch.nn as nn
import torch.optim as optim
import optuna

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, dropout_rate, num_layers):
        super().__init__()
        layers = []
        dims = [input_dim] + [hidden_dim] * num_layers + [1]
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i+1]))
            if i < num_layers:
                layers.append(nn.ReLU())
                layers.append(nn.Dropout(dropout_rate))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x).squeeze(-1)

def pytorch_objective(trial):
    hidden_dim = trial.suggest_int('hidden_dim', 32, 256, log=True)
    dropout_rate = trial.suggest_float('dropout_rate', 0.1, 0.5)
    num_layers = trial.suggest_int('num_layers', 1, 4)
    lr = trial.suggest_float('lr', 1e-5, 1e-2, log=True)
    batch_size = trial.suggest_categorical('batch_size', [16, 32, 64, 128])
    weight_decay = trial.suggest_float('weight_decay', 1e-6, 1e-3, log=True)

    model = SimpleNet(input_dim=10, hidden_dim=hidden_dim,
                      dropout_rate=dropout_rate, num_layers=num_layers)
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    criterion = nn.MSELoss()

    for epoch in range(50):
        # Training on a batch
        # ... training loop ...
        val_loss = 0.0  # Calculate validation loss

        trial.report(val_loss, epoch)
        if trial.should_prune():
            raise optuna.TrialPruned()

    return val_loss

Integration with Keras / TensorFlow

import optuna
import tensorflow as tf
from tensorflow import keras

def keras_callback(study, trial):
    """Callback for pruning with Keras."""
    if trial.should_prune():
        raise optuna.TrialPruned()

def keras_objective(trial):
    units = trial.suggest_int('units', 32, 256, log=True)
    dropout = trial.suggest_float('dropout', 0.1, 0.5)
    lr = trial.suggest_float('lr', 1e-4, 1e-2, log=True)
    optimizer_name = trial.suggest_categorical('optimizer', ['adam', 'rmsprop', 'sgd'])

    if optimizer_name == 'adam':
        optimizer = keras.optimizers.Adam(learning_rate=lr)
    elif optimizer_name == 'rmsprop':
        optimizer = keras.optimizers.RMSprop(learning_rate=lr)
    else:
        momentum = trial.suggest_float('momentum', 0.8, 0.99)
        optimizer = keras.optimizers.SGD(learning_rate=lr, momentum=momentum)

    model = keras.Sequential([
        keras.layers.Dense(units, activation='relu', input_shape=(20,)),
        keras.layers.Dropout(dropout),
        keras.layers.Dense(units // 2, activation='relu'),
        keras.layers.Dropout(dropout / 2),
        keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=['accuracy'])

    model.fit(X_train, y_train, epochs=50, batch_size=32,
              validation_data=(X_val, y_val), verbose=0,
              callbacks=[optuna.integration.TFKerasPruningCallback(trial, 'val_accuracy')])

    return model.evaluate(X_val, y_val, verbose=0)[1]

keras_study = optuna.create_study(direction='maximize',
                                  pruner=optuna.pruners.MedianPruner())
keras_study.optimize(keras_objective, n_trials=50)

Optuna Configuration Hyperparameters

Configuring an Optuna study relies on several key parameters that determine the behavior and efficiency of the optimization:

Choosing the Sampler

The choice of sampler significantly influences performance:

• TPESampler (default): excellent quality/price ratio, suitable for the majority of cases. Works well with mixed spaces (continuous, integer, categorical).
• CmaEsSampler: based on the CMA-ES (Covariance Matrix Adaptation Evolution Strategy) algorithm, particularly effective for continuous spaces of moderate dimension (up to ~50 dimensions).
• RandomSampler: purely random search, useful as a baseline or for quick initial exploration.
• NSGAIISampler: designed for multi-objective optimization, uses a non-dominated genetic algorithm.

Choosing the Pruner

• MedianPruner: median-based pruning, robust and simple. Good default choice.
• HyperbandPruner: a variant of the Hyperband algorithm, particularly effective for iterative models (neural networks, gradient boosting).
• PatientPruner: more tolerant, only prunes if performance is significantly below the reference. Useful when performance fluctuates.
• ThresholdPruner: prunes if performance exceeds a predefined fixed threshold.
• PercentilePruner: a variant of MedianPruner using a configurable percentile.

Advantages and Limitations of Optuna

Advantages

1. Superior efficiency: The TPE Sampler converges significantly faster than grid search and random search, especially in high-dimensional spaces. Published benchmarks often show improvements of 2 to 10 times over random methods.
2. Intuitive and flexible API: The define-by-run approach is natural to use and allows complex conditional search spaces without additional configuration. The Python syntax is elegant and expressive.
3. Built-in pruning: Automatic pruning of mediocre trials saves considerable computation time, often 30 to 50% of the total budget. The system is flexible with multiple strategies to choose from.
4. Rich visualization: The built-in visualization ecosystem (optuna.visualization) offers powerful tools for analyzing results, understanding hyperparameter importance, and diagnosing search problems.
5. Extensive integrations: Optuna integrates natively with scikit-learn, PyTorch, TensorFlow/Keras, XGBoost, LightGBM, CatBoost, Hugging Face Transformers, MLflow, and many other libraries.
6. Persistence and resumption: Studies can be saved in SQLite, PostgreSQL, or MySQL databases, allowing optimization to be resumed after an interruption or results to be shared across multiple machines.
7. Multi-objective optimization: Optuna natively supports optimization with multiple simultaneous objectives through the NSGA-II algorithm, producing a Pareto front of optimal solutions.
8. Active community: Developed by Preferred Networks with a significant open-source community, Optuna benefits from regular updates and quality documentation.

Limitations

1. Learning curve: Although the API is intuitive, mastering the different samplers, pruners, and advanced parameters requires a non-negligible time investment.
2. Initial computation time: The first 10 to 20 trials are random and do not yet benefit from Bayesian intelligence. For extremely expensive models, even these initial trials can represent a significant budget.
3. Less performant native distributed parallelism: Although Optuna supports parallelism, concurrency management can become complex at scale, requiring an external database to synchronize trials across workers.
4. Sensitivity to sampler hyperparameters: TPE itself has hyperparameters (n_startup_trials, gamma, weights) that influence its performance. Poor tuning can reduce efficiency.
5. Less suited to very vast spaces: For search spaces exceeding 100 dimensions, Bayesian methods like TPE can suffer from the curse of dimensionality and lose efficiency.

Use Cases

1. Deep Learning Model Optimization

The most classic use case for Optuna is optimizing the hyperparameters of a deep neural network: learning rate, hidden layer sizes, dropout rate, batch size, L2 regularization weight, and optimizer choice. With built-in pruning, Optuna can halve the time needed to find the best architecture.

2. Model Selection and Tuning for Kaggle Competition

In the competitive Kaggle context, every fraction of a point of accuracy matters. Optuna allows simultaneous optimization of multiple algorithms (XGBoost, LightGBM, CatBoost) with their specific hyperparameters, using pruning to quickly eliminate underperforming configurations and focus on the most promising approaches. Multi-objective optimization also allows finding the best trade-off between precision and inference time.

3. Tuning Complex scikit-learn Pipelines

For traditional machine learning pipelines combining preprocessing (imputation, encoding, normalization) and modeling, Optuna allows end-to-end optimization of the entire pipeline. The define-by-run nature allows dynamically choosing the type of preprocessing based on the data, then adjusting the model hyperparameters accordingly, creating rich and expressive conditional search spaces.

4. Multi-Objective Optimization for Production Deployment

In production, the goal is not just the best accuracy, but the best trade-off between accuracy, inference latency, model size, and memory consumption. Optuna natively supports multi-objective optimization with NSGA-II, producing a Pareto front of solutions that allows engineers to choose the model best suited to specific production constraints.

See Also

• Mastering Geometry with Python: Creating and Analyzing Concave Triangles
• Exploring Polar Polygons with Python: Complete Guide and Tips for Developers