GP training with derivatives at different input points than the function observations

[ankushaggarwal]

I am trying to create a GP that is trained on derivatives and function values, but at different input points. You have already shared an example for GP regression with derivatives (https://docs.gpytorch.ai/en/latest/examples/08_Advanced_Usage/Simple_GP_Regression_Derivative_Information_1d.html), however, as far as I understand, it assumes that the function values and derivative observations are provided at the same input points.

Is there a way to relax this assumption? I suspect it might be similar to the Hadamard Multitask GP Regression? Do I need to implement a new kernel or can I use the existing ones?

[gpleiss]

I am trying to create a GP that is trained on derivatives and function values, but at different input points.

Can you describe the problem in more detail? I'm a bit confused by the setup.

[ankushaggarwal]

Sorry that my question wasn't clear. The example for GP regression with derivatives uses two tasks at the same inputs. That is, train_y contains the function values and derivatives at the same values of train_x

n = 50

train_x = torch.linspace(lb, ub, n).unsqueeze(-1)
train_y = torch.stack([
    torch.sin(2*train_x) + torch.cos(train_x),
    -torch.sin(train_x) + 2*torch.cos(2*train_x)
], -1).squeeze(1)

train_y += 0.05 * torch.randn(n, 2)

However, I need something like the following

n1 = 50 #function values
train_x1 = torch.linspace(lb, ub, n1).unsqueeze(-1)
train_y1 = torch.sin(2*train_x1) + torch.cos(train_x) + 0.05 * torch.randn(n1, 2)

n2=20 #derivative values at different x locations
train_x2 = torch.linspace(lb, ub, n1).unsqueeze(-1)
train_y2 = -torch.sin(train_x) + 2*torch.cos(2*train_x) + 0.05 * torch.randn(n2, 2)

How can I train a GP using both train_y1 and train_y2?

[gpleiss]

I see. You probably need to implement a new kernel then. The derivative kernels currently assume that there are d + 1 outputs per input.

[ankushaggarwal]

Thanks @gpleiss That is what I suspected. Do you have any suggestions to start? I am wondering if I can follow Hamadard example ? What changes will I need to make?

[gpleiss]

You could get a really simple implementation using gpytorch.kernels.RBFKernelGrad and indexing:

# x =[x1 , x2]
# We have function observations at x1 (size n1 x d)
# We have derivative observations at x2 (size n2 x d)

# Assume self.covar_module = gpytorch.kernels.RBFKernelGrad()
full_kernel = self.covar_module(x)
# this is a ((n1 + n2) + (n1 x n2)) x ((n1 + n2) + (n1 x n2)) kernel
# 

index = torch.cat([
    torch.arange(n1, device=x.device),  # index of function covar information for x1
    torch.arange(n2, device=x.device) + 2 * n1 + n2,  # index of deriv covar information for x2
])
return full_kernel[..., index, index]

[ankushaggarwal]

Thanks, @gpleiss I tried to follow that and have written this code. However it does not behave as expected. There are huge oscillations using the index that you suggested. I tried another option, which although gives smoother behavior, its fit to data is rather poor and it gives NaN for confidence_region.

I am pasting the code below. Could you please help me find where the error is?

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

lb, ub = 0.0, 5*math.pi
n1 = 50 #function values
train_x1 = torch.linspace(lb, ub, n1)
train_y1 = torch.sin(2*train_x1) + torch.cos(train_x1) + 0.05 * torch.randn(train_x1.size())

n2=20 #derivative values at different x locations
train_x2 = torch.linspace(lb, ub, n2)
train_y2 = -torch.sin(train_x2) + 2*torch.cos(2*train_x2) + 0.05 * torch.randn(train_x2.size())

train_x = torch.cat([train_x1 , train_x2])
train_y = torch.cat([train_y1,train_y2])

class GPModelWithDerivatives(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPModelWithDerivatives, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.base_kernel = gpytorch.kernels.RBFKernelGrad()
        self.covar_module = gpytorch.kernels.ScaleKernel(self.base_kernel)

    def forward(self, x1, x2):
        x = torch.cat([x1,x2])
        n1 = x1.shape[0]
        n2 = x2.shape[0]
        option_index = 1 #1 for the suggestion by @gpleiss, 2 for (value, gradient, value, gradient, ...) format
        if option_index==1:
            index = torch.cat([
                torch.arange(n1, device=x.device),  # index of function covar information for x1
                torch.arange(n2, device=x.device) + 2 * n1+n2,  # index of deriv covar information for x2
            ])
        else:
            index = torch.cat([
                torch.arange(0,2*n1,2, device=x.device),  # index of function covar information for x1
                torch.arange(1,2*n2,2, device=x.device) + 2 * n1,  # index of deriv covar information for x2
            ])
        mean_x = self.mean_module(x)
        full_kernel = self.covar_module(x)
        covar_x = full_kernel[..., index,:][...,:,index]
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPModelWithDerivatives((train_x1,train_x2), train_y, likelihood)

# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
training_iter = 2 if smoke_test else 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_iter):
    optimizer.zero_grad()
    output = model(train_x1,train_x2)
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (
        i + 1, training_iter, loss.item(),
        model.covar_module.base_kernel.lengthscale.item(),
        model.likelihood.noise.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))

n11,n22=100,100
test_x1 = torch.linspace(lb, ub, n11)
test_x2 = torch.linspace(lb, ub, n22)

# Make predictions
with torch.no_grad(), gpytorch.settings.max_cg_iterations(50):
    predictions = likelihood(model(test_x1,test_x2))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

