Integrating Variational Multi-output GPs with the GP-LVM

[AndreaBraschi]

Hi all,

I have been trying to integrate the variational multi-output GPs with the GPLVM. I have been using the inter coregionalization model from Bonilla et al. (2008) to establish correlations across tasks, and, at the moment I get an acceptable data fit, but the correlations across tasks and the latent variables are not what I expect.

Basically, the covariance structure of the GP prior should be a triple kronecker product between the inter-task covariance, covariance of observed inputs and covariance of latent variables. The noise model should also be a triple kronecker between identity matrix, task noises and latent variable noises.

I've tried to adapt the Multitask kernel and the MultiTaskGaussianLikelihood to a triple Kronecker. Here below you can find the plots that show the data fit after 500 iterations and the code. Although it's a long code, you can just run it and you'll see a bunch of plots that will show which will show much better what I'm talking about.

I've been also thinking of using two separate sets of inducing points, one for the observed input domain and the other one for the latent space. However, I'm not really sure what would be the best way to then combine the means of the resulting 2 approximate posterior q(f).

I am first uploading the classes that handle the Kernel, the likelihood and the variational strategy which are needed to run the main code.

Any suggestion or food for thoughts will be much appreciated.

Forgive me for the extremely long code, but I've had to try to change lots of the source classes!!!

Noise Model

from typing import Any
import torch
from torch import Tensor

from gpytorch.constraints import GreaterThan
from gpytorch.distributions import base_distributions
from gpytorch.lazy import (
    ConstantDiagLazyTensor,
    DiagLazyTensor,
    KroneckerProductDiagLazyTensor,
    KroneckerProductLazyTensor,
    LazyEvaluatedKernelTensor,
    RootLazyTensor,
    lazify,
)
from gpytorch.likelihoods import Likelihood, _GaussianLikelihoodBase


class _MultitaskGaussianLikelihoodBase(_GaussianLikelihoodBase):

    def __init__(self, num_tasks, noise_covar, rank=0, task_correlation_prior=None, batch_shape=torch.Size()):

        super().__init__(noise_covar=noise_covar)
        self.num_tasks = num_tasks
        self.rank = rank


    def _shaped_noise_covar(self, shape, add_noise=True, *params, **kwargs):
        if not self.has_task_noise:
            noise = ConstantDiagLazyTensor(self.noise, diag_shape=shape[-2] * self.num_tasks)
            return noise

        if self.rank == 0:
            task_noises = self.raw_task_noises_constraint.transform(self.raw_task_noises)
            task_var_lt = DiagLazyTensor(task_noises)
            latent_noise = self.raw_latent_noise_constraint.transform(self.raw_latent_noise)
            latent_var_lt = DiagLazyTensor(latent_noise)
            dtype, device = task_noises.dtype, task_noises.device

        else:
            task_noise_covar_factor = self.task_noise_covar_factor
            task_var_lt = RootLazyTensor(task_noise_covar_factor)
            dtype, device = task_noise_covar_factor.dtype, task_noise_covar_factor.device

        shape = torch.Size([101])
        eye_lt = ConstantDiagLazyTensor(torch.ones(*shape[:-2], 1, dtype=dtype, device=device), diag_shape=shape)

        if add_noise and self.has_global_noise:
            noise = ConstantDiagLazyTensor(self.noise, diag_shape=task_var_lt.shape[-1])
            task_var_lt = task_var_lt + noise
            noise = ConstantDiagLazyTensor(self.noise, diag_shape=latent_var_lt.shape[-1])
            latent_var_lt = latent_var_lt + noise

        covar_kron_lt = KroneckerProductLazyTensor(eye_lt, task_var_lt, latent_var_lt)

        return covar_kron_lt

    def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> base_distributions.Normal:
        noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
        noise = noise.reshape(*noise.shape[:-1], *function_samples.shape[-2:])
        return base_distributions.Independent(base_distributions.Normal(function_samples, noise.sqrt()), 1)


