Implement variance reduction in SLQ logdet backward pass.

[gpleiss]

Based on "Reducing the Variance of Gaussian Process Hyperparameter Optimization with Preconditioning" by Wenger et al., 2021.

When using iterative methods (i.e. CG/SLQ) to compute the log determinant, the forward pass currently computes:
logdet K \approx logdet P + SLQ( P^{-1/2} K P^{-1/2} ),
where P is a preconditioner, and SLQ is a stochastic estimate of the log determinant. If the preconditioner is a good approximation of K, then this forward pass can be seen as a form of variance reduction.

In this PR, we apply this same variance reduction strategy to the backward pass. We compute the backward pass as:
d logdet(K)/dtheta \approx d logdet(P)/dtheta + d SLQ/dtheta

TODOs:

• Implement pivoted cholesky as a torch.autograd.Function, so that we can compute backward passes through it.
• Redo inv_quad_logdet function to apply variance reduction in the forward and backward passes.

[gpleiss]

@jacobrgardner @JonathanWenger ready for review

[gpleiss]

(actually ready for review now. I just fixed broken tests.)

[jacobrgardner]

@gpleiss something seems off about how this is computing the preconditioner log determinant now. We're still computing it efficiently using the QR decomposition in the init_cache methods on AddedDiagLT, and then it looks like we discard that efficient computation and call logdet on precond_lt which could be as bad as calling Cholesky on the full n x n preconditioner matrix, right? PsdSumLazyTensor doesn't override logdet?

Entirely possible I just missed the relevant code here...

[jacobrgardner]

Ideally, this would be the logdet value we'd return in the forward pass:

gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py

Line 147 in 70b79c2

self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()

[jacobrgardner]

Looks solid to me now 👍

[gpleiss]

I just profiled this PR on the KeOps example notebook - just to double check. It is just as fast as what is on is on master.