Gpytorch: gradients of Kernels / backward functions

[cherepanovic]

here is the backward function of rbf_covariance from rbf_covariance.py. I am confused what happens in that function, could you give me a hand regrading this formula, thanks!

    @staticmethod
    def backward(ctx, grad_output):
        d_output_d_input = ctx.saved_tensors[0]
        lengthscale_grad = grad_output * d_output_d_input
return None, None, lengthscale_grad, None

second question, how can I use the rbf in order to get gradients w.r.t inputs of GP (see [1]), that means that GP stays unchangeable.

[1] - https://stats.stackexchange.com/questions/373446/computing-gradients-via-gaussian-process-regression

[KeAWang]

Regarding your first question:

d_output_d_input is computed and saved during the forward pass when gradients are required to save on memory: https://github.com/cornellius-gp/gpytorch/blob/f6388f3be558f7b21d227803218faf0beae7d050/gpytorch/functions/rbf_covariance.py#L20. It computes the derivative of exp(-||x - x'||/(2sigma^2)) with respect to sigma and saved it for the backward pass. Then in backward, we use the chain rule to get the gradient with respect to the output.

Note that the custom Function is only run when using a single lengthscale. We use autograd when there's one lengthscale per dimenson.

[cherepanovic]

has a lengthscale according to [2] it will be the eye matrix, isn't?

[2] - https://gpytorch.readthedocs.io/en/latest/kernels.html#gpytorch.kernels.RBFKernel

is that [1] derivation implemented as a backward function?

[1] - https://stats.stackexchange.com/questions/373446/computing-gradients-via-gaussian-process-regression

[KeAWang]

No the derivative of k_RBF(x1, x2, lengthscale) with respect to the lengthscale is given by the elementwise product of the pairwise distance matrix and the kernel matrix divided by the lengthscale times 2, i.e. https://www.wolframalpha.com/input/?i=derivative+of+exp(-a%2F(2s%5E2))+with+respect+to+s

[KeAWang]

Regarding your question about getting the gradients of a GP, you can do something like this after training your model

# Get into evaluation (predictive posterior) mode
model.eval()
likelihood.eval()

# Test points are regularly spaced along [0,1]
# Make predictions by feeding model through likelihood
with gpytorch.settings.fast_pred_var():
    test_x = torch.linspace(0, 1, 51, requires_grad=True)
    observed_pred = likelihood(model(test_x))
    dydtest_x = torch.tensor([torch.autograd.grad(observed_pred.mean[i], test_x, retain_graph=True)[0][i] for i, _ in enumerate(test_x)])

with torch.no_grad():
    # Initialize plot
    f, ax = plt.subplots(1, 1, figsize=(4, 3))

    # Get upper and lower confidence bounds
    lower, upper = observed_pred.confidence_region()
    # Plot training data as black stars
    ax.plot(train_x.numpy(), train_y.numpy(), 'k*')
    # Plot predictive means as blue line
    ax.plot(test_x.detach().numpy(), observed_pred.mean.detach().numpy(), 'b')
    # Plot derivative
    ax.plot(test_x.detach().numpy(), dydtest_x.numpy(), 'r')
    # Plot real derivative
    ax.plot(test_x.detach().numpy(), 2 * math.pi * torch.cos(2 * math.pi * test_x).detach().numpy(), 'g')
    ax.set_ylim([-7, 7])
    ax.legend(['Observed Data', 'Mean',  'Estimated Derivative', 'True Derivative'])

which is what the stackexchange answer shows.

However, this loops over the test_x elements since autograd.grad doesn't support jacobian computation. I could see this becoming unwieldy for large N or for multi-output GPs. I wonder if there a faster way of computing this? @jacobrgardner

[KeAWang]

Oh actually for this special case since the derivative of the GP at each point depends only on the corresponding test point, you could do something like this instead

dydtest_x_fast = torch.autograd.grad(observed_pred.mean.sum(), test_x)[0]

Speed difference:

[cherepanovic]

thank a lot for your help!

[cherepanovic]

@KeAWang

observed_pred = likelihood(model(test_x))
torch.autograd.grad(observed_pred.mean.sum(), test_x)[0]

I am a little confused, which backward function of the observed_pred will be called?

thanks!

[KeAWang]

The exact backward function will depend on how observed_pred was computed. You shouldn't have to worry about this though since most of the time the backward function is automatically generated by autograd.

[cherepanovic]

@KeAWang

The exact backward function will depend on how observed_pred was computed

sure

You shouldn't have to worry about this though since most of the time the backward function is automatically generated by autograd.

does it not exist a backward function which was implemented by GPytorch creators? How does autograd generate a backward automatically?

