Repository: cornellius-gp/gpytorch

Issue:
Option to not compute the variances
In some use cases, we only ever care about the predictive mean of the GP. It would save a lot of computation time if there were an option to not compute the variances. Currently we have a hack to turn off variance computation at https://github.com/cornellius-gp/gpytorch/commit/34912d666664c34b0e420dafbd7b424445a52bc0

Maybe we can un-hackify this and make it an option for all GPs. 
Option to not compute the variances
In some use cases, we only ever care about the predictive mean of the GP. It would save a lot of computation time if there were an option to not compute the variances. Currently we have a hack to turn off variance computation at https://github.com/cornellius-gp/gpytorch/commit/34912d666664c34b0e420dafbd7b424445a52bc0

Maybe we can un-hackify this and make it an option for all GPs. 


Discussion:
We could make a LazyTensor for the variance computations, but it might be a little complicated.
Alternatively, we could have a feature flag (or option passed into ExactGP models) that returns a ZeroLazyTensor for the multivariate normal and skips covariance computation. 
We could make a LazyTensor for the variance computations, but it might be a little complicated.
Alternatively, we could have a feature flag (or option passed into ExactGP models) that returns a ZeroLazyTensor for the multivariate normal and skips covariance computation. 

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
@@ -4,7 +4,9 @@
 
 from .. import settings
 from ..distributions import MultivariateNormal
-from ..lazy import InterpolatedLazyTensor, LazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor
+from ..lazy import (
+    InterpolatedLazyTensor, LazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor, ZeroLazyTensor
+)
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.memoize import cached
 
@@ -305,6 +307,12 @@ def exact_predictive_covar(self, test_test_covar, test_train_covar):
         """
         from ..distributions import MultivariateNormal
 
+        if settings.fast_pred_var.on():
+            self._last_test_train_covar = test_train_covar
+
+        if settings.skip_posterior_variances.on():
+            return ZeroLazyTensor(*test_test_covar.size())
+
         if settings.fast_pred_var.off():
             train_train_covar = self.likelihood(
                 MultivariateNormal(torch.zeros(1), self.train_train_covar), self.train_inputs
@@ -312,15 +320,12 @@ def exact_predictive_covar(self, test_test_covar, test_train_covar):
             test_train_covar = test_train_covar.evaluate()
             train_test_covar = test_train_covar.transpose(-1, -2)
             covar_correction_rhs = train_train_covar.inv_matmul(train_test_covar).mul(-1)
-            res = test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
-            return res
+            return test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
 
-        self._last_test_train_covar = test_train_covar
         precomputed_cache = self.covar_cache
-
-        covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache, test_train_covar)
-        res = test_test_covar + RootLazyTensor(covar_inv_quad_form_root).mul(-1)
-        return res
+        covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache,
+                                                                                   test_train_covar)
+        return test_test_covar + RootLazyTensor(covar_inv_quad_form_root).mul(-1)
 
 
 @register_prediction_strategy(InterpolatedLazyTensor)
diff --git a/gpytorch/settings.py b/gpytorch/settings.py
--- a/gpytorch/settings.py
+++ b/gpytorch/settings.py
@@ -113,6 +113,16 @@ class check_training_data(_feature_flag):
     _state = True
 
 
+class skip_posterior_variances(_feature_flag):
+    """
+    Whether or not to skip the posterior covariance matrix when doing an ExactGP
+    forward pass. If this is on, the returned gpytorch MultivariateNormal will have a
+    ZeroLazyTensor as its covariance matrix. This allows gpytorch to not compute
+    the covariance matrix when it is not needed, speeding up computations.
+    """
+    _state = False
+
+
 class detach_test_caches(_feature_flag):
     """
     Whether or not to detach caches computed for making predictions. In most cases, you will want this,


Test patch:
diff --git a/test/examples/test_simple_gp_regression.py b/test/examples/test_simple_gp_regression.py
--- a/test/examples/test_simple_gp_regression.py
+++ b/test/examples/test_simple_gp_regression.py
@@ -134,6 +134,48 @@ def test_posterior_latent_gp_and_likelihood_without_optimization_cuda(self):
         if torch.cuda.is_available():
             self.test_posterior_latent_gp_and_likelihood_without_optimization(cuda=True)
 
+    def test_gp_posterior_mean_skip_variances_fast_cuda(self):
+        if torch.cuda.is_available():
+            train_x, test_x, train_y, _ = self._get_data(cuda=True)
+            likelihood = GaussianLikelihood()
+            gp_model = ExactGPModel(train_x, train_y, likelihood)
+
+            gp_model.cuda()
+            likelihood.cuda()
+
+            # Compute posterior distribution
+            gp_model.eval()
+            likelihood.eval()
+
+            with gpytorch.settings.skip_posterior_variances(True):
+                mean_skip_var = gp_model(test_x).mean
+            mean = gp_model(test_x).mean
+            likelihood_mean = likelihood(gp_model(test_x)).mean
+
+            self.assertTrue(torch.allclose(mean_skip_var, mean))
+            self.assertTrue(torch.allclose(mean_skip_var, likelihood_mean))
+
+    def test_gp_posterior_mean_skip_variances_slow_cuda(self):
+        if torch.cuda.is_available():
+            train_x, test_x, train_y, _ = self._get_data(cuda=True)
+            likelihood = GaussianLikelihood()
+            gp_model = ExactGPModel(train_x, train_y, likelihood)
+
+            gp_model.cuda()
+            likelihood.cuda()
+
+            # Compute posterior distribution
+            gp_model.eval()
+            likelihood.eval()
+
+            with gpytorch.settings.fast_pred_var(False):
+                with gpytorch.settings.skip_posterior_variances(True):
+                    mean_skip_var = gp_model(test_x).mean
+                mean = gp_model(test_x).mean
+                likelihood_mean = likelihood(gp_model(test_x)).mean
+            self.assertTrue(torch.allclose(mean_skip_var, mean))
+            self.assertTrue(torch.allclose(mean_skip_var, likelihood_mean))
+
     def test_posterior_latent_gp_and_likelihood_with_optimization(self, cuda=False):
         train_x, test_x, train_y, test_y = self._get_data(cuda=cuda)
         # We're manually going to set the hyperparameters to something they shouldn't be
@@ -173,7 +215,8 @@ def test_posterior_latent_gp_and_likelihood_with_optimization(self, cuda=False):
         # Test the model
         gp_model.eval()
         likelihood.eval()
-        test_function_predictions = likelihood(gp_model(test_x))
+        with gpytorch.settings.skip_posterior_variances(True):
+            test_function_predictions = likelihood(gp_model(test_x))
         mean_abs_error = torch.mean(torch.abs(test_y - test_function_predictions.mean))
 
         self.assertLess(mean_abs_error.item(), 0.05)