[Bug]: botorch_community VBLLModel posterior doesn't work with singel value tensor

[CarloMittelbach]

What happened?

VBBLs posterior method seems to convert mean & variance to scalar value, when input tensor is single value.
This isn't compatible with torch.diag_embed() also used in the posterior function.

Please provide a minimal, reproducible example of the unexpected behavior.

import torch
import traceback
from botorch_community.models.vblls import VBLLModel


def main():
    torch.set_default_dtype(torch.float64)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model = VBLLModel(in_features=1, out_features=1).to(device)
    test_X_single = torch.tensor([[0.5]], device=device)
    try:
        posterior = model.posterior(test_X_single)
    except IndexError as e:
        traceback.print_exc()
        
    test_X_multi = torch.tensor([[0.25], [0.75]], device=device)
    try:
        posterior_multi = model.posterior(test_X_multi)
        print(f"Worked as expected!")
    except Exception as e:
        print(f"You won't get here: {e}")


if __name__ == "__main__":
    main()

Please paste any relevant traceback/logs produced by the example provided.

Ouput:

Traceback (most recent call last):
  File "/home/carlo/uni/SoSe25/automahfic/automahfic/bo/surrogates/vbbl/min_example_failure.py", line 13, in main
    posterior = model.posterior(test_X_single)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carlo/uni/SoSe25/automahfic/venv/lib/python3.12/site-packages/botorch_community/models/vblls.py", line 492, in posterior
    cov = torch.diag_embed(variance)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got -2)

Multi worked as expected!

BoTorch Version

0.14.0

Python Version

3.12.3

Operating System

Linux Ubuntu 24.04

(Optional) Describe any potential fixes you've considered to the issue outlined above.

in vbbls.py VBBLModel

    def posterior(
        self,
        X: Tensor,
        output_indices=None,
        observation_noise=None,
        posterior_transform=None,
    ) -> Posterior:
        if X.dim() > 3:
            raise ValueError(f"Input must have at most 3 dimensions, got {X.dim()}.")

        # Determine if the input is batched
        batched = X.dim() == 3

        if not batched:
            N, D = X.shape
            B = 1
        else:
            B, N, D = X.shape
            X = X.reshape(B * N, D)

        posterior = self.model(X).predictive
        # Extract mean and variance
        mean = posterior.mean.squeeze()
        variance = posterior.variance.squeeze(-1) # make sure variance isn't converted to scalar for single value input
        cov = torch.diag_embed(variance)

        K = self.num_outputs
        mean = mean.reshape(B, N * K)

        # Cov must be `(B, N*K, N*K)`
        cov = cov.reshape(B, N, K, B, N, K)
        cov = torch.einsum("bnkbrl->bnkrl", cov)  # (B, N, K, N, K)
        cov = cov.reshape(B, N * K, N * K)

        # Remove fake batch dimension if not batched
        if not batched:
            mean = mean.squeeze(0)
            cov = cov.squeeze(0)

        # pass as MultivariateNormal to GPyTorchPosterior
        mvn_dist = MultivariateNormal(mean, cov)
        post_pred = GPyTorchPosterior(mvn_dist)
        return BLLPosterior(
            posterior=post_pred, model=self, X=X, output_dim=self.num_outputs
        )

making sure, that variance stays a tensor seems to fix the issue

Pull Request

Yes

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[hvarfner]

Hi @CarloMittelbach ,

Thanks for reporting! Not knowing this part of the codebase by heart, I think the culprit is the squeeze here, which lacks a squeeze dimension and thus squeezes all of them in the case when the input is of length 1. This should probably be a squeeze(dim=-1).

I'll have to familiarize myself with the code a bit to make sure this is indeed the issue, so hopefully you can make due with not passing in inputs of length 1 for now!

@brunzema , can you confirm whether this is the issue?

[brunzema]

Hi @CarloMittelbach , hi @hvarfner ,

great catch, thanks for spotting and reporting this! I was able to reproduce this and squeeze(dim=-1) for the mean and var should fix it. @CarloMittelbach, would you like to put together a quick PR for this, as mentioned in your bug report? If not, I’m happy to take care of it over the weekend—whichever works best for you.

