Announcing Optuna 3.6

mamu
Optuna
Published in
8 min readMar 18, 2024

--

Introduction

We have released v3.6, the latest version of the black-box optimization framework, Optuna. This release includes various new features, refactoring, and bug fixes. In this blog, we will be covering the highlights of v3.6 and many feature improvements.

TL;DR

  • Support for various new algorithms such as the Wilcoxon pruner, Gaussian process (GP) based sampler, and PED-ANOVA importance evaluator
  • Implement various improvements related to Optuna’s quality, such as stricter verification logic for FrozenTrial, refactoring the Dashboard, and migration of Integration

Wilcoxon Pruner

Before Optuna v3.6, Optuna pruners assumed machine learning-like applications. In such applications, the learning curve could be used for early stopping bad trials, but pruners always had to consider the possibility that slow-starting trials can actually end up better.

We introduce a different type of pruner `WilcoxonPruner` in Optuna v3.6, that works well in applications where the objective function is the mean/median, etc. of independently evaluated scores of many inputs and the evaluations can be stopped early for bad parameters.

This pruner can be used for objective functions like mean performance of a heuristic algorithm (like simulated annealing and SAT solver), the k-fold cross validation score of a machine learning model, and even mean performance of a prompt given to a large language model.

In such cases, the evaluation score for each input can be considered independent, and `WilcoxonPruner` can prune trials more aggressively by using Wilcoxon signed-rank test.

The following example demonstrates how to use `WilcoxonPruner`:

import optuna
import numpy as np

# We minimize the mean evaluation loss over all the problem instances.
def evaluate(param, instance):
return (param - problem_instance) ** 2

problem_instances = np.linspace(-1, 1, 100)

def objective(trial):
# Sample a parameter.
param = trial.suggest_float("param", 0, 1)

# Evaluate performance of the parameter.
results = []
# For best results, shuffle the evaluation order in each trial.
instance_ids = np.random.permutation(len(problem_instances))
for instance_id in instance_ids:
loss = evaluate(param, problem_instances[instance_id])
results.append(loss)

# Report loss together with the instance id.
# You need to pass the same id for the same instance.
trial.report(loss, instance_id)

if trial.should_prune():
# Return the current predicted value instead of raising `TrialPruned`.
# This is a workaround to tell the sampler to utilize the pruned trials.
return sum(results) / len(results)

return sum(results) / len(results)

# Higher p_threshold means trials are pruned more aggressively.
study = optuna.create_study(pruner=optuna.pruners.WilcoxonPruner(p_threshold=0.1))
study.optimize(objective, n_trials=100)

For a more realistic example, please see this tutorial of optimizing hyperparameters for simulated annealing.

Light-Weight GP-Based Sampler

Optuna v3.6 natively supports Gaussian process-based Bayesian optimization (GP-BO) through `optuna.samplers.GPSampler`. While we had previously supported GP-BO by `optuna.integration.BoTorchSampler`, it had limitations mainly coming from heavily relying on BoTorch. The newly introduced `GPSampler` has the following advantages over `BoTorchSampler`:

  • Better performance for discrete/mixed search spaces. `BoTorchSampler` basically only worked well on continuous search spaces, and the new sampler finally got to fix this issue.
  • Faster computation. `BoTorchSampler` was slow due to heavy abstractions in the internal implementation. In the new sampler, we kept the internal design simple, and the computation is ~5x faster than `BoTorchSampler`.
  • Lighter dependency. `BoTorchSampler` had a large list of transitive dependencies, including BoTorch, GPyTorch, Pyro, and so on. This often caused installation problems, especially for restrictive environments. The new sampler only depends on PyTorch and SciPy, so we expect it to work on more environments.
  • Ability to explicitly tell the sampler that the objective function is deterministic. Other Optuna samplers can repeatedly suggest the same parameters when the search space is fully discrete, and has been an issue for a while. The new `GPSampler` has a flag `deterministic_objective` and when set to `True`, the sampler takes that into account when fitting the surrogate model (thus the chance the same parameters are sampled repeatedly is minimized).

The basic usage of `GPSampler` is simple: just specify it when creating a new study.

study = optuna.create_study(sampler=optuna.samplers.GPSampler())

`BoTorchSampler` is still supported in `optuna_integration` module, and it has certain features (multi-objective optimization, `consider_running_trials` option, and constrained optimization) that `GPSampler` does not currently support. If you need those features, you can use `BoTorchSampler` by installing

$ pip install optuna-integration botorch

Speeding up Importance Evaluation with PED-ANOVA

Optuna has traditionally used a method called f-ANOVA to evaluate importance, but it has become a problem that the computation time becomes very long as the number of parameters and trials increases. Against this background, the Optuna Dashboard was using a Cython implementation of f-ANOVA.

In 2023, a faster algorithm than f-ANOVA, PED-ANOVA, was presented at the international conference for AI, IJCAI*1. This release introduces PED-ANOVA to Optuna as one of the solutions to resolve the issue of slow importance evaluation as the number of trials increases.

It can be used by passing the instantiated `optuna.importance.PedAnovaImportanceEvaluator` as shown below to `optuna.visualization.plot_param_importances`.

import optuna

def objective(trial) -> float:
x1 = trial.suggest_float("x1", -5, 5)
x2 = trial.suggest_float("x2", -5, 5)
return x1 ** 2 + x2 ** 2 / 1000

study = optuna.create_study()
study.optimize(objective, n_trials=30)

