Skip to content

Cache pair-wise distances to speed up GPSampler#6244

Merged
nabenabe0928 merged 7 commits intooptuna:masterfrom
nabenabe0928:enhance/cache-pair-wise-distance
Sep 17, 2025
Merged

Cache pair-wise distances to speed up GPSampler#6244
nabenabe0928 merged 7 commits intooptuna:masterfrom
nabenabe0928:enhance/cache-pair-wise-distance

Conversation

@nabenabe0928
Copy link
Copy Markdown
Contributor

@nabenabe0928 nabenabe0928 commented Aug 13, 2025

Motivation

This PR speeds up GPSampler by skipping the pair-wise distance calculations in the kernel parameter fitting.
Please note that this PR depends on:

Important

7% speedup by this PR (240 sec to 223 sec) in the benchmark below 🎉

Description of the changes

  • Cache the pair-wise distances
import optuna


def objective(trial: optuna.Trial) -> float:
    return sum(trial.suggest_float(f"x{i}", -5, 5)**2 for i in range(5))


sampler = optuna.samplers.GPSampler(seed=0)
study = optuna.create_study(sampler=sampler)
study.optimize(objective, n_trials=300)
print((study.trials[-1].datetime_complete - study.trials[0].datetime_start).total_seconds())

@nabenabe0928 nabenabe0928 added the enhancement Change that does not break compatibility and not affect public interfaces, but improves performance. label Aug 13, 2025
@github-actions
Copy link
Copy Markdown
Contributor

This pull request has not seen any recent activity.

@github-actions github-actions bot added the stale Exempt from stale bot labeling. label Aug 21, 2025
@nabenabe0928 nabenabe0928 removed the stale Exempt from stale bot labeling. label Aug 22, 2025
@nabenabe0928 nabenabe0928 marked this pull request as ready for review August 29, 2025 09:34
@nabenabe0928
Copy link
Copy Markdown
Contributor Author

@sawa3030 @kAIto47802
Could you review this PR?

optuna/_gp/gp.py Outdated
else:
if X2 is None:
X2 = self._X_train
d2 = (X1[..., None, :] - X2[..., None, :, :]) ** 2
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit confused about this part. Would it be more straightforward to write it as d2 = (X1[..., :, None, :] - X2[..., None, :, :]) ** 2?

Copy link
Copy Markdown
Contributor Author

@nabenabe0928 nabenabe0928 Sep 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for pointing it out:)
X1[..., :, None, :] - X2[..., None, :, :] first requires X1.shape[-1] == X2.shape[-1] and len(X1.shape) >= 2 and len(X2.shape) >= 2.

The newer version loosens the requirements to X1.shape[-1] == X2.shape[-1] and len(X1.shape) >= 1 and len(X2.shape) >= 2.

It essentially makes the following line look cleaner:

# Original (x must be 2 dimensions or more, requiring [..., None, :].)
cov_fx_fX = self.kernel(x[..., None, :])[..., 0, :]
# New (x can be 1 dimension, eliminating the need of [..., None, :] and flitering afterwards.)
cov_fx_fX = self.kernel(x)

Copy link
Copy Markdown
Collaborator

@sawa3030 sawa3030 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the explanation. LGTM

@github-actions
Copy link
Copy Markdown
Contributor

This pull request has not seen any recent activity.

@github-actions github-actions bot added the stale Exempt from stale bot labeling. label Sep 14, 2025
@nabenabe0928 nabenabe0928 removed the stale Exempt from stale bot labeling. label Sep 15, 2025
Copy link
Copy Markdown
Collaborator

@kAIto47802 kAIto47802 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR! Basically LGTM, leaving a minor comment.

Co-authored-by: kAIto47802 <115693559+kAIto47802@users.noreply.github.com>
@nabenabe0928 nabenabe0928 added this to the v4.6.0 milestone Sep 17, 2025
Copy link
Copy Markdown
Member

@y0z y0z left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@nabenabe0928 nabenabe0928 merged commit ad9abaf into optuna:master Sep 17, 2025
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement Change that does not break compatibility and not affect public interfaces, but improves performance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants