[Bug] High GPU memory consumption of ```ProductKernel``` compared to its equivalent variant

🐛 Bug

I want to apply different kernels on different dimensions for a specific problem, but, I am running out of GPU memory when using GPyTorch. Narrowing down the problem to its simplest version:

• Using products of multiple RBFKernel on each dimension takes more GPU memory compared to a single RBFKernel with active_dims argument.
• Memory consumption difference increases with more dimensions.

To reproduce

import torch
import gpytorch


class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, kernel):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = kernel

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


likelihood = gpytorch.likelihoods.GaussianLikelihood().cuda()

n_observations = 10000
dims = 3

train_x = torch.rand(n_observations, dims).cuda()
train_y = (
    torch.sin(train_x[:, 0] * 5)
    + torch.sin(train_x[:, 1] * 10)
    + torch.sin(train_x[:, 2] * 15)
).cuda()

### Version 1
# K = None
# for i in range(dims):
#     if K is None:
#         K = gpytorch.kernels.RBFKernel(ard_num_dims=1, active_dims=[i])
#     else:
#         K = K * gpytorch.kernels.RBFKernel(ard_num_dims=1, active_dims=[i])

### Version 2
K = gpytorch.kernels.RBFKernel(ard_num_dims=dims, active_dims=[0, 1, 2])

model = ExactGPModel(
    train_x, train_y, likelihood, kernel=gpytorch.kernels.ScaleKernel(K)
).cuda()
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

model.train()
likelihood.train()

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

training_iter = 100
for i in range(training_iter):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    print(
        i + 1,
        training_iter,
        loss.item(),
        model.covar_module.base_kernel.lengthscale,
        model.likelihood.noise.item(),
    )
    optimizer.step()

Expected Behavior

For ### Version 1 and ### Version 2 kernels, GPU should consume the same amount of GPU memory but currently, it consumes 5069 MB for ### Version 1 and 3161 MB for ### Version 2.

System information

• GPyTorch 1.5.1
• PyTorch 1.9.0+cu102
• Ubuntu 18.04

[gpleiss]

Are x1 and x2 integer category values? You can use IndexKernel, and set the covar_factor parameter to be zero, and set the raw_var factor to be 1.

https://docs.gpytorch.ai/en/stable/kernels.html#indexkernel

The IndexKernel is very memory efficient, and is designed to be used as part of a ProductKernel.

[gpleiss]

Sure!