BAxUS tutorial

[LeoIV]

Motivation

This pull request adds a tutorial for BAxUS (published at NeurIPS 2022), a novel BO method for high-dimensional problems. BAxUS allows for sample-efficient BO by using nested random embeddings that preserve observations.

Including BAxUS as a tutorial in Botorch helps other researchers to build on our work and makes it accessible for practitioners.

Have you read the Contributing Guidelines on pull requests?

Yes.

Test Plan

This PR does not add any core functionality, therefore no tests were added.

Related PRs

No new functionality.

[dme65]

Hi Leonard, thank you for the PR! I'll review this in the next few days.

[codecov]

Codecov Report

Merging #1559 (c132685) into main (e5f4d14) will not change coverage.
The diff coverage is n/a.

❗ Current head c132685 differs from pull request most recent head 6c1e66b. Consider uploading reports for the commit 6c1e66b to get more accurate results

@@            Coverage Diff            @@
##              main     #1559   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          170       170           
  Lines        14610     14610           
=========================================
  Hits         14610     14610           

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[dme65]

Thank you for contributing this tutorial and sorry for the delayed review! This looks great to me overall.

On a high level there is plenty of overlap between this tutorial and the TuRBO and SCBO tutorials. This is outside the scope of this diff, but we should probably get some of the core trust region functionality checked into BoTorch so we can cut some of the code in these tutorials. cc @Balandat .

• Can you run isort and flake8 on the notebook? Some of the lines are quite long which makes the code hard to read.
• Why are you using a domain of [-5, 15]^500? This is a bit non-standard for Branin which is usually optimized over [-5, 10] x [0, 15]. You can do something similar to https://github.com/pytorch/botorch/blob/main/tutorials/saasbo.ipynb in which case you can avoid having to define DummyBranin.
• Why are you learning the GP hyperparameters using Adam instead of calling fit_gpytorch_mll which will use L-BFGS-B? We generally find that fit_gpytorch_mll works better and you can see how this is done in https://botorch.org/tutorials/turbo_1.
• Can you rename X_turbo, Y_turbo, etc. to X_baxus, Y_baxus, … ?
• Dummy Branin -> Embedded Branin

[LeoIV]

Thank you for the review! I agree that there's a lot of overlap, but I think it's better to use similar code than to implement it differently if we cannot share code between the tutorials.

• I switched to fit_gpytorch_mll. I changed it originally because I ran into errors with that before. I guess the reason is that when I increase the target space, all observations are constrained to a linear subspace. But I switched back and have yet to face problems.
• I now use [-5, 10] and [0, 15] (same method as Saasbo)
• Renamed "Dummy" to ``Embedded'' and X_turbo to `X_baxus` etc.
• I ran isort and black on the notebook

Please let me know if I need to make additional changes!

[facebook-github-bot]

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

[Balandat]

thanks, there are a few places in which things can be cleaned up but overall this lgtm.

[Balandat]

Thanks @LeoIV for making the edits. @dme65 is currently out - let's wait for him to get back to merge this in.

[facebook-github-bot]

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

[saitcakmak]

Looks like the model fitting occasionally fails in the smoke test mode.
First attempt failed: https://github.com/pytorch/botorch/actions/runs/4239527520/attempts/1
Re-run succeeded: https://github.com/pytorch/botorch/actions/runs/4239527520/attempts/2

Would you mind setting the seed so that this doesn't lead to CI failures? If you have other ideas for having it deterministically succeed, that'd also be great.

[LeoIV]

Thanks, I fixed the seed. When the embedding is increased, the previous observations lie in a lower-dimensional subspace of that higher-dimensional embedding. This seems to cause problems with LBFGS in some cases. What I do in my work is to catch the ModelFittingError and revert to Adam-based MLE. I could add that but it would increase the notebook length by ~10-20 lines.

[saitcakmak]

What I do in my work is to catch the ModelFittingError and revert to Adam-based MLE. I could add that but it would increase the notebook length by ~10-20 lines.

I think that could be helpful for someone else that's using this tutorial to optimize their problem and runs into a similar issue. Let's add it with a simple comment explaining why this is needed.

[LeoIV]

I added the fallback optimization and a comment about that. Also, I added a context manager to suppress the input data validation. The data is normalized but to [-1, 1] instead of [0, 1] due to the choice of the embedding.

[facebook-github-bot]

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

[facebook-github-bot]

@dme65 merged this pull request in 23d1a3b.