Performance & runtime improvements to info-theoretic acquisition functions (1/N) #2748

hvarfner · 2025-02-18T16:43:32Z

A series of improvements directed towards improving the performance of PES & JES, as well as their MultiObj counterparts.

Motivation

As pointed out by @SebastianAment in this paper, the BoTorch variant of JES, and to an extent PES, is brutally slow an suspiciously ill-performing. To bring them up to their potential, I've added a series of performance improvements:

1. Improvement to get_optimal_samples and optimal_posterior_samples: As this is an integral part of their efficiency, I've added suggestions (similar to sample_around_best) to optimize_posterior_samples.
Marginal runtime improvement in acquisition optimization (sampling time practically unchanged):

Substantial performance improvement: .

2. Added initializer to acquisition funcction optimization: Similar to KG, ES methods have sensible suggestions for acquisition function optimization in the form of the sampled optima. This drastically reduces the time of acquisition function optimization, which could on occasion take 30+ seconds when num_restarts was large >4.

Benchmarking INC

2b. Multi-objective support for initializer: By re-naming arguments of the multi-objective variants, we get consistency and support for MO variants.

3. Enabled gradient-based optimization for PES: The current implementation contains a while-loop which forces the gradients to be recursively computed. This commonly causes NaN gradients, which is why the recommended option is "with_grad": False in the tutorial. One detach() alleviates this issue, enabling gradient-based optimization.

NOTE: this has NOT been ablated, since the non-grad optimization is extremely computationally demanding.

Test Plan

Unit tests and benchmarking.

Related PRs

First of a couple!

Bonus: while benchmarking, I had issues repro'ing the LogEI performance initially. I found that sample_around_best made LogEI worse on Mich5. All experiments are otherwise a repro of the settings used in the LogEI paper.

facebook-github-bot · 2025-02-18T16:50:25Z

@sdaulton has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

sdaulton · 2025-02-18T16:50:35Z

Thanks! It seems like sample_around_best could definitely lead to the AF optimization getting stuck in a local optima

codecov · 2025-02-18T16:59:15Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.99%. Comparing base (9a7c517) to head (75afd87).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2748   +/-   ##
=======================================
  Coverage   99.99%   99.99%           
=======================================
  Files         203      203           
  Lines       18690    18689    -1     
=======================================
- Hits        18689    18688    -1     
  Misses          1        1

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

hvarfner · 2025-02-18T17:05:04Z

@sdaulton for sure! I currently observe similar things for JES, but I'm not sure whether the found points are actually higher in acquisition function value or not (for either LogEI or JES)

sdaulton · 2025-02-18T17:34:20Z

That would be interesting to see

saitcakmak

Hi Carl! This seems like a decent improvement. Just a few comments in-line

botorch/utils/sampling.py

saitcakmak · 2025-02-20T17:26:31Z

botorch/utils/sampling.py

@@ -1008,13 +1012,17 @@ def optimize_posterior_samples(
            negate the objective or otherwise transform the output.
        return_transformed: A boolean indicating whether to return the transformed
            or non-transformed samples.
+        suggested_points


Let's complete the docstring here

saitcakmak · 2025-02-20T17:29:06Z

botorch/utils/sampling.py

+        bounds=bounds, n=round(raw_samples * frac_random), q=1
+    ).squeeze(-2)
+    if suggested_points is not None:
+        from botorch.optim.initializers import sample_truncated_normal_perturbations


