Repository: cornellius-gp/gpytorch

Issue:
Flag to turn off logdet computations for MLLs
For exact and variational GPs, the marginal log likelihood requires computing the log determinant of the prior covariance matrix. This isn't a super expensive operation (since we get most of the terms for free from a single call to `linear_cg`). However, we still have to make a few calls to `torch.symeig`, which can add a bit of a performance hit.

If we are simply using the MLL as a loss function to optimize hyperparameters (and not for Gibbs sampling, model selection, etc.), then we don't actually care about the logdet value (just its derivatives, which we get for free from `linear_cg`). Therefore, we could avoid these `torch.symeig` calls.

I propose a `gpytorch.beta_features.skip_logdet` feature flag that skips the logdet computation (but still computes the correct gradient, obviously). It will be off by default.


Discussion:

Resolving patch:
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -95,7 +95,7 @@ def forward(self, *args):
         # Perform solves (for inv_quad) and tridiagonalization (for estimating log_det)
         rhs = torch.cat(rhs_list, -1)
         t_mat = None
-        if self.log_det:
+        if self.log_det and not settings.skip_logdet_forward.on():
             solves, t_mat = linear_cg(
                 matmul_closure,
                 rhs,
@@ -109,17 +109,17 @@ def forward(self, *args):
             solves = linear_cg(
                 matmul_closure,
                 rhs,
-                n_tridiag=num_random_probes,
+                n_tridiag=0,
                 max_iter=settings.max_cg_iterations.value(),
                 preconditioner=self.preconditioner,
             )
 
         # Final values to return
-        log_det_term = torch.empty(0, dtype=self.dtype, device=self.device)
-        inv_quad_term = torch.empty(0, dtype=self.dtype, device=self.device)
+        log_det_term = torch.zeros(lazy_tsr.batch_shape, dtype=self.dtype, device=self.device)
+        inv_quad_term = torch.zeros(lazy_tsr.batch_shape, dtype=self.dtype, device=self.device)
 
         # Compute log_det from tridiagonalization
-        if self.log_det:
+        if self.log_det and not settings.skip_logdet_forward.on():
             if torch.any(torch.isnan(t_mat)).item():
                 log_det_term = float("nan")
             else:
diff --git a/gpytorch/settings.py b/gpytorch/settings.py
--- a/gpytorch/settings.py
+++ b/gpytorch/settings.py
@@ -143,6 +143,32 @@ class num_trace_samples(_value_context):
     _global_value = 10
 
 
+class skip_logdet_forward(_feature_flag):
+    """
+    .. warning:
+
+        ADVANCED FEATURE. Use this feature ONLY IF you're using
+        `gpytorch.mlls.MarginalLogLikelihood` as loss functions for optimizing
+        hyperparameters/variational parameters.  DO NOT use this feature if you
+        need accurate estimates of the MLL (i.e. for model selection, MCMC,
+        second order optimizaiton methods, etc.)
+
+    This feature does not affect the gradients returned by
+    :meth:`gpytorch.distributions.MultivariateNormal.log_prob`
+    (used by `gpytorch.mlls.MarginalLogLikelihood`).
+    The gradients remain unbiased estimates, and therefore can be used with SGD.
+    However, the actual likelihood value returned by the forward
+    pass will skip certain computations (i.e. the logdet computation), and will therefore
+    be improper estimates.
+
+    If you're using SGD (or a varient) to optimize parameters, you probably
+    don't need an accurate MLL estimate; you only need accurate gradients. So
+    this setting may give your model a performance boost.
+    """
+
+    _state = False
+
+
 class terminate_cg_by_size(_feature_flag):
     """
     If set to true, cg will terminate after n iterations for an n x n matrix.


Test patch:
diff --git a/test/functions/test_inv_quad_log_det.py b/test/functions/test_inv_quad_log_det.py
--- a/test/functions/test_inv_quad_log_det.py
+++ b/test/functions/test_inv_quad_log_det.py
@@ -86,6 +86,25 @@ def test_inv_quad_log_det_many_vectors(self):
         self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
         self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
 
+    def test_inv_quad_log_det_many_vectors_improper(self):
+        # Forward pass
+        actual_inv_quad = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
+        actual_log_det = self.mat_clone.logdet()
+        with gpytorch.settings.num_trace_samples(1000), gpytorch.settings.skip_logdet_forward(True):
+            non_lazy_tsr = NonLazyTensor(self.mat)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vecs, log_det=True)
+        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
+        self.assertAlmostEqual(res_log_det.item(), torch.zeros(non_lazy_tsr.batch_shape), places=1)
+
+        # Backward
+        actual_inv_quad.backward()
+        actual_log_det.backward()
+        res_inv_quad.backward(retain_graph=True)
+        res_log_det.backward()
+
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
     def test_inv_quad_only_many_vectors(self):
         # Forward pass
         res = NonLazyTensor(self.mat).inv_quad(self.vecs)
@@ -164,6 +183,36 @@ def test_inv_quad_log_det_many_vectors(self):
         self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
         self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
 
+    def test_inv_quad_log_det_many_vectors_improper(self):
+        # Forward pass
+        actual_inv_quad = (
+            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
+            .sum(2)
+            .sum(1)
+        )
+        actual_log_det = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
+        with gpytorch.settings.num_trace_samples(2000), gpytorch.settings.skip_logdet_forward(True):
+            non_lazy_tsr = NonLazyTensor(self.mats)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vecs, log_det=True)
+
+        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
+        self.assertEqual(res_log_det.shape, actual_log_det.shape)
+        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
+        self.assertLess(torch.max(res_log_det.abs()).item(), 1e-1)
+
+        # Backward
+        inv_quad_grad_output = torch.randn(5, dtype=torch.float)
+        log_det_grad_output = torch.randn(5, dtype=torch.float)
+        actual_inv_quad.backward(gradient=inv_quad_grad_output)
+        actual_log_det.backward(gradient=log_det_grad_output)
+        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
+        res_log_det.backward(gradient=log_det_grad_output)
+
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
     def test_inv_quad_only_many_vectors(self):
         # Forward pass
         res = NonLazyTensor(self.mats).inv_quad(self.vecs)