[Feature] Variational Bayesian last layer models as surrogate models

[brunzema]

Motivation

This PR adds variational Bayesian last layers (VBLLs) [1], which demonstrated very promising results in the context of BO in our last paper [2], to BoTorch. The goal is to provide a BoTorch-compatible implementation of VBLL surrogates for standard use cases (single-output models), making them accessible to the community as quickly as possible. This PR does not yet contain all the features discussed in [2] such as the continual learning. If there is the interest to also add the continual learning, I am happy to add them down the line!

The VBLLs can be used in standard acquisition functions such as (log)EI but are especially nice for Thompson sampling as the Thompson sample of a Bayesian last layer model is a differentiable standard feed forward neural network which is useful for (almost) global optimization of the sample for the next query location.

Implementation details

This PR adds the implementation to the community folders--also here, if there is a large interest in the model, I am happy to help merge them into the main part of the repo. The added files of this PR are the following

botorch_community
|-- acquisition
|   |-- bll_thompson_sampling.py # TS for Bayesian last layer models
|-- models
|   |-- vblls.py # BoTorch wrapper for VBLLs
|-- posteriors
|   |-- bll_posterior.py # Posterior class for Bayesian last layer models
notebooks_community
|-- vbll_thompson_sampling.ipynb # Tutorial on how use the VBLL model
test_community
|-- models
|   |-- test_vblls.py # test for the VBLL models functionality (backbone freezing for feature reuse, etc)

The current implementation build directly on the VBLL repo, which is actively maintained and depends only on PyTorch. Using this repo allows improvements—e.g., better variational posterior initialization—to be directly beneficial for BO.

Have you read the Contributing Guidelines on pull requests?

Yes.

Test Plan

The PR does not change any functionality of the current code base. The core functionality of the VBLLs should be covered by test_vblls.py. Let me know if further tests are required.

Related PRs

This PR does not change functionality and I did not see any PRs regarding last layer models in BoTorch. Maybe this implementation can useful also for other BLLs.

References

[1] P. Brunzema, M. Jordahn, J. Willes, S. Trimpe, J. Snoek, J. Harrison. Bayesian Optimization via Continual Variational Last Layer Training. International Conference on Learning Representations (ICLR), 2025.

[2] J. Harrison, J. Willes, J. Snoek. Variational Bayesian Last Layers. International Conference on Learning Representations (ICLR), 2024.

[eytan]

Edit: max also responded. As he mentions, dependencies are a pain. Since this isn’t much code it makes sense to copy over and we can always reconsider if there is substantial new functionality and more extensive unit tests. Thanks for this contribution, Paul! This was a very intriguing paper and it’s great to see it in botch. We typically don’t take on new external dependencies but if it’s pure PyTorch that should be OK,, as long as it’s possible for your package to have some sort of unit or integration test that ensures any future changes to your package do not cause problems for the botch implementation/ functionality demonstrated in the nb.

…