class MultitaskGaussianLikelihood(_MultitaskGaussianLikelihoodBase):

    def __init__(
        self,
        num_tasks,
        num_x,
        rank=0,
        task_prior=None,
        batch_shape=torch.Size(),
        noise_prior=None,
        noise_constraint=None,
        has_global_noise=True,
        has_task_noise=True,
    ):

        super(Likelihood, self).__init__()
        if noise_constraint is None:
            noise_constraint = GreaterThan(1e-4)

        if not has_task_noise and not has_global_noise:
            raise ValueError(
                "At least one of has_task_noise or has_global_noise must be specified. "
                "Attempting to specify a likelihood that has no noise terms."
            )

        if has_task_noise:
            if rank == 0:
                self.register_parameter(
                    name="raw_task_noises", parameter=torch.nn.Parameter(torch.zeros(*batch_shape, num_tasks))
                )
                self.register_constraint("raw_task_noises", noise_constraint)

                self.register_parameter(
                    name="raw_latent_noise", parameter=torch.nn.Parameter(torch.zeros(*batch_shape, num_x))
                )
                self.register_constraint("raw_latent_noise", noise_constraint)

                if noise_prior is not None:
                    self.register_prior("raw_task_noises_prior", noise_prior, lambda m: m.task_noises)
                if task_prior is not None:
                    raise RuntimeError("Cannot set a `task_prior` if rank=0")
            else:
                self.register_parameter(
                    name="task_noise_covar_factor",
                    parameter=torch.nn.Parameter(torch.randn(*batch_shape, num_tasks, rank)),
                )
                if task_prior is not None:
                    self.register_prior("MultitaskErrorCovariancePrior", task_prior, lambda m: m._eval_covar_matrix)
        self.num_tasks = num_tasks
        self.rank = rank

        if has_global_noise:
            self.register_parameter(name="raw_noise", parameter=torch.nn.Parameter(torch.zeros(*batch_shape, 1)))
            self.register_constraint("raw_noise", noise_constraint)
            if noise_prior is not None:
                self.register_prior("raw_noise_prior", noise_prior, lambda m: m.noise)

        self.has_global_noise = has_global_noise
        self.has_task_noise = has_task_noise

    @property
    def noise(self):
        return self.raw_noise_constraint.transform(self.raw_noise)

    @noise.setter
    def noise(self, value):
        self._set_noise(value)

    @property
    def task_noises(self):
        if self.rank == 0:
            return self.raw_task_noises_constraint.transform(self.raw_task_noises)
        else:
            raise AttributeError("Cannot set diagonal task noises when covariance has ", self.rank, ">0")

    @task_noises.setter
    def task_noises(self, value):
        if self.rank == 0:
            self._set_task_noises(value)
        else:
            raise AttributeError("Cannot set diagonal task noises when covariance has ", self.rank, ">0")

    @property
    def latent_noise(self):
        if self.rank == 0:
            return self.raw_latent_noise_constraint.transform(self.raw_latent_noise)
        else:
            raise AttributeError("Cannot set diagonal latent noises when covariance has ", self.rank, ">0")

    @latent_noise.setter
    def latent_noise(self, value):
        if self.rank == 0:
            self._set_latent_noise(value)
        else:
            raise AttributeError("Cannot set diagonal task noises when covariance has ", self.rank, ">0")

    def _set_noise(self, value):
        self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))

    def _set_task_noises(self, value):
        self.initialize(raw_task_noises=self.raw_task_noises_constraint.inverse_transform(value))

    def _set_latent_noise(self, value):
        self.initialize(raw_latent_noise=self.raw_latent_noise_constraint.inverse_transform(value))

    @property
    def task_noise_covar(self):
        if self.rank > 0:
            return self.task_noise_covar_factor.matmul(self.task_noise_covar_factor.transpose(-1, -2))
        else:
            raise AttributeError("Cannot retrieve task noises when covariance is diagonal.")

    @task_noise_covar.setter
    def task_noise_covar(self, value):
        # internally uses a pivoted cholesky decomposition to construct a low rank
        # approximation of the covariance
        if self.rank > 0:
            with torch.no_grad():
                self.task_noise_covar_factor.data = lazify(value).pivoted_cholesky(rank=self.rank)
        else:
            raise AttributeError("Cannot set non-diagonal task noises when covariance is diagonal.")

    def _eval_covar_matrix(self):
        covar_factor = self.task_noise_covar_factor
        noise = self.noise
        D = noise * torch.eye(self.num_tasks, dtype=noise.dtype, device=noise.device)
        return covar_factor.matmul(covar_factor.transpose(-1, -2)) + D

