Gpytorch: [Feature Request] Warped Kernels

[fonnesbeck]

It would be great to allow for more flexibility in GPyTorch's covariance functions by allowing warped kernels to be implemented; for example, when processes are not homogeneous along a particular dimension. I realize one could simply warp the inputs, but the parameters of the warping function may not be known ahead of time, so it would be nice to have a kernel with the associated warping hyperparameters. Has anyone thought of adding this, or is it easy for an end user to subclass in order to realize it?

[jacobrgardner]

You mention parametric warpings specifically. Do you have examples of warpings you'd like to apply that couldn't be applied as PyTorch modules before the kernel is applied? In general, the following is a valid forward method for a GPyTorch model:

def forward(self, x):
    x = self.some_module(x)
    covar_x = self.some_kernel(x)
    mean_x = self.some_mean(x)
    return MultivariateNormal(mean_x, covar_x)

where self.some_module is any torch.nn.Module (or gpytorch.Module if you would like to use gpytorch.constraints on the parameters) that you would like learned through the marginal likelihood of a GP.

[Balandat]

If you wanted to do sth like the beta CDF warping a la Snoek et al you could also add parameters to the gpytorch module and register priors using register_prior for the distributional parameters, this should automatically get you the right MLL (this would be a MAP estimate though, and not the fully Bayesian treatment as in the paper).

[fonnesbeck]

What would need to be implemented in the Module subclass? Just forward and backward? Is there an example of a custom transformation like this anywhere?

[jacobrgardner]

You wouldn't need to implement a backward method for a Module, just the forward method. In general the primary example of this we have in the repo is deep kernel learning, but you can do pretty much anything you can do with standard pytorch.

Here are two example Modules, one using built in PyTorch methods and one using a custom definition:

# Module 1: Make a torch.nn.Sequential of built on ops
# Here, we use a linear layer to reduce the input dimensionality from 10 to 5
# so that the GP operates on a 5 dimensional input.
my_module1 = torch.nn.Sequential(
    torch.nn.Linear(10, 5),
    torch.nn.ReLU(),
)

# Example 2: essentially the same thing, but defined manually without the use of torch.nn.Linear
class MyWarping(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # Let's learn a matrix that maps from input_dim to output_dim (this is basically a linear layer)
        self.warping_mat = torch.nn.Parameter(torch.randn(input_dim, output_dim))

    def forward(self, x):
         # Assume: x is n x input_dim
        projected_x = x.matmul(self.warping_mat)  # Project using our warping mat
        projected_x = torch.nn.functional.relu(projected_x)  # Pass through a ReLU for this toy example

        # Now do something silly like pass the data through a sine function 50% of the time. 
        # The point is you can really do whatever you want here.
        if torch.rand(1).item() < 0.5:
            projected_x = torch.sin(projected_x)
        return projected_x

[fonnesbeck]

That's very helpful, but putting anything like this in my model raises the RuntimeError that complains about not having retain_graph=True added to backward(). I assume the module is generically instantiated in the model __init__, correct?

[jacobrgardner]

The retain_graph=True error you are getting is probably coming from a bug somewhere else in your code. I am able to drop the following in to the simple GP example notebook:

# Example 2: essentially the same thing, but defined manually without the use of torch.nn.Linear
class MyWarping(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # Let's learn a matrix that maps from input_dim to output_dim (this is basically a linear layer)
        self.warping_mat = torch.nn.Parameter(torch.randn(input_dim, output_dim))

    def forward(self, x):
         # Assume: x is n x input_dim
        projected_x = x.matmul(self.warping_mat)  # Project using our warping mat
        return projected_x


# We will use the simplest form of GP model, exact inference
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        self.warping = MyWarping(1, 1)

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

# initialize likelihood and model
likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(train_x, train_y, likelihood)

And it trains just fine. In general, this integration with other torch Modules is a fairly well tested feature of GPyTorch.

In order to debug your backward error, I'd need to see more of your code.

[fonnesbeck]

I've got something very similar; in some ways simpler, but I'm trying to separately warp two different input features. Here's the warping function as a Module:

class CustomWarp(torch.nn.Module):

    def __init__(self):
        super(CustomWarp, self).__init__()
        self.a = torch.nn.Parameter(torch.zeros(1))
        self.b = torch.nn.Parameter(torch.zeros(1))
        self.c = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return x + (self.a * torch.tanh(self.b * (x - self.c)))

And here is the application in my GP model.

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 = gpytorch.means.ConstantMean(prior=gpytorch.priors.NormalPrior(0, 0.2))
        self.covar_module = (gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.MaternKernel(nu=5/2, active_dims=[0])) # Horizontal location 
                        + gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.MaternKernel(nu=5/2, active_dims=[1])) # Vertical location 
                        + gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.RBFKernel(active_dims=[2])) # Time
                            )

        self.warp_xdim = CustomWarp()
        self.warp_zdim = CustomWarp()

    def forward(self, x):
        x[:, 0] = self.warp_xdim(x[:, 0])
        x[:, 1] = self.warp_zdim(x[:, 1])
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