On Fri, Feb 21, 2025 at 9:57 AM Paul Brunzema ***@***.***> wrote: Motivation This PR adds variational Bayesian last layers (VBLLs) [1], which demonstrated pre very promising results in the context of BO in our last paper [2], to BoTorch. The goal is to provide a BoTorch-compatible implementation of VBLL surrogates for standard use cases (single-output models), making them accessible to the community as quickly as possible. This PR does not yet contain all the features discussed in [2] such as the continual learning. If there is the interest to also add the continual learning, I am happy to add them down the line! The VBLLs can be used in standard acquisition functions such as (log)EI but are especially nice for Thompson sampling as the Thompson sample of a Bayesian last layer model is a differentiable standard feed forward neural network which is useful for (almost) global optimization of the sample for the next query location. Implementation details This PR adds the implementation to the community folders--also here, if there is a large interest in the model, I am happy to help merge them into the main part of the repo. The added files of this PR are the following botorch_community |-- acquisition | |-- bll_thompson_sampling.py # TS for Bayesian last layer models |-- models | |-- vblls.py # BoTorch wrapper for VBLLs |-- posteriors | |-- bll_posterior.py # Posterior class for Bayesian last layer models notebooks_community |-- vbll_thompson_sampling.ipynb # Tutorial on how use the VBLL model test_community |-- models | |-- test_vblls.py # test for the VBLL models functionality (backbone freezing for feature reuse, etc) The current implementation build directly on the VBLL repo <https://github.com/VectorInstitute/vbll>, which is actively maintained and depends only on PyTorch. Using this repo allows improvements—e.g., better variational posterior initialization—to be directly beneficial for BO. Have you read the Contributing Guidelines on pull requests <https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests> ? Yes. Test Plan The PR does not change any functionality of the current code base. The core functionality of the VBLLs should be covered by test_vblls.py. Let me know if further tests are required. Related PRs This PR does not change functionality and I did not see any PRs regarding last layer models in BoTorch. Maybe this implementation can useful also for other BLLs. References [1] P. Brunzema, M. Jordahn, J. Willes, S. Trimpe, J. Snoek, J. Harrison. Bayesian Optimization via Continual Variational Last Layer Training <https://arxiv.org/abs/2412.09477>. International Conference on Learning Representations (ICLR), 2025. [2] J. Harrison, J. Willes, J. Snoek. Variational Bayesian Last Layers <https://arxiv.org/abs/2404.11599>. International Conference on Learning Representations (ICLR), 2024. ------------------------------ You can view, comment on, or merge this pull request online at: #2754 Commit Summary - f599471 <f599471> add vbll surrogate and notebook - 6b82184 <6b82184> update notebook and base implementation - 1b0f899 <1b0f899> update optim config - 2606ac3 <2606ac3> add tests for vbll model - 873163f <873163f> update tutorial for vblls File Changes (6 files <https://github.com/pytorch/botorch/pull/2754/files>) - *A* botorch_community/acquisition/bll_thompson_sampling.py <https://github.com/pytorch/botorch/pull/2754/files#diff-e1d4fb223629e80cd1469b9e3965e95cbde8c7cd3aabc2a588e7a859c75fd8d4> (133) - *A* botorch_community/models/vblls.py <https://github.com/pytorch/botorch/pull/2754/files#diff-7c675eabda341230ee44f210363a77bc8ec06e8a1f5d905bcec7a614eef78502> (418) - *A* botorch_community/posteriors/__init__.py <https://github.com/pytorch/botorch/pull/2754/files#diff-c655d31ec372b021b283dd2eb9469ffe897917b48cb753120689904c029241d2> (4) - *A* botorch_community/posteriors/bll_posterior.py <https://github.com/pytorch/botorch/pull/2754/files#diff-a45c3f4d1096e4c276c5269ed330a44650361f816d5e5830eb29d7a0b38adf94> (54) - *A* notebooks_community/vbll_thompson_sampling.ipynb <https://github.com/pytorch/botorch/pull/2754/files#diff-4faf6dfab200cc0e3b077731b3438146ce836eecef404a848be04a27aa5b3020> (803) - *A* test_community/models/test_vblls.py <https://github.com/pytorch/botorch/pull/2754/files#diff-17cdea2ea479a1ee1616e238b6b1a155df7a92b3223d281ce7f37d4beb9cda61> (186) Patch Links: - https://github.com/pytorch/botorch/pull/2754.patch - https://github.com/pytorch/botorch/pull/2754.diff — Reply to this email directly, view it on GitHub <#2754>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAAW34NCIDINOF67WZXXLGL2Q45FJAVCNFSM6AAAAABXTNRYUOVHI2DSMVQWIX3LMV43ASLTON2WKOZSHA3DSMRWGIYDIMQ> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[Balandat]

Thanks for putting this up, @brunzema. The notebook looks great, and I plan to review this PR in more detail over the next day or two.

