Reduce the time complexity of HSSP 2d from O(NK^2 log K) to O((N - K)K)

[nabenabe0928]

Motivation

As the bottleneck of MO-TPE is HSSP, I enhanced the time complexity of HSSP for 2D from $O(NK^2 \log K)$ to $O((N-K)K)$ where $N$ is the number of trials to be considered and $K$ is the subset size.
I further speeded up the implementation by using numpy vectorization.

Description of the changes

• The algorithm of HSSP for 2d setups, and
• The speedup by numpy vectorization.

Another breaking change is in the reproducibility when study has inf, -inf, and nan.
However, this is actually a bug in the original implementation because the original implementation gets nan for inf hypervolume contribution.
It happened because hypervolume contributions were incrementally calculated by doing inf - inf = nan.
Meanwhile, the new implementation correctly calculates inf hypervolume contributions.
You can check such behavior using:

import numpy as np


solutions = [[-np.inf, np.inf], [0, 1], [1, 0], [np.inf, -np.inf]]

In this example, the reference point is [np.inf, np.inf], so the hypervolumes for solutions[0] and solutions[3] cannot be computed because they exhibit nan from (np.inf - np.inf) * np.inf.
On the other hand, the hypervolumes for solutions[1] and solutions[2] are inf.

Verification Code

import time

import numpy as np

from optuna._hypervolume.hssp import _solve_hssp


def _is_pareto_front_2d(unique_lexsorted_loss_values: np.ndarray) -> np.ndarray:
    n_trials = unique_lexsorted_loss_values.shape[0]
    cummin_value1 = np.minimum.accumulate(unique_lexsorted_loss_values[:, 1])
    on_front = np.ones(n_trials, dtype=bool)
    on_front[1:] = cummin_value1[1:] < cummin_value1[:-1]
    return on_front


def _is_pareto_front_for_unique_sorted(unique_lexsorted_loss_values: np.ndarray) -> np.ndarray:
    (n_trials, n_objectives) = unique_lexsorted_loss_values.shape
    if n_objectives == 1:
        on_front = np.zeros(len(unique_lexsorted_loss_values), dtype=bool)
        on_front[0] = True
        return on_front
    elif n_objectives == 2:
        return _is_pareto_front_2d(unique_lexsorted_loss_values)
    else:
        return _is_pareto_front_nd(unique_lexsorted_loss_values)


def is_pareto_front(loss_values: np.ndarray, assume_unique_lexsorted: bool = True) -> np.ndarray:
    if assume_unique_lexsorted:
        return _is_pareto_front_for_unique_sorted(loss_values)

    unique_lexsorted_loss_values, order_inv = np.unique(loss_values, axis=0, return_inverse=True)
    on_front = _is_pareto_front_for_unique_sorted(unique_lexsorted_loss_values)
    return on_front[order_inv]


def _solve_hssp_2d(
    rank_i_loss_vals: np.ndarray,
    rank_i_indices: np.ndarray,
    subset_size: int,
    reference_point: np.ndarray,
) -> np.ndarray:
    # The time complexity is O(subset_size * rank_i_loss_vals.shape[0]).
    assert rank_i_loss_vals.shape[-1] == 2 and subset_size <= rank_i_loss_vals.shape[0]
    n_trials = rank_i_loss_vals.shape[0]
    order = np.argsort(rank_i_loss_vals[:, 0])
    sorted_loss_vals = rank_i_loss_vals[order]
    rectangle_diagonal_points = np.repeat(reference_point[np.newaxis, :], n_trials, axis=0)

    subset_indices = np.zeros(subset_size, dtype=int)
    for i in range(subset_size):
        contribs = np.prod(rectangle_diagonal_points - sorted_loss_vals, axis=-1)
        max_index = np.argmax(contribs)
        subset_indices[i] = rank_i_indices[order[max_index]]
        rectangle_diagonal_points[: max_index + 1, 0] = np.minimum(
            sorted_loss_vals[max_index, 0], rectangle_diagonal_points[: max_index + 1, 0]
        )
        rectangle_diagonal_points[max_index:, 1] = np.minimum(
            sorted_loss_vals[max_index, 1], rectangle_diagonal_points[max_index:, 1]
        )

        keep = np.ones(n_trials - i, dtype=bool)
        keep[max_index] = False
        # Remove the chosen point.
        order = order[keep]
        rectangle_diagonal_points = rectangle_diagonal_points[keep]
        sorted_loss_vals = sorted_loss_vals[keep]

    return subset_indices


