Converting torch mean and var tensors into multi-output posterior objectives for qNoisyExpectedHypervolumeImprovement

[exs-hkenlay]

I'm attempting to use qNoisyExpectedHypervolumeImprovement acquisition function with a PyTorch model which outputs a mean and var tensor describing a normal distribution for each objective.

Before describing my issue, I noticed the following observation which I think provides a clue as to why my solution is not working. In the multi-objective Bayesian optimisation tutorial, the sampler generates samples of shape samples x b x m from the posterior as seen below. I'm using 512 samples, a batch size of 32 and a model that takes a 4-dimensional input and outputs 2 objectives.

from botorch.models.gp_regression import FixedNoiseGP
from botorch.models.model_list_gp_regression import ModelListGP
from botorch.models.transforms.outcome import Standardize
from gpytorch.mlls.sum_marginal_log_likelihood import SumMarginalLogLikelihood
from botorch.sampling.normal import SobolQMCNormalSampler
import torch

def initialize_model(train_x, train_obj):
    # define models for objective and constraint
    models = []
    for i in range(train_obj.shape[-1]):
        train_y = train_obj[..., i : i + 1]
        models.append(
            FixedNoiseGP(
                train_x, train_y, train_y, outcome_transform=Standardize(m=1)
            )
        )
    model = ModelListGP(*models)
    mll = SumMarginalLogLikelihood(model.likelihood, model)
    return mll, model

_, model = initialize_model(torch.rand(32, 4), torch.rand(32, 2))
posterior = model.posterior(torch.rand(32, 4))
sampler = SobolQMCNormalSampler(torch.Size([512]))
print(sampler(posterior).shape) # outputs torch.Size([512, 32, 2])

Here I use a toy model that takes the first two dimensions as the mean and the next two as the variance. I try to use MultivariateNormal from gpytorch, but an extra dimension is added.

from gpytorch.distributions import MultivariateNormal
from botorch.posteriors.gpytorch import GPyTorchPosterior
from botorch.models.model import Model
import torch
from botorch.sampling.normal import SobolQMCNormalSampler

class Surrogate(Model):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        mean = x[:, :2]
        cov = x[:, 2:]
        return mean, cov

    def posterior(self, x: torch.tensor):
        mean, cov = self(x)
        cov = torch.diag_embed(cov) # (b, d) -> (b, d, d)
        distribution = MultivariateNormal(mean, cov)
        return GPyTorchPosterior(distribution)

surrogate = Surrogate()
sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
posterior = surrogate.posterior(torch.rand(32, 4))
print(sampler(posterior).shape) # outputs torch.Size([512, 32, 2, 1])

I then tried the MultitaskMultivariateNormal distribution because I saw here https://botorch.org/api/posteriors.html#module-botorch.posteriors.gpytorch that it should be used for multiple outputs. I use unsqueeze(1) corresponding to two tasks (task in gpytorch seems to correspond to output in botorch?). This changes the output dimension but still adds an extra dimension.

from gpytorch.distributions import MultitaskMultivariateNormal
from botorch.posteriors.gpytorch import GPyTorchPosterior
from botorch.models.model import Model
import torch
from botorch.sampling.normal import SobolQMCNormalSampler

class Surrogate(Model):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        mean = x[:, :2]
        cov = x[:, 2:]
        return mean, cov

    def posterior(self, x: torch.tensor):
        mean, cov = self(x)
        cov = torch.diag_embed(cov) # (b, d) -> (b, d, d)
        mean = mean.unsqueeze(1) # (b, d) -> (b, 1, d)
        distribution = MultitaskMultivariateNormal(mean, cov)
        return GPyTorchPosterior(distribution)

surrogate = Surrogate()
sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
posterior = surrogate.posterior(torch.rand(32, 4))
print(sampler(posterior).shape) # outputs torch.Size([512, 32, 1, 2])

when I use either of these surrogate models and pass them to the qNoisyExpectedHypervolumeImprovement acqusition function using this snippet

from botorch.acquisition.multi_objective.monte_carlo import (
    qNoisyExpectedHypervolumeImprovement,
)

acq_func = qNoisyExpectedHypervolumeImprovement(
    model=surrogate,
    ref_point=(0, 0),
    X_baseline = torch.rand(32, 4),
    sampler = sampler
)

I get an error at this line https://github.com/pytorch/botorch/blob/main/botorch/acquisition/multi_objective/monte_carlo.py#L580. The error I get is RuntimeError: The q-batch shape of the objective values does not agree with the q-batch shape of X. Got 2 and 32. This may happen if you used a one-to-many input transform but forgot to use a corresponding objective., I'm not quite sure what this means. The self.objective here is IdentityMCMultiOutputObjective by default which subclasses MCMultiOutputObjective. In the docstring for MCMultiOutputObjective it says the samples should be of shape sample_shape x batch_shape x q x m (which is the case for when I use MultitaskMultivariateNormal, but not for the ModelListGP used in the MOO tutorial).