Regarding the dependency on vbll: Right now it looks like the code only uses ~120 lines of pure torch code from the vbll repo (namely this
https://github.com/VectorInstitute/vbll/blob/main/vbll/layers/regression.py#L34-L149 plus the minimal function https://github.com/VectorInstitute/vbll/blob/main/vbll/utils/distributions.py#L94-L98). It seems questionable to me to take that dependency, especially since the vllb repo doesn't include unit tests and/or CI. We spend a nontrivial amount of time fixing issues with downstream dependencies, so we're really careful about adding them only if really necessary.

My preference would be to move the relevant pieces of the vbll code mentioned above into a helper module (and clearly attribute the source there, of course) so we can avoid the dependency for now. If we do end up expanding the functionality and use additional features from vbll, then I'd be happy to reconsider (provided the vbll repo adds proper unit tests and a CI setup).

[Balandat]

Unit tests are failing b/c vbll is not installed in the CI. This will not be necessary if we move the minimal code required in a helper module that we include here.

[Balandat]

oops, submitted prematurely. Here is the rest of the review.

[Balandat]

The TS samples shown are not the samples that were optimized to obtain the next candidate point. This initially caused some confusion; can you make the plots in a way so that this is consistent (i.e. that we show the set of TSs that were optimized to obtain the next candidate)?

[brunzema]

Mhm, yeah thats true.. I now removed the samples from the plot to avoid this. I could also include in BLLMaxPosteriorSampling to return the sampled functions but this use case seemed too specific. Do you think this is also relevant beyond visualization? Then, this might be nice imo

[Balandat]

It would be really nice to run both TS and logEI and then show these in a two-panel figure to have the comparison. This would also ensure that both code paths work properly

[brunzema]

I changed it to loop over both logEI and TS and then show the models for all iterations. I think the model "evolution" is very similar. I wanted to keep the main BO loop as simple as possible rather than optimizing in parallel with both acquisition function for easier reuse of the loop. If you think two-panel figures would be more interesting, I'm happy to make that change!

[brunzema]

@eytan Thank you for the nice comment and @Balandat thank you for the detailed review!

Just wanted to quickly give an update on the vblls: I talked to my collaborators and we are happy to move the relevant part (regression layer + some utils) to botorch. I will update this PR to incorporate all suggested changes within the next few days + include the vbll code 👍

[brunzema]

@Balandat Thank you again for the detailed review--I really appreciate it! I’ve updated the PR to address all your points, but please let me know if any further changes are needed. I’m happy to make any additional updates!

Biggest change is of course the added code from the vbll package, let me know if the way I know included it is ok (additional file + credit at the top).

[Balandat]

@brunzema thanks a lot for the updates - will review this within the next couple of days!

[codecov]

Codecov Report

Attention: Patch coverage is 68.71287% with 158 lines in your changes missing coverage. Please review.

Project coverage is 99.16%. Comparing base (da28b43) to head (d01e1fd).

Files with missing lines	Patch %	Lines
botorch_community/models/vbll_helper.py	66.28%	88 Missing ⚠️
...rch_community/acquisition/bll_thompson_sampling.py	0.00%	52 Missing ⚠️
botorch_community/models/vblls.py	93.54%	10 Missing ⚠️
botorch_community/posteriors/bll_posterior.py	78.37%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2754      +/-   ##
==========================================
- Coverage   99.99%   99.16%   -0.84%     
==========================================
  Files         202      206       +4     
  Lines       18520    19025     +505     
==========================================
+ Hits        18519    18866     +347     
- Misses          1      159     +158     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Balandat]

Thanks for the updates and in particular for avoiding another external dependency! I left a number of inline comments, here are some higher level ones:

1. Please rebase this on a recent version of master (the base of commit of this PR is pretty old)
2. Please address flake8 (incl. line length) and import sorting errors - you can install the pre-commit hooks to make sure that your code confirms to the standard (see https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pre-commit-hooks)
3. I spotted some some rather problematic numerical code in the vbll helpers, let's update that (highlighted inline), shall we?