If you do go ahead with the PR, it would be great to also include a small unit test in test_community/models/test_vblls.py, just to ensure this case is covered in future updates.

[CarloMittelbach]

Hi @brunzema, @hvarfner
sorry for the late reply, the last few days were quiet busy for me.
If its not to late, I would gladly create a PR in the next few days, however I'd have two more suggestions.

1. It would be nice, if the BBLPosterior implemented all functions of Posterior (even if its just means calling the corresponding function from GPyTorchPosterior) This would make using it more convenient.
   Something like that. (Docstrings would be added later)

class BLLPosterior(Posterior):
    def __init__(
        self,
        posterior: GPyTorchPosterior,
        model: AbstractBLLModel,
        X: Tensor,
        output_dim: int,
    ):
        """A posterior for Bayesian last layer models.

        Args:
            posterior: A posterior object from the feature extractor.
            model: A BLL model.
            X: Input data on which the posterior was computed.
            output_dim: Output dimension of the model.
        """
        super().__init__()
        self.posterior = posterior
        self.model = model
        self.output_dim = output_dim
        self.X = X
        self._is_mt = output_dim > 1

    @property
    def device(self) -> torch.device:
        return self.posterior.device

    @property
    def dtype(self) -> torch.dtype:
        return self.posterior.dtype

    @property
    def mean(self) -> Tensor:
        return self.posterior.mean

    @property
    def variance(self) -> Tensor:
        return self.posterior.variance

    
    def rsample(
        self,
        sample_shape: torch.Size | None = None,
    ) -> Tensor:
        n_samples = 1 if sample_shape is None else math.prod(sample_shape)
        samples_list = [self.model.sample()(self.X) for _ in range(n_samples)]
        samples = torch.stack(samples_list, dim=0)

        # reshape to [sample_shape, n, output_dim]
        sample_shape = torch.Size([1]) if sample_shape is None else sample_shape
        new_shape = sample_shape + samples.shape[-2:]
        return samples.reshape(new_shape)


    def quantile(self, value: Tensor) -> Tensor:
        return self.posterior.quantile(value=value)

    def _extended_shape(
        self,
        sample_shape: torch.Size = torch.Size(),
    ) -> torch.Size:
        return sample_shape + torch.Size([self.X.shape[-2], self.output_dim])

    
    @property
    def base_sample_shape(self) -> torch.Size:
        return self.posterior.base_sample_shape

    def rsample_from_base_samples(
        self,
        sample_shape: torch.Size,
        base_samples: Tensor,
    ) -> Tensor:
        return self.posterior.rsample_from_base_samples(
            sample_shape=sample_shape, base_samples=base_samples
        )

    @property
    def batch_range(self) -> torch.Size:
        return self.posterior.batch_range    

2. There seems to some problem, when using the model with specific acquisitionfunctions:

import torch
from botorch_community.models.vblls import VBLLModel
from botorch.optim import optimize_acqf
from botorch.acquisition import qLogExpectedImprovement
from automahfic.bo.candidates import gen_candidates
from botorch.sampling import SobolQMCNormalSampler

def main():
    torch.set_default_dtype(torch.float64)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    best_value = 0.0
    model = VBLLModel(in_features=1, out_features=1).to(device)
    
    sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
    acq_function = qLogExpectedImprovement(model=model, best_f=best_value, sampler=sampler)

    X_new, _ = optimize_acqf(
            acq_function=acq_function,
            q=1,
            bounds = torch.tensor([[0.0], [1.0]]),
            num_restarts=1,
            raw_samples=512,
            gen_candidates=gen_candidates,
    )

if __name__ == "__main__":
    main()

(this example assumes you use the update BBLPosterior, if not you have to register it)

When running this code you get the following error:

venv/lib/python3.12/site-packages/botorch_community/models/vblls.py", line 494, in posterior
    raise ValueError(f"Input must have at most 3 dimensions, got {X.dim()}.")
ValueError: Input must have at most 3 dimensions, got 4.

