Independent observation noise in MultitaskGaussianLikelihood

[madsendennis]

Hi, I'm new to gpytorch and trying to convert some code from Sklearn to hopefully speed up some processes.
I am working in a 3D setting - so with num_task=3.

I would like to provide an uncertainty to the 3D deformations that are observerved, ideally the full covariance matrix, but to begin with just I*sigma^2.

In Sklearn I do so by:

gpr = GaussianProcessRegressor(kernel=kernel, alpha=np.array([1.0, noise**2]))

In gpytorch I tried to do

noise_deformation_point_0 = torch.tensor([1.0, 1.0, 1.0]).to(device)
noise_deformation_point_1 = torch.tensor([noise**2, noise**2, noise**2]).to(device)
combined_noise = torch.stack([noise_deformation_point_0, noise_deformation_point_1]).reshape(-1)

likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(
    num_tasks=3,
    rank=0,
    # noise_constraint = gpytorch.constraints.GreaterThan(1.0),
    has_global_noise=False,
    has_task_noise=True,
    ).to(device)

likelihood.task_noises = combined_noise

But it only works if I set

likelihood.task_noises = torch.tensor([1.0, 1.0, 1.0]).to(device)

Which to me seems that the same task noise is given to each observation.
I also tried the FixedNoiseGaussianLikelihood, which also doesn't give the output I am looking for.

likelihood = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(noise = torch.tensor([1.0, 1.0, 1.0, noise**2, noise**2, noise**2]), learn_additional_noise=False).to(device)

For context, I created a small toy example where I define a 2D plane.
I then take a corner point and the center point as the places I observe some deformation.
And the deformation being a vector along the Z axis.

I then compute the posterior model and add the deformation to the original points:

i, j = np.indices((board_size, board_size))
points = np.stack((i, j, np.ones_like(i) * board_size/2.0), axis=-1).reshape(-1, 3)

pos_points = [points[0], points[int(board_size**2/2-board_size/2)]]

deformation_vector = torch.tensor([0, 0, board_size], dtype=torch.float32).to(device)
deformation_point_0 = torch.tensor(pos_points[0], dtype=torch.float32).to(device)
deformation_point_1 = torch.tensor(pos_points[1], dtype=torch.float32).to(device)
... 
predicted_distribution = likelihood(model(points_tensor))
predicted_mean = predicted_distribution.mean
new_points = points_tensor + predicted_mean

This here is my model in gpytorch:

class MultitaskGPModel(gpytorch.models.ExactGP):
  def __init__(self, train_x, train_y, likelihood):
      super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
      self.mean_module = gpytorch.means.MultitaskMean(
          gpytorch.means.ZeroMean(), num_tasks=3
      )
      self.covar_module = gpytorch.kernels.MultitaskKernel(
          gpytorch.kernels.RBFKernel(), num_tasks=3
      )

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

In Sklearn I get the expected output. The corner point has constant noise 1.0, so it is the same over the 3 runs. The center point is computed with noise 1, 2 and 5:

In gpytorch I currently get:

I see that there are some similarities to this discussion #2117
Ideally, I would also like a full covariance matrix per observation, but to begin with, a single variance term, just as in Sklearn, would also be fine; I just cannot figure out where to properly define this.