Introducing Timeline Plot: a new feature of Optuna v3.2

The newest version of Optuna, Optuna v3.2, an open source software tool for black box optimization, was released at the end of May. This article introduces one of its new features, Timeline Plot.

Timeline Plot is a type of plot that displays each trial as a horizontal bar, with the time on the x-axis and the trial number on the y-axis, as shown in the figure below.

hands-on

Here is an example for a hands-on exercise: let us use the sample classification dataset provided by scikit-learn to perform machine learning with LightGBM. The Python source code is as follows:

import lightgbm as lgb
import numpy as np
import optuna
import sklearn.datasets
import sklearn.metrics
from sklearn.model_selection import train_test_split

data, target = sklearn.datasets.load_breast_cancer(return_X_y=True)
train_x, val_x, train_y, val_y = train_test_split(
    data, target, test_size=0.25
)


def objective(trial):
    dtrain = lgb.Dataset(train_x, label=train_y)
    param = {
        "objective": "binary",
        "metric": "binary_logloss",
        "boosting_type": "gbdt",
        "lambda_l1": trial.suggest_float("lambda_l1", 1e-8, 10.0, log=True),
        "lambda_l2": trial.suggest_float("lambda_l2", 1e-8, 10.0, log=True),
        "num_leaves": trial.suggest_int("num_leaves", 2, 256),
        "feature_fraction": trial.suggest_float("feature_fraction", 0.4, 1.0),
        "bagging_fraction": trial.suggest_float("bagging_fraction", 0.4, 1.0),
        "bagging_freq": trial.suggest_int("bagging_freq", 1, 7),
        "min_child_samples": trial.suggest_int("min_child_samples", 5, 100),
    }
    gbm = lgb.train(param, dtrain)
    preds = gbm.predict(val_x)
    pred_labels = np.rint(preds)
    return sklearn.metrics.accuracy_score(val_y, pred_labels)


study = optuna.create_study(direction="maximize")
study.optimize(objective, n_jobs=4, n_trials=50)
fig = optuna.visualization.plot_timeline(study)
fig.show()

Executing the above source code will output the following figure.

From the Timeline Plot, it can be observed that the computations for trials after the 15th trial took relatively longer time. Optuna recognizes that values of certain parameters of LightGBM that take longer to calculate tend to result in higher accuracy, so such trials are likely to increase in number from the middle of the search process. Specifically, increasing the value of min_child_samples parameter can lengthen the computation time.

• As an aside, in the aforementioned code, the dataset was partitioned into two subsets (for training and validation) and then optimized across 50 trials. Were this to be increased to, for instance, 1000 trials, the validation accuracy would likely improve; however, the generalization performance for test data may deteriorate. This potential outcome indicates that the hyperparameters could overfit the validation dataset. To mitigate this issue, a new feature dubbed the ‘Optuna Terminator’, which was introduced in Optuna v3.2. Optuna Terminator can be utilized to halt the optimization early, hence curtailing overfitting and simultaneously reducing computation time.

Situations where the Timeline Plot is particularly effective

Monitoring trials in distributed environments

Timeline Plot is particularly effective in visualizing issues that often arise in distributed environments. When running distributed optimization in a shared cluster, Optuna workers are often preempted, causing the computation to be forcibly terminated midway. If Optuna is not notified about the worker being preempted, the status of the trial that was being worked on at the time of preemption will remain ‘Running’ indefinitely. Of course, there could also be cases where it genuinely remains ‘Running’ indefinitely, such as if there is a mistake in the implementation of the evaluation function causing an infinite loop. Such mistakes are easy to overlook in a distributed environment, but they can be clearly visualized using a Timeline Plot.

The Timeline Plot function can accept an ongoing study as an argument. This feature enables quick detection of any anomalous trials, facilitating the initiation of bug fixes or introducing solutions to these problems, such as implementing a ‘heartbeat’ to detect preemption.

An example of visualizing anomalous Trials. The trial indicated by the yellow arrow persists in the ‘RUNNING’ state. In this scenario, even though initial parallel computations launched two independent trials, the throughput was reduced by half for all subsequent trials, following the one that became stalled.

Checking trial details interactively

Timeline Plot can visualize delays or failures on a per-trial basis. Furthermore, in the case of the plotly version, you can hover the mouse cursor over the horizontal bars representing each trial to view the parameters and evaluation function values of that trial as hover text. This feature would be of great help in quickly eliminating difficult bugs, such as ones causing infinite loops only when a parameter has a certain value.

This is a screenshot of the Timeline Plot. The states of trials are displayed according to the sequence of trial numbers. The red bubble in the center represents hover text, which reveals the parameter and evaluation function value (shown as null due to an error) for a single trial when you hover the mouse cursor over it.

Timeline Plot in Optuna Dashboard

Timeline Plot is available in Optuna Dashboard v0.10.0, as shown in the screenshot below. As an example of the functionality that can be leveraged through this combination, if a trial that remains in the ‘Running’ status exists, it can be easily and promptly identified and removed using the dashboard.

Conclusions

In this article, I introduced the Timeline Plot added in Optuna v3.2. The Timeline Plot can play a crucial role in troubleshooting in distributed environments and tuning machine learning hyperparameters that require trial and error. It is particularly useful for promptly identifying the presence of anomalous trials and parameters of noteworthy trials.

In addition to this, other useful features were added in Optuna v3.2. The Optuna Terminator is a feature for early stopping optimization, helping to reduce computation time and prevent overfitting.

By using new features for Optuna, you can expect improvements in processing speed for optimization tasks in machine learning, as well as the discovery of better solutions.