How to use Batch Independent Multioutput GPs with Inducing Points?

[famura]

Hi there,

I want to use batch-independent GPs (see this tutorial) with the inducing points method (see this tutorial). I managed to do so for the fully-correlated multi-task setting using the MultitaskKernel, however, fail at doing this in the described case.

Here is my minimal example script which is almost equal to the first tutorial above. Both versions I tried failed with the same (output-dim based?) error. I have the feeling that I just missed one detail.

import math
import torch
import gpytorch
from matplotlib import pyplot as plt


train_x = torch.linspace(0, 1, 200)

train_y = torch.stack([
    torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
], -1)

class BatchIndependentMultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([2]))

        # v1 -- RuntimeError: grad can be implicitly created only for scalar outputs
        base_covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2])),
            batch_shape=torch.Size([2])
        )
        self.covar_module = gpytorch.kernels.InducingPointKernel(
            base_covar_module,
            inducing_points=train_x[::10].clone(), # every 10th
            likelihood=likelihood
        )

        # v2 -- RuntimeError: grad can be implicitly created only for scalar outputs
        # base_covar_module = gpytorch.kernels.InducingPointKernel(
        #     gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2])),
        #     inducing_points=train_x[::10].clone(),  # every 10th
        #     likelihood=likelihood
        # )
        # self.covar_module = gpytorch.kernels.ScaleKernel(
        #     base_covar_module,
        #     batch_shape = torch.Size([2])
        # )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(
            gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        )


likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=2)
model = BatchIndependentMultitaskGPModel(train_x, train_y, likelihood)

# this is for running the notebook in our testing framework
import os
training_iterations = 50


# Find optimal model hyperparameters
model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)  # Includes GaussianLikelihood parameters

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))
    optimizer.step()

# Set into eval mode
model.eval()
likelihood.eval()

# Initialize plots
f, (y1_ax, y2_ax) = plt.subplots(1, 2, figsize=(12, 6))

# Make predictions
with torch.no_grad(), gpytorch.settings.fast_pred_var():
    test_x = torch.linspace(0, 1, 51)
    predictions = likelihood(model(test_x))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

# This contains predictions for both tasks, flattened out
# The first half of the predictions is for the first task
# The second half is for the second task

# Plot training data as black stars
y1_ax.plot(train_x.detach().numpy(), train_y[:, 0].detach().numpy(), 'k*')
# Predictive mean as blue line
y1_ax.plot(test_x.numpy(), mean[:, 0].numpy(), 'b')
# Shade in confidence
y1_ax.fill_between(test_x.numpy(), lower[:, 0].numpy(), upper[:, 0].numpy(), alpha=0.5)
y1_ax.set_ylim([-3, 3])
y1_ax.legend(['Observed Data', 'Mean', 'Confidence'])
y1_ax.set_title('Observed Values (Likelihood)')

# Plot training data as black stars
y2_ax.plot(train_x.detach().numpy(), train_y[:, 1].detach().numpy(), 'k*')
# Predictive mean as blue line
y2_ax.plot(test_x.numpy(), mean[:, 1].numpy(), 'b')
# Shade in confidence
y2_ax.fill_between(test_x.numpy(), lower[:, 1].numpy(), upper[:, 1].numpy(), alpha=0.5)
y2_ax.set_ylim([-3, 3])
y2_ax.legend(['Observed Data', 'Mean', 'Confidence'])
y2_ax.set_title('Observed Values (Likelihood)')

plt.show()

[famura]

I created issue #2384 since there is no response here.

[gpleiss]

The inducing point tutorial you linked to was for SVGP, but here you're using SGPR?

If you want to do batched independent GPs, you just need to supply the batch_shape parameter to all of the kernel/mean/likelihood/variational objects that you create.

This tutorial should cover exactly what you're looking for: https://docs.gpytorch.ai/en/stable/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html

[famura]

Thanks @gpleiss for your reply. In the end I just missed that example. When I found it, I closed the related issue but forgot to close this discussion too. Sorry for that.