Merge pull request #1717 from AustinT/sgpr-diag-correction

Add setting to turn off diagonal correction for InducingPointKernel.

Author: jacobrgardner

gpytorch/kernels/inducing_point_kernel.py

@@ -5,6 +5,7 @@
 
 import torch
 
+from .. import settings
 from ..distributions import MultivariateNormal
 from ..lazy import DiagLazyTensor, LowRankRootAddedDiagLazyTensor, LowRankRootLazyTensor, MatmulLazyTensor, delazify
 from ..mlls import InducingPointKernelAddedLossTerm
@@ -61,7 +62,7 @@ def _get_covariance(self, x1, x2):
             covar = LowRankRootLazyTensor(k_ux1.matmul(self._inducing_inv_root))
 
             # Diagonal correction for predictive posterior
-            if not self.training:
+            if not self.training and settings.sgpr_diagonal_correction.on():
                 correction = (self.base_kernel(x1, x2, diag=True) - covar.diag()).clamp(0, math.inf)
                 covar = LowRankRootAddedDiagLazyTensor(covar, DiagLazyTensor(correction))
         else:

gpytorch/settings.py

@@ -640,6 +640,19 @@ class prior_mode(_feature_flag):
     _default = False
 
 
+class sgpr_diagonal_correction(_feature_flag):
+    """
+    If set to true, during posterior prediction the variances of the InducingPointKernel
+    will be corrected to match the variances of the exact kernel.
+
+    If false then no such correction will be performed (this is the default in other libraries).
+
+    (Default: True)
+    """
+
+    _default = True
+
+
 class skip_logdet_forward(_feature_flag):
     """
     .. warning:

test/kernels/test_inducing_point_kernel.py

@@ -48,3 +48,14 @@ def test_kernel_output(self):
             output = model.likelihood(model(test_x))
             _ = output.mean + output.variance  # Compute something to break through any lazy evaluations
             self.assertTrue(ps_mock.called)
+
+        # Check whether changing diagonal correction makes a difference (ensuring that cache is cleared)
+        model.train()
+        model.eval()
+        with gpytorch.settings.sgpr_diagonal_correction(True), torch.no_grad():
+            output_mean_correct = model(test_x).mean
+        model.train()
+        model.eval()
+        with gpytorch.settings.sgpr_diagonal_correction(False), torch.no_grad():
+            output_mean_no_correct = model(test_x).mean
+        self.assertNotAlmostEqual(output_mean_correct.sum().item(), output_mean_no_correct.sum().item())