Motivation

The multi-start optimization of the acquisition function in the GPSampler is a significant performance bottleneck, accounting for 50-60% of its total execution time. The current implementation uses 10 initial points (n_local_search=10) and performs optimization from each point sequentially.

This pull request aims to significantly speed up this multi-start optimization by batching the evaluation of the acquisition function and its gradient. Benchmark results show that this change can reduce the execution time by approximately 40-50% without compromising the optimization accuracy of the objective function.

Description of the Changes

Previous Sequential Approach

The previous implementation runs an optimization loop for each starting point individually.

# Previous pseudo-code
def multistart_optimize(acqf, x0_list, ...):
    ...
    min_fvals = ...
    min_xs = ...
    for i, x0 in enumerate(x0_list):
        # optimize() performs iterative optimization for a single initial point x0
        min_fval, min_x = optimize(acqf, x0, ...)
        min_fvals[i] = min_fval
        min_xs[i] = min_x
    best_idx = np.argmin(min_fvals)
    return min_xs[best_idx], min_fvals[best_idx]

Proposed Batched Approach

The proposed approach executes each optimization step for all initial points in a single batch.

Evaluations of the acquisition function (acqf) and its gradient (acqf_grad) are performed at once using vectorized operations on torch.Tensor. The iterative optimization from each starting point is managed cooperatively using greenlet. This enables a cooperative, iterative workflow. In each iteration, the function values for all starting points are evaluated as a single batch. Subsequently, each optimization is advanced by one step to determine the next set of points to be evaluated, which then forms the input for the following batch evaluation cycle.

# New pseudo-code
def multistart_optimize(acqf, x0_list, ...):
    ...
    # optimize_batched manages optimization for multiple initial points x0_list
    min_xs, min_fvals = optimize_batched(acqf, x0_list, ...)
    best_idx = np.argmin(min_fvals)
    return min_xs[best_idx], min_fvals[best_idx]

For backward compatibility, a fallback mechanism has been implemented. If greenlet fails to import, the logic will revert to the original sequential execution.

Dependencies

This change utilizes greenlet, which is already included as a dependency of sqlalchemy, a core Optuna dependency. Therefore, this PR does not introduce any new library dependencies.

Performance Evaluation (Benchmark)

To validate the effectiveness of the proposed method, we compared the performance in the following three modes:

1. This PR (Batched): The implementation in this PR with batch evaluation enabled.
2. original: The current implementation on the master branch.
3. This PR (Fallback): The implementation in this PR with batch evaluation intentionally disabled to activate the sequential fallback logic. This helps confirm that the new implementation's overhead is acceptable and that the implementation is functioning correctly.

Experimental Setup

• CPU: Apple M4 (10-core)
• Trials: 300 trials per study
• Seed: 42, 43, 44 (results are averaged over three runs)
• Objective Functions:
  • rotated_ellipsoid (10 dimensions)
  • f6 from optunahub/benchmarks/bbob (20 dimensions)
  • wfg (2 objectives, 3 dimensions, k=1) from optunahub/benchmarks/wfg

Results

The following plot show that the hypervolume values remain consistent across all modes, confirming that the change does not negatively impact search performance.

Analysis

• The implementation with batched evaluation is approximately 1.8x to 2.0x faster than the current implementation.
• The best objective values obtained are comparable across all modes. The observed variations are likely due to changes in the order of floating-point operations and do not indicate a degradation in the sampler's search performance.
• The fallback mode is as fast as the original implementation, suggesting that the structural overhead of the proposed changes is minimal.

Objective Value Progression:

f6 (D=20, seed=42)

The following tables show that the objective values remain consistent across all modes, confirming that the change does not negatively impact search performance.

rotated_ellipsoid (D=10, seed=42)

import json
import os
from itertools import product

import numpy as np
import optunahub

import optuna

np.random.seed(42)

NUM_CONTINUOUS_PARAMS = 10
ROTATION_MATRIX, _ = np.linalg.qr(np.random.rand(NUM_CONTINUOUS_PARAMS, NUM_CONTINUOUS_PARAMS))

CONDITIONING = 10
WEIGHTS = np.array([CONDITIONING**(i / (NUM_CONTINUOUS_PARAMS - 1)) for i in range(NUM_CONTINUOUS_PARAMS)])

def rotated_ellipsoid(trial):
    xs = np.array([trial.suggest_float(f"x_{i}", -1, 1) for i in range(NUM_CONTINUOUS_PARAMS)])
    rotated_xs = ROTATION_MATRIX @ xs
    return np.sum(WEIGHTS * (rotated_xs**2))

bbob = optunahub.load_module("benchmarks/bbob")
f6 = bbob.Problem(function_id=6, dimension=20, instance_id=1)

def execute_benchmark(
    mode: str,
    objective_type: str,
    n_trials: int,
    seed: int,
    results_file="results.jsonl",
    output_dir="output",
):
    sampler = optuna.samplers.GPSampler(seed=seed)
    name = f"{objective_type}_{seed}_{mode}_{n_trials}"
    log_file = f"{name}.jsonl"
    study = optuna.create_study(study_name=name, sampler=sampler)

    study.optimize(func=f6 if objective_type == "f6" else rotated_ellipsoid, n_trials=n_trials)
    print(study.best_trial.params, study.best_trial.value)
    elapsed = (study.trials[-1].datetime_complete - study.trials[0].datetime_start).total_seconds() # type: ignore
    print(f"{mode} took {elapsed:f} seconds. ")

    result = {
        "objective_type": objective_type,
        "seed": seed,
        "mode": mode,
        "elapsed": round(elapsed, 2),
        "n_trials": n_trials,
        "best_value": study.best_trial.value,
    }

    with open(os.path.join(output_dir, results_file), "a") as f:
        f.write(json.dumps(result) + "\n")

seeds = [42,43,44]
n_trials = 300

# To test different modes, you need to edit the source code of the GPSampler directly.
# This script runs with the batched evaluation enabled by default.
modes: list[str] = [
    "This PR (Batched)",
    # "This PR (Fallback)",
    # "original",
]

objective_types = [
    "rotated_ellipsoid",
    "f6",
]

for seed, mode, objective_type in product(seeds, modes, objective_types):
    execute_benchmark(
        mode=mode,
        objective_type=objective_type,
        n_trials=n_trials,
        seed=seed,
    )

Tests

Unit tests must be added.