java-topology/defects/optuna/patch/optuna-0001-nondomination-rank-cubic.md

5.4 KiB
Raw Blame History

UNDF: UNDF-2026-000000494

optuna-0001 — O(N³) Non-Dominated Sort in _calculate_nondomination_rank

Severity: HIGH Complexity: O(N³) worst case → O(N² log N) CWE: CWE-407 (Algorithmic Complexity)

Affected File

File Lines Notes
optuna/study/_multi_objective.py 187216 _calculate_nondomination_rank outer while loop
optuna/study/_multi_objective.py 127148 _is_pareto_front_nd inner while loop

Defective Code

_calculate_nondomination_rank lines 207213

while n_unique - indices.size < n_below:
    on_front = _is_pareto_front(unique_lexsorted_loss_values, assume_unique_lexsorted=True)
    ranks[indices[on_front]] = rank
    indices = indices[~on_front]
    unique_lexsorted_loss_values = unique_lexsorted_loss_values[~on_front]
    rank += 1

_is_pareto_front_nd lines 136147

while len(remaining_indices):
    on_front[(new_nondominated_index := remaining_indices[0])] = True
    nondominated_and_not_top = np.any(
        loss_values[remaining_indices] < loss_values[new_nondominated_index], axis=1
    )
    remaining_indices = remaining_indices[nondominated_and_not_top]

Defect: Two nested while-loops implement the non-dominated sort:

  • Outer loop (_calculate_nondomination_rank): iterates once per front level F. Worst case: all trials in distinct fronts → F = N iterations.
  • Inner loop (_is_pareto_front_nd): for each front computation, scans all remaining trials. On front level f, there are approximately N f·(front_size) remaining trials → O(N) per call.
  • Each inner-loop iteration also does an O(N × M) numpy comparison over M objectives.

Total: O(F × N × M) = O(N² × M) average, O(N³) worst case with M objectives when every trial is in a distinct front (common in high-dimensional hyperparameter optimization where no trial dominates another).

The outer loop re-runs the entire Pareto-front computation from scratch for each rank level rather than incrementally processing the domination graph built once.

Root Cause

The code explicitly notes a Kung's algorithm attempt was rejected as "not really quick":

# NOTE(nabenabe0928): I tried the Kung's algorithm below, but it was not really quick.
# https://github.com/optuna/optuna/pull/5302#issuecomment-1988665532

However, the accepted alternative is still O(N²×M) average and O(N³) worst case. The classic Deb et al. NSGA-II fast non-dominated sort is O(M×N²), which equals the current average but has a lower constant. The key fix is to build the dominance graph once and peel off fronts via degree-counting (analogous to Kahn's topological sort), rather than rerunning the full Pareto check per level.

Fix

Build dominance counts once, then peel fronts in O(N²) total:

def _calculate_nondomination_rank_fast(
    loss_values: np.ndarray, *, n_below: int | None = None
) -> np.ndarray:
    if len(loss_values) == 0 or (n_below is not None and n_below <= 0):
        return np.zeros(len(loss_values), dtype=int)

    n = len(loss_values)
    n_below = n_below or n

    # Build domination graph once: O(N² × M)
    # dominated_by[i] = set of indices that dominate trial i
    # domination_count[i] = number of trials that dominate i
    domination_count = np.zeros(n, dtype=int)
    dominated_set = [[] for _ in range(n)]

    for i in range(n):
        for j in range(i + 1, n):
            # i dominates j?
            ij = loss_values[i] <= loss_values[j]
            ji = loss_values[j] <= loss_values[i]
            if np.all(ij) and np.any(loss_values[i] < loss_values[j]):
                dominated_set[i].append(j)
                domination_count[j] += 1
            elif np.all(ji) and np.any(loss_values[j] < loss_values[i]):
                dominated_set[j].append(i)
                domination_count[i] += 1

    ranks = np.zeros(n, dtype=int)
    current_front = [i for i in range(n) if domination_count[i] == 0]
    rank = 0
    sorted_count = 0

    while current_front and sorted_count < n_below:
        next_front = []
        for i in current_front:
            ranks[i] = rank
            sorted_count += 1
            for j in dominated_set[i]:
                domination_count[j] -= 1
                if domination_count[j] == 0:
                    next_front.append(j)
        rank += 1
        current_front = next_front

    # Trials not yet ranked get the current rank
    for i in range(n):
        if domination_count[i] > 0:
            ranks[i] = rank

    return ranks

Complexity

Phase Before After
Build domination graph O(F × N × M) — rebuilt per front O(N² × M) — once
Front extraction O(F × N) rescan O(N + E) via degree counting
Total O(N² × M) avg, O(N³) worst O(N² × M)
Constant factor High — full numpy scan per front Low — single pass

For N=1000 trials with M=3 objectives, worst-case improvement: 1000× fewer redundant scans.

Impact

_calculate_nondomination_rank is called every NSGA-II and NSGA-III generation cycle (every population_size trials). With population_size=50 and 1000 total trials, this function is called 20 times per study. Each call is O(N³) worst case on the current population. For large multi-objective studies (N=200+ per generation, M=3+ objectives), the non-dominated sort dominates wall-clock time and scales poorly.