Variational Multitask GP with correlated outputs

Hi,

I noticed in the tutorials that you can create an exact multi-task GP with correlated outputs:

https://gpytorch.readthedocs.io/en/latest/examples/03_Multitask_Exact_GPs/Multitask_GP_Regression.html

Is it possible to do this with a variational GP? I have tried the following which doesn’t work:

class MultitaskSVGP(gpytorch.models.ApproximateGP):
    def __init__(self, inducing_points, input_dim, output_dim):
        batch_shape = torch.Size([output_dim])

        variational_distribution = CholeskyVariationalDistribution(inducing_points.size(-2), 
                                                                   batch_shape=batch_shape)

        variational_strategy = MultitaskVariationalStrategy(
                                VariationalStrategy(self, inducing_points, 
                                                   variational_distribution, 
                                                   learn_inducing_locations=True), 
                                num_tasks=output_dim)
        
        super().__init__(variational_strategy)
        
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=output_dim
        )
        
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.RBFKernel(ard_num_dims=input_dim), num_tasks=output_dim, rank=1
        )

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

[gpleiss]

Right - so we’re working on a feature that will add ICM and LCM-type variational multitask models. At the moment this kind of example won’t work - because the MultitaskVariationalStrategy assumes it’s getting a batch of indepenent prior GPs rather than a multi-output GP.

For now, if you try the SVGP multitask example you can learn a multi-output SVGP model. It doesn’t assume any correlations

[gpleiss]

@dinithins There isn’t a straightforward way to exploit Kronecker structure for variational models. This is because - even though the prior has Kronecker structure - there isn’t any immediately obvious structure that can be assigned to the variational distribution.

The variational multitask GP implementation - which implements the linear model of coregionalization - is actually quite efficient though!