Ax: Modular `BoTorchModel`: combining BoTorch components like `Model`, `AcquisitionFunction`, `Objective` etc. in Ax

[nwollek]

Hello Ax team,

You did a great job so far with the project researching the topic of bayesian optimization (BO)!
I have a feature proposal that I already discussed with Max @Balandat beforehand and would like to further explain it here to your whole team.
This issue is currently preventing that I use your library for certain features or at all.

My use case:
Implementing a software suite for parameter optimization, I also integrated your library as BO backend (via the developer API), next to other BO implementations.
One important feature of my software is to let the user control, which acquisition function (ACF) to use for BO (EI, POI, UCB).
Furthermore, the trade-off between exploration and exploitation (via xi/kappa in the ACF) should be adjustable while optimizing.
Last but not least, the user should be able to understand the decisions of the underlying BO process.
To accomplish this, the ACF should be rendered as a plot after every step.
As you can see, I would like to leverage existing BO implementations and integrate them into a whole user experience for a specific domain.

The problem:
While some other BO libraries offer an easy to use interface for accessing and adjusting the ACF (while keeping a certain abstraction), it seems to be quite laborious with Ax and BoTorch.
First of all, the initialization functions of your three APIs (Loop, Service, Developer) do not offer switching to another one of the popular acquisition functions (POI, UCB) and specifying the trade-off.
It might be technically possible to do so, but it requires much effort and lot of code on a very low layer, which again requires to understand the whole Ax code on that low layer.
Trying to evaluate the ACF for certain parametrizations, I have to read through multiple abtraction layers of your library, trying to understand all conversions (for example normalizing parameter spaces and later converting them back, which seems not to be documented anywhere).
Afterwards I have to manually call several functions on all those layers and apply all conversions to reconstruct the ACF to be able to evaluate it.
As a result, too much time needs to be invested just to re-implement all steps your BO optimizer is doing anyway in its main functions.
This also leads to potential mistakes due to wrong or missing conversions, everytime a user of your library wants to access the ACF.

The proposal:
I assume that the average user of Ax/BoTorch would not like to bother with digging to the lowest layer of the Ax/BoTorch framework just to be able to do something that basic for BO like choosing the ACF or evaluating the ACF for a specific parametrization.
I really appreciate the advanced features of your framework and it would be a shame to not use it because of this deficit regarding the ACF.
For the use case I explained in the beginning, I propose the following features to solve the problem:

1. Possibility to specifiy which of the implemented ACFs to use (for example via a simple identification string) on initialization of the optimizer as well as switching it dynamically in the optimization process

2. Possibility to control the trade-off (exploitation vs. exploration) on initialization as well as adjusting it dynamically in the optimization process

3. Possibility to evaluate the very identical ACF the BO optimization loop uses for a given parametrization without building it manually

I know those things might be partially doable already by many highly specific function calls on a low layer of your framework, but the corresponding research and coding overhead stands in no relation to the very basic functionality achieved by it.

Those three features are, in my opinion, very essential to BO and would further enhance your existing framework.
Also they would increase the user base as other users like me would not have to switch to other libraries just to be able to render an ACF.
Other features of your library are already easy to use on a high abstraction layer, so I assume providing the proposed functionality would match the existing framework greatly.

I hope that I expressed everything in an understandable way and would like to hear about your opinion on that!

Thanks in advance!

[lena-kashtelyan]

Hi, @nwollek! We are currently working on something that might cover at least part of this feature request: an improved and modular interface between Ax and BoTorch, which will allow for mixing and matching of surrogate models and acquisition functions. We can put the work in its current (unfinished) state up as a PR in the next few days, and you could give us feedback? How does that sound?

[nwollek]

Hello Lena,

Thanks for your fast answer!
I would definitely take a quick look into the PR, but have to admit that I don't have time for a full review.
So you could create it in the next days with respect to your time schedule and I will check it out sometime soon.

[lena-kashtelyan]

Hi, @nwollek! On Friday, @EricZLou added our work on modular BoTorchModel in Ax to our open-source repo, in https://github.com/facebook/Ax/pull/374. The PR contains a draft tutorial that explains the new setup and how to use it: https://github.com/facebook/Ax/blob/master/tutorials/botorch_modular.ipynb.

As you take a look at the new functionality, please bear in mind that both the code and the tutorial are incomplete (active work underway to add the missing pieces), and we'll be happy to take into account any feedback you might have!

cc @Balandat

[nwollek]

Hey @lena-kashtelyan, I just took a quick look at the tutorial draft and it seems to be going into the right direction.
I am looking forward to see further developments and tutorials explaining how to use those classes for steering the optimizer between exploration and exploitation on a high abstraction layer. I think this issue should be kept open until then. Keep me updated!

[lena-kashtelyan]

@nwollek, agreed! I'll keep the issue open and post updates here.

[lena-kashtelyan]

New and improved tutorial: https://github.com/facebook/Ax/blob/master/tutorials/Setup_and_Usage_of_BoTorch_Models_in_Ax.ipynb. We continue the work on this and there will be a more polished version of the tutorial and the new abstractions in the near future, but the current version should be good to work with!

[lena-kashtelyan]

Closing this issue as the feature request is now shipped; it is still in beta mode, so please open any follow-up issues if you have questions or concerns about how to use it, @nwollek!

[lena-kashtelyan]

Will keep this open for better discoverability

[lena-kashtelyan]

Also marking this as "in progress" as we are planning more work on this setup in the near future, after which we will publish two tutorials on it on the website and open a designated issue for any questions or feedback.