Hadamard Multitask GP revisited

[prockenschaub]

Hi,

I am currently trying to implement a Hadamard Multitask GP model. The tutorial in the docs was a great help in getting something up and running quickly, and the modular nature of gpytorch is a joy to work with!

That being said, what I ultimately ended up with feels a bit hacky. Unlike the tutorial, I would like my GP to support different means and noise variables per task. I also wanted to wrap the kernel into a subclass of gpytorch.kernels.Kernel to make it more readable and to easily allow for multiple kernels with gpytorch.kernels.AdditiveKernel. Similar to #654, I ran into difficulties communicating task_indices to all model parts that needed them.

I thought I'd share what I ended up with here

1. for anyone else looking into this in the future and
2. for pointers on anything that I might have missed or ideas on how this might be achieved more elegantly

Mean

As in #654, the main problem is that gpytorch.means.Mean.__call__() only takes x as the input. The solution for me was to reserve x[..., 0] for the task indices.

class HadamardMean(MultitaskMean):
    def __init__(self, base_means, num_tasks):
        super(HadamardMean, self).__init__(base_means, num_tasks)

    def forward(self, input):
        i, x = input[..., [0]], input[..., 1:]
      
        torch.stack([sub_mean(x) for sub_mean in self.base_means], dim=-1)
        return means.gather(dim=-1, index=i.long()).squeeze(-1)

It would also have been possible to override the __call__() function, as in #654. However, this would not been as straightforward for other objects like likelihoods, as we'll see below.

Kernel

Since gpytorch.kernels.Kernel.__call__() takes a **params argument, the Kernel could in theory be passed separate task_indices1 and task_indices2 parameters. However, to be compatible with the rest, I again used a part of x to carry the task indices.

class HadamardKernel(Kernel):
    def __init__(self, base_kernel, num_tasks, rank=1):
        super(HadamardKernel, self).__init__()
        self.num_tasks = num_tasks
        self.base_kernel = base_kernel
        self.task_kernel = IndexKernel(num_tasks, rank)

    def forward(self, input1, input2, **params):
        i1, x1 = input1[..., 0], input1[..., 1:]
        i2, x2 = input2[..., 0], input2[..., 1:]

        covar_x = self.base_kernel(x1, x2, **params)
        covar_i = self.task_kernel(i1, i2)
        return covar_x.mul(covar_i)

Likelihood

Here is where things get hacky. In order to get a one noise parameter per task, I think I need to create my own Noise class, which will need to know the task indices i.

class HadamardHomoskedasticNoise(_HomoskedasticNoiseBase):
    def __init__(self, noise_prior=None, noise_constraint=None, num_tasks=1):
        super().__init__(noise_prior, noise_constraint, torch.Size(), num_tasks)

    def forward(self, *params, shape=None, noise=None, i=None, **kwargs):
        # Note: removed batching and additional checks/logic for simplicity
        
        # For each observation, pick the noise indicated by i
        noise = self.noise
        noise_diag = noise.expand(shape[0], len(noise)).contiguous()
        noise_diag = noise_diag.gather(-1, i).squeeze(-1)
        return DiagLazyTensor(noise_diag)

In order to make sure that i is always passed to the noise, I also need to subclass the likelihood, relying on the data to be passed as the first element in *args. The reason for this is that although we have **kwargs in the call to the likelihood, this was the only quick way that I could find to pass indices during a call to the gpytorch.mlls.ExactMarginalLogLikelihood (which does not pass on **kwargs to the likelihood) while simultaneously allowing me to use gpytorch.models.exact_prediction_strategies.DefaultPredictionStrategy (which I did not feel comfortable amending without breaking anything) and which otherwise would throw an error here:

gpytorch/gpytorch/models/exact_prediction_strategies.py

Line 60 in e64cb6f

mvn = self.likelihood(train_prior_dist, train_inputs)

My implementation of the likelihood looks like this

