Repository: cornellius-gp/gpytorch (python)

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.

Pull request description:
Option to skip LazyTensor.logdet forward pass
This is going to be a not-highly-advertised feature.

This PR introduces `gpytorch.settings.skip_logdet_forward`. Setting this to True makes `LazyTensor.logdet` return a zero tensor, but one that does the correct thing on the backward pass.

This allows advanced users (like us 😁) to get much faster performance for hyperparameter optimization. If we're using a single-order method - we don't actually care about the logdet value - we just care about its derivatives. This settings will allow us to skip unnecessary forward-pass computation.

Fixes #396 

Patch:
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index fe3147d2..7c0082e9 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -95,7 +95,7 @@ class InvQuadLogDet(Function):
         # 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 @@ class InvQuadLogDet(Function):
             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
index 2de721c3..dcff9ede 100644
--- 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
index 487b6492..085b111b 100644
--- a/test/functions/test_inv_quad_log_det.py
+++ b/test/functions/test_inv_quad_log_det.py
@@ -86,6 +86,25 @@ class TestInvQuadLogDetNonBatch(unittest.TestCase):
         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 @@ class TestInvQuadLogDetBatch(unittest.TestCase):
         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)