evaluator = optuna.importance.PedAnovaImportanceEvaluator()
optuna.visualization.plot_param_importances(
study, evaluator=evaluator
)

Here are the results of the speed benchmark. We compared the implemented PED-ANOVA, Optuna’s f-ANOVA, and the Cython implementation of f-ANOVA with various numbers of parameters and trials. The Cython implementation of f-ANOVA is faster than Optuna’s f-ANOVA, but it takes more than a minute when the number of parameters or trials is large. On the other hand, PED-ANOVA operates within a second in all settings, proving its speed.

*1: Watanabe et al., PED-ANOVA: Efficiently Quantifying Hyperparameter Importance in Arbitrary Subspaces. Proceedings of International Joint Conference on Artificial Intelligence 2023 (IJCAI’23). URL: https://www.ijcai.org/proceedings/2023/488

Stricter Verification Logic for FrozenTrial

Optuna has a class called `FrozenTrial`. This class is a type of `Trial` class, which represents a single attempt at optimization and is mainly used by users outside the objective function. The entire series of optimization processes is managed by the `Study` class, while the `FrozenTrial` represents a single trial not tied to a `Study`.

A `FrozenTrial` can be freely created with the `optuna.create_trial` function, and can be inserted into a `Study` with the `Study.add_trial` method. This feature can be used to visualize only the necessary trials, or to pre-insert known superior parameters and their evaluation values into the `Study` to provide hints to the algorithm (as in the code example below).

However, the previous `create_trial` function was also capable of creating abnormal `FrozenTrial`s that did not conform to the rules assumed in Optuna’s internal implementation. Better early detection of such cases was decided upon in this update by doing a more detailed verification of `FrozenTrial`. This is expected to reduce the occurrence of troublesome problems such as runtime errors and unintended degradation of optimization performance.

import optuna

def objective(trial: optuna.Trial) -> float:
x = trial.suggest_float("x", -1.0, 1.0)
return x**2.0

study = optuna.create_study()
trial = optuna.trial.create_trial(
params={"x": -0.2},
distributions={
"x": optuna.distributions.FloatDistribution(-1.0, 1.0),
},
value=0.04,
)
study.add_trial(trial)
study.enqueue_trial({"x": 0.1})
study.optimize(objective, n_trials=20)
print(f"Best value: {study.best_value} (params: {study.best_params})")

Refactoring the Optuna Dashboard

Our web dashboard, Optuna Dashboard, used to implement visualization features in the front-end using the JavaScript library Plotly. This was a reimplementation of Optuna’s Python visualization features, allowing for high interactivity but raising concerns about maintainability and consistency as more components were added.

To mitigate this issue, we have tentatively added an option to use Optuna’s visualization features on the back-end from Optuna Dashboard. This can be enabled by toggling on “Use Plotlypy” from the settings button at the bottom left of the dashboard, allowing users to get consistent output with Optuna’s visualization functions. Also, the calculations required for plotting the graph are done on the back-end, which can be beneficial in terms of lightening the calculations on the client side when running Optuna Dashboard on a server.

Currently, this option only provides simple functionalities and minimal interactivity. We would appreciate your feedback or any issues and PRs based on your experience using this feature.

Migration to Optuna Integration

Optuna provides a set of modules to conveniently use Optuna in combination with third-party libraries. Originally, these modules were introduced as `optuna.integration` in the Optuna core, but since they are not essential for the core operation, a partial migration to a separate package, `optuna-integration`, was started from v3.2 to reduce codebase bloating and burden on CI in the core. With the release of v3.6, we are pleased to announce that the migration of all modules to `optuna-integration` has been completed.

To use the modules that were migrated in v3.6, you will need to install `optuna-integration` as follows.

pip install optuna-integration

On the other hand, to minimize the impact on existing user codes in the migration, `optuna.integration.XXX` (where XXX is any module) remains as a thin wrapper referring to `optuna_integration.XXX`. As long as `optuna-integration` is installed in your environment, the traditional code that uses `optuna.integration` will work correctly without any error. An example is shown below.

# It is also possible to import from optuna_integration
from optuna_integration import OptunaSearchCV

# The following, as usual, is also OK, but will result in an error in an environment without optuna-integration.
from optuna.integration import OptunaSearchCV

Conclusion

Our committers are working diligently every day to make Optuna the best black-box optimization framework in the world. We are improving existing features, releasing various new features, and running ambitious projects such as implementation in Rust. Since Optuna is constantly evolving, we are always looking for new contributors. If you’re interested, please don’t hesitate to reach out to our committers on GitHub or participate in development events.

Over the next few weeks, we will be publishing blogs about experimental initiatives taken internally by committers, introducing new features, and showcasing application examples. We have some exciting content prepared, so stay tuned! Here’s a sneak peek:

  • Introduction of the Wilcoxon Pruner: its ideas, algorithm, and use cases
  • Details on the lightweight Gaussian process-based sampler
  • Application example of Wilcoxon pruner to the heuristic contest
  • Prototyping a fast and low-dependency Optuna using Rust

Contributors

This release was made possible by the authors and the people who participated in the reviews and discussions.

@Alnusjaponica, @DanielAvdar, @HarshitNagpal29, @HideakiImamura, @SimonPop, @adjeiv, @buruzaemon, @c-bata, @contramundum53, @dheemantha-bhat, @eukaryo, @gen740, @hrntsm, @knshnb, @nabenabe0928, @not522, @nzw0301, @porink0424, @ryota717, @shahpratham, @toshihikoyanase, @y0z

--

--