Trouble Reproducing Variational Nearest Neighbors GP Performance

[kevinli1324]

Hi all, I've been trying to reproduce the results from the "Variational Nearest Neighbor Gaussian Process" paper using the code from the VNNGP tutorial. Unfortunately, I can't seem to get results similar to those achieved in the paper. My training loss tends to stay high after hundreds of epochs, and the test RMSEs don't reach near the levels of the paper/that of SVGP. I've fiddled around with the batchsizes/#nearest neighbors but nothing seems to have made a difference. Can anyone who has had success with VNNGP give pointers to what might be going wrong?

I've listed the code I'm using below with the elevators dataset. The code follows the tutorial nearly identically.

# load packages
import tqdm
import math
import torch
import gpytorch
from matplotlib import pyplot as plt

import urllib.request
import os
from scipy.io import loadmat
from math import floor


# load adata
print('Downloading \'elevators\' UCI dataset...')
urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1jhWL3YUHvXIaftia4qeAyDwVxo6j1alk', '../elevators.mat')


data = torch.Tensor(loadmat('../elevators.mat')['data'])
data = data[torch.randperm(data.shape[0])]
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0].clamp_min(1e-6)) - 1
y = data[:, -1]
y = y.sub(y.mean()).div(y.std())


train_n = int(floor(0.8 * len(X)))
train_x = X[:train_n, :].contiguous()
train_y = y[:train_n].contiguous()

test_x = X[train_n:, :].contiguous()
test_y = y[train_n:].contiguous()

if torch.cuda.is_available():
    train_x, train_y, test_x, test_y = train_x.cuda(), train_y.cuda(), test_x.cuda(), test_y.cuda()

# define GP model
from gpytorch.models import ApproximateGP
from gpytorch.variational.nearest_neighbor_variational_strategy import NNVariationalStrategy


class GPModel(ApproximateGP):
    def __init__(self, inducing_points, likelihood, k=256, training_batch_size=256):

        m, d = inducing_points.shape
        self.m = m
        self.k = k

        variational_distribution = gpytorch.variational.MeanFieldVariationalDistribution(m)

        if torch.cuda.is_available():
            inducing_points = inducing_points.cuda()

        variational_strategy = NNVariationalStrategy(self, inducing_points, variational_distribution, k=k,
                                                     training_batch_size=training_batch_size)
        super(GPModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=d)

        self.likelihood = likelihood

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

    def __call__(self, x, prior=False, **kwargs):
        if x is not None:
            if x.dim() == 1:
                x = x.unsqueeze(-1)
        return self.variational_strategy(x=x, prior=False, **kwargs)

k = 256
training_batch_size = 64

likelihood = gpytorch.likelihoods.GaussianLikelihood()
# Note: one should use full training set as inducing points!
model = GPModel(inducing_points=train_x, likelihood=likelihood, k=k, training_batch_size=training_batch_size)

if torch.cuda.is_available():
    likelihood = likelihood.cuda()
    model = model.cuda()

#optimize  model
num_epochs = 500
num_batches = model.variational_strategy._total_training_batches


model.train()
likelihood.train()


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

# Our loss object. We're using the VariationalELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0), beta = .1)


epochs_iter = tqdm.notebook.tqdm(range(num_epochs), desc="Epoch")
for epoch in epochs_iter:
    minibatch_iter = tqdm.notebook.tqdm(range(num_batches), desc="Minibatch", leave=False)

    for i in minibatch_iter:
        optimizer.zero_grad()
        output = model(x=None)
        # Obtain the indices for mini-batch data
        current_training_indices = model.variational_strategy.current_training_indices
        # Obtain the y_batch using indices. It is important to keep the same order of train_x and train_y
        y_batch = train_y[...,current_training_indices]
        loss = -mll(output, y_batch)
        minibatch_iter.set_postfix(loss=loss.item())
        loss.backward()
        optimizer.step()
    epochs_iter.set_postfix(loss = loss.item())

Thanks!

[gpleiss]

cc/ @LuhuanWu

[LuhuanWu]

Hi, the reported behaviors may be related to the beta=0.1 KL regularization factor you set in
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0), beta = .1)

The experiments presented in the paper were run with standard ELBO setting, i.e. beta=1.0. I ran the code with beta=1.0 and the results look good to me. I also changed to training_batch_size to 256 to be consistent with the paper. Here are the results under a few different configurations.

beta = 1.0:

ntrain = 80%, nepochs=300, lr=1e-2, training_batch_size=256: test_rmse = 0.3816, test_nll = 0.5160
ntrain = 80%, nepochs=500, lr=1e-2, training_batch_size=256: test_rmse = 0.3789, test_nll = 0.5012 
ntrain = 80%, nepochs=500, lr=5e-3, training_batch_size=256: test_rmse = 0.3808, test_nll = 0.4856

beta = 0.1:

ntrain = 80%, nepochs=500, lr=5e-3, training_batch_size=256: test_rmse = 0.4466, test_nll = 0.7541

For reference, the results in the paper are test_rmse = .373 ± .004, and test_nll=.463 ± .020 (which uses different training-valid-test split, uses adaptive learning rate schedule, and performs hyperparameter search, see Appendix D.3.1. TRAINING DETAILS in the paper for details).

Here are the codes for reproducing results above, including an additional evaluation part:

# load packages
import tqdm
import math
import torch
import gpytorch
from matplotlib import pyplot as plt

import urllib.request
import os
from scipy.io import loadmat
from math import floor
from gpytorch import settings
from math import ceil


# load adata