Based on this it seems the issue either comes from how the objective function is used in qNoisyExpectedHypervolumeImprovement or that the Posterior objects I'm using include an empty dimension. I guess its the later because qParEGO in the tutorial uses a posterior which samples to shape samples x b x m too. Is this the correct approach to go from mean and var tensors to a Posterior type object? Any help would be appreciated, thanks.

System Info

Please provide information about your setup, including

• BoTorch Version: 0.9.2
• GPyTorch Version: 1.11
• PyTorch Version: 2.0.0
• Computer OS: Amazon Linux

[saitcakmak]

Hi @hkenlay. Here's a toy example that shows various tensor shapes that we expect when working with qNEHVI. I used a SingleTaskGP rather than ModelListGP as the tutorial does, though this will not affect the shapes at all.

import torch
from botorch.models.gp_regression import SingleTaskGP
from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement

torch.set_default_dtype(torch.float64)
train_X = torch.rand(10, 4)
train_Y = torch.randn(10, 2)
model = SingleTaskGP(train_X, train_Y).eval()
qNEHVI = qNoisyExpectedHypervolumeImprovement(
    model=model,
    ref_point=torch.zeros(2),
    X_baseline=train_X,
)
test_X = torch.rand(64, 8, 4)  # batch_shape=64 x q=8 x d=4
posterior = model.posterior(test_X)
samples = qNEHVI.get_posterior_samples(posterior)  # default sampler has sample_shape=128. Samples has shape sample_shape=128 x batch_shape=64 x q=8 x m=2
acq_vals = qNEHVI(test_X)  # has shape batch_shape=64
qNEHVI.sampler.base_samples.shape  # This is sample_shape=128 x 1 (we're re-using same base samples across batch_shape) x q=18 (=10 from train_X + 8 from test_X) x m=2

In the multi-objective Bayesian optimisation tutorial, the sampler generates samples of shape samples x b x m from the posterior as seen below.

This is the case because the posterior does not have a batch shape in your example. The acquisition functions always use an explicit batch shape (via @t_batch_mode_transform). If you update the code, you'll get the batch shape in there.

posterior = model.posterior(torch.rand(batch, 32, 4))
sampler = SobolQMCNormalSampler(torch.Size([512]))
print(sampler(posterior).shape) # Will output torch.Size([512, batch, 32, 2])

I am not fully sure whether you should be using MultivariateNormal or MultiTaskMultivariateNormal. This will depend on the covariance your model produces. You should check to make sure the model produces samples with the correct correlation.

For example, SingleTaskGP uses MultivariateNormal regardless of the number of outputs. For multi-output use cases, it internally represents the inputs in a batched manner, so you have something like this:

    def forward(self, x: Tensor) -> MultivariateNormal:
        # x is batch x m x q x d
        mean_x = self.mean_module(x)  # batch x m x q
        covar_x = self.covar_module(x)  # batch x m x q x q
        return MultivariateNormal(mean_x, covar_x)

[exs-hkenlay]

Hello @saitcakmak, thank you for your detailed reply, I really appreciate your help. Although I think I understand your example, I'm still struggling to get my own toy example working. The first question I have from working through your example and putting debug flags throughout BoTorch source code is the following:

In the docstring of the Model base class it states the input of the posterior method should be X, a tensor of shape b x q x d. However, in this line of the acquisition function self.X_baseline is fed to the posterior method, which is of shape r x d (where I guess r can be thought of as a t-batch dimension). Is it the case that the posterior method should also support an input tensor X of shape b x d?

Thank you for your suggestion to implement the batch_size class variable method, I added this but ran into the same error. I believe the issue I'm running into is occurring before the lack of batch_size is a problem, otherwise the NotImplementedError would be raised from the Model base class. The acquisition function forward is decorated by t_batch_mode_transform but it seems the error I see is called before the forward is called. In the __init__ of the qNoisyExpectedHypervolumeImprovement object the self._set_cell_bounds method is called. The code then fails in the line obj = self.objective(samples, X=self.X_baseline). It seems likely to me that this is because samples is the wrong shape, presumably because I've constructed the posterior object incorrectly.

It seems like the SingleTaskGP forward uses a MultivariateNormal, but the posterior combines the outputs into a MultiTaskMultivariateNormal. This behaviour is inherited from the base class.

For my use case, I have a neural network which takes as input a batch of size b x d and outputs a mean tensor and a variance tensor both of size b x m which represent b x m Gaussian distributions which are assumed to be independent (both across the batch dimension and the objective dimension). Here b is a generic dimension and could be a q-batch or a t-batch (since independence is assumed across all input and output dimensions).

