[Bug] Possibly inconsistent likelihood calls with input transforms

[ruicoelhopedro]

🐛 Bug

While tinkering around with heteroskedastic noise models (in the follow-up of #861), I've noticed a possible issue with input transforms on likelihood calls.

In train mode, the likelihood calls receive the X in the original input space. However, in eval mode, both the train and test points are transformed, so the likelihood takes the X in the transformed input space. Is this behaviour expected?

Below is an MWE with a wrapper around a FixedNoiseGaussianLikelihood to print the arguments for each call, which shows the different input spaces. Removing the input_transform from the model fixes the difference.

To reproduce

** Code snippet to reproduce **

import torch
from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from botorch.models.transforms import Normalize


class WrapperLikelihood(FixedNoiseGaussianLikelihood):

    def __call__(self, *args, **kwargs):
        for arg in args:
            print(arg)
        return super().__call__(*args, **kwargs)


X_train = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.float64)
Y_train = torch.tensor([[-1.0], [0.0], [1.0]], dtype=torch.float64)
Yvar_train = torch.full_like(Y_train, 0.1)

model = SingleTaskGP(
    X_train,
    Y_train,
    train_Yvar=Yvar_train,
    likelihood=WrapperLikelihood(Yvar_train.squeeze(-1)),
    input_transform=Normalize(d=X_train.shape[-1]),
)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll(mll)

print("\nTraining complete\n")

x = torch.tensor([[1.5]], dtype=torch.float64)
posterior = model.posterior(x, observation_noise=True)

** Stack trace/error message **

MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)
MultivariateNormal(loc: torch.Size([3]))
tensor([[1.],
        [2.],
        [3.]], dtype=torch.float64)

Training complete

MultivariateNormal(loc: torch.Size([3]))
[tensor([[0.0000],
        [0.5000],
        [1.0000]], dtype=torch.float64)]
MultivariateNormal(loc: torch.Size([3]))
[tensor([[0.0000],
        [0.5000],
        [1.0000]], dtype=torch.float64)]
MultivariateNormal(loc: tensor([-0.5816], dtype=torch.float64, grad_fn=<ViewBackward0>))
tensor([[0.2500]], dtype=torch.float64)

Expected Behavior

The likelihood to be evaluated for the same input space in both modes.

System information

• BoTorch Version 0.11.0
• GPyTorch Version 1.11
• PyTorch Version 2.3.0
• Linux

[saitcakmak]

Hi @ruicoelhopedro. Thanks for reporting! This appears to be due to a missing input transform call here: https://www.internalfb.com/code/fbsource/[414006f6ab1a]/fbcode/pytorch/botorch/botorch/optim/closures/model_closures.py?lines=178 We'll put up a fix