Multiple input adaptation for the SVGP Multitask GP Regression example

[jfaulhab]

Hi,

I am looking for a way to adapt your provided example on Variational GPs with Multiple Outputs (linear dependencies between the outputs) to multiple-input dimensions. How do I need to change the model class MultitaskGPModel if I want to use, let´s say, 2 input dimensions (instead of one) and 4 tasks/output dimensions?

Many thanks in advance for your help and your great work!

Here is some code and the corresponding error message:

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

train_x = torch.stack([torch.linspace(0, 1, 100),
                     torch.linspace(0,-1,100)
                     ], -1)

train_y = torch.stack([
  torch.linspace(0,-10,100),
  torch.linspace(0,-5,100),
  torch.linspace(0,-10,100),
  torch.linspace(0,5,100)
], -1)

num_latents = 3
num_tasks = 4


class MultitaskGPModel(gpytorch.models.ApproximateGP):
  def __init__(self):
      # Let's use a different set of inducing points for each latent function
      inducing_points = torch.rand(num_latents, 16, 1)

      # We have to mark the CholeskyVariationalDistribution as batch
      # so that we learn a variational distribution for each task
      variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
          inducing_points.size(-2), batch_shape=torch.Size([num_latents])
      )

      # We have to wrap the VariationalStrategy in a LMCVariationalStrategy
      # so that the output will be a MultitaskMultivariateNormal rather than a batch output
      variational_strategy = gpytorch.variational.LMCVariationalStrategy(
          gpytorch.variational.VariationalStrategy(
              self, inducing_points, variational_distribution, learn_inducing_locations=True
          ),
          num_tasks=4,
          num_latents=3,
          latent_dim=-1
      )

      super().__init__(variational_strategy)

      # The mean and covariance modules should be marked as batch
      # so we learn a different set of hyperparameters
      self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
      self.covar_module = gpytorch.kernels.ScaleKernel(
          gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_latents])),
          batch_shape=torch.Size([num_latents])
      )

  def forward(self, x):
      # The forward function should be written as if we were dealing with each output
      # dimension in batch
      mean_x = self.mean_module(x)
      covar_x = self.covar_module(x)
      return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

model = MultitaskGPModel()
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)

model.train()
likelihood.train()

optimizer = torch.optim.Adam([
  {'params': model.parameters()},
  {'params': likelihood.parameters()},
], lr=0.1)

# Our loss object. We're using the VariationalELBO, which essentially just computes the ELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))

# We use more CG iterations here because the preconditioner introduced in the NeurIPS paper seems to be less
# effective for VI.
epochs_iter = tqdm.tqdm(range(100), desc="Epoch")
for i in epochs_iter:
  # Within each iteration, we will go over each minibatch of data
  optimizer.zero_grad()
  output = model(train_x)
  loss = -mll(output, train_y)
  epochs_iter.set_postfix(loss=loss.item())
  loss.backward()
  optimizer.step()

The code above is giving me following error message:

Traceback (most recent call last):
File ".../example.py", line 82, in
output = model(train_x)
File ".../gpytorch/models/approximate_gp.py", line 81, in call
return self.variational_strategy(inputs, prior=prior, **kwargs)
File ".../gpytorch/variational/lmc_variational_strategy.py", line 124, in call
function_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
File ".../gpytorch/variational/variational_strategy.py", line 168, in call
return super().call(x, prior=prior, **kwargs)
File ".../gpytorch/variational/_variational_strategy.py", line 129, in call
**kwargs,
File " .../gpytorch/module.py", line 28, in call
outputs = self.forward(*inputs, **kwargs)
File ".../gpytorch/variational/variational_strategy.py", line 89, in forward
full_inputs = torch.cat([inducing_points, x], dim=-2)
RuntimeError: Sizes of tensors must match except in dimension 1. Got 1 and 2 in dimension 2 (The offending index is 1)

[jfaulhab]

I came up with a "reasonable" solution - Changing the Initialization inputs and the inducing points works for model training. Feel free to add some remarks. Here is the code:

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

train_x = torch.stack([torch.linspace(0, 1, 100),
                       torch.linspace(0,-1,100)
                       ], -1)

train_y = torch.stack([
    torch.linspace(0,-10,100),
    torch.linspace(0,-5,100),
    torch.linspace(0,-10,100),
    torch.linspace(0,5,100)
], -1)

num_latents = 2
num_tasks = 4
num_inputs = 2

class MultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self, num_inputs, num_latents, num_tasks, num_inducing_points):
        inducing_points = torch.rand(num_latents, num_inducing_points, num_inputs)

        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_latents])
        )

        variational_strategy = gpytorch.variational.LMCVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=num_tasks,
            num_latents=num_latents,
            latent_dim=-1
        )

        super().__init__(variational_strategy)

        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_latents])),
            batch_shape=torch.Size([num_latents])
        )

    def forward(self, x):
        # The forward function should be written as if we were dealing with each output
        # dimension in batch
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

model = MultitaskGPModel(num_inputs = 2, num_tasks = num_tasks, num_latents = num_latents, num_inducing_points = 1000)
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)

model.train()
likelihood.train()

optimizer = torch.optim.Adam([
    {'params': model.parameters()},
    {'params': likelihood.parameters()},
], lr=0.1)

mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))

epochs_iter = tqdm.tqdm(range(100), desc="Epoch")
loss_ = []
for i in epochs_iter:
    # Within each iteration, we will go over each minibatch of data
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    epochs_iter.set_postfix(loss=loss.item())
    loss_.append(loss.item())
    loss.backward()
    optimizer.step()

plt.plot(loss_)
plt.show()

[gpleiss]

The issue is here:

inducing_points = torch.rand(num_latents, 16, 1)

The inducing points need to have the same dimensionality as your inputs. So it should be:

inducing_points = torch.rand(num_latents, 16, 2)