The main goal here is to somehow go from these tensors to a Posterior object that I can then feed into the qNEHVI acquisition function. I'm starting to think the gpytorch Posteriors are not appropriate for non-GP models, is this correct? Is there a Posterior object I can use for a univariate Gaussian per output (each with its own mean and variance), or do I need to write my own?

[saitcakmak]

In the docstring of the Model base class it states the input of the posterior method should be X, a tensor of shape b x q x d. However, in this line of the acquisition function self.X_baseline is fed to the posterior method, which is of shape r x d (where I guess r can be thought of as a t-batch dimension). Is it the case that the posterior method should also support an input tensor X of shape b x d?

In BoTorch, we typically express the shapes of inputs to a model as batch x q x d, with the implicit assumption that the batch might be empty (so, q x d is also generally supported). We separate q from the rest of the batch (also referred to as t-batch) since the q-batch typically represents inputs that get jointly evaluated to produce a full-covariance multi-output posterior, whereas the t-batches typically behave independently. When the posterior is called with a q x d-dim tensor, this is equivalent to calling it with a 1 x q x d-dim tensor with the exception of the additional batch (of size 1) in the output. If you were to instead use a q x 1 x d-dim tensor, this would be very much like calling it in a for loop with 1 x 1 x d-dim tensors, which will produce independent posteriors for each input.

For my use case, I have a neural network which takes as input a batch of size b x d and outputs a mean tensor and a variance tensor both of size b x m which represent b x m Gaussian distributions which are assumed to be independent (both across the batch dimension and the objective dimension). Here b is a generic dimension and could be a q-batch or a t-batch (since independence is assumed across all input and output dimensions).

The main goal here is to somehow go from these tensors to a Posterior object that I can then feed into the qNEHVI acquisition function. I'm starting to think the gpytorch Posteriors are not appropriate for non-GP models, is this correct? Is there a Posterior object I can use for a univariate Gaussian per output (each with its own mean and variance), or do I need to write my own?

If your setup does not produce full-covariance posteriors, the effort to fit it into a GPyTorch MVN is probably unnecessary. One option for you would be to use PyTorch MVN, which should represent independent batched posteriors quite easily. We have a TorchPosterior class that you can pair with the PyTorch MVN objects: https://github.com/pytorch/botorch/blob/main/botorch/posteriors/torch.py#L21C33-L21C33

[saitcakmak]

Actually, you don't even need an MVN if it is all independent. You can just use PyTorch Normal distribution: https://pytorch.org/docs/stable/distributions.html#normal

[exs-hkenlay]

Thank you for your help and patience @saitcakmak. I took the PyTorch route as you suggested and managed to get a working solution. I needed to implement some functionality from the Posterior API which are not implemented by TorchPosterior. I also needed to pass a sampler to qNoisyExpectedHypervolumeImprovement because my new posterior class was not registered. For the benefit of anyone else who stumbles across this issue the full code is:

from typing import Tuple

import torch
from botorch.acquisition.multi_objective.monte_carlo import (
    qNoisyExpectedHypervolumeImprovement,
)
from botorch.models.model import Model
from botorch.posteriors.torch import TorchPosterior
from botorch.sampling.normal import SobolQMCNormalSampler
from torch import Tensor
from torch.distributions import Normal


class Surrogate(Model):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, x: torch.Tensor):
        mean = x[..., :2]
        var = x[..., 2:] ** 2
        return mean, var

    def posterior(self, X: torch.Tensor):
        mean, var = self.forward(X)
        return NormalPosterior(mean, var)


class NormalPosterior(TorchPosterior):
    def __init__(
        self, mean: torch.Tensor, var: torch.Tensor, jitter: float = 10e-6
    ) -> None:
        self.mean = mean
        self.var = var
        self.jitter = jitter
        super().__init__(Normal(mean, var ** (1 / 2) + jitter))

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

    @property
    def batch_range(self) -> Tuple[int, int]:
        return 0, -1

    def rsample_from_base_samples(
        self,
        sample_shape: torch.Size,
        base_samples: Tensor,
    ) -> Tensor:
        return self.loc + base_samples * self.scale


model = Surrogate()

qNEHVI = qNoisyExpectedHypervolumeImprovement(
    model=model,
    ref_point=torch.zeros(2),
    X_baseline=torch.zeros(1, 4),
    sampler=SobolQMCNormalSampler(torch.Size([512])),
)
test_X = torch.rand(64, 8, 4)  # b=64 x q=8 x d=4
posterior = model.posterior(test_X)
samples = qNEHVI.get_posterior_samples(posterior)  # samples = 512, b=64, q=8, m=2
acq_vals = qNEHVI(test_X)  # b=64

[saitcakmak]

Happy to help & thanks for sharing your solution! I'll move this over to a discussion for future discoverability