Implementing a GP Adapter

[prockenschaub]

I am trying to implement a Gaussian Process Adapter as described in Li and Marlin (2016). The basic idea is to use a GP to convert a dataset of N irregularly sampled time-series into a fixed time grid of size NxT, which is then coupled with a black-box model (e.g., RNN) for classification/regression.

A key difference in this setup is that the loss is not calculated on the "training data" (which are the irregularly sampled key-value-pairs) via a marginal log-likelihood but on the posterior predictions using a standard prediction loss (e.g., MSE or cross-entropy). That means model training moves from gpytorch's gp_model.train() to gp_model.eval().

I believe this currently does not work with an gpytorch.models.ExactGP out of the box because of the caches used in the prediction strategy. As the kernel parameters still change, we must recalculate the train-train covariance at each iteration, precluding the use of caches. I have tried to do so by subclassing DefaultPredictionStrategy and forcing it to recalculate the cache each time (making "cache" a misnomer that's just left for compatibility):

class AdapterPredictionStrategy(DefaultPredictionStrategy):
    def __init__(self, train_inputs, train_prior_dist, train_labels, likelihood):
        super(AdapterPredictionStrategy, self).__init__(train_inputs, train_prior_dist, train_labels, likelihood)

    @property
    def covar_cache(self):
        train_train_covar = self.lik_train_train_covar
        train_train_covar_inv_root = delazify(train_train_covar.root_inv_decomposition().root)
        return train_train_covar_inv_root

    @property
    def mean_cache(self):
        mvn = self.likelihood(self.train_prior_dist, self.train_inputs)
        train_mean, train_train_covar = mvn.loc, mvn.lazy_covariance_matrix

        train_labels_offset = (self.train_labels - train_mean).unsqueeze(-1)
        mean_cache = train_train_covar.evaluate_kernel().inv_matmul(train_labels_offset).squeeze(-1)

        return mean_cache

This custom prediction strategy is then set manually for each iteration of the training cycle.

for iter in range(num_iter):
    optimizer.zero_grad()
    
    adapter.train()
    train_prior_dist = gp(irreg_x)
    adapter.prediction_strategy = AdapterPredictionStrategy(
        train_inputs=irreg_x,
        train_prior_dist=train_prior_dist,
        train_labels=irreg_v,
        likelihood=likelihood,
    )

    adapter.eval()
    with gpytorch.settings.detach_test_caches(False): 
        reg_v_mc = likelihood(gp(reg_x)).rsample()
    reg_y_hat = predictor(reg_v_mc.unsqueeze(-1)).view(-1)

    loss = mse_loss(reg_y_hat, reg_y)
    loss.backward()
    optimizer.step()

When running the above on a toy example adapted from the Exact GP Tutorial, I mostly get sensible-ish results, although often with confidence limits that are a bit too narrow.

However, sometimes I get degenerate results like these

The full code for this example can be found here

Click to expand

import torch
import torch.nn.functional as F
import math
import matplotlib.pyplot as plt
import gpytorch
from gpytorch.models.exact_prediction_strategies import DefaultPredictionStrategy
from gpytorch.lazy import delazify

import warnings
warnings.filterwarnings('ignore', '.*triangular_solve is deprecated.*', )


# Training data is 100 irregularly spaced times in [0,1] at which a predictor is
# measured
irreg_x = torch.rand(100).sort().values

# True function of the predictor is sin(2*pi*x) with Gaussian noise
def predictor_value(x):
   return torch.sin(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2
irreg_v = predictor_value(irreg_x)

# The regression outcome of interest is measured at evenly spaced times
reg_x = torch.linspace(0, 1, 11)

# The value of the predictor at those time is not known, only the regression 
# outcome is observed but it depends linearly on the predictor
reg_y = 1. + 0.5 * predictor_value(reg_x) + torch.randn(reg_x.size()) * 0.5


# GP model copied more or less directly from ExactGP tutorial
class GPAdapter(gpytorch.models.ExactGP):
   def __init__(self, train_x, train_y, likelihood):
       super(GPAdapter, self).__init__(train_x, train_y, likelihood)
       self.mean_module = gpytorch.means.ConstantMean()
       self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel())

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


# Prediction Strategy that does not use caches
class AdapterPredictionStrategy(DefaultPredictionStrategy):
   def __init__(self, train_inputs, train_prior_dist, train_labels, likelihood):
       super(AdapterPredictionStrategy, self).__init__(train_inputs, train_prior_dist, train_labels, likelihood)

   @property
   def covar_cache(self):
       train_train_covar = self.lik_train_train_covar
       train_train_covar_inv_root = delazify(train_train_covar.root_inv_decomposition().root)
       return train_train_covar_inv_root

   @property
   def mean_cache(self):
       mvn = self.likelihood(self.train_prior_dist, self.train_inputs)
       train_mean, train_train_covar = mvn.loc, mvn.lazy_covariance_matrix

       train_labels_offset = (self.train_labels - train_mean).unsqueeze(-1)
       mean_cache = train_train_covar.evaluate_kernel().inv_matmul(train_labels_offset).squeeze(-1)

       return mean_cache


# initialize likelihood and model
likelihood = gpytorch.likelihoods.GaussianLikelihood()
adapter = GPAdapter(irreg_x, irreg_v, likelihood)
predictor = torch.nn.Linear(1, 1)

# Set optimizer and losses (MSE +/- MLL)
optimizer = torch.optim.Adam(list(adapter.parameters()) + list(predictor.parameters()), lr=0.1)
mse = F.mse_loss
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, adapter)

