Gpytorch: Multitask GP syntactic sugar

[gpleiss]

Multitask GPs can be written in GPyTorch as is (see #208), but right now it is extremely ugly.

Proposed interface:

• Introduce a MultitaskKernel model, which wraps a standard kernel function and returns the Kronecker product between an index (task) kernel and the supplied input kernel.
• Mean functions take in an n_tasks argument, which defaults as None. If n_tasks=d, means will return an n x d output.
• Multiple outputs (for train_y) should be expressed in an n x d matrix.

Necessary updates:

• We introduce a syntactic sugar for GaussianRandomVariables. If the mean is an n x d matrix and the covariance is an nd x nd LazyVariable, we will figure things out
  
  
  
  • var, std, etc. will return n x d matrices
  
  
  • sample will return an n x d x s tensor
  
  
  • Explicitly evaluating the covariance might return an n x d x n x d tensor, or maybe a nd x nd tensor? Thoughts? (Probably won't happen too often, but probably a good corner case to catch.)
  
  
• inv_quad_log_det is updated to handle this new syntactic sugar.

The hadamard multitask case (one task per input) will remain essentially the same.

cc @Balandat @darbour @jacobrgardner

[jacobrgardner]

One potential issue I see is a conflict of the multi dimensional label idea with the batch GP training we did recently (e.g. see the batch GP regression unit test). Does a x b labels mean a GPs on b data points or some number of tasks?

While the former can be viewed as multitask with identity task covariance, it's an example of Hadamard multitask, not Kronecker, so this is a collision of meaning that needs to be resolved.

[jacobrgardner]

Maybe labels should always be n x 1 and then b x n x t means b batches of n data on t tasks?

[Balandat]

I think having labels always be either n x t in the non-batch case (iirc this is the sklearn interface) and b x n x t in the batch case makes a lot of sense. We could automatically convert single-dimensional inputs to n x 1 add issue a deprecation warning, which should cover the most common use cases. We can improve documentation for the batch case so this is clear to people.

[gpleiss]

I was always annoyed that targets had to be n x 1 in GPy :joy: but that might be the most straightforward.

Additionally, a pattern that we have elsewhere in the code is that all batch operations must be 3-dimensional tensors. What this means is that batch targets must be b x n x 1, but non-batch targets could be either n or n x 1.

Not sure if this is a good pattern or not?

[jacobrgardner]

If you allow both, how do you distinguish a x b as batch vs multitask?

[gpleiss]

a x b would be multitask. Batch would specifically require a c x a x b 3d tensor.

(I think?) this is what's currently done for batch inputs, right? a x b inputs are interpreted as a b-dimensional input, and you would need a c x a x b input for batch mode?

[gpleiss]

The pattern that I know exists for matmul, inv_matmul, and a lot of other LazyVariable operations is that all batch operations require 3D tensors.

E.g. a batch of vectors for an MVM must be represented as a b x n x 1 Tensor. However, a vector for a (non-batch) MVM can be represented either as a n x 1 matrix or an n vector.

Again, not sure if this pattern is too confusing or not. It might make things easier if we required everything to be a matrix (non-batch mode) or a 3D tensor (batch mode)l

[jacobrgardner]

I think this would be okay. So basically a x b would be multitask but a x b x 1 would be batch.

[jacobrgardner]

Another way of looking at it is that all tensors have a default implicit batch size of 1, rather than explicitly requiring all tensors to have a task dimension of 1.

[Balandat]

So I think internally (down in the Lazy variables etc.) we should always using b x n x t tensors to avoid any ambiguities. I think it's fine to infer a output dimension of 1 if 1-dim data is passed in, to default to batch size 1, and to require explicitly specifying the output dimension when using batch mode.

[Balandat]

I'm going to play around a bit with this augmented GaussianRandomVariable, this will be useful for implementing the batch acquisition functions over multiple outcomes.

cc @rajkumarkarthik, @darbour

[Balandat]

Here is something quick an dirty that takes the n x t mean and nt x nt covar and gets the job done at least for sampling:

class MultiOutputGaussianRandomVariable(RandomVariable):
    def __init__(self, mean, covar):
        """
        Constructs a multi-output multivariate Gaussian random variable, based on mean and covariance
        Can be multi-output multivariate, or a batch of multi-output multivariate Gaussians

        Passing a matrix mean corresponds to a multi-output multivariate Gaussian
        Passing a matrix mean corresponds to a batch of multivariate Gaussians

        Params:
        - mean (Tensor: matrix n x t, or batch matrix b x n x t) mean of Gaussian distribution
        - covar (Tensor: matrix nt x nt, or batch matrix b x nt x nt) covariance of Gaussian distribution
        """
        super(MultiOutputGaussianRandomVariable, self).__init__(mean, covar)
        if not torch.is_tensor(mean) and not isinstance(mean, LazyVariable):
            raise RuntimeError("The mean of a GaussianRandomVariable must be a Tensor or LazyVariable")

        if not torch.is_tensor(covar) and not isinstance(covar, LazyVariable):
            raise RuntimeError("The covariance of a GaussianRandomVariable must be a Tensor or LazyVariable")

        if mean.ndimension() not in {2, 3}:
            raise RuntimeError("mean should be a matrix or a batch matrix (batch mode)")

        if not isinstance(covar, LazyVariable):
            covar = NonLazyVariable(covar)

        self._mean = mean
        self._covar = covar
        self._nbatch = mean.shape[0] if mean.ndimension() == 3 else None
        self._n = mean.shape[-2]
        self._t = mean.shape[-1]

    def covar(self):
        raise NotImplementedError("Explicit construction of covar not implemented")

    def mean(self):
        return self._mean

    def representation(self):
        return self._mean, self._covar

    def sample(self, n_samples):
        samples = (
            self
            ._covar
            .zero_mean_mvn_samples(n_samples)
            .view(self._t, self._n, n_samples)
            .transpose(0, 1)
            .contiguous()
            .add(self._mean.unsqueeze(-1))
        )
        return samples

    def var(self):
        return self._covar.diag().view(self._n, self._t)