class HadamardGaussianLikelihood(_GaussianLikelihoodBase):
    def __init__(self, num_tasks, noise_prior=None, noise_constraint=None, **kwargs):
        noise_covar = HadamardHomoskedasticNoise(
            noise_prior=noise_prior, noise_constraint=noise_constraint, num_tasks=num_tasks
        )
        super().__init__(noise_covar=noise_covar)

    @property
    def noise(self):
        return self.noise_covar.noise

    @noise.setter
    def noise(self, value):
        self.noise_covar.initialize(noise=value)

    @property
    def raw_noise(self):
        return self.noise_covar.raw_noise

    @raw_noise.setter
    def raw_noise(self, value):
        self.noise_covar.initialize(raw_noise=value)

    def __call__(self, input, *args, **kwargs):
        # Extract the task identifiers from the second column of the inputs
        # to pass on to the evaluation of self.noise_covar
        xi = args[0][0]
        i = xi[..., [0]].long()
        
        # Conditional
        if torch.is_tensor(input):
            return super().__call__(input, i=i, *args, **kwargs)
        # Marginal
        elif isinstance(input, MultivariateNormal):
            return self.marginal(input, i=i, *args, **kwargs)
        # Error
        else:
            raise RuntimeError(
                "Likelihoods expects a MultivariateNormal or Normal input to make marginal predictions, or a "
                "torch.Tensor for conditional predictions. Got a {}".format(input.__class__.__name__)
            )

Putting it all together

Putting all of the above together in a subclass of ExactGP, I can fit a slightly extended tutorial. The full code for this can be found below (Python 3.7, pytorch==1.11.0, gpytorch==1.6.0)

Click to see full example

# Implementation of Hadamard Multitask GP based on 
# https://docs.gpytorch.ai/en/stable/examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.html

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

from gpytorch.models import ExactGP
from gpytorch.means import ConstantMean, MultitaskMean
from gpytorch.kernels import Kernel, IndexKernel, MaternKernel, AdditiveKernel
from gpytorch.likelihoods import _GaussianLikelihoodBase
from gpytorch.likelihoods.noise_models import _HomoskedasticNoiseBase
from gpytorch.distributions import MultivariateNormal
from gpytorch.lazy import DiagLazyTensor

class HadamardGP(ExactGP):
    """The base class for a Hadamard multitask GP regression to be used in 
    conjunction with exact inference.

    Args: 
        num_tasks (int): number of tasks fitted by the model 
        num_kernels (int, optional): number of kernels to fit; kernels are 
            combined additively
        rank (int, optional): rank of the inter-task correlation
    """
    def __init__(self, train_inputs, train_targets, likelihood, num_tasks, num_kernels=1, rank=1):
        super(HadamardGP, self).__init__(train_inputs, train_targets, likelihood)
        self.mean_module = HadamardMean(ConstantMean(), num_tasks)
        self.covar_module = AdditiveKernel(*[
            HadamardKernel(MaternKernel(), num_tasks, rank)
            for _ in range(num_kernels)
        ])

    def forward(self, input):
        mean = self.mean_module(input)
        covar = self.covar_module(input)
        return MultivariateNormal(mean, covar)


class HadamardMean(MultitaskMean):
    """Mean function for a Hadamard Multitask GP with one mean per task

    Args:
        base_means (:obj:`list` or :obj:`gpytorch.means.Mean`): If a list, each mean is applied to the data.
            If a single mean (or a list containing a single mean), that mean is copied `t` times.
        num_tasks (int): Number of tasks. If base_means is a list, this should equal its length.
    """
    def __init__(self, base_means, num_tasks):
        super(HadamardMean, self).__init__(base_means, num_tasks)

    def forward(self, input):
        """
        Evaluate the mean in self.base_means corresponding to each element of 
        the input data, and return as an n-vector of means
        """
        i, x = input[..., [0]], input[..., 1:]
        
        # Get means at x for each possible task and then gather the right one 
        # for each row based on the task number i.
        means = torch.cat(
            [sub_mean(x).unsqueeze(-1) for sub_mean in self.base_means], 
            dim=-1
        )
        return means.gather(dim=-1, index=i.long()).squeeze(-1)


