How to use Hyperopt or Optuna for Bayesian hyperparameter optimization on Databricks

Summary

• Hyperopt and Optuna use Bayesian optimization with Bayesian estimation to efficiently explore hyperparameter space and find good parameters with fewer trials.
• Hyperopt's Tree-structured Parzen Estimator (TPE) models the relationship between hyperparameters and performance, guiding the search toward promising regions.
• Optuna adds pruning, which stops poorly-performing trials early, cutting wall-clock time by stopping clearly bad configurations mid-training.
• Both libraries integrate with Databricks clusters, running trials in parallel across workers, and log results to MLflow for comparison.

How to use Hyperopt or Optuna for Bayesian hyperparameter optimization on Databricks

Hyperopt and Optuna use Bayesian optimization with probabilistic models to efficiently explore hyperparameter space and find good parameters with fewer trials. Hyperopt's Tree-structured Parzen Estimator (TPE) models the relationship between hyperparameters and performance. Optuna adds pruning, which stops poorly-performing trials early. Both libraries integrate with Databricks clusters, running trials in parallel across workers, and log results to MLflow for comparison and reproducibility.

Why Bayesian optimization improves on grid and random search

Grid search tests all combinations of a predefined set of values. This misses good parameter combinations between grid points. Random search tests random combinations without guidance. Bayesian optimization models the relationship between hyperparameters and model performance using a probabilistic model, then uses that model to decide which parameters to test next. This concentrates the search on the most promising regions.

Key capabilities to evaluate

Capability	Why it matters
Tree-structured Parzen Estimator (TPE)	Hyperopt's algorithm; models good vs. bad hyperparameter combinations efficiently.
Optuna pruning	Stops trials that are clearly underperforming without waiting for them to finish.
Distributed trial execution	Run multiple trials on different workers simultaneously, reducing total time.
MLflow logging	Each trial logs hyperparameters, metrics, and model; query MLflow to find the best trial.
Cross-validation per trial	Evaluate each candidate using k-fold cross-validation on training data to avoid overfitting.

How Databricks approaches Bayesian optimization with Hyperopt

Define an objective function that trains a model and returns loss. Hyperopt calls this function with different hyperparameters, learning as it goes:

from hyperopt import hp, fmin, tpe, Trials, STATUS_OK
import mlflow

def objective(params):
    with mlflow.start_run(nested=True):
        learning_rate = params['learning_rate']
        max_depth = int(params['max_depth'])
        n_estimators = int(params['n_estimators'])

        model = xgb.XGBClassifier(
            learning_rate=learning_rate,
            max_depth=max_depth,
            n_estimators=n_estimators
        )
        scores = cross_val_score(model, X_train, y_train, cv=5, scoring='roc_auc')

        mlflow.log_params({
            'learning_rate': learning_rate,
            'max_depth': max_depth,
            'n_estimators': n_estimators
        })
        mlflow.log_metric('cv_roc_auc', scores.mean())

        return {
            'loss': 1 - scores.mean(),
            'status': STATUS_OK
        }

space = {
    'learning_rate': hp.loguniform('learning_rate', -3, 0),
    'max_depth': hp.quniform('max_depth', 2, 10, 1),
    'n_estimators': hp.quniform('n_estimators', 50, 500, 10)
}

trials = Trials()
best = fmin(
    fn=objective,
    space=space,
    algo=tpe.suggest,
    max_evals=100,
    trials=trials,
    verbose=1
)

Hyperopt uses TPE to propose new hyperparameters, calls your objective function, observes the loss, and updates its model. After many trials, the best hyperparameters are available in best.
For Optuna with pruning:

import optuna
from optuna.pruners import MedianPruner
import mlflow

def objective(trial):
    learning_rate = trial.suggest_float('learning_rate', 0.001, 1, log=True)
    max_depth = trial.suggest_int('max_depth', 2, 10)
    n_estimators = trial.suggest_int('n_estimators', 50, 500, step=10)

    model = xgb.XGBClassifier(
        learning_rate=learning_rate,
        max_depth=max_depth,
        n_estimators=n_estimators
    )
    scores = cross_val_score(model, X_train, y_train, cv=5, scoring='roc_auc')

    trial.report(scores.mean(), step=0)
    if trial.should_prune():
        raise optuna.TrialPruned()

    with mlflow.start_run(nested=True):
        mlflow.log_params({
            'learning_rate': learning_rate,
            'max_depth': max_depth,
            'n_estimators': n_estimators
        })
        mlflow.log_metric('cv_roc_auc', scores.mean())

    return scores.mean()

study = optuna.create_study(
    direction='maximize',
    pruner=MedianPruner()
)
study.optimize(objective, n_trials=100)
best_trial = study.best_trial

Optuna's pruner stops trials that are clearly underperforming, saving compute time.

Use cases

A fraud detection team has labeled transactions and wants to tune an XGBoost classifier. Testing many hyperparameter combinations sequentially would take considerable time. They use Hyperopt on a Databricks cluster with multiple workers. Hyperopt runs trials in parallel. After a reasonable duration, many trials have completed. Hyperopt has explored combinations and identified good hyperparameters. The best trial achieves strong cross-validation performance. The team checks MLflow to see which hyperparameters were best.
An e-commerce company uses Optuna to tune a random forest for recommendation ranking. With pruning enabled, Optuna stops clearly bad trials early. Fewer trials run to completion than without pruning, cutting tuning time significantly. The best model uses settings that give strong ranking performance. They log all trials to MLflow and can see which hyperparameters had the most impact.

FAQs

How do I choose between Hyperopt and Optuna?

Both are excellent. Hyperopt's TPE is well-established and proven. Optuna has better pruning and a cleaner API. Try both on your problem; they're similar in performance.

Can I parallelize Hyperopt or Optuna across a cluster?

Yes. Use Hyperopt with a Spark backend or Optuna's distributed engine. Each trial runs on a different worker, and trials communicate results back to the central optimizer.

What if the search space is very large?

Use prior knowledge to constrain the search space. For example, if you know learning rates should be between 0.001 and 0.1, specify that range. A smaller space converges faster.