Is this a local import, because it leads to cyclical dependencies? If so, we could move sample_truncated_normal_perturbations under utils (assuming it doesn't depend other code in optim).

I think so! let me check

botorch/utils/sampling.py

saitcakmak · 2025-02-20T17:33:44Z

botorch/utils/sampling.py

+    candidate_set = draw_sobol_samples(
+        bounds=bounds, n=round(raw_samples * frac_random), q=1
+    ).squeeze(-2)
+    if suggested_points is not None:


If suggested_points is None, we end up with a candidate_set of size smaller than raw_samples. Should we make sure we always use raw_samples points?

Ah, good catch. Absolutely

saitcakmak · 2025-02-20T17:48:24Z

botorch/utils/sampling.py

+    weights = (
+        candidate_queries - candidate_queries.mean(dim=-1, keepdim=True)
+    ) / candidate_queries.std(dim=-1, keepdim=True)
+    eta = options.get("eta", 2.0)
+    weights = torch.exp(eta * weights)
+
+    # weights can be more than 2D, which is not supported by torch.multinomial
+    # the argsort picks out the indices that are nonzero, i.e. those that are drawn
+    # (without replacement, so we will always have num_restarts nonzero ones)
+    idx = (
+        Multinomial(num_restarts, probs=weights)
+        .sample()
+        .argsort(descending=True)[..., :num_restarts]
+    )


This seems like a very similar logic to initialize_q_batch. Would it make sense to re-use that?

It is extremely similar, but torch.multinomial does not support 3D multinomial sampling which is why I re-used the logic (see the comment), but I initially tried initialize_q_batch.

Literally all the duplicated logic is to accomodate fully Bayesian models (otherwise I could have just used ThompsonSampling and optimize_acqf, but the extra dim that fully Bayesian models generate really throws it off.

code, reshuffling of other sampling methods (that don't take an acqf)

improve performance and runtime of PES/JES

saitcakmak · 2025-02-21T21:22:34Z

botorch/optim/initializers.py

-        idcs = batched_multinomial(
-            weights=weights.permute(*range(1, len(batch_shape) + 1), 0), num_samples=n
-        ).permute(-1, *range(len(batch_shape)))
+    idcs = boltzmann_sample(


Cool, this deduplicated a bunch of repeated logic here.

botorch/utils/sampling.py

facebook-github-bot · 2025-02-21T21:30:43Z

@sdaulton has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

facebook-github-bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Feb 18, 2025

hvarfner mentioned this pull request Feb 20, 2025

Performance & runtime improvements to info-theoretic acquisition functions (2/N) - AcqOpt initializer #2751

Open

saitcakmak requested changes Feb 20, 2025

View reviewed changes

hvarfner force-pushed the optimize_posterior_samples branch 2 times, most recently from 520aad7 to 43744fa Compare February 21, 2025 13:06

hvarfner mentioned this pull request Feb 21, 2025

Performance & runtime improvements to info-theoretic acquisition functions (0/N) - Restructuring of sampling methods #2753

Open

hvarfner force-pushed the optimize_posterior_samples branch 5 times, most recently from d5d9f9d to 3ae01cb Compare February 21, 2025 15:17

hvarfner added 2 commits February 21, 2025 16:18

Boltzmann sampling function added in utils/sampling to remove duplicate

14eddf6

code, reshuffling of other sampling methods (that don't take an acqf)

Improvements of get_optimal_samples and optimize_posterior_samples to

f4f01bf

improve performance and runtime of PES/JES

hvarfner force-pushed the optimize_posterior_samples branch from 3ae01cb to f4f01bf Compare February 21, 2025 15:22

saitcakmak approved these changes Feb 21, 2025

View reviewed changes

Update botorch/utils/sampling.py

75afd87

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Performance & runtime improvements to info-theoretic acquisition functions (1/N) #2748

Performance & runtime improvements to info-theoretic acquisition functions (1/N) #2748

hvarfner commented Feb 18, 2025

facebook-github-bot commented Feb 18, 2025

sdaulton commented Feb 18, 2025

codecov bot commented Feb 18, 2025 •

edited

Loading

hvarfner commented Feb 18, 2025

sdaulton commented Feb 18, 2025

saitcakmak left a comment

saitcakmak Feb 20, 2025

saitcakmak Feb 20, 2025

hvarfner Feb 21, 2025

saitcakmak Feb 20, 2025

hvarfner Feb 21, 2025

saitcakmak Feb 20, 2025

hvarfner Feb 21, 2025 •

edited

Loading

saitcakmak Feb 21, 2025

facebook-github-bot commented Feb 21, 2025

Performance & runtime improvements to info-theoretic acquisition functions (1/N) #2748

Are you sure you want to change the base?

Performance & runtime improvements to info-theoretic acquisition functions (1/N) #2748

Conversation

hvarfner commented Feb 18, 2025

Motivation

Test Plan

Related PRs

facebook-github-bot commented Feb 18, 2025

sdaulton commented Feb 18, 2025

codecov bot commented Feb 18, 2025 • edited Loading

Codecov Report

hvarfner commented Feb 18, 2025

sdaulton commented Feb 18, 2025

saitcakmak left a comment

Choose a reason for hiding this comment

saitcakmak Feb 20, 2025

Choose a reason for hiding this comment

saitcakmak Feb 20, 2025

Choose a reason for hiding this comment

hvarfner Feb 21, 2025

Choose a reason for hiding this comment

saitcakmak Feb 20, 2025

Choose a reason for hiding this comment

hvarfner Feb 21, 2025

Choose a reason for hiding this comment

saitcakmak Feb 20, 2025

Choose a reason for hiding this comment

hvarfner Feb 21, 2025 • edited Loading

Choose a reason for hiding this comment

saitcakmak Feb 21, 2025

Choose a reason for hiding this comment

facebook-github-bot commented Feb 21, 2025

codecov bot commented Feb 18, 2025 •

edited

Loading

hvarfner Feb 21, 2025 •

edited

Loading