[Balandat]

maybe rename this to discrete_inputs to make it clear what this refers to.

[Balandat]

Suggested change

	raise ValueError("Model must be an instance of AbstractBLLModel")
	raise ValueError("Model must be an instance of AbstractBLLModel, is {type(model)}")

[Balandat]

Suggested change

	self.lb = torch.zeros(self.model.num_inputs, dtype=torch.float64).cpu()
	self.ub = torch.ones(self.model.num_inputs, dtype=torch.float64).cpu()
	self.lb = torch.zeros(self.model.num_inputs, dtype=torch.float64, device=torch.device("cpu))
	self.ub = torch.ones(self.model.num_inputs, dtype=torch.float64, device=torch.device("cpu))

[Balandat]

Should we warn here or error out? Seems like this is not really a use case we need to / should support?

[brunzema]

yeah makes sense, changed it to raise an error

[Balandat]

Same question as above - why ignore the X_cand arg in that case? Should we instead just define two different instances of AbstractBLLModel, one for discrete one for non-discrete use cases?

Also, what about the mixed setting in which only some of the dimensions of the input are discrete and others are continuous? Seems like that case could be supported reasonably straightforwardly but is not here (which I guess is fine for a first implementation).

[brunzema]

Happy to include this down the line! Seems very relevant for incorporating some one-hot encoding-like inputs + some parameters in R^d. I am mainly thinking about some autoML applications. Would it be ok to include this in a later PR?

[Balandat]

Ugh, this is pretty bad. You're calling inv on the same tensor twice. Also, you already have the Cholesky decomposition here, why are you calling inv in the first place? You can do this with a simple triangular solve.

Suggested change

	return tp(torch.linalg.inv(self.scale_tril)) @ torch.linalg.inv(self.scale_tril)
	Eye = torch.eye(self.scale_tril.shape[-1], device=self.scale_tril.device, dtype=self.scale_tril.dtype)
	W = torch.linalg.solve_triangular(self.scale_tril, Eye, upper=False)
	return tp(W) @ W

[brunzema]

uff, thank you! Will also push this to the vbll repo

[Balandat]

should use solve_triangular here as well

[Balandat]

seems like the torch.distributions.LowRankMultivariateNormal does this just fine, why overwrite?

[brunzema]

Actually not super sure if there was a deeper reason here, I will double check!

[Balandat]

in-pace updates can be problematic for backprop, suggest avoiding that

Suggested change

	total_elbo += self.regularization_weight * (wishart_term - kl_term)
	total_elbo = total_elbo + self.regularization_weight * (wishart_term - kl_term)

[Balandat]

is this still needed if you use from __future__ import annotations in all the modules to do postponed type evaluations? We should generally use that.

[Balandat]

@brunzema checking in here, anything needed to get this over the finish line? Seems like we're very close.

[brunzema]

@Balandat no, the delay is fully on me, sorry! Everything is pretty much done, flake8 is adressed and also with the import sorting + from __future__ import annotations I no longer require the not so nice type_checking condition 👍 I will make a fresh rebase tomorrow and a final pass and hopefully in the morning your time, I have submitted the updated PR :)

[Balandat]

Excellent!

[brunzema]

@Balandat I updated the PR. Not sure why I thought the importing was fixed yesterday—I still had the type checking in place. I noticed that some files in the main repo also use it, so it should be fine?

That said, I think a cleaner approach (which I’ve implemented in the updated PR) is to extract the abstract Bayesian last-layer (BLL) class and use it as a parent class/interface. This should also make it more extensible for future BLL models.

botorch_community
|-- acquisition
|   |-- bll_thompson_sampling.py # AbstractBLL to specify "interface"
|-- models
|   |-- blls.py # Define AbstractBLL model
|   |-- vblls.py # VBLL inherit from AbstractBLL
|-- posteriors
|   |-- bll_posterior.py # AbstractBLL to specify "interface"
...