feat: Improve Acquisition Function optimization in discrete and mixed search spaces

[Sohambasu07]

Issue

• Optimization of acquisition function in NePS BO was rather slow and capped at 30 combinations for search spaces with categorical hyperparameters (Limit on categoricals in GP #219).

Additions and Deletions

• Removed maximum_allowed_categorical_combinations which previously set a cap on number of allowed categorical combinations in the search space.
• Removed optimize_acqf_mixed()
• Added WrappedAcquisition module in neps.optimizers.acquisition.wrapped_acquisition.
• Added optimize_acqf_discrete_local_search() for purely categorical and mixed search spaces.

What remains the same

1. Purely numerical search spaces:
   • Remains unchanged.
2. Mixed or purely categorical search spaces with number of categorical combinations < num_restarts:
   • Generate all possible categorical combinations.
   • Use optimize_acqf_mixed using all possible categorical combinations as fixed_features.
   • NOTE: The reason behind still keeping this is that optimize_acqf_discrete_local_search which we now use for search spaces with large number of categoricals fails to generate num_restarts number of minimum candidates if the total number of categorical combinations is lower than num_restarts. Since num_restarts is usually very low (defaults to 20 in NePS GP), this will not be very computationally expensive.

What has changed:

For search spaces with a high number of categorical dimensions (and n_combos > num_restarts), the following changes have been introduced:

1. Purely categorical search spaces (:

   • Replace optimize_acqf_mixed() with optimize_discrete_local_search() which scales better with increasing number of categorical dimensions and combinations.

2. Mixed search spaces:

   • For mixed search spaces, we first remove the need to generate all possible categorical combinations, which further reduces a lot of computational overhead.
   • Instead we randomly select one combination per iteration of the optimization.
   • We then perform a sequential 2-step optimization:
     • First, keeping the randomly selected categorical combination fixed, we optimize over the continuous features using optimize_acqf().
     • Next, we wrap the acquisition function inside WrappedAcquisition to keep the best seen values of the continuous features fixed and optimize only over the categoricals using optimize_acqf_discrete_local_search(). Finally, we merge the best seen values of both the numerical and categorical features into a single tensor, perform one forward pass over the acquisition function and return the candidate and the score.

Tests

• Tested the optimizers - PriMO, BO and PriorbandBO on the JAHSBench CIFAR10 task.

[nastaran78]

Can you also share your observation regarding performance improvement using optimize_acqf_discrete_local_search for mixed and categorical space over optimize_acqf_mixed?

[nastaran78]

Shall we rename it to something conveying the mission of this class better?
maybe FixedNeumericalAcquisition

[Sohambasu07]

I kept it as WrappedAcquisition because it's pretty much what it does. I was a bit conflicted too since fixing the numericals is not the main thing it does, but rather optimizing over only the categorical space.

[Sohambasu07]

Can you also share your observation regarding performance improvement using optimize_acqf_discrete_local_search for mixed and categorical space over optimize_acqf_mixed?

So, previously the number of allowed categorical combinations were 30, because the complexity of the outer loop in botorch's optimize_acqf_mixed increases with the number of fixed_features i.e., categorical combinations, since it just iterates over all of them. So NePS BO previously didn't allow search spaces with more than 30 categorical combinations at all. So, even if we removed the cap, running something like Jahsbench with ~10,000 categorical combinations or HW-GPT-Bench with ~10E7 categorical choices would be practically infeasible.

I would say that the performance improvement is that this PR makes it possible to optimize over any large search space within seconds (about 10 seconds per BO loop).