using SVDKL with GNN

[alonsocampana]

Hi,
I'm trying to perform gaussian process regression using GNNs (implemented in torch geometric), and Gpytorch.
It's a single output regression problem. I've adapted the code in the tutorial, but the learning is extremely slow, and probably it's because my implementation has issues. I'm using the initial graph-features generated by the GNN as inducing points, and I suspect the issue could be related to this, but I don't know how to do it differently.

class GaussianProcessLayer(gpytorch.models.ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(inducing_points.size(0))
        variational_strategy = gpytorch.variational.VariationalStrategy(self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(GaussianProcessLayer, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

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

class DKLModel(gpytorch.Module):
    def __init__(self, feature_extractor, inducing_points, grid_bounds=(-10., 10.)):
        super(DKLModel, self).__init__()
        self.feature_extractor = feature_extractor
        self.gp_layer = GaussianProcessLayer(inducing_points)
        self.grid_bounds = grid_bounds

        # This module will scale the NN features so that they're nice values
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(self.grid_bounds[0], self.grid_bounds[1])

    def forward(self, x):
        features = self.feature_extractor(x)
        # features = self.scale_to_bounds(features) # if on predictions are 0, if off it converges to 0.
        res = self.gp_layer(features)
        return res

feature_extractor = EmbedderModel() # A GNN mapping graphs to a vector.
model = DKLModel(feature_extractor, feature_extractor(X_batch)) # This is what I think is wrong
likelihood = gpytorch.likelihoods.GaussianLikelihood() # regression
if torch.cuda.is_available():
    model = model.cuda()
    likelihood = likelihood.cuda()
optimizer = SGD([
{'params': model.feature_extractor.parameters(), 'weight_decay': 1e-4},
{'params': model.gp_layer.hyperparameters(), 'lr': lr * 0.01},
{'params': model.gp_layer.variational_parameters()},
{'params': likelihood.parameters()},
], lr=lr, momentum=0.9, nesterov=True, weight_decay=0)
scheduler = MultiStepLR(optimizer, milestones=[0.5 * n_epochs, 0.75 * n_epochs], gamma=0.1)
mll = gpytorch.mlls.VariationalELBO(likelihood,
                                    model.gp_layer,
                                    num_data=len(train_loader.dataset))

mse_loss = nn.MSELoss()
device = torch.device("cuda")
model.to(device)
metrics = torchmetrics.MetricCollection([torchmetrics.PearsonCorrCoef(),
                           torchmetrics.MeanSquaredError(),               
                           torchmetrics.MeanAbsoluteError(),
                           torchmetrics.SpearmanCorrCoef()]).to(device)
loss_epoch = []
for epoch in range(1000):
    metrics.reset()
    model.train()
    likelihood.train()
    losses = []
    with gpytorch.settings.num_likelihood_samples(32):
        for data in train_loader:
            if torch.cuda.is_available():
                data, target = data.cuda(), data["y"].cuda()
            optimizer.zero_grad()
            output = model(data)
            loss = -mll(output, target).sum()
            loss.backward()
            losses += [loss.item()]
            optimizer.step()

Any tip on how should I face this problem?
EDIT: My learning rate was very low, now it learns something in training, but in test all the predictions are almost the same.