dataset_path = None # fill in your dataset path 
dataset = "elevators"
data = torch.Tensor(loadmat(os.path.join(dataset_path, f"{dataset}.mat"))['data'])
data = data[torch.randperm(data.shape[0])]
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0].clamp_min(1e-6)) - 1
y = data[:, -1]
y = y.sub(y.mean()).div(y.std())

train_n_percent = 0.64 
test_n_percent = 0.20 
train_n = int(floor(train_n_percent * len(X)))
test_n = int(floor(test_n_percent * len(X)))

train_x = X[:train_n, :].contiguous()
train_y = y[:train_n].contiguous()

#test_x = X[train_n:, :].contiguous()
#test_y = y[train_n:].contiguous()
test_x = X[-test_n:, :].contiguous()
test_y = y[-test_n:].contiguous()

if torch.cuda.is_available():
    train_x, train_y, test_x, test_y = train_x.cuda(), train_y.cuda(), test_x.cuda(), test_y.cuda()

# define GP model
from gpytorch.models import ApproximateGP
from gpytorch.variational.nearest_neighbor_variational_strategy import NNVariationalStrategy


class GPModel(ApproximateGP):
    def __init__(self, inducing_points, likelihood, k=256, training_batch_size=256):

        m, d = inducing_points.shape
        self.m = m
        self.k = k

        variational_distribution = gpytorch.variational.MeanFieldVariationalDistribution(m)

        if torch.cuda.is_available():
            inducing_points = inducing_points.cuda()

        variational_strategy = NNVariationalStrategy(self, inducing_points, variational_distribution, k=k,
                                                     training_batch_size=training_batch_size)
        super(GPModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=d)

        self.likelihood = likelihood

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

    def __call__(self, x, prior=False, **kwargs):
        if x is not None:
            if x.dim() == 1:
                x = x.unsqueeze(-1)
        return self.variational_strategy(x=x, prior=False, **kwargs)

k = 256
training_batch_size = 256 #64

likelihood = gpytorch.likelihoods.GaussianLikelihood()
# Note: one should use full training set as inducing points!
model = GPModel(inducing_points=train_x, likelihood=likelihood, k=k, training_batch_size=training_batch_size)

if torch.cuda.is_available():
    likelihood = likelihood.cuda()
    model = model.cuda()

#optimize  model
num_epochs = 500
num_batches = model.variational_strategy._total_training_batches


model.train()
likelihood.train()


optimizer = torch.optim.Adam(model.parameters(), lr=5e-3)

# Our loss object. We're using the VariationalELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0), beta=1.0)


############
# training #
############

#epochs_iter = tqdm.notebook.tqdm(range(num_epochs), desc="Epoch")
epochs_iter = tqdm.tqdm(range(num_epochs), desc="Epoch")
for epoch in epochs_iter:
    #minibatch_iter = tqdm.notebook.tqdm(range(num_batches), desc="Minibatch", leave=False)
    minibatch_iter = tqdm.tqdm(range(num_batches), desc="Minibatch", leave=False)


    for i in minibatch_iter:
        optimizer.zero_grad()
        output = model(x=None)
        # Obtain the indices for mini-batch data
        current_training_indices = model.variational_strategy.current_training_indices
        # Obtain the y_batch using indices. It is important to keep the same order of train_x and train_y
        y_batch = train_y[...,current_training_indices]
        loss = -mll(output, y_batch)
        minibatch_iter.set_postfix(loss=loss.item())
        loss.backward()
        optimizer.step()
    epochs_iter.set_postfix(loss = loss.item())


##############
# prediction #
##############
eval = True


def get_prediction_stats(model, likelihood_type, x, y):
    def sample_batch_indices(N, mbs):
        num_batches = int(ceil(N / mbs))
        idx = torch.randperm(N)
        return [idx[i * mbs: (i + 1) * mbs] for i in range(num_batches)]

    with torch.no_grad():
        def predict(model, likelihood_type, test_x, test_y, mbs=256):
            indices = sample_batch_indices(test_x.size(0), mbs)
            test_ll = 0.0
            test_mse = 0.0
            for idx in indices:
                test_x_batch = test_x[idx]
                test_y_batch = test_y[idx]

                observed_pred = model.likelihood(model(test_x_batch))  # (likelihood_nsamples, bsz)
                if likelihood_type == 'gaussian' or likelihood_type == 'studentt':
                    diff = torch.pow(observed_pred.mean - test_y_batch, 2)
                    diff = diff.sum(dim=-1) / test_x.size(0) # sum over bsz and scaling
                    diff = diff.mean() # average over likelihood_nsamples
                    test_mse += diff
                elif likelihood_type == 'bernoulli':
                    diff = (observed_pred.mean.gt(0) != test_y_batch.bool())
                    assert diff.shape == test_y_batch.shape
                    test_mse += diff.sum() / test_x.size(0)

                batch_ll = model.likelihood.log_marginal(observations=test_y_batch.float(), function_dist=model(test_x_batch))
                assert batch_ll.shape == (len(test_x_batch), ), batch_ll.shape
                test_ll += batch_ll.sum() / test_x.size(0)

            return test_ll.item(), math.sqrt(test_mse.item())

        ll, rmse = predict(model, likelihood_type, x, y)
        nll = -ll
    return rmse, nll

if eval:
    model.eval()
    likelihood.eval()

    jitter = 1e-3
    likelihood_type = 'gaussian'

    with settings.cholesky_jitter(float=jitter, double=jitter):
        test_rmse, test_nll = get_prediction_stats(model, likelihood_type, test_x, test_y)
    
    print(f"test_rmse = {test_rmse:.4f}, test_nll = {test_nll:.4f}")

I also noticed that the google drive link in the notebook is not working. You can download elevators.mat here.