if __name__ == "__main__":
    rng = np.random.RandomState()
    for _ in range(1000):
        Y = rng.normal(size=(10000, 2)) + 10.0
        Y = np.unique(Y, axis=0)
        on_front = is_pareto_front(loss_values=Y)
        Y_pareto = Y[on_front]
        reference_point = np.max(Y, axis=0) * 1.1
        n_sols = Y_pareto.shape[0]

        start = time.time()
        ans = _solve_hssp(Y_pareto.copy(), np.arange(n_sols), int(n_sols * 0.3), reference_point.copy())
        runtime1 = time.time() - start

        start = time.time()
        output = _solve_hssp_2d(Y_pareto.copy(), np.arange(n_sols), int(n_sols * 0.3), reference_point.copy())
        runtime2 = time.time() - start
        print(runtime1 / runtime2)
        assert np.all(ans == output)

The figure below shows the solutions ($v_i$ where $v_{1,1} \leq v_{2,1} \leq v_{3,1} \leq v_{4,1}$ and $v_{1,2} \geq v_{2,2} \geq v_{3,2} \geq v_{4,2}$. Recall that $\{v_i\}$ is a non-dominated set) and their diagonal points ($d_i$).
Note that I defined $d_i$ (pink dots) so that the hypervolume contribution (the gray rectangular) of the $i$-th solution becomes $|v_i - d_i |^2$ given a current state.
In the first figure, we have not picked any point.
In the second figure, we picked $v_2$ as a member of the subset.
Then the diagonal point for $v_1$ becomes $d_1 \leftarrow [\min(v_{2,1}, d_{1,1}), d_{1, 2}]$ and the diagonal points for $v_3, v_4$ become $d_3 \leftarrow [d_{3,1}, \min(v_{2, 2}, d_{3,2})]$ and $d_4 \leftarrow [d_{4,1}, \min(v_{2, 2}, d_{4,2})]$.
In the third figure, we picked $v_3$ as another member of the subset.
Then the diagonal point for $v_1$ becomes $d_1 \leftarrow [\min(v_{3,1}, d_{1,1}), d_{1, 2}]$ and the diagonal point for $v_4$ becomes $d_4 \leftarrow [d_{4,1}, \min(v_{3, 2}, d_{4,2})]$.

In principle, when $v_i$ is picked, $d_j$ will be updated as $d_j \leftarrow [\min(v_{i,1}, d_{j,1}), d_{j,2}]$ for $j < i$ and $d_j \leftarrow [d_{j,1}, \min(v_{i,2}, d_{j,2})]$ for $j > i$.

These updates can be done by $O(N)$ and we need to repeat it $K$ times, so the time complexity is $O(NK)$.

[nabenabe0928]

Benchmarking using MO-TPE:

Benchmarking Code

from __future__ import annotations

import time

import optuna


def objective2d(trial: optuna.Trial) -> tuple[float, float]:
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    return x ** 2 + y ** 2, (x - 2) ** 2 + (y - 2) ** 2


if __name__ == "__main__":
    start = time.time()
    # optuna.logging.set_verbosity(optuna.logging.CRITICAL)
    n_trials = 1000
    sampler = optuna.samplers.TPESampler(seed=42)
    study = optuna.create_study(sampler=sampler, directions=["minimize"]*2)
    study.optimize(objective2d, n_trials=n_trials)
    print(f"n_obj=2, {study.sampler}", time.time() - start)

The results are here:

This PR	Original
53	121

Note that the unit of the table is seconds.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 89.66%. Comparing base (181d65f) to head (75d7b6d).
Report is 90 commits behind head on master.

❗ Current head 75d7b6d differs from pull request most recent head 94a5770. Consider uploading reports for the commit 94a5770 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5346      +/-   ##
==========================================
+ Coverage   89.52%   89.66%   +0.13%     
==========================================
  Files         194      210      +16     
  Lines       12626    13309     +683     
==========================================
+ Hits        11303    11933     +630     
- Misses       1323     1376      +53     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[eukaryo]

@gen740 @contramundum53 Could you review this PR?

[nabenabe0928]

@contramundum53 I unassigned you as discussed online for now!

[github-actions]

This pull request has not seen any recent activity.

[HideakiImamura]

@eukaryo Could you review this PR?

[eukaryo]

I have interpreted HSSP as referring to the Hypervolume Subset Selection Problem.

[nabenabe0928]

I have interpreted HSSP as referring to the Hypervolume Subset Selection Problem.

Yes, HSSP is the Hypervolume Subset Selection Problem.

[eukaryo]

LGTM. I have reviewed the algorithm and your implementation. The logic is sound and the code is well-written.
While the inherent complexity of the algorithm naturally leads to challenging code, your implementation is clear and relatively easy to follow.

[github-actions]

This pull request has not seen any recent activity.

[eukaryo]

@not522 Could you review this PR?

[github-actions]

This pull request has not seen any recent activity.

[github-actions]

This pull request has not seen any recent activity.