[Question] Proper use of `Normalize` for inputs (possible bug?)

[lukasfro]

Issue description

I'm struggling to understand how to properly use the input normalization for the SingleTaskGP (presumably the same issue arises for different models). As far as I know, the default kernel hyper-parameters are optimized for normalized inputs and standardized targets. Hence, I would expect to obtain reasonable posterior predictions when the inputs are normalized but somewhat unreasonable length scales for unnormalized inputs. However, I observed the exact opposite when running the small code example below. I'm unsure if I'm just using Normalize in a wrong manner, if the observed behaviour is intended, or if I'm just completely off track. Any help is much appreciated!

Code example

import torch
import matplotlib.pyplot as plt

from botorch.models.transforms import Normalize, Standardize
from botorch.models import SingleTaskGP
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.fit import fit_gpytorch_scipy

import math


def f(x):
    return torch.sin(math.pi * x)


def train_gp(x, y, input_transform):
    gp = SingleTaskGP(
        train_X=x,
        train_Y=y,
        outcome_transform=Standardize(m=1),
        input_transform=input_transform,
    )
    mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
    mll.train()
    fit_gpytorch_scipy(mll)
    mll.eval()

    return gp


def plot_gp(gp, x_plot, color, label):
    p = gp.posterior(x_plot)
    m = p.mean.squeeze().detach()
    s = p.variance.sqrt().squeeze().detach()

    plt.plot(x_plot, m, color, label=label)
    plt.fill_between(x_plot.squeeze(), m + s, m - s, color=color, alpha=0.3)


# Problem set up
d = 1
bounds = torch.tensor([[-2.0], [2.0]])
train_X = bounds[0] + (bounds[1] - bounds[0]) * torch.rand((10, 1))
train_Y = f(train_X)
train_Y += 0.1 * torch.randn_like(train_Y)
eval_X = torch.linspace(bounds[0].item(), bounds[1].item(), 500).unsqueeze(-1)

# Train 3 models with different input normalization
gp1 = train_gp(train_X, train_Y, Normalize(d=d, bounds=bounds))
gp2 = train_gp(train_X, train_Y, Normalize(d=d, bounds=bounds, transform_on_eval=False))
gp3 = train_gp(train_X, train_Y, None)

# Visualize
plt.figure()
plt.plot(eval_X, f(eval_X), "k")
plt.plot(train_X, train_Y, "ro")
plot_gp(gp1, eval_X, "C0", label="Normalize")
plot_gp(gp2, eval_X, "C1", label="Normalize(transform_on_eval=False)")
plot_gp(gp3, eval_X, "C2", label="No Normalization")
plt.legend()
plt.show()

Exemplary Output

System Info

Please provide information about your setup, including

• BoTorch Version: 0.5.0
• GPyTorch Version: 1.5.0
• PyTorch Version: 1.8.1
• Computer OS: macOS Big Sur (11.3.1)

[saitcakmak]

Hi @lukasfro. The issue here is that you're not training your models, so the GP hyper-parameters are set to some arbitrary initial values. If you do

gp1 = train_gp(train_X, train_Y, Normalize(d=d, bounds=bounds))
mll = ExactMarginalLogLikelihood(gp1.likelihood, gp1)
fit_gpytorch_model(mll)

you should get a plot that looks much more similar to what you expect. Regarding the correct use of Normalize, you should use it as in gp1 with the bounds argument being optional (it is inferred if not given). Note that training the model via fit_gpytorch_model is necessary for the input transforms to work correctly (this will store the training inputs as transformed). The transform_on_eval=False argument would be appropriate if you intended to train the model with inputs from original scale ([-2, 2]^d) and evaluate it with inputs from the standardized scale ([0, 1]^d).

[lukasfro]

Hi @saitcakmak, thanks for your quick reply! With your recommended solution it works fine. Just to fully understand the issue here: why does the training inside the train_gp function has no effect? Or is the issue that I'm using fit_gpytorch_scipy instead?

[saitcakmak]

Ah, looks like I completely missed that you were training the model there. Short answer is yes, it is because you're using fit_gpytorch_scipy. fit_gpytorch_model internally uses fit_gpytorch_scipy to optimize the mll. Once the optimization is done, it calls _set_transformed_inputs, which updates the model training inputs with the transformed ones.

If you were curious about why this is necessary: We apply input transforms in model.forward(...) in train mode and in model.posterior(...) in eval mode. Because of how things work on gpytorch end, applying the transforms in model.posterior(...) means we need to have the model.train_inputs stored as transformed (otherwise we would be mixing transformed and non-transformed inputs). fit_gpytorch_model handles this by calling _set_transformed_inputs, which is not done by fit_gpytorch_scipy. Hopefully, this won't be necessary in the future (cornellius-gp/gpytorch#1652) but this is how things work at the moment to make sure we can support the one-to-many transforms introduced with 0.5 release (#820, #827).

[lukasfro]

Wow, thanks for the detailed reply. I very much appreciate it! Guess I'll be using fit_gpytorch_model from now on. I was wondering already why this is not yet a part of GPyTorch, but now it's clear to me.