Independent Stochastic Variational GP Regression

[luisroque]

Hello,

I was wondering if anyone could provide insight on the expected difference between using Stochastic Variational GP Regression and Exact GP Regression with the IndependentModelList container in GPyTorch. I have datasets with a large number of time series that I would like to be able to fit using both approaches. I have done some testing and it appears that using the Stochastic Variational GP Regression is actually slower and uses even more memory than the Exact GP Regression.

My datasets often have a number of time series between 300 and 5000 and each time series is around 300 points. Could this be because I am not defining a big enough number of time points? I have also noticed that for larger datasets, the Stochastic Variational GP Regression often gives an output that the matrix is not positive definite. Is there a specific reason for this? Thank you for any insight you can provide.

Below you a toy setup that I tried to build to make it more clear.

Creating synthetic data:

train_x = torch.reshape(torch.linspace(0, 100, 100), (-1,1))
train_y = torch.sin(train_x) + torch.normal(0, 1, size=(100, 500))
train_y = train_y.numpy()

scaler = StandardScaler()
scaler.fit(train_y)

scaled_train_y = scaler.transform(train_y)
scaled_train_y = torch.from_numpy(scaled_train_y)

Defining ExactGP and using the IndependentModelList container:

# Define the GP model
class GPRegressionModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel()) + gpytorch.kernels.LinearKernel() + gpytorch.kernels.PeriodicKernel()

    def forward(self, x):
        mean = self.mean_module(x)
        covar = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean, covar)
    
models = []
likelihoods = []
for i in range(train_y.shape[1]):
    models.append(GPRegressionModel(train_x, scaled_train_y[:,i], GaussianLikelihood()))
    likelihoods.append(models[i].likelihood)
model = gpytorch.models.IndependentModelList(*models)
likelihood = gpytorch.likelihoods.LikelihoodList(*likelihoods)

Training:

start_time = time.perf_counter()
loss_fn = SumMarginalLogLikelihood(likelihood, model)

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

model.train()
likelihood.train()

training_iterations = 100
for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(*model.train_inputs)
    loss = -loss_fn(output, model.train_targets)
    loss.backward()
    optimizer.step()

end_time = time.perf_counter()
elapsed_time = end_time - start_time
print(f'RAM usage: {process.memory_info().rss / 1024**3:.2f} GB')
print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))
print(f'Elapsed time: {elapsed_time:.2f} seconds')

I get the following output:

RAM usage: 1.86 GB
Iter 100/100 - Loss: 1.461
Elapsed time: 351.72 seconds

And one time series looks like this:

For the SVGP I defined it as:

class SvgGP(gpytorch.models.ApproximateGP):
    def __init__(self, train_x, train_y, inducing_points):
        variational_distribution = CholeskyVariationalDistribution(
            inducing_points.size(0)
        ).to(torch.float64)
        variational_strategy = VariationalStrategy(
            self,
            inducing_points=inducing_points,
            variational_distribution=variational_distribution,
            learn_inducing_locations=True,
        ).to(torch.float64)
        
        super().__init__(variational_strategy)
        self.train_inputs = train_x
        self.train_targets = train_y
        
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(ard_num_dims=1)
        )
        self.likelihood = gpytorch.likelihoods.GaussianLikelihood()
        
        
    def forward(self, x):
        if x.requires_grad:
            x = x.detach()
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

I also created a class SumVariationalELBO that following the same logic as the marginal likelihood one:

class SumVariationalELBO(MarginalLogLikelihood):
    
    def __init__(self, likelihood, model, num_data, mll_cls=VariationalELBO):
        super().__init__(model.likelihood, model)
        self.mlls = ModuleList([mll_cls(mdl.likelihood, mdl, num_data) for mdl in model.models])

    def forward(self, outputs, targets, *params):
        if len(params) == 0:
            sum_mll = sum(mll(output, target) for mll, output, target in zip(self.mlls, outputs, targets))
        else:
            sum_mll = sum(
                mll(output, target, *iparams)
                for mll, output, target, iparams in zip(self.mlls, outputs, targets, params)
            )
        return sum_mll.div_(len(self.mlls))

Then, I created the inducing points:

inducing_points = torch.linspace(0, 100, 50)

And run the exact same training procedure; the only change was the mll.

In terms of time and memory the output was the following:

RAM usage: 1.82 GB
Iter 100/100 - Loss: 1.452
Elapsed time: 367.84 seconds

And one time series looks like:

[gpleiss]

I have done some testing and it appears that using the Stochastic Variational GP Regression is actually slower and uses even more memory than the Exact GP Regression... My datasets often have a number of time series between 300 and 5000 and each time series is around 300 points. Could this be because I am not defining a big enough number of time points?

I'm not too surprised by this finding, given your data. Exact GP inference is very efficient when you are conditioning on only 100 datapoints, so you won't see any speedups from any scalable method.

I have also noticed that for larger datasets, the Stochastic Variational GP Regression often gives an output that the matrix is not positive definite. Is there a specific reason for this?

This sounds like some numerical instability. You could try training these models in double precision (i.e. torch.set_default_dtype(torch.double), though this will use more memory.