Repository: cornellius-gp/gpytorch

Issue:
[Bug] `gpytorch.settings.deterministic_probes` computes an extra Lanczos factorization
# 🐛 Bug

In preconditioned CG mode, we make the following lining calls for each training iteration:

```
LinAlg (Verbose) - DEBUG - Running Pivoted Cholesky on a torch.Size([5000, 5000]) RHS for 15 iterations.

# these two calls shouldn't exist
LinAlg (Verbose) - DEBUG - Running Lanczos on a torch.Size([5000, 5000]) matrix with a torch.Size([5000, 1]) RHS for 100 iterations.
LinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([9, 9]).

LinAlg (Verbose) - DEBUG - Running CG on a torch.Size([5000, 11]) RHS for 1000 iterations (tol=1). Output: torch.Size([5000, 11]).
LinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([10, 4, 4]).

LinAlg (Verbose) - INFO - Iter 1/50 - Loss: 0.808   lengthscale: 0.693   noise: 0.693
```

This Lanczos/symeig calls come from this line in [inv_quad_logdet](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/functions/_inv_quad_log_det.py#L98)

```python
            else:  # When preconditioning, probe vectors must be drawn from N(0, P)
                if precond_lt.size()[-2:] == torch.Size([1, 1]):
                    covar_root = precond_lt.evaluate().sqrt()
                else:
                    covar_root = precond_lt.root_decomposition().root
```

Before the `gpytorch.settings.deterministic_probes` option was added (#929), the code looked like this:

```python
            else:  # When preconditioning, probe vectors must be drawn from N(0, P)
                probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)
                # ...
```

This is much more efficient, as it does not perform a root decomposition.

@jacobrgardner there is an obvious easy fix: we don't call `root_decomposition` unless we're using deterministic probes. However, I think this makes the `deterministic_probes` setting really undesirable - the slight variance reductions gains we get from it come at the cost of an extra linalg routine.

Thoughts?


Discussion:
@gpleiss So, the original reason we added deterministic probes was I think for @Balandat because it maybe seemed to work much better when using L-BFGS as an optimizer, which sometimes doesn't handle stochasticity particularly well? Now I think there might be better solutions to this. I certainly don't think we should ever use deterministic probes with Adam.
> which sometimes doesn't handle stochasticity particularly well

At least standard L-BFGS-B (not the more advanced stochastic quasi-second order methods) will generally break without deterministic probes. I agree they should never be used for Adam.

We're not actively using determinstic probes currently, but it would be good to keep the setting around if it doesn't hurt since we're like to revisit this. So for now can we just put calling `root_decomposition` behind the deterministic probe setting? 
Sounds good! I'll put up a PR.

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
@@ -92,12 +92,15 @@ def forward(
                     probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)
                     probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)
             else:  # When preconditioning, probe vectors must be drawn from N(0, P)
-                if precond_lt.size()[-2:] == torch.Size([1, 1]):
-                    covar_root = precond_lt.evaluate().sqrt()
-                else:
-                    covar_root = precond_lt.root_decomposition().root
-
                 if settings.deterministic_probes.on():
+                    # NOTE: calling precond_lt.root_decomposition() is expensive
+                    # because it requires Lanczos
+                    # We don't have any other choice for when we want to use deterministic probes, however
+                    if precond_lt.size()[-2:] == torch.Size([1, 1]):
+                        covar_root = precond_lt.evaluate().sqrt()
+                    else:
+                        covar_root = precond_lt.root_decomposition().root
+
                     warnings.warn(
                         "Deterministic probes will currently work only if you aren't training multiple independent"
                         " models simultaneously.",
@@ -116,13 +119,6 @@ def forward(
 
                     probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))
                 else:
-                    base_samples = torch.randn(
-                        *precond_lt.batch_shape,
-                        covar_root.size(-1),
-                        num_random_probes,
-                        dtype=precond_lt.dtype,
-                        device=precond_lt.device,
-                    )
                     probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)
                 probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()
                 probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)


Test patch: