Gpytorch: gradients didn't backward to the end

[pinle-go]

I don't have much experience with the Gaussian process and I found this repo very helpful! I met one problem about the GPModel module.

When the input of the Gaussian process model comes from some embedding layers (which are part of the model), the embedding layers' weights have no gradients after the loss backward.

Here's an example.

import gpytorch
import torch
from gpytorch.kernels import RBFKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.means import ConstantMean
from gpytorch.random_variables import GaussianRandomVariable
from torch import nn
from torch.autograd import Variable


class LatentFunction(gpytorch.AdditiveGridInducingPointModule):
    def __init__(self):
        super(LatentFunction, self).__init__(grid_size=100, grid_bounds=[(-10, 10)], n_components=2)
        self.mean_module = ConstantMean(constant_bounds=[-1e-5, 1e-5])
        self.covar_module = RBFKernel(log_lengthscale_bounds=(-5, 6))
        self.register_parameter('log_outputscale', nn.Parameter(torch.Tensor([0])), bounds=(-5, 6))

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        covar_x = covar_x.mul(self.log_outputscale.exp())
        latent_pred = GaussianRandomVariable(mean_x, covar_x)
        return latent_pred


class GPRegressionModel(gpytorch.GPModel):
    def __init__(self):
        super(GPRegressionModel, self).__init__(GaussianLikelihood())
        self.latent_function = LatentFunction()

    def forward(self, x):
        return self.latent_function(x)


if __name__ == '__main__':
    n = 10
    embs = nn.Embedding(10, 2)
    train_x = (torch.rand(n) * 10).type(torch.LongTensor)
    train_x = Variable(train_x)
    train_x = embs(train_x)
    train_y = (train_x.data[:, :1] - train_x.data[:, 1:]).norm(p=2, dim=1)
    train_y = Variable(train_y)

    model = GPRegressionModel()
    output = model(train_x)
    loss = -model.marginal_log_likelihood(output, train_y)
    loss.backward()
    print('Embedding has no gradients?', embs.weight.grad is None)

    if train_x.grad:
        train_x.grad.zero_()
    model = nn.Linear(2, 1)
    output = model(train_x)
    loss = nn.functional.mse_loss(output, train_y)
    loss.backward()

    print('Linear has no gradients?', embs.weight.grad is None, '\tnorm:', embs.weight.grad.norm().data[0])

The output is

Embedding has no gradients? True
Linear has no gradients? False
norm: 1.098832130432129

I noticed it extends the "torch.nn.Module" class so I expected it has some similar properties like other "torch.nn.Module" subclasses (such as Linear Layer or Conv Layer). But looks like it doesn't.

I don't know if it's the Gaussian process's property or some "bugs". Correct me if there are some trivial mistakes.

Many thanks!

[gpleiss]

Hmmm you should be getting gradients! I'll look into this.

[gpleiss]

Alright, I reproduced your issue, and I think I have pinpointed the problem. Some of our interpolation code is not computing gradients, which then prevents embeddings from getting gradients.

Hopefully we'll have a fix for this soon.

[gpleiss]

Okay - we have a fix for this on the alpha_release branch.

The interface has changed slightly on this branch, so you'll probably have to update your model a bit.
Could you test out your model on the alpha_release branch and see if it works?

[gpleiss]

Actually - I ported your example to work on the alpha_release branch:

import gpytorch
import torch
from gpytorch.kernels import RBFKernel, AdditiveGridInterpolationKernel
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.means import ConstantMean
from gpytorch.random_variables import GaussianRandomVariable
from torch import nn
from torch.autograd import Variable


class GP(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GP, self).__init__(train_x, train_y, likelihood)
        self.embedding = nn.Embedding(10, 2)
        self.mean_module = ConstantMean(constant_bounds=[-1e-5, 1e-5])
        self.base_covar_module = RBFKernel(log_lengthscale_bounds=(-5, 6))
        self.covar_module = AdditiveGridInterpolationKernel(self.base_covar_module,
                                                            grid_size=100, grid_bounds=[(-10, 10)], n_components=2)
        self.register_parameter('log_outputscale', nn.Parameter(torch.Tensor([0])), bounds=(-5, 6))

    def forward(self, x):
        embs = self.embedding(x)
        mean_embs = self.mean_module(embs)
        covar_embs = self.covar_module(embs)
        covar_embs = covar_embs.mul(self.log_outputscale.exp())
        latent_pred = GaussianRandomVariable(mean_embs, covar_embs)
        return latent_pred


class GPRegressionModel(gpytorch.Module):
    def __init__(self, train_x, train_y):
        super(GPRegressionModel, self).__init__()
        self.likelihood = GaussianLikelihood()
        self.gp = GP(train_x, train_y, self.likelihood)

    def forward(self, x):
        return self.likelihood(self.gp(x))


if __name__ == '__main__':
    n = 10
    train_x = (torch.rand(n) * 10).type(torch.LongTensor)
    train_x = Variable(train_x)
    train_y = torch.randn(n)
    train_y = Variable(train_y)

    model = GPRegressionModel(train_x.data, train_y.data)
    output = model(train_x)
    loss = -model.gp.marginal_log_likelihood(model.likelihood, output, train_y)
    loss.backward()
    print('Embedding has gradients!', model.gp.embedding.weight.grad)

"""
Output:
('Embedding has gradients!', Variable containing:
 0.0000  0.0000
 0.0000  0.0000
 0.0000  0.0000
 0.0175 -0.0341
 0.0000  0.0000
-0.0136  0.0151
-0.0094  0.0722
 0.0290 -0.0124
 0.0907  0.0171
-0.0657 -0.1519
[torch.FloatTensor of size 10x2]
)
"""

Note that the nn.Embeddings layer has to live inside the GP module.