class HadamardKernel(Kernel):
    """Kernel function for a Hadamard Multitask GP of the form 

    K_x(x_1, x_2) \times K_i(i_1, i_2)
    
    where x denotes locations (e.g., in time) and i denotes a task identifier.

    Args:
        base_kernel (:obj:`gpytorch.kernels.Kernel): the base class for the location kernel K_x, 
        num_tasks (int): number of tasks denoting size of task covariance K_i
        rank (int): rank of the inter-task correlation 
    """
    def __init__(self, base_kernel, num_tasks, rank):
        super(HadamardKernel, self).__init__()
        self.num_tasks = num_tasks
        self.base_kernel = base_kernel
        
        if rank is None:
            rank = num_tasks
        self.task_covar_module = IndexKernel(num_tasks, rank)

    def forward(self, input1, input2, **params):
        i1, x1 = input1[..., 0], input1[..., 1:]
        i2, x2 = input2[..., 0], input2[..., 1:]

        # Get input-input covariance
        covar_x = self.base_kernel(x1, x2, **params)
        # Get task-task covariance
        covar_i = self.task_covar_module(i1, i2)
        # Multiply the two together to get the covariance we want
        return covar_x.mul(covar_i)


class HadamardGaussianLikelihood(_GaussianLikelihoodBase):
    r"""
    Likelihood for a Hadamard multitask GP regression. Assumes a different 
    homoskedastic noise for each task i

    p(y_i \mid f_i) = f_i + \epsilon_i, \quad \epsilon_i \sim \mathcal N (0, \sigma_i^2)

    where :math:`\sigma_i^2` is the noise parameter of task i.

    .. note::
        Does not currently allow for batched training. 

    :param num_tasks: The number of tasks in the multitask GP.
    :type num_tasks: int
    :param noise_prior: Prior for noise parameter :math:`\sigma^2`.
    :type noise_prior: ~gpytorch.priors.Prior, optional
    :param noise_constraint: Constraint for noise parameter :math:`\sigma^2`.
    :type noise_constraint: ~gpytorch.constraints.Interval, optional
    
    :var torch.Tensor noise: :math:`\sigma_i^2` parameters (noise)
    """
    def __init__(self, num_tasks, noise_prior=None, noise_constraint=None, **kwargs):
        noise_covar = HadamardHomoskedasticNoise(
            noise_prior=noise_prior, noise_constraint=noise_constraint, num_tasks=num_tasks
        )
        super().__init__(noise_covar=noise_covar)

    @property
    def noise(self):
        return self.noise_covar.noise

    @noise.setter
    def noise(self, value):
        self.noise_covar.initialize(noise=value)

    @property
    def raw_noise(self):
        return self.noise_covar.raw_noise

    @raw_noise.setter
    def raw_noise(self, value):
        self.noise_covar.initialize(raw_noise=value)

    def __call__(self, input, *args, **kwargs):
        if not args:
            raise ValueError(
                "The first element of *args must be a list of the training" 
                "inputs."
            )
        
        # Extract the task identifiers from the first column of the inputs
        # to pass on to the evaluation of self.noise_covar
        xi = args[0][0]
        i = xi[..., [0]].long()
        
        # Conditional
        if torch.is_tensor(input):
            return super().__call__(input, i=i, *args, **kwargs)
        # Marginal
        elif isinstance(input, MultivariateNormal):
            return self.marginal(input, i=i, *args, **kwargs)
        # Error
        else:
            raise RuntimeError(
                "Likelihoods expects a MultivariateNormal or Normal input to make marginal predictions, or a "
                "torch.Tensor for conditional predictions. Got a {}".format(input.__class__.__name__)
            )