Kernel

import torch
from gpytorch.kernels import Kernel, IndexKernel
from gpytorch.lazy import lazify, LazyEvaluatedKernelTensor, KroneckerProductLazyTensor

from torch import Tensor

class TripleKronecker(Kernel):

    def __init__(
            self,
    data_covar_mat: Tensor,
    task_covar_module: IndexKernel,
    latent_covar_module: Kernel,
    num_tasks: int,
    data_pts: int,
    ):

        super().__init__()
        self.data_covar_mat = data_covar_mat
        self.task_covar_module = task_covar_module
        self.latent_covar_module = latent_covar_module
        self.num_tasks = num_tasks
        self.data_pts = data_pts

    def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):

        covar_tasks = self.task_covar_module.covar_matrix
        covar_data = self.data_covar_mat

        covar_latent = lazify(self.latent_covar_module.forward(x1, x2))

        if covar_data.ndim > 2:
            batches = torch.Size([covar_data.shape[0], covar_data.shape[1]])
            covar_tasks = covar_tasks.expand(*batches, *covar_tasks.shape[0:])
            covar_latent = covar_latent.expand(*batches, *covar_latent.shape[-2:])
            full_covar = KroneckerProductLazyTensor(covar_tasks, covar_data, covar_latent).sum(0).sum(0)
        else:
            full_covar = KroneckerProductLazyTensor(covar_tasks, covar_data, covar_latent)

        return full_covar.diag() if diag else full_covar

    def num_outputs_per_input(self, x1, x2):
        return self.num_tasks * self.data_pts

Variational Strategy

import torch
import gpytorch
from gpytorch.kernels import IndexKernel
from gpytorch.module import Module
import gpytorch.utils.broadcasting
from gpytorch.lazy import (
    KroneckerProductLazyTensor,
    RootLazyTensor
)
from gpytorch.variational import _VariationalStrategy,

class ICM_VariationalStrategy(_VariationalStrategy):

    def __init__(
        self,
        base_variational_strategy: gpytorch.variational.VariationalStrategy,
        num_tasks: int,
        N: int,
        rank: int
    ):
        Module.__init__(self)
        self.base_variational_strategy = base_variational_strategy
        self.batch_shape = self.base_variational_strategy._variational_distribution.batch_shape
        self.num_tasks = num_tasks
        self.N = N
        self.rank = rank

        # Kernels
        self.task_covar_module = IndexKernel(
            num_tasks=num_tasks, batch_shape=torch.Size([]), rank=rank, prior=None
        )

    @property
    def prior_distribution(self):
        return self.base_variational_strategy.prior_distribution

    @property
    def variational_distribution(self):
        return self.base_variational_strategy.variational_distribution

    @property
    def variational_params_initialized(self):
        return self.base_variational_strategy.variational_params_initialized

    def kl_divergence(self):
        return super().kl_divergence().sum(0).sum(0)

    def __call__(self, time, prior=False, **kwargs):

        # sample from q(X) to get X's across all latent dim
        x = self.base_variational_strategy.model.sample_latent_variable()
        covar_x = self.base_variational_strategy.model.covar_module(x)

        # Get an approximate posterior for each task
        q_f = self.base_variational_strategy(time, latent=False)
        covar_f = q_f.covariance_matrix

        factors = self.task_covar_module.covar_factor
        factors = factors.expand(*torch.Size([q_f.mean.shape[0]]), *factors.shape[-2:])
        latent_mean = factors @ q_f.mean

        num_tasks = self.num_tasks
        mean_reshaped = torch.zeros(len(x), len(time) * num_tasks)
        k = len(time)
        j = 0
        time_stamps = len(time)
        for i in range(num_tasks):
            mean_reshaped[:, j:k] = latent_mean[:, i, :]
            j = j + time_stamps
            k = j + time_stamps
        #.cuda()
        latent_latent_mean = covar_x @ mean_reshaped.cuda()

        # Covariance
        full_covar_module = TripleKronecker(covar_f, self.task_covar_module,
                                            self.base_variational_strategy.model.covar_module,
                                            num_tasks, len(time))
        full_covar = full_covar_module(x)

        # Distribution
        distribution = gpytorch.distributions.MultitaskMultivariateNormal(latent_latent_mean.T, full_covar)

        return distribution

