Repository: cornellius-gp/gpytorch

Issue:
[Bug] New release broke SGPR predictive variance?
# 🐛 Bug

Hi guys and thanks a lot for your work.

The predictive variance in my SGPR use cases is broken, approximately since the new release. This happens in both single output and batch mode, it doesn't throw an error but the values for the predictive variance are abnormally large and often negative. Probably a bug that slipped unnoticed through the release, maybe in the InducingPointKernel.

Thanks for your help!

## To reproduce

```python
import time
import urllib.request
import gpytorch
import torch
from gpytorch.distributions import MultivariateNormal
from gpytorch.kernels import ScaleKernel, RBFKernel, InducingPointKernel
from gpytorch.means import ConstantMean
from scipy.io import loadmat

if __name__ == '__main__':
    # Run GPyTorch SGPR + independent multioutputs example: approximate
    # https://docs.gpytorch.ai/en/v1.2.1/examples/02_Scalable_Exact_GPs/SGPR_Regression_CUDA.html
    # https://github.com/cornellius-gp/gpytorch/issues/1043
    print('Downloading \'elevators\' UCI dataset...')
    urllib.request.urlretrieve(
        'https://drive.google.com/uc?export=download&id=1jhWL3YUHvXIaftia4qeAyDwVxo6j1alk',
        '../elevators.mat')
    output_size = 1
    input_size = 18
    nb_inducing_points = 500
    data = torch.Tensor(loadmat('../elevators.mat')['data'])
    X = data[:, :-1]
    X = X - X.min(0)[0]
    X = 2 * (X / X.max(0)[0]) - 1
    X = X[:, :input_size]
    y = data[:, -1]
    # MAKE MULTIOUTPUT DATA
    y = y.reshape(-1, 1)
    y = y.repeat(1, output_size)
    input_size = X.shape[1]
    train_x = X[:10000, :].contiguous()
    train_y = y[:10000].contiguous()
    test_x = X[10000:20000, :].contiguous()
    test_y = y[10000:20000].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()
    # CONVERT TO BATCH GP
    train_x = train_x.repeat(output_size, 1, 1)
    train_y = train_y.transpose(-2, -1)
    test_x = test_x.repeat(output_size, 1, 1)
    test_y = test_y.transpose(-2, -1)


    class GPRegressionModel(gpytorch.models.ExactGP):
        def __init__(self, train_x, train_y, likelihood):
            super(GPRegressionModel, self).__init__(train_x, train_y,
                                                    likelihood)
            self.mean_module = ConstantMean(
                batch_shape=torch.Size([output_size]))
            self.base_covar_module = ScaleKernel(RBFKernel(
                batch_shape=torch.Size([output_size])),
                batch_shape=torch.Size([output_size]))
            inducing_points = train_x[:, :nb_inducing_points, :]
            self.covar_module = InducingPointKernel(
                self.base_covar_module,
                inducing_points=inducing_points,
                likelihood=likelihood)

        def forward(self, x):
            mean_x = self.mean_module(x)
            covar_x = self.covar_module(x)
            return MultivariateNormal(mean_x, covar_x)


    likelihood = gpytorch.likelihoods.GaussianLikelihood(
        batch_shape=torch.Size([output_size]))
    model = GPRegressionModel(train_x, train_y, likelihood)
    if torch.cuda.is_available():
        model = model.cuda()
        likelihood = likelihood.cuda()
    # Train
    training_iterations = 10
    model.train()
    likelihood.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
    for i in range(training_iterations):
        start = time.time()
        # Zero backprop gradients
        optimizer.zero_grad()
        # Get output from model
        output = model(train_x)
        # Calc loss and backprop derivatives
        loss = -mll(output, train_y).sum()
        loss.backward()
        end = time.time()
        print('Iter %d/%d - Loss: %.3f' % (
            i + 1, training_iterations, loss.item()), 'in', str(end - start))
        optimizer.step()
        torch.cuda.empty_cache()
    model.eval()
    likelihood.eval()
    with gpytorch.settings.max_preconditioner_size(10), torch.no_grad():
        with gpytorch.settings.max_root_decomposition_size(
                30), gpytorch.settings.fast_pred_var():
            preds = model(test_x)
    print(torch.mean(preds.covariance_matrix))
    print(torch.mean(preds.variance))
```

## Expected Behavior

With version 1.3.1 of GPyTorch, outputs `tensor(3.6882e-05, grad_fn=<MeanBackward0>)`, `tensor(0.0162, grad_fn=<MeanBackward0>)`.
With version 1.4.0, outputs `tensor(-2916.0591)` then hangs. 

## System information

- GPyTorch version 1.4.0
- PyTorch version 1.7.1
- Mac OS Big Sur



Discussion:
And while I'm at it, before the new release I used to initialize the value of a MultitaskGaussianLikelihood with
```python
likelihood.noise_covar.noise = torch.tensor([0.04])
likelihood.noise = torch.tensor([1e-4])
```
but now this throws an error saying `likelihood.noise_covar.noise` does not exist and `likelihood.noise` is the wrong size for `num_tasks > 1`. Any idea how I am supposed to set the value of the MultitaskGaussianLikelihood now?

Thanks!
@monabf I'm looking into the SGPR issue now - seems like something was broken by the release.

Can you please open up a separate issue for the initialization code?
Sure!
Thanks a lot for your quick response!

Resolving patch:
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -667,19 +667,24 @@ class SGPRPredictionStrategy(DefaultPredictionStrategy):
     def covar_cache(self):
         # Here, the covar_cache is going to be the inverse of K_{XX} + \sigma^2 I
         # This is easily computed using Woodbury
