Custom GP as linear combination of two GPs

[fonnesbeck]

I'm looking to model a 2D function as a sum of two Gaussian processes, where the second one is only applied with the active elements of a mask. I was hoping the following would work, where train_x is comprised of xy coordinates in the first two columns and binary mask values in the third.

class GPRegressionModel(gpytorch.models.ExactGP):

    def __init__(self, train_x, train_y, likelihood):
        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean()

        self.base_covar_module_0 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_0 = InducingPointKernel(self.base_covar_module_0, inducing_points=train_x[:200, :], likelihood=likelihood)

        self.base_covar_module_1 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_1 = InducingPointKernel(self.base_covar_module_1, inducing_points=train_x[:200, :], likelihood=likelihood)

    def forward(self, x):
        mean_0 = self.mean_module(x)
        covar_0 = self.covar_module_0(x)
        covar_1 = self.covar_module_1(x)
        return MultivariateNormal(mean_0, covar_0) + x[:, 2]* MultivariateNormal(torch.zeros(x.shape[0]).to(device), covar_1)

In other words, some elements are just the first GP and others are the sum of the two GPs. However, it does not permit me to multiply a MultivariateNormal with a Tensor in the last statement, so I'm clearly doing this wrong. How do I mask out the values as needed to accomplish this? Some sort of affine transformation elsewhere?

[fonnesbeck]

I have not made any headway on this. I assume it is not possible?

[gpleiss]

Sorry for the slow reply @fonnesbeck . Is there a reason you want to use InducingPointKernel? Do you need that scalability?

[fonnesbeck]

I appreciate the reply, thanks! I may eventually need to scale it pretty big (hundreds of thousands of points), but it seems to work fine without the sparse approximation, yes.

[gpleiss]

Rather than multiplying the MVN objects, you need to multiply the kernels together instead.

class GPRegressionModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean()

        self.base_covar_module_0 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_0 = InducingPointKernel(self.base_covar_module_0, inducing_points=train_x[:200, :], likelihood=likelihood)

        self.base_covar_module_1 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_1 = InducingPointKernel(self.base_covar_module_1, inducing_points=train_x[:200, :], likelihood=likelihood)

    def forward(self, x):
        mean_0 = self.mean_module(x)
        covar_0 = self.covar_module_0(x)
        covar_1 = self.covar_module_1(x) * (x[:, 2].unsqueeze(-1) @ x[:, 2].unsqueeze(-2))
        return MultivariateNormal(mean_0, covar_0 + covar_1)

A slightly more efficient version would be:

class GPRegressionModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean()

        self.base_covar_module_0 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_0 = InducingPointKernel(self.base_covar_module_0, inducing_points=train_x[:200, :], likelihood=likelihood)

        self.base_covar_module_1 = ScaleKernel(RBFKernel(ard_num_dims=2, active_dims=[0,1]))
        self.covar_module_1 = InducingPointKernel(self.base_covar_module_1, inducing_points=train_x[:200, :], likelihood=likelihood)

    def forward(self, x):
        mean_0 = self.mean_module(x)
        covar_0 = self.covar_module_0(x)
        covar_1 = self.covar_module_1(x) * gpytorch.lazy.RootLazyTensor(x[:, 2].unsqueeze(-1))
        return MultivariateNormal(mean_0, covar_0 + covar_1)

[fonnesbeck]

One quick follow-up to this: after fitting the model, if I am interested in extracting inferences from one of the covariance functions, what's the best practice for doing this? Do I export that kernel and build a new model with it, or modify forward to only use the kernel of interest when predicting? Seems like something pretty common, so if there is already an example out there for this, feel free to just point me there.

[gpleiss]

if I am interested in extracting inferences from one of the covariance functions, what's the best practice for doing this?

What inference in particular? Do you want to see how much each additive component contributes to the posterior?

[fonnesbeck]

Exactly, yes. I'm looking to extract a surface representing the additive effect of a group that is represented by the indicator.

[gpleiss]

@fonnesbeck unfortunately, I can't think of an easy way to do this at the moment with our code :/

One (unsatisfying) solution would be to use an additive variational model. It is often easier in GPyTorch to perform these sorts of inferences with our VI implementation.

@jacobrgardner - any suggestions?

[jacobrgardner]

I would just define something like:

...
def forward(self, x, use_covar1=True, use_covar2=True):
    ...
    # Just use if statements to decide whether to return `covar1`, `covar2`, or `covar1 + covar2`.

Then calling e.g. model(x, use_covar1=True, use_covar2=False) should work just fine in GPyTorch.