[Docs] - Supplying gradients from an external, non-torch function

[ghutchis]

🚀 Feature Request

We've been using Bayesian optimization (with GPyOpt) for molecular geometries. I'd like to migrate to BoTorch and as a non-expert, I'm hitting some walls with the current docs.

Yet... I'd like to connect to external codes for evaluation functions and gradients to supply them to the optimization loop (i.e., rather than using auto-diff from a torch model).

I suspect I'm just missing some docs..

Motivation

Is your feature request related to a problem? Please describe.

Many optimization problems use external codes. Sometimes they have well-defined, known gradients w.r.t variables.

It would be great to have a tutorial with evaluation function and gradient function for gradient-assisted optimization loops.

Pitch

Describe the solution you'd like

An example class like the synthetic test functions code, but with a gradient method.

Describe alternatives you've considered

I've used simple functions without gradients, but optimization is slow.

Are you willing to open a pull request? (See CONTRIBUTING)

I'd be happy to write out a notebook tutorial and contribute if someone can point me in the right direction.

[Balandat]

Thanks for your interest in BoTorch. Let me just make sure I fully understand the request:

You evaluating some black box function f (e.g. some metric returned by a molecular simulation), and in addition to the metric you also get gradient information grad(f) w.r.t. the input parameters?

I don't think we have any current examples for this in BoTorch, but the underlying GPyTorch models do have the ability to use gradient information, see e.g. this tutorial.

Making this work basically amounts to hooking such a model that uses derivatives into the BoTorch Model API, which should be relatively straightforward. Internally, we'd just have to pick out the function value to optimize (assuming your objective is on the function value and does not include the derivative). Then auto-differentiating the model prediction internally should work without issues.

I'd be happy to cook up a minimal example for you to start off from.

[ghutchis]

You evaluating some black box function f (e.g. some metric returned by a molecular simulation), and in addition to the metric you also get gradient information grad(f) w.r.t. the input parameters?

Yes, exactly.

Internally, we'd just have to pick out the function value to optimize (assuming your objective is on the function value and does not include the derivative).

Right. The objective is to minimize the energy function. You can separately calculate the energy and the gradient of the energy w.r.t parameters.

Your approach makes sense - basically, I'd want something like the MultiTaskGPyTorchModel - let's say for an example that i'm changing the O-H bond lengths and the H-O-H angle in water - computing energies and gradients given those parameters.

[Balandat]

Great. I have a basic version working. It's just some plumbing that needs to be done. I'll put up a notebook once It's in a decent shape.

One potential concern I have with the current gpytorch derivative-enabled GP setup is scalability - largely b/c how the indexing into the returned multi-task MVN is done. How many data points do you typically work with in your applications?

[ghutchis]

Awesome. I'm not concerned with scalability. If there are more than 50-100 points, something is very wrong.

[Balandat]

What's a typical dimensionality of the parameter / search space?

[ghutchis]

For now, I'd be content with simpler models (e.g. water is 2 or 3 parameters) - so ~30 parameters is fine.

There's been some success with the general approach (e.g., https://doi.org/10.1021/acs.jctc.0c00257). I think we can place fairly strong prior beliefs on each parameter so I'm working on a custom mean - which presumably improves performance.

(I also don't think there's been much effort to try different acquisition functions, which is why BoTorch is interesting.)

[Balandat]

OK so let me at least share some partial results here: It turns out there is some issue with backpropagating through the acquisition function computed on the posterior of the derviative-enabled GP. That will take a bit of debugging on our end to figure out what's going on there.

But, fwiw, here is a notebook for how to set up a model for this that, once that bug has been squashed, should provide a model for your use case. We can clean that up and merge it in since it's more generally useful to a broader audience

demo_deriv_enabled_BO.ipynb.txt

[ghutchis]

Awesome, many thanks! I'll give it a try on an analytical form of the water energy surface first - but I'd guess it's generally useful to a wide range of users.

[Balandat]

I found the backprop issue: #638. Once that fix is merged in (imminent), things should be good to go (ish - pending some scalability optimizations).

Here is the updated notebook that shows how to optimize UCB on top of this:
demo_deriv_enabled_BO.ipynb.txt

[Balandat]

@ghutchis Did you get anywhere with this starting from my example above?

I'm going to close out this issue since the basic question should be resolved. Happy go deeper on this in the form of a Discussion topic.

[ollie-bell]

Apologies to comment on a closed thread. I just wanted to say thank you for that example notebook, as I too am starting on a journey on Bayesian optimization of functions where gradient information is available. Do you have any intuition as to how easy/hard it would be to implement a "derivative-enabled" acquisition function, such as the one described here https://arxiv.org/pdf/1703.04389.pdf. Thanks

[saitcakmak]

@acse-ogb119, you may find this notebook useful: https://docs.gpytorch.ai/en/stable/examples/08_Advanced_Usage/Simple_GP_Regression_Derivative_Information_1d.html

[Balandat]

Do you have any intuition as to how easy/hard it would be to implement a "derivative-enabled" acquisition function, such as the one described here https://arxiv.org/pdf/1703.04389.pdf

So AFAICT d-KG is really just KG on top of a model that includes derivative observations. On a high level, there shouldn't be any issue with implementing this - we essentially just have to make sure that the fantasize API works for this Kernel, and then compute the posterior mean of the function (not including the predicted derivatives). It could create some scalability issues.

In a sense the UCB function I used in the example nb above is already "derivative enabled" in that it uses the model including derivative observations.