# Define training parameters
num_iter = 100     # Number of training iterations (=optimizer steps on data)
num_samples = 10    # Number of Monte Carlo samples to draw at regular times
incl_mll = False    # Flag whether to use MSE+MLL loss (True) or MSE alone (False)

predictor.train()

# Start training iterations
for iter in range(num_iter):
   optimizer.zero_grad()
   
   # Get the train-train covariance for the prediction strategy
   adapter.train()
   irreg_prior_dist = adapter(irreg_x)
   adapter.prediction_strategy = AdapterPredictionStrategy(
       train_inputs=irreg_x,
       train_prior_dist=irreg_prior_dist,
       train_labels=irreg_v,
       likelihood=likelihood,
   )

   # Use prediction strategy to make backpropagatable predictions 
   # at regular grid tims
   adapter.eval()
   with gpytorch.settings.detach_test_caches(False):
       mc_size = torch.Size((num_samples, ))
       outc_v_mc = likelihood(adapter(reg_x)).rsample(mc_size)
   output = predictor(outc_v_mc.unsqueeze(-1)).view(-1)
   target = reg_y.repeat(num_samples) # repeat to match multiple mc samples of reg_v

   loss = mse(output, target) # prediction loss (=MSE)
   if incl_mll:
       loss -= mll(irreg_prior_dist, irreg_v) # "reconstruction" loss (=MLL)
   loss.backward(retain_graph=True)
   
   print('Iter %d/%d - loss: %.3f  mean:  %.3f  scale: %.3f  ls: %.3f  noise: %.3f' % (
       iter, num_iter, loss.item(),
       adapter.mean_module.constant.item(),
       adapter.covar_module.outputscale.item(),
       adapter.covar_module.base_kernel.lengthscale.item(),
       adapter.likelihood.noise.item()
   ))
   optimizer.step()

   # Plot each 10 iterations to see how the GP evolves
   if iter % 10 == 0:
       # Set into eval mode
       adapter.eval()

       # Test points are regularly spaced along [0,1]
       # Make predictions by feeding model through likelihood
       with torch.no_grad(), gpytorch.settings.fast_pred_var():
           test_x = torch.linspace(0, 1, 51)
           observed_pred = likelihood(adapter(test_x))

       with torch.no_grad():
           # Initialize plot
           f, ax = plt.subplots(1, 1, figsize=(4, 3))

           # Get upper and lower confidence bounds
           lower, upper = observed_pred.confidence_region()
           # Plot observed predictor values and times
           ax.plot(irreg_x.numpy(), irreg_v.numpy(), 'k*')
           # Plot predictive means as blue line
           ax.plot(test_x.numpy(), observed_pred.mean.numpy(), 'b')
           # Shade between the lower and upper confidence bounds
           ax.fill_between(test_x.numpy(), lower.numpy(), upper.numpy(), alpha=0.5)
           ax.set_ylim([-3, 3])
           ax.legend(['Observed Data', 'Mean', 'Confidence'])

I therefore have a couple of open questions left:

• Does the above implementation seem generally correct? Is it sufficient to remove the caches or are there any other things that I missed?
• Would degenerate results generally be expected for this setup (e.g., because of bad initialisations) and would we generally expect it to underestimate the variance (e.g., because of the generally lower number of prediction test points or because it is not it's main objective to estimate the variance well)? Would a joint loss of MSE + alpha * MLL fix this by explicitly incentivising a good fit of the predictors?
• Can this be used with gpytorch.settings.fast_pred_samples(True) and gpytorch.settings.fast_pred_samples(True)?

I'd be super grateful for any feedback, comments, and insights you might have, particularly with regards to gpytorch's internals!

Edits: typos in the code and text