Am I applying the warp functions incorrectly?

[jacobrgardner]

@fonnesbeck Yeah, there are two issues with your code:

1. In general, initializing parameters to zero can be problematic, because you'll often get zero gradients back. I would change torch.zeros(1) -> torch.rand(1) in all cases.

2. More relevantly to the specific problem you are encountering, the challenge is that you are modifying x in place with tensors that require grad:

        x[:, 0] = self.warp_xdim(x[:, 0])
        x[:, 1] = self.warp_zdim(x[:, 1])

This is problematic for two reasons. First, since we pass your model the training data, you are actually overwriting your own training data, which you probably did not intend to do. Second, since the training data now requires grad, you are creating the "double backwards" problem you are encountering.

One simple solution is to add an x = x.clone() to the top of your forward method before calling your warping module on it. That way, you'll only modify a copy of the data in place rather than the actual training data.

In total to fix these two problems, here's a fixed version of your code:

class CustomWarp(torch.nn.Module):

    def __init__(self):
        super(CustomWarp, self).__init__()
        self.a = torch.nn.Parameter(torch.rand(1))
        self.b = torch.nn.Parameter(torch.rand(1))
        self.c = torch.nn.Parameter(torch.rand(1))

    def forward(self, x):
        return x + (self.a * torch.tanh(self.b * (x - self.c)))

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 = gpytorch.means.ConstantMean(prior=gpytorch.priors.NormalPrior(0, 0.2))
        self.covar_module = (gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.MaternKernel(nu=5/2, active_dims=[0])) # Horizontal location 
                        + gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.MaternKernel(nu=5/2, active_dims=[1])) # Vertical location 
                        + gpytorch.kernels.ScaleKernel(
           gpytorch.kernels.RBFKernel(active_dims=[2])) # Time
                            )

        self.warp_xdim = CustomWarp()
        self.warp_zdim = CustomWarp()

    def forward(self, x):
        x = x.clone()
        x[:, 0] = self.warp_xdim(x[:, 0])
        x[:, 1] = self.warp_zdim(x[:, 1])
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPRegressionModel(train_x, train_y, likelihood)

[fonnesbeck]

Yes, that did the trick. I think I will go ahead and implement the Beta CDF warping as @Balandat suggests.

EDIT: this is complicated by the fact that the incomplete beta function is not available in PyTorch!

I will close this issue; hopefully it will be useful for others. Thanks again.

[fonnesbeck]

So, I did manage to get a Beta CDF warping function implemented, using SciPy's incomplete beta.

from scipy.stats import beta

class BetaCDFWarp(gpytorch.Module):

    def __init__(self, a=None, b=None):
        super(BetaCDFWarp, self).__init__()
        a = torch.rand(1) if a is None else torch.tensor([a])
        b = torch.rand(1) if b is None else torch.tensor([b])

        self.register_parameter(name="a_raw", parameter=torch.nn.Parameter(a))
        self.register_constraint("a_raw", gpytorch.constraints.Positive())
        self.register_prior("a_prior", gpytorch.priors.GammaPrior(1, 1), lambda: self.a)

        self.register_parameter(name="b_raw", parameter=torch.nn.Parameter(b))
        self.register_constraint("b_raw", gpytorch.constraints.Positive())
        self.register_prior("b_prior", gpytorch.priors.GammaPrior(1, 1), lambda: self.b)

    @property
    def a(self):
        return self.a_raw_constraint.transform(self.a_raw)

    @property
    def b(self):
        return self.b_raw_constraint.transform(self.b_raw)

    def forward(self, x):
        a = self.a.detach()
        b = self.b.detach()
        warp_numpy = beta(a, b).cdf(x)
        return torch.from_numpy(warp_numpy)

I think I have the constraints and priors set up correctly, but when I run it as a warping function, the estimated beta parameters are always pushed towards zero (a_raw and b_raw end up being around -4).

[Balandat]

The issue here is that you're detaching the parameters and use them in the scipy distribution, so you're not getting any gradients w.r.t. a and b from the MLL (the gradient you do get is from the prior). So the result of the fitting will be that you just pick the mode of the prior.

To set this up correctly you need to be able to differentiate through beta(a, b).cdf(x), which means you should use the Beta distribution in pytorch. Unfortunately the cdf for this isn't implemented at this point so you may have to implement that (ideally this can be done upstream).