Main code

import numpy as np
import gpytorch
import math
from matplotlib import pyplot as plt
from gpytorch.models.gplvm.latent_variable import *
from gpytorch.models.gplvm import BayesianGPLVM
from gpytorch.means import ZeroMean
from gpytorch.kernels import ScaleKernel, MaternKernel, RBFKernel
from gpytorch.distributions import MultivariateNormal
from gpytorch.priors import NormalPrior
from multitask_likelihood import MultitaskGaussianLikelihood
from gpytorch.variational import TrilNaturalVariationalDistribution
from gpytorch.variational import VariationalStrategy
from tqdm.notebook import trange


# function for Linear PCA
def lin_PCA(Data, latent_dim):
    U, S, V = torch.pca_lowrank(Data, q=latent_dim)
    return U, S, V, torch.nn.Parameter(torch.matmul(Data, V[:, :latent_dim]))

# class for variational multi-task bayesian GP-LVM
class MultiTask_GPLVM(BayesianGPLVM):
    def __init__(self, N: int, num_tasks: int, time_stamps: int,
                 n_inducing: list, x_pca: torch.Tensor,
                 latent_dim: int, time: torch.Tensor, rank: int):

        self.batch_shape = torch.Size([N, rank])
        self.time = time
        self.latent_dim = latent_dim

        n_induc_time = n_inducing[0]
        n_induc_latent = n_inducing[1]

        self.inducing_inputs = torch.randn(n_induc_time, 1)
        #
        variational_dist = TrilNaturalVariationalDistribution(n_induc_time, batch_shape=self.batch_shape)

        q_f = ICM_VariationalStrategy(VariationalStrategy(
            self, self.inducing_inputs, variational_dist,
            learn_inducing_locations=True), num_tasks, N, rank=rank
        )

        # prior over X
        x = x_pca
        prior_mu = torch.zeros_like(x)
        prior_sigma = torch.ones_like(x)
        X_prior = NormalPrior(prior_mu, prior_sigma)

        # Define Variational Distribution q(X)
        X = VariationalLatentVariable(N, num_tasks * time_stamps, latent_dim, x, X_prior)

        super().__init__(X, q_f)

        # time
        self.t_module = ZeroMean()
        self.t_covar_module = ScaleKernel(MaternKernel())

        # LATENT SPACE
        self.mean_module = ZeroMean(ard_num_dims=latent_dim)
        self.covar_module = ScaleKernel(
            RBFKernel(ard_num_dims=latent_dim),
            )
        # inducing inputs
        induc_latent = torch.randn(n_induc_latent,)
        self.latent_inducing = induc_latent.expand(*torch.Size([latent_dim]), *induc_latent.shape[-2:])
        # variational distribution
        latent_variational_dist = TrilNaturalVariationalDistribution(n_induc_latent, batch_shape=torch.Size([]))
        # variational strategy: how inducing points are integrated out
        self.latent_q_f = VariationalStrategy(self, self.latent_inducing.T, latent_variational_dist,
                                              learn_inducing_locations=True, latent=True)


    def forward(self, x, latent=False):
        if latent:
            mean = self.mean_module(x)
            covar = self.covar_module(x)
            dist = MultivariateNormal(mean, covar)
            return dist
        else:
            mean = self.t_module(x)
            covar = self.t_covar_module(x)
            dist = MultivariateNormal(mean, covar)
            return dist