# Plot training data as black stars
y1_ax.plot(train_x[:n1].detach().numpy(), train_y[:n1].detach().numpy(), 'k*')
# Predictive mean as blue line
y1_ax.plot(test_x1.detach().numpy(), mean[:n11].detach().numpy(), 'b')
# Shade in confidence
y1_ax.fill_between(test_x1.detach().numpy(), lower[:n11].detach().numpy(), upper[:n11].detach().numpy(), alpha=0.5)
y1_ax.legend(['Observed Values', 'Mean', 'Confidence'])
y1_ax.set_title('Function values')

# Plot training data as black stars
y2_ax.plot(train_x[n1:].detach().numpy(), train_y[n1:].detach().numpy(), 'k*')
# Predictive mean as blue line
y2_ax.plot(test_x2.detach().numpy(), mean[n11:].detach().numpy(), 'b')
# Shade in confidence
y2_ax.fill_between(test_x2.detach().numpy(), lower[n11:].detach().numpy(), upper[n11:].detach().numpy(), alpha=0.5)
y2_ax.legend(['Observed Derivatives', 'Mean', 'Confidence'])
y2_ax.set_title('Derivatives')

plt.show()

Result with option_index=1

Result with option_index=2

[ankushaggarwal]

I think I have figured out a solution, but I have to comment out an error to make it work. Here is my attempt

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

lb, ub = 0.0, 5*math.pi
n1 = 50 #function values
train_x1 = torch.linspace(lb, ub, n1)#.unsqueeze(-1)
train_y1 = torch.sin(2*train_x1) + torch.cos(train_x1) + 0.05 * torch.randn(train_x1.size())

n2=20 #derivative values at different x locations
train_x2 = torch.linspace(lb, ub, n2)#.unsqueeze(-1)
train_y2 = -torch.sin(train_x2) + 2*torch.cos(2*train_x2) + 0.05 * torch.randn(train_x2.size())

train_x = torch.cat([train_x1 , train_x2])
train_y = torch.cat([train_y1,train_y2])

train_index = torch.tensor([True]*n1+[False]*n2)

class GPModelWithDerivatives(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPModelWithDerivatives, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.base_kernel = gpytorch.kernels.RBFKernelGrad()
        self.covar_module = gpytorch.kernels.ScaleKernel(self.base_kernel)

    def forward(self, x, index):
        index = index.squeeze()
        index = torch.stack((index,~index)).T.reshape(-1)
        mean_x = self.mean_module(x)
        full_kernel = self.covar_module(x)
        covar_x = full_kernel[..., index,:][...,:,index]
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPModelWithDerivatives((train_x,train_index), train_y, likelihood)

# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
training_iter = 2 if smoke_test else 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_iter):
    optimizer.zero_grad()
    output = model(train_x,train_index)
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (
        i + 1, training_iter, loss.item(),
        model.covar_module.base_kernel.lengthscale.item(),
        model.likelihood.noise.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))

n11,n22=200,200
test_x1 = torch.linspace(lb, ub, n11)
test_x2 = torch.linspace(lb, ub, n22)
test_x = torch.cat([test_x1 , test_x2])
test_index = torch.tensor([True]*n11+[False]*n22)

# Make predictions
with torch.no_grad(), gpytorch.settings.max_cg_iterations(50):
    predictions = likelihood(model(test_x,test_index))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

# Plot training data as black stars
y1_ax.plot(train_x[:n1].detach().numpy(), train_y[:n1].detach().numpy(), 'k*')
# Predictive mean as blue line
y1_ax.plot(test_x1.detach().numpy(), mean[:n11].detach().numpy(), 'b')
# Shade in confidence
y1_ax.fill_between(test_x1.detach().numpy(), lower[:n11].detach().numpy(), upper[:n11].detach().numpy(), alpha=0.5)
y1_ax.legend(['Observed Values', 'Mean', 'Confidence'])
y1_ax.set_title('Function values')

# Plot training data as black stars
y2_ax.plot(train_x[n1:].detach().numpy(), train_y[n1:].detach().numpy(), 'k*')
# Predictive mean as blue line
y2_ax.plot(test_x2.detach().numpy(), mean[n11:].detach().numpy(), 'b')
# Shade in confidence
y2_ax.fill_between(test_x2.detach().numpy(), lower[n11:].detach().numpy(), upper[n11:].detach().numpy(), alpha=0.5)
y2_ax.legend(['Observed Derivatives', 'Mean', 'Confidence'])
y2_ax.set_title('Derivatives')

plt.show()

However, at the slicing covar_x = full_kernel[..., index,:][...,:,index] it raises an error File "/opt/anaconda3/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1732, in __getitem__ "or your custom LazyTensor.".format(self.__class__.__name__, expected_shape, res.shape) RuntimeError: NonLazyTensor.__getitem__ failed! Expected a final shape of size torch.Size([70, 140]), got torch.Size([70, 70]). This is a bug with GPyTorch, or your custom LazyTensor.

If I comment out the raise RuntimeError in the lazy_tensor.py file, the code runs fine and gives expected result. I think this might be a bug. I am starting a new post regarding this.

[gpleiss]

Thanks for the bug report. For now you can also replace that line with covar_x = full_kernel.evaluate_kernel()[..., index,:][...,:,index], and that should also work.

[ankushaggarwal]

I get the same error, unfortunately. In fact, after using the evaluate_kernel, I get the error on the first indexing operation.
That is, full_kernel.evaluate_kernel()[..., index,:] raises an error, while full_kernel[..., index,:] did not.