Is there a multivariate implementation of GaussianSymmetrizedKLKernel?

[brucejwittmann]

I would like to train a GP on the output of a VAE. The input "x" can in this case thus be thought of as a multivariate gaussian with a diagonal covariance matrix. From what I can see in the documentation for "GP Regression with Uncertain Inputs", only 1-dimensional inputs are currently accepted?

I looked into the docstring for the distance function of GaussianSymmetrizedKLKernel, and it looks like higher dimensional distributions can be passed in by placing the means in the first half of the input space and the log-variances in the second half. I tried this with my example and the model trains; however, come inference time I rapidly run out of GPU memory (>160 GB required for predicting against 7000 examples in a batch). If I use my same programmatic infrastructure to make predictions with a standard kernel (e.g., RBF) using just the mean as my features, then I have no problems. This leads me to believe that I am missing something in my understanding of how to implement GaussianSymmetrizedKLKernel for higher-dimensional inputs. Are there any examples out there for using higher-dimensional distributions as inputs?

[wjmaddox]

Could you post a code example?

It's been a while but I've tested out GaussianSymmetrizedKLKernel with several dimensional inputs (as shown in its unit tests).

[brucejwittmann]

Sure! See below. I will get an OOM error when I try to call observed_pred.confidence_region() in the last line. I've copied the trace below the code block.

from gpytorch.models import ExactGP
from gpytorch.kernels import GaussianSymmetrizedKLKernel, ScaleKernel
from gpytorch.means import ConstantMean

# Get dummy data. This represents a VAE output.
train_x_distributional = torch.rand(1, 768, 2)
test_x_distributional = torch.rand(7000, 768, 2)
train_y = torch.rand(1)

# Reshape so that first half of input is mean, second half is log-var
train_x_distributional = torch.flatten(torch.transpose(train_x_distributional, -2, -1), 
                                       start_dim = -2)
test_x_distributional = torch.flatten(torch.transpose(test_x_distributional, -2, -1), 
                                       start_dim = -2)

# Set training iterations
training_iter = 50

# Everything below this comment is a direct copy from documentation
class ExactGPModel(ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean()
        self.covar_module = ScaleKernel(GaussianSymmetrizedKLKernel())

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

# initialize likelihood and model
likelihood = gpytorch.likelihoods.GaussianLikelihood().to("cuda")
model = ExactGPModel(train_x_distributional, train_y, likelihood).to("cuda")

# Find optimal model hyperparameters
model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.25)  # Includes GaussianLikelihood parameters

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

for i in range(training_iter):
    # Zero gradients from previous iteration
    optimizer.zero_grad()
    # Output from model
    output = model(train_x_distributional.to("cuda"))
    # Calc loss and backprop gradients
    loss = -mll(output, train_y.to("cuda"))
    loss.backward()
    optimizer.step()
    
# Make predictions
model.eval()
likelihood.eval()

with torch.no_grad():
    observed_pred = likelihood(model(test_x_distributional.to("cuda")))
    
# Get confidence region
observed_pred.confidence_region()

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-19-62ecaf84a917> in <module>
     46 
     47 # Get confidence region
---> 48 observed_pred.confidence_region()

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/distributions/multivariate_normal.py in confidence_region(self)
     99             the confidence region.
    100         """
--> 101         std2 = self.stddev.mul_(2)
    102         mean = self.mean
    103         return mean.sub(std2), mean.add(std2)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/distributions/multivariate_normal.py in stddev(self)
    241     def stddev(self):
    242         # self.variance is guaranteed to be positive, because we do clamping.
--> 243         return self.variance.sqrt()
    244 
    245     @property

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/distributions/multivariate_normal.py in variance(self)
    247         if self.islazy:
    248             # overwrite this since torch MVN uses unbroadcasted_scale_tril for this
--> 249             diag = self.lazy_covariance_matrix.diag()
    250             diag = diag.view(diag.shape[:-1] + self._event_shape)
    251             variance = diag.expand(self._batch_shape + self._event_shape)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/lazy/sum_lazy_tensor.py in diag(self)
     93 
     94     def diag(self):
---> 95         return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/lazy/sum_lazy_tensor.py in <genexpr>(.0)
     93 
     94     def diag(self):
---> 95         return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/lazy/sum_lazy_tensor.py in diag(self)
     93 
     94     def diag(self):
---> 95         return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/lazy/sum_lazy_tensor.py in <genexpr>(.0)
     93 
     94     def diag(self):
---> 95         return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     57         kwargs_pkl = pickle.dumps(kwargs)
     58         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
---> 59             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
     60         return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
     61 

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in diag(self)
    253         x2 = self.x2
    254 
--> 255         res = super(Kernel, self.kernel).__call__(
    256             x1, x2, diag=True, last_dim_is_batch=self.last_dim_is_batch, **self.params
    257         )

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     28 
     29     def __call__(self, *inputs, **kwargs):
---> 30         outputs = self.forward(*inputs, **kwargs)
     31         if isinstance(outputs, list):
     32             return [_validate_module_outputs(output) for output in outputs]

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/kernels/scale_kernel.py in forward(self, x1, x2, last_dim_is_batch, diag, **params)
     90 
     91     def forward(self, x1, x2, last_dim_is_batch=False, diag=False, **params):
---> 92         orig_output = self.base_kernel.forward(x1, x2, diag=diag, last_dim_is_batch=last_dim_is_batch, **params)
     93         outputscales = self.outputscale
     94         if last_dim_is_batch:

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/kernels/distributional_input_kernel.py in forward(self, x1, x2, diag, *args, **kwargs)
     31 
     32     def forward(self, x1, x2, diag=False, *args, **kwargs):
---> 33         negative_covar_func = -self.distance_function(x1, x2)
     34         res = negative_covar_func.div(self.lengthscale).exp()
     35 

~/anaconda3/envs/SchemaModel/lib/python3.8/site-packages/gpytorch/kernels/gaussian_symmetrized_kl_kernel.py in _symmetrized_kl(dist1, dist2, eps)
     28     dist2_var = eps + dist2_logvar.exp()
     29 
---> 30     var_ratio12 = dist1_var / dist2_var
     31     # log_var_ratio12 = var_ratio12.log()
     32     # note that the log variance ratio cancels because of the summed KL.

RuntimeError: CUDA out of memory. Tried to allocate 140.19 GiB (GPU 0; 11.78 GiB total capacity; 635.92 MiB already allocated; 9.71 GiB free; 786.00 MiB reserved in total by PyTorch)

[wjmaddox]

Sorry for the slow response.

This is a) real and b) more of an implementation goof-up on my end as the code is ignoring the need to get out the self-KL (aka zero) when diag = True.

In general, there's probably a more memory efficient implementation, but I'm not entirely sure how to refactor it at the moment.

[wjmaddox]

In this setting, a workaround (as you're only looking for the predictive variances, which are conditionally independent) is to batch the test time variances in the following manner:

pred_list = []
for i in range(0, 7000, 500):
    with torch.no_grad():
        observed_pred = likelihood(model(test_x_distributional[i:(i+500)].to("cuda")))
        pred_list.append(observed_pred.confidence_region())

predictive_c_region = [
    torch.cat([x[0] for x in pred_list], dim=0),
    torch.cat([x[1] for x in pred_list], dim=0),
]

[brucejwittmann]

Thanks for the help! I ended up implementing something similar with a torch DataLoader. That fixed all problems.