class HadamardHomoskedasticNoise(_HomoskedasticNoiseBase):
    r"""
    Noise for a Hadamard multitask GP regression with a different homoskedastic 
    noise for each task i:
    """
    def __init__(self, noise_prior=None, noise_constraint=None, num_tasks=1):
        super().__init__(noise_prior, noise_constraint, torch.Size(), num_tasks)

    def forward(self, *params, shape=None, noise=None, i=None, **kwargs):
        # Note: removed batching and additional checks/logic for simplicity
        
        # For each observation, pick the noise indicated by i
        noise = self.noise
        noise_diag = noise.expand(shape[0], len(noise)).contiguous()
        noise_diag = noise_diag.gather(-1, i).squeeze(-1)
        return DiagLazyTensor(noise_diag)



if __name__ == '__main__':
    torch.random.manual_seed(42)

    # Generate x, i, y for 2 tasks
    x1 = torch.rand((50, 1))
    x2 = torch.rand((50, 1))

    i1 = torch.full(x1.shape, dtype=torch.long, fill_value=0)
    i2 = torch.full(x2.shape, dtype=torch.long, fill_value=1)

    y1 = torch.sin(x1 * (2 * math.pi)) + torch.randn(x1.size()) * 0.2
    y2 = torch.cos(x2 * (2 * math.pi)) + torch.randn(x2.size()) * 0.4

    # Combine into a common input (i, x) and target (y)
    input = torch.cat((torch.cat([i1, i2]), torch.cat([x1, x2])), dim=-1)
    target = torch.cat([y1, y2]).view(-1)

    # Setup likelihood an model
    likelihood = HadamardGaussianLikelihood(num_tasks=2)
    model = HadamardGP(input, target, likelihood, num_tasks=2)


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

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

    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    for i in range(50):
        optimizer.zero_grad()
        output = model(input)
        loss = -mll(output, target, [input])
        loss.backward()
        print('Iter %d/50 - Loss: %.3f' % (i + 1, loss.item()))
        optimizer.step()

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

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

    # Test points every 0.02 in [0,1]
    test_x = torch.linspace(0, 1, 51)[:, None]
    test_i1 = torch.full((test_x.shape[0],1), dtype=torch.long, fill_value=0)
    test_i2 = torch.full((test_x.shape[0],1), dtype=torch.long, fill_value=1)

    test_input1 = torch.cat((test_i1, test_x), dim=-1)
    test_input2 = torch.cat((test_i2, test_x), dim=-1)

    # Make predictions - one task at a time
    # We control the task we cae about using the indices

    # The gpytorch.settings.fast_pred_var flag activates LOVE (for fast variances)
    # See https://arxiv.org/abs/1803.06058
    with torch.no_grad(), gpytorch.settings.fast_pred_var():
        observed_pred_y1 = likelihood(model(test_input1), [test_input1])
        observed_pred_y2 = likelihood(model(test_input2), [test_input2])

    # Define plotting function
    def ax_plot(ax, train_y, train_x, rand_var, title):
        # Get lower and upper confidence bounds
        lower, upper = rand_var.confidence_region()
        # Plot training data as black stars
        ax.plot(train_x.numpy(), train_y.numpy(), 'k*')
        # Predictive mean as blue line
        ax.plot(test_x.numpy(), rand_var.mean.detach().numpy(), 'b')
        # Shade in confidence
        ax.fill_between(test_x.view(-1).numpy(), lower.detach().numpy(), upper.detach().numpy(), alpha=0.5)
        ax.set_ylim([-3, 3])
        ax.legend(['Observed Data', 'Mean', 'Confidence'])
        ax.set_title(title)

    # Plot both tasks
    ax_plot(y1_ax, y1, x1, observed_pred_y1, 'Observed Values (Likelihood)')
    ax_plot(y2_ax, y2, x2, observed_pred_y2, 'Observed Values (Likelihood)')

I am sure this can be done better, so any pointers are greatly appreciated.