+        # K_{XX} + \sigma^2 I = R R^T + \sigma^2 I
+        #                     = \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )
         train_train_covar = self.lik_train_train_covar.evaluate_kernel()
 
         # Get terms needed for woodbury
-        root = train_train_covar._lazy_tensor.root_decomposition().root
-        inv_diag = train_train_covar._diag_tensor.inverse()
+        root = train_train_covar._lazy_tensor.root_decomposition().root.evaluate()  # R
+        inv_diag = train_train_covar._diag_tensor.inverse()  # \sigma^{-2}
 
         # Form LT using woodbury
-        ones = torch.tensor(1.0, dtype=inv_diag.dtype, device=inv_diag.device)
-        chol_factor = (root.transpose(-1, -2) @ root).add_diag(ones).cholesky().evaluate()
-        woodbury_term = torch.triangular_solve(
-            inv_diag.diag().unsqueeze(-2) * root.evaluate().transpose(-1, -2), chol_factor, upper=False
-        )[0]
-        inverse = AddedDiagLazyTensor(MatmulLazyTensor(woodbury_term.transpose(-1, -2), -woodbury_term), inv_diag)
+        ones = torch.tensor(1.0, dtype=root.dtype, device=root.device)
+        chol_factor = lazify(root.transpose(-1, -2) @ (inv_diag @ root)).add_diag(ones)  # (I + \sigma^{-2} R^T R)^{-1}
+        woodbury_term = inv_diag @ torch.triangular_solve(
+            root.transpose(-1, -2), chol_factor.cholesky().evaluate(), upper=False
+        )[0].transpose(-1, -2)
+        # woodbury_term @ woodbury_term^T = \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T \sigma^{-2}
+
+        inverse = AddedDiagLazyTensor(inv_diag, MatmulLazyTensor(-woodbury_term, woodbury_term.transpose(-1, -2)))
+        # \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )
         return inverse
 
     def get_fantasy_strategy(self, inputs, targets, full_inputs, full_targets, full_output, **kwargs):
@@ -700,4 +705,6 @@ def exact_prediction(self, joint_mean, joint_covar):
 
     def exact_predictive_covar(self, test_test_covar, test_train_covar):
         train_train_precision = self.covar_cache
-        return (test_train_covar @ (train_train_precision @ test_train_covar.transpose(-1, -2))) + test_test_covar
+        test_train_covar = test_train_covar.evaluate()  # Evaluation is more efficient than lazy evaluation
+        res = test_test_covar - (test_train_covar @ (train_train_precision @ test_train_covar.transpose(-1, -2)))
+        return res


Test patch:
diff --git a/test/examples/test_sgpr_regression.py b/test/examples/test_sgpr_regression.py
--- a/test/examples/test_sgpr_regression.py
+++ b/test/examples/test_sgpr_regression.py
@@ -17,6 +17,8 @@
 from gpytorch.utils.warnings import NumericalWarning
 from torch import optim
 
+from gpytorch.test.base_test_case import BaseTestCase
+
 
 # Simple training data: let's try to learn a sine function,
 # but with SGPR
@@ -50,7 +52,7 @@ def forward(self, x):
         return MultivariateNormal(mean_x, covar_x)
 
 
-class TestSGPRRegression(unittest.TestCase):
+class TestSGPRRegression(unittest.TestCase, BaseTestCase):
     def setUp(self):
         if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
             self.rng_state = torch.get_rng_state()
@@ -101,45 +103,14 @@ def test_sgpr_mean_abs_error(self):
 
         self.assertLess(mean_abs_error.squeeze().item(), 0.05)
 
-    def test_sgpr_fast_pred_var(self):
-        # Suppress numerical warnings
-        warnings.simplefilter("ignore", NumericalWarning)
-
-        train_x, train_y, test_x, test_y = make_data()
-        likelihood = GaussianLikelihood()
-        gp_model = GPRegressionModel(train_x, train_y, likelihood)
-        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
-
-        # Optimize the model
-        gp_model.train()
-        likelihood.train()
-
-        optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
-        for _ in range(50):
-            optimizer.zero_grad()
-            output = gp_model(train_x)
-            loss = -mll(output, train_y)
-            loss.backward()
-            optimizer.step()
-
-        for param in gp_model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-
-        # Test the model
-        gp_model.eval()
-        likelihood.eval()
-
-        with gpytorch.settings.max_preconditioner_size(5), gpytorch.settings.max_cg_iterations(50):
-            with gpytorch.settings.fast_pred_var(True):
-                fast_var = gp_model(test_x).variance
-                fast_var_cache = gp_model(test_x).variance
-                self.assertLess(torch.max((fast_var_cache - fast_var).abs()), 1e-3)
-
-            with gpytorch.settings.fast_pred_var(False):
-                slow_var = gp_model(test_x).variance
-
-        self.assertLess(torch.max((fast_var_cache - slow_var).abs()), 1e-3)
+        # Test variances
+        test_vars = likelihood(gp_model(test_x)).variance
+        self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
+        self.assertGreater(test_vars.min().item() + 0.05, likelihood.noise.item())
+        self.assertLess(
+            test_vars.max().item() - 0.05,
+            likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
+        )
 
     def test_sgpr_mean_abs_error_cuda(self):
         # Suppress numerical warnings
@@ -179,6 +150,15 @@ def test_sgpr_mean_abs_error_cuda(self):
 
             self.assertLess(mean_abs_error.squeeze().item(), 0.02)
 
+            # Test variances
+            test_vars = likelihood(gp_model(test_x)).variance
+            self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
+            self.assertGreater(test_vars.min().item() + 0.05, likelihood.noise.item())
+            self.assertLess(
+                test_vars.max().item() - 0.05,
+                likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
+            )
+
 
 if __name__ == "__main__":
     unittest.main()