GPyTorch prediction for fixed hyperparameter massively underestimates variances

[theorashid]

Hi there,

I have an exact GP model trying to fit some spatially structured data. The domain of X is around -2 to 2. The y are centred around zero and Gaussian (residuals from a model) and range from around -200 to 200 (sd 73.89).

I struggled to fit this ExactGP with Matern-52 kernel model using MLE through GPyTorch, finding that although the lengthscale converged, the noise and variance. I fitted the model using PyMC's GP with MCMC (faster than Pyro for sampling) to obtain a posterior, finding mean posterior hyperparameter values of lengthscale=0.01, gp_outputscale=75.15 and likelihood noise=34.64, which are sensible values given the sd of the data.

I then fixed these values of the hyperparameters in GPyTorch to these values.

device = "cuda" if torch.cuda.is_available() else "cpu"

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(
        self,
        train_x: torch.Tensor,
        train_y: torch.Tensor,
        likelihood: gpytorch.likelihoods.Likelihood,
    ) -> None:
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.MaternKernel(nu=2.5)
        )

    def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultivariateNormal:
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(X_train, y_train, likelihood).to(device) # conditioned on X_train, y_train
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model).to(device)

model.likelihood.noise_covar.noise = 34.
model.covar_module.outputscale = 62.
model.covar_module.base_kernel.lengthscale = 0.01

model.train()
with torch.no_grad():
    output = model(X_train)
    loss = -mll(output, y_train)
    
model.eval()
likelihood.eval()
with torch.no_grad(), gpytorch.settings.fast_pred_samples():
    test_pred = model(X_test)
    test_samples = test_pred.sample(sample_shape=torch.Size([1000]))

However, these samples are extremely low variance (results below), far less that would be expected from those values of likelihood noise and outputscale.

Question: Is GPyTorch doing any scaling of the data under the hood? (This might also explain why GPyTorch hyperparameter tuning led to the noise and outputscale increasing and increasing).

Best,
Theo

PyMC MCMC samples credible interval width...

are much much larger than GPyTorch credible interval width for values of outputscale and noise based on the MCMC results.

PS: I'm currently working on a reproducible example on simulated data and will update when possible. Posting here before in case there's an obvious answer.

[jacobrgardner]

Your code isn't using the likelihood noise at all in predicting. model(test_x) gives you predictions about the noiseless latent function f, while likelihood(model(test_x)) gives you predictions about y, the noisy version of f.

Try fixing that first and see what happens?

[theorashid]

The intervals from GPyTorch are still too small when adding likelihood(model(test_x)).

But, I got sensible CrI in GPJax with the code at the bottom. I would like to use GPyTorch if possible because it has some really useful features for the task I'm working on.

I have a feeling it is to do with the scale of the variables. The docs say noise is $\sigma^2$, so maybe I should be doing 34. ** 2? The docs also say outputscale is estimated on a log scale, but this isn't clear to me what I should be doing to get the equivalent of the GPJax behaviour.

model.likelihood.noise_covar.noise = 34.
model.covar_module.outputscale = 62.
model.covar_module.base_kernel.lengthscale = 0.01

GPJax code

kern = gpx.kernels.Matern52(active_dims=[0, 1], lengthscale=0.01, variance=62.0**2)
meanf = gpx.mean_functions.Zero()
prior = gpx.gps.Prior(mean_function=meanf, kernel=kern)

likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=34.0)
posterior = prior * likelihood

[jacobrgardner]

If 34 is the standard deviation, which I'm assuming from the GPJax variable name, you'll need to square it for sure. noise in gpytorch refers to the variance of the distribution.

It also looks like you need to be squaring the outputscale. ScaleKernel is outputscale * k(x, x').

[theorashid]

It's still a bit low but squaring the terms and using likelihood(model(X_test)) gets it closer. Not sure why it's not matching GPJax or PyMC but this is the correct maths so I've marked as answer. Thank you