Spatiotemporal GP exploiting Kronecker Structure

[ekinugurel]

I have a multi-task problem where I would like to model two correlated tasks as a function of temporal variables and spatial features (10 and 2 dimensions, respectively). This is loosely related to #1048, #1590, and #1429.

Independently, I found the covariance structure that works best for spatial and temporal features. Neither is fit with a single base kernel, but rather a composite kernel with either additive or product structure. I now want to take the Kronecker product of these two (potentially nonstationary) kernel functions. Is this doable?

I tried the SKIP method but ran into issues while trying to use more than one base kernel (and did not see multi-task support). Is there a way to do this within the SKIP framework?

I first tried this which resulted in RuntimeError: The covariance of a MultitaskMultivariateNormal must be a Tensor or LinearOperator.

from gpytorch.kernels import ScaleKernel, RBFKernel as RBF, MaternKernel as MAT, \
    RQKernel as RQ, PeriodicKernel as PER

class SpatiotemporalMTGPRegressor(gpytorch.models.ExactGP):
    def __init__(self, X, y, mean=gpytorch.means.ConstantMean(), likelihood=None, num_tasks=2, rank=1):
        if likelihood is None:
            likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
        
        super().__init__(X, y, likelihood)
        self.mean = gpytorch.means.MultitaskMean(mean, num_tasks=num_tasks)
        self.covar_space  = gpytorch.kernels.MultitaskKernel(
            ScaleKernel(MAT(nu=2.5, active_dims=[0,1])) + ScaleKernel(RQ(active_dims=[0,1])) \
                + ScaleKernel(MAT(nu=2.5, active_dims=[0,1])) + ScaleKernel(RQ(active_dims=[0,1])), 
                num_tasks=num_tasks, rank=rank)
        self.covar_time = gpytorch.kernels.MultitaskKernel(
            ScaleKernel(RQ(active_dims=[2,3,4,5,6,7,8,9,10,11])*PER(active_dims=[2])) \
                + ScaleKernel(RQ(active_dims=[2,3,4,5,6,7,8,9,10,11])*PER(active_dims=[2])),
                num_tasks=num_tasks, rank=rank)
        self.likelihood = likelihood
    
    def forward(self, x):
        mean_x = self.mean(x)
        covar_x = self.covar_space * self.covar_time
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

The following works (and is the only way I know of doing this non-independent multitask), but I can't tell if we're exploiting the Kronecker structure in the process.

class MTGPRegressor(gpytorch.models.ExactGP):
    def __init__(self, X, y, kernel, mean=ConstantMean(), likelihood=None, num_tasks=2, rank=1):
        if likelihood is None:
            likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
        
        super().__init__(X, y, likelihood)
        self.mean = MultitaskMean(mean, num_tasks=num_tasks)
        self.covar_module = MultitaskKernel(
            ScaleKernel(MAT(nu=2.5, active_dims=[0,1])) + ScaleKernel(RQ(active_dims=[0,1])) \
                + ScaleKernel(MAT(nu=2.5, active_dims=[0,1])) + ScaleKernel(RQ(active_dims=[0,1])) \
                * ScaleKernel(RQ(active_dims=[2,3,4,5,6,7,8,9,10,11])*PER(active_dims=[2])) \ 
                + ScaleKernel(RQ(active_dims=[2,3,4,5,6,7,8,9,10,11])*PER(active_dims=[2])), 
                num_tasks=num_tasks, rank=rank)
        self.likelihood = likelihood
    
    def forward(self, x):
        mean_x = self.mean(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

[Turakar]

To inspect what solve strategy linear_operator uses, you can use a debugger and step into the code - I am not aware of any better approach. Your entry point can be the log_prob() function in MultivariateNormal.

linear_operator supports two solve strategies: Modified conjugate gradients (paper) and Cholesky decomposition. It switches from Cholesky to CG based on the size of the LinearOperator. This can be controlled via gpytorch.settings.max_cholesky_size(...).

For solving, there are two ways the kronecker structure can be exploited by linear_operator:

1. If your outermost LinearOperator is a KroneckerProductLinearOperator, it will exploit the Kronecker structure directly and compute the solve with each component individually, to combine them before return (KroneckerProductLinearOperator._solve()). The individual solves may use CG, Cholesky, or some existing decomposition (RootLinearOperator, which powers MultitaskKernel, already has a decomposition).
2. If your outermost LinearOperator is not a KroneckerProductLinearOperator, as it would be the case in your product setting and in Better support for missing labels #2288), the Kronecker structure can be exploited for the matrix multiplication happening during a CG solve (KroneckerProductLinearOperator._matmul()). Most likely, linear_operator won't choose a Cholesky decomposition as Kronecker matrices are usually large.

To combine two kernels via the product, you might want to have a look at ProductKernel.