# number of inducing inputs
n_inducing = [25, 5]

# data
x = torch.linspace(0, 1, 101)

y = torch.stack([
    torch.sin(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    torch.sin(x * (3 * math.pi)) + torch.randn(x.size()) * 0.2,
    torch.cos(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    torch.sin(x * (2 * math.pi)) + 2 * torch.cos(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    -torch.cos(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    -torch.cos(x * (3 * math.pi)) + torch.randn(x.size()) * 0.2,
    torch.cos(x * (2 * math.pi)) + 2 * torch.sin(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    -torch.sin(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
    torch.sin(x * (2 * math.pi)) + 2 * torch.sin(x * (2 * math.pi)) + torch.randn(x.size()) * 0.2,
], -1)

N = 3
tasks = 3

reshape_y = torch.zeros(N, len(x), tasks)
j = 0
k = 3
for i in range(N):
    reshape_y[i, :, :] = y[:, j:k]
    j = j + 3
    k = j + 3
train_y = torch.zeros(N, len(x) * tasks)
j = 0
k = len(x)
for i in range(tasks):
    train_y[:, j:k] = reshape_y[:, :, i]
    j = j + len(x)
    k = j + len(x)

# initialise latent variables with linear PCA
latent_dim = 3
U, S, V, X_init = lin_PCA(train_y, latent_dim)


# Model and Likelihood7
num_tasks = tasks
Model = MultiTask_GPLVM(N, tasks, len(x), n_inducing, X_init, latent_dim, x, rank=3)
Likelihood = MultitaskGaussianLikelihood(num_tasks=num_tasks, num_x=N)

# NGD and Adam optimisers
NGD = gpytorch.optim.NGD([
    {'params': Model.variational_parameters()},
], num_data=train_y.shape[-1], lr=0.1)

Adam = torch.optim.Adam([
    {'params': Model.hyperparameters()},
    {'params': Likelihood.parameters()}],
     lr=0.01)

# objective function: Variational ELBO
mll = gpytorch.mlls.VariationalELBO(Likelihood, Model, num_data=train_y.shape[-1])

iterator = trange(500, leave=True)
training_iter = 500
loss_list = []

if torch.cuda.is_available():
    training_x = x.cuda()
    training_y = train_y.cuda()
    Model.cuda()
    Likelihood.cuda()

# set training mode
Model.train()
Likelihood.train()
for i in iterator:
    Adam.zero_grad()
    NGD.zero_grad()
    output = Model(training_x)
    loss = -mll(output, training_y.T).sum()
    loss_list.append(loss.item())
    loss.backward()

    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iter, loss.item()))

    Adam.step()
    NGD.step()

plt.figure()
full_covar = output.covariance_matrix
plt.imshow(full_covar.cpu().detach().numpy(), origin='lower')
plt.colorbar()

plt.figure()
B = Model.variational_strategy.task_covar_module.covar_matrix
plt.imshow(B.cpu().detach().numpy(), origin='lower')
plt.colorbar()


plt.figure()
# latent space and K_xx ARD dimensions
inv_lengthscale = 1 / Model.covar_module.base_kernel.lengthscale.cpu()
values, indices = torch.topk(Model.covar_module.base_kernel.lengthscale.cpu(), k=2, largest=False)
plt.bar(np.arange(latent_dim), height=inv_lengthscale.detach().numpy().flatten())
plt.title('Inverse Lengthscale with ARD kernel', fontsize='small')


plt.figure()
X = Model.sample_latent_variable()
plt.suptitle('Latent Space Covariance Matrix')
plt.imshow(Model.covar_module(X).cpu().detach().numpy(), origin='lower')
plt.colorbar()

plt.show()