Some debugging shows, that optimize_acqf internally creates Tensors with a Dimension > 3, which are passed into the posterior function.
This seems so be happening in candidates.py.
One solution would be, to allow for more than 3 dimensions and reshape accordingly in the posterior function.
I'm currently working on a solution, which isnt finished yet.

[hvarfner]

@CarloMittelbach thanks for working on this!

A closer adherence to the general posterior is certainly better! If a method matches the method of Posterior exactly, you don't have to add it at all since the superclass variant will just be used instead.

Generally, the BoTorch posteriors allow for an arbitrary number of input dimensions. As long as squeezes and unsqueezes are counted backwards (i.e. unsqueeze(-2), squeeze(-3) instead of unsqueeze(0), squeeze(1)) and you adhere to tensor broadcasting rules, this should hopefully work out nicely.

Optimize_acqf creates an additional dim (the q-dim at dimension -2) which allows for batch optimization. This is intentional, but I am not sure why the tensor is 3-dimensional prior to this. My guess is that it shouldn't be, so you should look there to try and find the bug.

So, to recap:

• Posteriors should ideally be able to handle arbitrary-dim tensors, as long as the dimension is greater than one and the last dimension matches the input dimensionality D. As long as the tensor is appended from the front, this should be okay.
• optimize_acqf adds a dimension at position -2

If you're unsure, feel free to ask more or run the VBLL model parallel to the SingleTaskGP - that one works as intended with all shapes. Getting all of this right is not necessary, but we'd like to have working functionality for the most common cases!

[brunzema]

Hi, fully agree that it would be great to extend the VBLLs to also handle acquisition functions besides the analytical ones and super cool that you are working on this @CarloMittelbach ! Some thoughts about what one might need to be careful about when extending the VBLLs base functionality to the batch setting:

Sampling from a VBLL (or in general a BLL) is fundamentally different compared to sampling from a GP as the predictive at different input location is not correlated. I'm not entirely sure on all the technical details on how the MC acquisition functions work under the hood but my rough understanding (and please @hvarfner @Balandat correct me if I am wrong) is that you have your candidate points $X \in R^{q \times d}$ and you evaluate your posterior GP on these candidates to obtain the joint MVN $f(X) \sim \mathcal{N}(\mu(X), \Sigma(X))$ over the candidates. To then approximate, e.g., the intractable batch (q-)EI $EI(X) = \mathbb{E}_q[max(max(f(X)) - y^*, 0)]$ , you draw samples from this joint distribution and compute the empirical average as

$$\hat{EI}(X) = \frac{1}{S} \sum_{s=1}^S max(max(f_s(X)) - y^*, 0), \text{ where } f_s(X) = [f_s(x_1), f_s(x_2), \ldots, f_s(x_q)] \text{ is the s-th sample from the joint MVN}$$

To actually optimize the candidate points, you do the sampling above using the fancy reparameterization trick, i.e., sampling from $f$ using the Cholesky factorization of $\Sigma(X)$, i.e., $f_s(X) = \mu(X) + L \epsilon_s$ where $L L^T = \Sigma(X)$.

Now with VBLLs, your predictions at different input locations are independent, i.e., $f(X) \sim N (\mu(X), \Sigma_d(X))$ where $\Sigma_d(X)$ is diagonal. Instead of a joint posterior over $f$ you actually have joint posterior over $W$ (weights in the last layer) and this posterior has a full covariance. So for VBLLs you actually want to generate your posterior samples as $f_s(X) = \phi_\theta(X)^\top w_s$ where $w_s$ is a sample from $W \sim N(\mu_w, \Sigma_w)$ and then use them for the estimator of the q-EI acquisition function. Also the base samples are no longer realizations $\epsilon_s$ but $w_s$. Now the nice thing is that you actually don't need the reparameterization here because the $f_s$ from a VBLL are directly differentiable (they are just standard MLPs) so the optimization problem might become a bit easier.

That all being said, I think most of that should be taken care of by the current rsample which essentially mimics a GP-like interface (though, given your findings, this may also need some adjustment given a fixed posterior method). However, if we would to add, e.g., a base sampler, we should be mindful about the differences mentioned above.

---

...Sorry for the long message, but I thought it might be useful for future reference as well....