[KeAWang]

For most functions you don't have to write a custom backward because pytorch automatically gets the gradient from the computational graph. However sometimes we write a custom backward function for gpytorch to save on memory/computation time.

[cherepanovic]

@KeAWang

pytorch automatically gets the gradient from the computational graph

how does it happen in this case? Since it is a derivation of a kernel (among other things), I mean it is something specific.

I thought, it is an abligation to write a backward function for own functions...

[KeAWang]

You only need to write a backward function for a custom torch.autograd.Function, e.g. https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html.

You don't need to write a custom backward if you just use pytorch functions, e.g. https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html.

Take a look at our PolynomialKernel for example: https://github.com/cornellius-gp/gpytorch/blob/f6388f3be558f7b21d227803218faf0beae7d050/gpytorch/kernels/polynomial_kernel.py#L83. We use Pytorch functions to compute the kernel so we don't need to write a custom backward. We wrote a custom backward for the RBFKernel and the MaternKernel because we wanted to speed up the computation and reduce memory usage for these two commonly used kernels. However, it's not necessary.

[cherepanovic]

If I would stack some pytorch NN models and would have GP between these NN, I don't have to think about the loss and backwards functions, autograd would do its work for the updates?

again, GP w.r.t. inputs would be done by autograd automatically?

Thanks a lot!

[jacobrgardner]

@cherepanovic Yes, you can forward propagate through a GP just like you would any PyTorch module, and when you call backward on a scalar you'll get gradients with respect to any tensors that require grad that were involved in the computation of said scalar. In general, there are lots of great PyTorch tutorials on autograd mechanisms.

In your specific example, there are a few minor GPyTorch specific issues to keep in mind. First, when you backpropagate through a GP posterior in GPyTorch, you'll want to be conscious of the fact that we compute caches for test time computation that you'll want to clear each time through the model (since these caches explicitly assume the parameters aren't changing). Second, as the output of a GP is a distribution and not a tensor, your arrow going from the GP to the second NN will actually need to be some operation that gives you a tensor, like sampling from the GP posterior.

Here's a full example that does something like what you want:

### Step 1: Define the GP. Here, we assume the GP takes the first NN as input and represents the full NN -> GP part of the model.
class GPRegressionModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, feature_extractor):
        # Batch shape 5 means we actually want 5 GPs in this "layer"
        batch_shape = torch.Size([5])

        train_x = train_x.expand(*batch_shape, *train_x.shape)
        train_y = train_y.expand(*batch_shape, *train_y.shape)

        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)


        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_shape=batch_shape), batch_shape=batch_shape)

        self.feature_extractor = feature_extractor

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


### Step 2: Define the GP -> NN part of the model. Forward here is made slightly complicated by the fact that when training
#   a model like this, we need to be sure to clear the GP test time caches each time.
class GPNN(torch.nn.Module):
    def __init__(self, gp, postprocess_nn):
        super().__init__()
        self.gp = gp
        self.postprocess_nn = postprocess_nn

    def forward(self, x):
        if self.training:
            # The next three lines are required to clear the GP test time caches since the GP parameters will change
            # each time
            with gpytorch.settings.detach_test_caches(False):
                self.gp.train()
                self.gp.eval()
                gp_output = self.gp(x)
        else:
            # If we aren't in training mode, we don't expect the GP parameters to change each iteration
            # so we don't need to clear the caches.
            gp_output = self.gp(x)

        # f_samples will be 10 x 5 x n in this example
        f_samples = gp_output.rsample(torch.Size([10]))

        # Transpose to be 10 x n x 5
        f_samples = f_samples.transpose(-2, -1)

        output = self.postprocess_nn(f_samples)
        output = output.mean(0).squeeze(-1)  # Average over GP sample dimension

        return output

### Step 3: Plug the pieces together: We define a simple feature extractor, the GP, and the postprocessing NN and combine them.
feature_extractor = torch.nn.Sequential(
    torch.nn.Linear(train_x.size(-1), 10),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 10),
    torch.nn.ReLU(),
)

gp = GPRegressionModel(
    train_x,
    train_y,
    gpytorch.likelihoods.GaussianLikelihood(),
    feature_extractor=feature_extractor
)

postprocess_nn = torch.nn.Sequential(
    torch.nn.Linear(5, 5),   # 5 because the GP layer has 5 GP outputs
    torch.nn.ReLU(),
    torch.nn.Linear(5, 1),
)

model = GPNN(gp, postprocess_nn)


# Compute mean squared error through this model and look at a gradient
mse = (model(train_x) - train_y).pow(2).mean()
mse.backward()
print(feature_extractor[0].weight.grad)