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]

Is this dep used somewhere? I don't see any import of it in the PR

[Sohambasu07]

It's not used in this PR. We use this for some of the multi-objective algorithms, especially Multi-objective ASHA, where I'd forgotten to include this before.

[nastaran78]

Suggested change

self.fixed_numericals = fixed_numericals

[nastaran78]

I believe it modifies the actual acquisition function that we have, we may want to copy the value?

[nastaran78]

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