Repository: cornellius-gp/gpytorch
Pull request #454 (2019-01-02 13:15:52)

Cache solves for fast variational
This PR introduces `CachedCGLazyTensor`, which caches the solves that are performed by CG. This improves the speed of variational models. Variational models now call CG 1 time per ELBO computation rather than 3 times.


Patch:
diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
index 6afc1696..480d0896 100644
--- a/gpytorch/distributions/multivariate_normal.py
+++ b/gpytorch/distributions/multivariate_normal.py
@@ -225,7 +225,7 @@ except ImportError:
 @register_kl(MultivariateNormal, MultivariateNormal)
 def kl_mvn_mvn(p_dist, q_dist):
     q_mean = q_dist.loc
-    q_covar = q_dist.lazy_covariance_matrix.add_jitter()
+    q_covar = q_dist.lazy_covariance_matrix
 
     p_mean = p_dist.loc
     p_covar = p_dist.lazy_covariance_matrix
@@ -236,7 +236,7 @@ def kl_mvn_mvn(p_dist, q_dist):
         # right now this just catches if root_p_covar is a DiagLazyTensor,
         # but we may want to be smarter about this in the future
         root_p_covar = root_p_covar.evaluate()
-    inv_quad_rhs = torch.cat([root_p_covar, mean_diffs.unsqueeze(-1)], -1)
+    inv_quad_rhs = torch.cat([mean_diffs.unsqueeze(-1), root_p_covar], -1)
     log_det_p_covar = p_covar.log_det()
     trace_plus_inv_quad_form, log_det_q_covar = q_covar.inv_quad_log_det(inv_quad_rhs=inv_quad_rhs, log_det=True)
 
diff --git a/gpytorch/functions/__init__.py b/gpytorch/functions/__init__.py
index ab5eac3d..9d8c0ae8 100644
--- a/gpytorch/functions/__init__.py
+++ b/gpytorch/functions/__init__.py
@@ -85,19 +85,39 @@ def matmul(mat, rhs):
     return mat.matmul(rhs)
 
 
-def inv_matmul(mat, rhs):
-    """
-    Computes a linear solve with several right hand sides.
+def inv_matmul(mat, right_tensor, left_tensor=None):
+    r"""
+    Computes a linear solve (w.r.t :attr:`mat` = :math:`A`) with several right hand sides :math:`R`.
+    I.e. computes
+
+    ... math::
+
+        \begin{equation}
+            A^{-1} R,
+        \end{equation}
+
+    where :math:`R` is :attr:`right_tensor` and :math:`A` is :attr:`mat`.
+
+    If :attr:`left_tensor` is supplied, computes
+
+    ... math::
+
+        \begin{equation}
+            L A^{-1} R,
+        \end{equation}
+
+    where :math:`L` is :attr:`left_tensor`. Supplying this can reduce the number of
+    CG calls required.
 
     Args:
-        - mat (matrix nxn) - Matrix to solve with
-        - rhs (matrix nxk) - rhs matrix or vector
+        - :obj:`torch.tensor` (n x k) - Matrix :math:`R` right hand sides
+        - :obj:`torch.tensor` (m x n) - Optional matrix :math:`L` to perform left multiplication with
 
     Returns:
-        - matrix nxk - (mat)^{-1} rhs
+        - :obj:`torch.tensor` - :math:`A^{-1}R` or :math:`LA^{-1}R`.
     """
     from ..lazy import lazify
-    return lazify(mat).inv_matmul(rhs)
+    return lazify(mat).inv_matmul(right_tensor, left_tensor)
 
 
 def inv_quad(mat, tensor):
diff --git a/gpytorch/functions/_inv_matmul.py b/gpytorch/functions/_inv_matmul.py
index 3ccf8fac..76fd2d58 100644
--- a/gpytorch/functions/_inv_matmul.py
+++ b/gpytorch/functions/_inv_matmul.py
@@ -2,34 +2,48 @@
 
 import torch
 from torch.autograd import Function
-from ..utils import linear_cg
 from .. import settings
 
 
 class InvMatmul(Function):
-    def __init__(self, representation_tree, preconditioner=None):
+    def __init__(self, representation_tree, preconditioner=None, has_left=False):
         self.representation_tree = representation_tree
         self.preconditioner = preconditioner
-
-    def forward(self, rhs, *matrix_args):
+        self.has_left = has_left
+
+    def forward(self, *args):
+        left_tensor = None
+        right_tensor = None
+        matrix_args = None
+        if self.has_left:
+            left_tensor, right_tensor, *matrix_args = args
+        else:
+            right_tensor, *matrix_args = args
+        orig_right_tensor = right_tensor
         lazy_tsr = self.representation_tree(*matrix_args)
-        matmul_closure = lazy_tsr._matmul
 
         self.is_vector = False
-        if rhs.ndimension() == 1:
-            rhs.unsqueeze_(-1)
+        if right_tensor.ndimension() == 1:
+            right_tensor = right_tensor.unsqueeze(-1)
             self.is_vector = True
 
         # Perform solves (for inv_quad) and tridiagonalization (for estimating log_det)
-        res = linear_cg(
-            matmul_closure, rhs, max_iter=settings.max_cg_iterations.value(), preconditioner=self.preconditioner
-        )
+        if self.has_left:
+            rhs = torch.cat([left_tensor.transpose(-1, -2), right_tensor], -1)
+            solves = lazy_tsr._solve(rhs, self.preconditioner)
+            res = solves[..., left_tensor.size(-2):]
+            res = left_tensor @ res
+        else:
+            solves = lazy_tsr._solve(right_tensor, self.preconditioner)
+            res = solves
 
         if self.is_vector:
-            res.squeeze_(-1)
-            rhs.squeeze_(-1)
+            res = res.squeeze(-1)
 
-        args = [res, rhs] + list(matrix_args)
+        if self.has_left:
+            args = [solves, left_tensor, orig_right_tensor] + list(matrix_args)
+        else:
+            args = [solves, orig_right_tensor] + list(matrix_args)
         self.save_for_backward(*args)
         if settings.memory_efficient.off():
             self._lazy_tsr = lazy_tsr
@@ -38,46 +52,52 @@ class InvMatmul(Function):
 
     def backward(self, grad_output):
         # Extract items that were saved
-        rhs_solves = self.saved_tensors[0]
-        rhs = self.saved_tensors[1]
-        matrix_args = self.saved_tensors[2:]
+        if self.has_left:
+            solves, left_tensor, right_tensor, *matrix_args = self.saved_tensors
+            left_solves = solves[..., :left_tensor.size(-2)]
+            right_solves = solves[..., left_tensor.size(-2):]
+        else:
+            right_solves, right_tensor, *matrix_args = self.saved_tensors
 
         # Get matrix functions
         if hasattr(self, "_lazy_tsr"):
             lazy_tsr = self._lazy_tsr
         else:
             lazy_tsr = self.representation_tree(*matrix_args)
-        matmul_closure = lazy_tsr._matmul
 
         # Define gradient placeholders
         arg_grads = [None] * len(matrix_args)
-        rhs_grad = None
+        left_grad = None
+        right_grad = None
         if any(self.needs_input_grad):
             # De-vectorize objects
             if self.is_vector:
-                rhs = rhs.unsqueeze(-1)
-                rhs_solves = rhs_solves.unsqueeze(-1)
+                right_tensor = right_tensor.unsqueeze(-1)
                 grad_output = grad_output.unsqueeze(-1)
 
-            # Compute self^{-1} grad_output
-            grad_output_solves = linear_cg(
-                matmul_closure,
-                grad_output,
-                max_iter=settings.max_cg_iterations.value(),
-                preconditioner=self.preconditioner,
-            )
-
-            # input_1 gradient
-            if any(self.needs_input_grad[1:]):
-                if lazy_tsr is not None:
-                    arg_grads = lazy_tsr._quad_form_derivative(grad_output_solves, rhs_solves.mul(-1))
-                else:
-                    arg_grads = (torch.matmul(grad_output_solves, rhs_solves.mul(-1).transpose(-1, -2)),)
-
-            # input_2 gradient
-            if self.needs_input_grad[0]:
-                rhs_grad = grad_output_solves
-                if self.is_vector:
-                    rhs_grad.squeeze_(-1)
-
-        return tuple([rhs_grad] + list(arg_grads))
+            if not self.has_left:
+                # Compute self^{-1} grad_output
+                left_solves = lazy_tsr._solve(grad_output, self.preconditioner)
+
+                if any(self.needs_input_grad[1:]):
+                    arg_grads = lazy_tsr._quad_form_derivative(left_solves, right_solves.mul(-1))
+                if self.needs_input_grad[0]:
+                    right_grad = left_solves
+                    if self.is_vector:
+                        right_grad.squeeze_(-1)
+
+                return tuple([right_grad] + list(arg_grads))
+
+            else:
+                left_solves = left_solves @ grad_output
+
+                if self.needs_input_grad[1]:
+                    left_grad = grad_output @ right_solves.transpose(-1, -2)
+                if any(self.needs_input_grad[2:]):
+                    arg_grads = lazy_tsr._quad_form_derivative(left_solves, right_solves.mul(-1))
+                if self.needs_input_grad[0]:
+                    right_grad = left_solves
+                    if self.is_vector:
+                        right_grad.squeeze_(-1)
+
+                return tuple([left_grad, right_grad] + list(arg_grads))
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 152ccd23..7f357b9a 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -4,7 +4,6 @@ import torch
 from torch.autograd import Function
 from ..utils.lanczos import lanczos_tridiag_to_diag
 from ..utils.stochastic_lq import StochasticLQ
-from ..utils import linear_cg
 from .. import settings
 
 
@@ -27,6 +26,8 @@ class InvQuadLogDet(Function):
         log_det=False,
         preconditioner=None,
         log_det_correction=None,
+        probe_vectors=None,
+        probe_vector_norms=None,
     ):
         if not (inv_quad or log_det):
             raise RuntimeError("Either inv_quad or log_det must be true (or both)")
@@ -41,6 +42,19 @@ class InvQuadLogDet(Function):
         self.preconditioner = preconditioner
         self.log_det_correction = log_det_correction
 
+        if (probe_vectors is None or probe_vector_norms is None) and log_det:
+            num_random_probes = settings.num_trace_samples.value()
+            probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)
+            probe_vectors.bernoulli_().mul_(2).add_(-1)
+            probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
+            if batch_shape is not None:
+                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)
+            probe_vectors = probe_vectors.div(probe_vector_norms)
+
+        self.probe_vectors = probe_vectors
+        self.probe_vector_norms = probe_vector_norms
+
     def forward(self, *args):
         """
         *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)
@@ -61,7 +75,6 @@ class InvQuadLogDet(Function):
 
         # Get closure for matmul
         lazy_tsr = self.representation_tree(*matrix_args)
-        matmul_closure = lazy_tsr._matmul
 
         # Collect terms for LinearCG
         # We use LinearCG for both matrix solves and for stochastically estimating the log det
@@ -69,19 +82,10 @@ class InvQuadLogDet(Function):
         num_random_probes = 0
         num_inv_quad_solves = 0
 
-        # Probe vector for lanczos quadrature (log_det estimation)
-        probe_vectors = None
-        probe_vector_norms = None
+        # RHS for log_det
         if self.log_det:
-            num_random_probes = settings.num_trace_samples.value()
-            probe_vectors = torch.empty(self.matrix_shape[-1], num_random_probes, dtype=self.dtype, device=self.device)
-            probe_vectors.bernoulli_().mul_(2).add_(-1)
-            probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
-            if self.batch_shape is not None:
-                probe_vectors = probe_vectors.expand(*self.batch_shape, self.matrix_shape[-1], num_random_probes)
-                probe_vector_norms = probe_vector_norms.expand(*self.batch_shape, 1, num_random_probes)
-            probe_vectors = probe_vectors.div(probe_vector_norms)
-            rhs_list.append(probe_vectors)
+            rhs_list.append(self.probe_vectors)
+            num_random_probes = self.probe_vectors.size(-1)
 
         # RHS for inv_quad
         self.is_vector = False
@@ -96,23 +100,10 @@ class InvQuadLogDet(Function):
         rhs = torch.cat(rhs_list, -1)
         t_mat = None
         if self.log_det and settings.skip_logdet_forward.off():
-            solves, t_mat = linear_cg(
-                matmul_closure,
-                rhs,
-                n_tridiag=num_random_probes,
-                max_iter=settings.max_cg_iterations.value(),
-                max_tridiag_iter=settings.max_lanczos_quadrature_iterations.value(),
-                preconditioner=self.preconditioner,
-            )
+            solves, t_mat = lazy_tsr._solve(rhs, self.preconditioner, num_tridiag=num_random_probes)
 
         else:
-            solves = linear_cg(
-                matmul_closure,
-                rhs,
-                n_tridiag=0,
-                max_iter=settings.max_cg_iterations.value(),
-                preconditioner=self.preconditioner,
-            )
+            solves = lazy_tsr._solve(rhs, self.preconditioner, num_tridiag=0)
 
         # Final values to return
         log_det_term = torch.zeros(lazy_tsr.batch_shape, dtype=self.dtype, device=self.device)
@@ -141,7 +132,7 @@ class InvQuadLogDet(Function):
         self.num_random_probes = num_random_probes
         self.num_inv_quad_solves = num_inv_quad_solves
 
-        to_save = list(matrix_args) + [solves, probe_vectors, probe_vector_norms]
+        to_save = list(matrix_args) + [solves, ]
         self.save_for_backward(*to_save)
 
         if settings.memory_efficient.off():
@@ -158,10 +149,8 @@ class InvQuadLogDet(Function):
         compute_log_det_grad = log_det_grad_output.abs().sum() and self.log_det
 
         # Get input arguments, and get gradients in the proper form
-        matrix_args = self.saved_tensors[:-3]
-        solves = self.saved_tensors[-3]
-        probe_vectors = self.saved_tensors[-2]
-        probe_vector_norms = self.saved_tensors[-1]
+        matrix_args = self.saved_tensors[:-1]
+        solves = self.saved_tensors[-1]
 
         if hasattr(self, "_lazy_tsr"):
             lazy_tsr = self._lazy_tsr
@@ -180,10 +169,10 @@ class InvQuadLogDet(Function):
         inv_quad_solves = None
         neg_inv_quad_solves_times_grad_out = None
         if compute_log_det_grad:
-            coef = 1.0 / probe_vectors.size(-1)
+            coef = 1.0 / self.probe_vectors.size(-1)
             probe_vector_solves = solves.narrow(-1, 0, self.num_random_probes).mul(coef)
-            probe_vector_solves.mul_(probe_vector_norms).mul_(log_det_grad_output)
-            probe_vectors = probe_vectors.mul(probe_vector_norms)
+            probe_vector_solves.mul_(self.probe_vector_norms).mul_(log_det_grad_output)
+            probe_vectors = self.probe_vectors.mul(self.probe_vector_norms)
         if self.inv_quad:
             inv_quad_solves = solves.narrow(-1, self.num_random_probes, self.num_inv_quad_solves)
             neg_inv_quad_solves_times_grad_out = inv_quad_solves.mul(inv_quad_grad_output).mul_(-1)
diff --git a/gpytorch/lazy/__init__.py b/gpytorch/lazy/__init__.py
index c148f4b1..e4bddb07 100644
--- a/gpytorch/lazy/__init__.py
+++ b/gpytorch/lazy/__init__.py
@@ -5,6 +5,7 @@ from .added_diag_lazy_tensor import AddedDiagLazyTensor
 from .batch_repeat_lazy_tensor import BatchRepeatLazyTensor
 from .block_lazy_tensor import BlockLazyTensor
 from .block_diag_lazy_tensor import BlockDiagLazyTensor
+from .cached_cg_lazy_tensor import CachedCGLazyTensor
 from .chol_lazy_tensor import CholLazyTensor
 from .constant_mul_lazy_tensor import ConstantMulLazyTensor
 from .diag_lazy_tensor import DiagLazyTensor
@@ -31,6 +32,7 @@ __all__ = [
     "BatchRepeatLazyTensor",
     "BlockLazyTensor",
     "BlockDiagLazyTensor",
+    "CachedCGLazyTensor",
     "CholLazyTensor",
     "ConstantMulLazyTensor",
     "DiagLazyTensor",
diff --git a/gpytorch/lazy/cached_cg_lazy_tensor.py b/gpytorch/lazy/cached_cg_lazy_tensor.py
new file mode 100644
index 00000000..c9c805d6
--- /dev/null
+++ b/gpytorch/lazy/cached_cg_lazy_tensor.py
@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+
+import torch
+import warnings
+from .lazy_tensor import LazyTensor
+from .. import settings
+
+
+class CachedCGLazyTensor(LazyTensor):
+    """
+    A LazyTensor wrapper that eagerly computes many CG calls in batch.
+    This maximizes CG parallelism for fast inference.
+    Used primarily for variational inference with GPs.
+
+    Args:
+        :attr:`base_lazy_tensor` (:class:`gpytorch.lazy.LazyTensor`):
+            the LazyTensor to wrap
+        :attr:`eager_rhss` (list of :class:`gpytorch.lazy.LazyTensor`):
+            list of right-hand sides with eagerly-computed solves
+        :attr:`solves` (list of :class:`gpytorch.lazy.LazyTensor`):
+            list of solves associated with :attr:`eager_rhss`
+        :attr:`probe_vectors` (:class:`gpytorch.lazy.LazyTensor`, optional):
+            normalized probe vectors (for computing logdet with SLQ)
+        :attr:`probe_vector_norms` (:class:`gpytorch.lazy.LazyTensor`, optional):
+            norms associated with :attr:`probe_vectors` that will return :attr:`probe_vectors`
+            to having identity covariance (for computing logdet with SLQ)
+        :attr:`probe_vector_solves` (:class:`gpytorch.lazy.LazyTensor`, optional):
+            solves associated with :attr:`probe_vectors` (for computing logdet with SLQ)
+        :attr:`probe_vector_tmats` (:class:`gpytorch.lazy.LazyTensor`, optional):
+            Lanczos tridiagonal matrices associated with :attr:`probe_vectors`
+            (for computing logdet with SLQ)
+    """
+
+    @classmethod
+    def precompute_terms(cls, base_lazy_tensor, eager_rhs, logdet_terms=True, include_tmats=True):
+        """
+        Computes the solves, probe vectors, probe_vector norms, probe vector solves, and probe vector
+        tridiagonal matrices to construct a CachedCGLazyTensor
+
+        Set logdet_terms to False if you are not going to compute the logdet of the LazyTensor
+        """
+        with torch.no_grad():
+            if logdet_terms:
+                # Generate probe vectors
+                num_random_probes = settings.num_trace_samples.value()
+                probe_vectors = torch.empty(
+                    base_lazy_tensor.matrix_shape[-1], num_random_probes, dtype=base_lazy_tensor.dtype,
+                    device=base_lazy_tensor.device
+                )
+                probe_vectors.bernoulli_().mul_(2).add_(-1)
+                probe_vectors = probe_vectors.expand(
+                    *base_lazy_tensor.batch_shape, base_lazy_tensor.matrix_shape[-1], num_random_probes
+                )
+                probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
+                probe_vectors = probe_vectors.div(probe_vector_norms)
+
+                # Compute solves
+                if include_tmats:
+                    all_solves, probe_vector_tmats, = base_lazy_tensor._solve(
+                        torch.cat([probe_vectors, eager_rhs], -1),
+                        preconditioner=base_lazy_tensor._preconditioner()[0],
+                        num_tridiag=probe_vectors.size(-1)
+                    )
+                else:
+                    all_solves = base_lazy_tensor._solve(
+                        torch.cat([probe_vectors, eager_rhs], -1),
+                        preconditioner=base_lazy_tensor._preconditioner()[0],
+                    )
+                    probe_vector_tmats = torch.tensor([])
+                probe_vector_solves = all_solves[..., :probe_vectors.size(-1)].detach()
+                solves = all_solves[..., probe_vectors.size(-1):]
+
+                return solves.detach(), probe_vectors.detach(), probe_vector_norms.detach(), \
+                    probe_vector_solves.detach(), probe_vector_tmats.detach()
+
+            else:
+                # Compute solves
+                solves = base_lazy_tensor._solve(eager_rhs, preconditioner=base_lazy_tensor._preconditioner()[0])
+                return solves.detach(), torch.tensor([]), torch.tensor([]), torch.tensor([]), torch.tensor([])
+
+    def __init__(
+        self, base_lazy_tensor, eager_rhss=[], solves=[], probe_vectors=torch.tensor([]),
+        probe_vector_norms=torch.tensor([]), probe_vector_solves=torch.tensor([]), probe_vector_tmats=torch.tensor([])
+    ):
+        super(CachedCGLazyTensor, self).__init__(
+            base_lazy_tensor, eager_rhss=eager_rhss, solves=solves, probe_vectors=probe_vectors,
+            probe_vector_norms=probe_vector_norms, probe_vector_solves=probe_vector_solves,
+            probe_vector_tmats=probe_vector_tmats,
+        )
+        self.base_lazy_tensor = base_lazy_tensor
+        self.eager_rhss = [eager_rhs.requires_grad_(False) for eager_rhs in eager_rhss]
+        self.solves = [solve.requires_grad_(False) for solve in solves]
+        self.probe_vectors = probe_vectors.requires_grad_(False)
+        self.probe_vector_norms = probe_vector_norms.requires_grad_(False)
+        self.probe_vector_solves = probe_vector_solves.requires_grad_(False)
+        self.probe_vector_tmats = probe_vector_tmats.requires_grad_(False)
+
+    @property
+    def requires_grad(self):
+        return self.base_lazy_tensor.requires_grad
+
+    @requires_grad.setter
+    def requires_grad(self, val):
+        self.base_lazy_tensor.requires_grad = val
+
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        return self.base_lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
+
+    def _getitem(self, *indices):
+        return self.base_lazy_tensor._getitem(*indices)
+
+    def _matmul(self, tensor):
+        return self.base_lazy_tensor._matmul(tensor)
+
+    def _probe_vectors_and_norms(self):
+        return self.probe_vectors, self.probe_vector_norms
+
+    def _quad_form_derivative(self, left_vecs, right_vecs):
+        return self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
+
+    def _solve(self, rhs, preconditioner, num_tridiag=None):
+        if num_tridiag:
+            probe_vectors = rhs[..., :num_tridiag].detach()
+            if torch.equal(probe_vectors, self.probe_vectors):
+                probe_vector_solves = self.probe_vector_solves
+                tmats = self.probe_vector_tmats
+            else:
+                if settings.debug.on():
+                    warnings.warn(
+                        "CachedCGLazyTensor did not recognize the supplied probe vectors for tridiagonalization."
+                    )
+                return super(CachedCGLazyTensor, self)._solve(rhs, preconditioner, num_tridiag=num_tridiag)
+
+        # Here we check to see what solves we've already performed
+        truncated_rhs = rhs[..., (num_tridiag or 0):]
+        for eager_rhs, solve in zip(self.eager_rhss, self.solves):
+            if torch.equal(truncated_rhs, eager_rhs):
+                if num_tridiag:
+                    return torch.cat([probe_vector_solves, solve], -1), tmats
+                else:
+                    return solve
+
+        if settings.debug.on():
+            warnings.warn(
+                "CachedCGLazyTensor had to run CG on a tensor of size {}. For best performance, this "
+                "LazyTensor should pre-register all vectors to run CG against.".format(rhs.shape)
+            )
+        return super(CachedCGLazyTensor, self)._solve(rhs, preconditioner, num_tridiag=num_tridiag)
+
+    def _size(self):
+        return self.base_lazy_tensor._size()
+
+    def _t_matmul(self, tensor):
+        return self.base_lazy_tensor._t_matmul(tensor)
+
+    def _transpose_nonbatch(self):
+        return self.base_lazy_tensor._transpose_nonbatch()
+
+    def detach_(self):
+        self.base_lazy_tensor.detach_()
+        return self
diff --git a/gpytorch/lazy/chol_lazy_tensor.py b/gpytorch/lazy/chol_lazy_tensor.py
index 48e29165..d57e831f 100644
--- a/gpytorch/lazy/chol_lazy_tensor.py
+++ b/gpytorch/lazy/chol_lazy_tensor.py
@@ -12,10 +12,7 @@ class CholLazyTensor(RootLazyTensor):
             chol = chol.evaluate()
 
         # Check that we have a lower triangular matrix
-        mask = torch.full((chol.size(-2), chol.size(-2)), fill_value=-1, dtype=chol.dtype, device=chol.device)
-        mask.tril_().add_(1)
-        if chol.ndimension() == 3:
-            mask.unsqueeze_(0)
+        mask = torch.ones(chol.shape[-2:], dtype=chol.dtype, device=chol.device).triu_(1)
         if torch.max(chol.mul(mask)).item() > 1e-3 and torch.equal(chol, chol):
             raise RuntimeError("CholLazyVaraiable should take a lower-triangular " "matrix in the constructor.")
 
diff --git a/gpytorch/lazy/diag_lazy_tensor.py b/gpytorch/lazy/diag_lazy_tensor.py
index 84ea0d44..e3cb92dc 100644
--- a/gpytorch/lazy/diag_lazy_tensor.py
+++ b/gpytorch/lazy/diag_lazy_tensor.py
@@ -120,8 +120,11 @@ class DiagLazyTensor(LazyTensor):
     def inverse(self):
         return DiagLazyTensor(self._diag.reciprocal())
 
-    def inv_matmul(self, tensor):
-        return self.inverse()._matmul(tensor)
+    def inv_matmul(self, right_tensor, left_tensor=None):
+        res = self.inverse()._matmul(right_tensor)
+        if left_tensor is not None:
+            res = left_tensor @ res
+        return res
 
     def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
 
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 881a7cd5..1ec97796 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -10,6 +10,7 @@ from ..functions._inv_matmul import InvMatmul
 from ..functions._inv_quad_log_det import InvQuadLogDet
 from ..functions._matmul import Matmul
 from ..functions._root_decomposition import RootDecomposition
+from ..utils import linear_cg
 from ..utils.broadcasting import _matmul_broadcast_shape
 from ..utils.memoize import cached
 from ..utils.qr import batch_qr
@@ -91,6 +92,19 @@ class LazyTensor(object):
         """
         raise NotImplementedError("The class {} requires a _matmul function!".format(self.__class__.__name__))
 
+    def _probe_vectors_and_norms(self):
+        return None, None
+
+    def _solve(self, rhs, preconditioner, num_tridiag=None):
+        return linear_cg(
+            self._matmul,
+            rhs,
+            n_tridiag=num_tridiag,
+            max_iter=settings.max_cg_iterations.value(),
+            max_tridiag_iter=settings.max_lanczos_quadrature_iterations.value(),
+            preconditioner=preconditioner
+        )
+
     def _size(self):
         """
         Returns the size of the resulting Tensor that the lazy tensor represents.
@@ -617,15 +631,36 @@ class LazyTensor(object):
         """
         return self.representation_tree()(*self.representation())
 
-    def inv_matmul(self, tensor):
+    def inv_matmul(self, right_tensor, left_tensor=None):
         """
-        Computes a linear solve (w.r.t self = :math:`K`) with several right hand sides :math:`M`.
+        Computes a linear solve (w.r.t self = :math:`A`) with several right hand sides :math:`R`.
+        I.e. computes
+
+        ... math::
+
+            \begin{equation}
+                A^{-1} R,
+            \end{equation}
+
+        where :math:`R` is :attr:`right_tensor` and :math:`A` is the LazyTensor.
+
+        If :attr:`left_tensor` is supplied, computes
+
+        ... math::
+
+            \begin{equation}
+                L A^{-1} R,
+            \end{equation}
+
+        where :math:`L` is :attr:`left_tensor`. Supplying this can reduce the number of
+        CG calls required.
 
         Args:
-            - :obj:`torch.tensor` (n x k) - Matrix :math:`M` right hand sides
+            - :obj:`torch.tensor` (n x k) - Matrix :math:`R` right hand sides
+            - :obj:`torch.tensor` (m x n) - Optional matrix :math:`L` to perform left multiplication with
 
         Returns:
-            - :obj:`torch.tensor` - :math:`K^{-1}M`
+            - :obj:`torch.tensor` - :math:`A^{-1}R` or :math:`LA^{-1}R`.
         """
         if not self.is_square:
             raise RuntimeError(
@@ -633,18 +668,24 @@ class LazyTensor(object):
                 "Got a {} of size {}.".format(self.__class__.__name__, self.size())
             )
 
-        if self.dim() == 2 and tensor.dim() == 1:
-            if self.shape[-1] != tensor.numel():
+        if self.dim() == 2 and right_tensor.dim() == 1:
+            if self.shape[-1] != right_tensor.numel():
                 raise RuntimeError(
                     "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
-                        self.shape, tensor.shape
+                        self.shape, right_tensor.shape
                     )
                 )
 
-        func = InvMatmul(self.representation_tree(), preconditioner=self._inv_matmul_preconditioner())
-        return func(tensor, *self.representation())
+        func = InvMatmul(
+            self.representation_tree(), preconditioner=self._inv_matmul_preconditioner(),
+            has_left=(left_tensor is not None)
+        )
+        if left_tensor is None:
+            return func(right_tensor, *self.representation())
+        else:
+            return func(left_tensor, right_tensor, *self.representation())
 
-    def inv_quad(self, tensor):
+    def inv_quad(self, tensor, reduce_inv_quad=True):
         """
         Computes an inverse quadratic form (w.r.t self) with several right hand sides.
         I.e. computes tr( tensor^T self^{-1} tensor )
@@ -658,7 +699,7 @@ class LazyTensor(object):
         Returns:
             - tensor - tr( tensor^T (self)^{-1} tensor )
         """
-        res, _ = self.inv_quad_log_det(inv_quad_rhs=tensor, log_det=False)
+        res, _ = self.inv_quad_log_det(inv_quad_rhs=tensor, log_det=False, reduce_inv_quad=reduce_inv_quad)
         return res
 
     def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
@@ -704,6 +745,7 @@ class LazyTensor(object):
         if inv_quad_rhs is not None:
             args = [inv_quad_rhs] + list(args)
 
+        probe_vectors, probe_vector_norms = self._probe_vectors_and_norms()
         inv_quad_term, log_det_term = InvQuadLogDet(
             representation_tree=self.representation_tree(),
             matrix_shape=self.matrix_shape,
@@ -714,6 +756,8 @@ class LazyTensor(object):
             log_det=log_det,
             preconditioner=self._preconditioner()[0],
             log_det_correction=self._preconditioner()[1],
+            probe_vectors=probe_vectors,
+            probe_vector_norms=probe_vector_norms,
         )(*args)
 
         if inv_quad_term.numel() and reduce_inv_quad:
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index 22301d84..984eba06 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -119,7 +119,7 @@ class ZeroLazyTensor(LazyTensor):
     def evaluate(self):
         return torch.zeros(*self.sizes)
 
-    def inv_matmul(self, tensor):
+    def inv_matmul(self, right_tensor, left_tensor=None):
         raise RuntimeError("ZeroLazyTensors are not invertible!")
 
     def inv_quad(self, tensor):
diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py
index 985a4f31..2b060837 100644
--- a/gpytorch/variational/variational_strategy.py
+++ b/gpytorch/variational/variational_strategy.py
@@ -2,8 +2,8 @@
 
 import math
 import torch
-from .. import beta_features
-from ..lazy import RootLazyTensor, PsdSumLazyTensor, DiagLazyTensor
+from .. import beta_features, settings
+from ..lazy import DiagLazyTensor, CachedCGLazyTensor, PsdSumLazyTensor, RootLazyTensor
 from ..module import Module
 from ..distributions import MultivariateNormal
 
@@ -87,17 +87,47 @@ class VariationalStrategy(Module):
         if torch.equal(x, inducing_points):
             return variational_dist
         else:
-            n_induc = inducing_points.size(-2)
+            num_induc = inducing_points.size(-2)
             full_inputs = torch.cat([inducing_points, x], dim=-2)
             full_output = self.model.forward(full_inputs)
             full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
 
-            test_mean = full_mean[..., n_induc:]
-            induc_mean = full_mean[..., :n_induc]
-            induc_induc_covar = full_covar[..., :n_induc, :n_induc].add_jitter()
-            induc_data_covar = full_covar[..., :n_induc, n_induc:]
-            data_data_covar = full_covar[..., n_induc:, n_induc:]
+            # Mean terms
+            test_mean = full_mean[..., num_induc:]
+            induc_mean = full_mean[..., :num_induc]
             var_dist_mean = variational_dist.mean
+            neg_mean_diff = (induc_mean - var_dist_mean).unsqueeze(-1)
+
+            # Covariance terms
+            induc_induc_covar = full_covar[..., :num_induc, :num_induc].evaluate_kernel().add_jitter()
+            induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()
+            data_data_covar = full_covar[..., num_induc:, num_induc:]
+            root_variational_covar = variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate()
+
+            # Cache the CG results
+            left_tensors = torch.cat([neg_mean_diff, root_variational_covar], -1)
+            with torch.no_grad():
+                eager_rhs = torch.cat([left_tensors, induc_data_covar], -1)
+                solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                    induc_induc_covar, eager_rhs.detach(), logdet_terms=self.training,
+                    include_tmats=(not settings.skip_logdet_forward.on())
+                )
+                eager_rhss = [
+                    eager_rhs.detach(), eager_rhs[..., left_tensors.size(-1):].detach(),
+                    eager_rhs[..., :left_tensors.size(-1)].detach()
+                ]
+                solves = [
+                    solve.detach(), solve[..., left_tensors.size(-1):].detach(),
+                    solve[..., :left_tensors.size(-1)].detach()
+                ]
+                if settings.skip_logdet_forward.on():
+                    eager_rhss.append(torch.cat([probe_vecs, left_tensors], -1))
+                    solves.append(torch.cat([probe_vec_solves, solve[..., :left_tensors.size(-1)]], -1))
+            induc_induc_covar = CachedCGLazyTensor(
+                induc_induc_covar, eager_rhss=eager_rhss, solves=solves, probe_vectors=probe_vecs,
+                probe_vector_norms=probe_vec_norms, probe_vector_solves=probe_vec_solves,
+                probe_vector_tmats=tmats,
+            )
 
             # Cache the prior distribution, for faster training
             if self.training:
@@ -105,18 +135,16 @@ class VariationalStrategy(Module):
                 self._prior_distribution_memo = prior_dist
 
             # Compute predictive mean/covariance
-            induc_data_covar = induc_data_covar.evaluate()
-            inv_product = induc_induc_covar.inv_matmul(induc_data_covar)
-            factor = variational_dist.lazy_covariance_matrix.root_decomposition().root.matmul(inv_product)
-            predictive_mean = torch.add(
-                test_mean, inv_product.transpose(-1, -2).matmul((var_dist_mean - induc_mean).unsqueeze(-1)).squeeze(-1)
-            )
-            predictive_covar = RootLazyTensor(factor.transpose(-2, -1))
+            inv_products = induc_induc_covar.inv_matmul(induc_data_covar, left_tensors.transpose(-1, -2))
+            predictive_mean = torch.add(test_mean, inv_products[..., 0, :].mul(-1))
+            predictive_covar = RootLazyTensor(inv_products[..., 1:, :].transpose(-1, -2))
 
+            # Compute a diagonal correction for the covariance. It is the diagonal of the Schur complement
+            # K_{data,data} - K_{data,induc} K_{induc,induc}^{-1} K_{induc,data}
             if beta_features.diagonal_correction.on():
-                fake_diagonal = (inv_product * induc_data_covar).sum(-2)
-                real_diagonal = data_data_covar.diag()
-                diag_correction = DiagLazyTensor((real_diagonal - fake_diagonal).clamp(0, math.inf))
+                # interp_data_data_var = diag(K_{data,induc} K_{induc,induc}^{-1} K_{induc,data})
+                interp_data_data_var = induc_induc_covar.inv_quad(induc_data_covar, reduce_inv_quad=False)
+                diag_correction = DiagLazyTensor((data_data_covar.diag() - interp_data_data_var).clamp(0, math.inf))
                 predictive_covar = PsdSumLazyTensor(predictive_covar, diag_correction)
 
             return MultivariateNormal(predictive_mean, predictive_covar)


Test patch:
diff --git a/test/examples/test_batch_svgp_gp_regression.py b/test/examples/test_batch_svgp_gp_regression.py
index 47ce6b62..58395025 100644
--- a/test/examples/test_batch_svgp_gp_regression.py
+++ b/test/examples/test_batch_svgp_gp_regression.py
@@ -71,7 +71,7 @@ class TestSVGPRegression(unittest.TestCase):
         model.train()
         likelihood.train()
         optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        for _ in range(150):
+        for _ in range(180):
             optimizer.zero_grad()
             output = model(train_x)
             loss = -mll(output, train_y)
@@ -106,7 +106,7 @@ class TestSVGPRegression(unittest.TestCase):
         model.train()
         likelihood.train()
         optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        for _ in range(150):
+        for _ in range(200):
             optimizer.zero_grad()
             output = model(train_x)
             loss = -mll(output, train_y)
diff --git a/test/examples/test_svgp_gp_regression.py b/test/examples/test_svgp_gp_regression.py
index 68c41654..91122e05 100644
--- a/test/examples/test_svgp_gp_regression.py
+++ b/test/examples/test_svgp_gp_regression.py
@@ -55,7 +55,7 @@ class TestSVGPRegression(unittest.TestCase):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
-    def test_regression_error(self):
+    def test_regression_error(self, skip_logdet_forward=False):
         train_x, train_y = train_data()
         likelihood = GaussianLikelihood()
         model = SVGPRegressionModel(torch.linspace(0, 1, 25))
@@ -65,13 +65,13 @@ class TestSVGPRegression(unittest.TestCase):
         model.train()
         likelihood.train()
         optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        optimizer.n_iter = 0
-        for _ in range(150):
-            optimizer.zero_grad()
-            output = model(train_x)
-            loss = -mll(output, train_y)
-            loss.backward()
-            optimizer.step()
+        with gpytorch.settings.skip_logdet_forward(skip_logdet_forward):
+            for _ in range(170):
+                optimizer.zero_grad()
+                output = model(train_x)
+                loss = -mll(output, train_y)
+                loss.backward()
+                optimizer.step()
 
         for param in model.parameters():
             self.assertTrue(param.grad is not None)
@@ -87,6 +87,9 @@ class TestSVGPRegression(unittest.TestCase):
         mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
         self.assertLess(mean_abs_error.item(), 1e-1)
 
+    def test_regression_error_skip_logdet_forward(self):
+        return self.test_regression_error(skip_logdet_forward=True)
+
     def test_regression_error_cuda(self):
         if torch.cuda.is_available():
             train_x, train_y = train_data(cuda=True)
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index fc2eaaf9..5afbe145 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -282,10 +282,43 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
 
-        test_vector = torch.randn(lazy_tensor.size(-1))
+        test_vector = torch.randn(lazy_tensor.size(-1), requires_grad=True)
+        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
         with gpytorch.settings.max_cg_iterations(200):
             res = lazy_tensor.inv_matmul(test_vector)
-        actual = evaluated.inverse().matmul(test_vector)
+        actual = evaluated.inverse().matmul(test_vector_copy)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        res.backward(gradient=grad)
+        actual.backward(gradient=grad)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_matmul_vector_with_left(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+
+        # We skip this test if we're dealing with batch LazyTensors
+        # They shouldn't multiply by a vec
+        if lazy_tensor.ndimension() > 2:
+            return
+
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = torch.randn(lazy_tensor.size(-1), requires_grad=True)
+        test_left = torch.randn(6, lazy_tensor.size(-1), requires_grad=True)
+        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
+        test_left_copy = test_left.clone().detach().requires_grad_(True)
+        with gpytorch.settings.max_cg_iterations(100):
+            res = lazy_tensor.inv_matmul(test_vector, test_left)
+        actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
 
         grad = torch.randn_like(res)
@@ -296,16 +329,49 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
                 self.assertLess(
                     ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
                 )
+        self.assertLess(
+            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
 
     def test_inv_matmul_matrix(self):
         lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
 
-        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
+        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
         with gpytorch.settings.max_cg_iterations(100):
             res = lazy_tensor.inv_matmul(test_vector)
-        actual = evaluated.inverse().matmul(test_vector)
+        actual = evaluated.inverse().matmul(test_vector_copy)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        res.backward(gradient=grad)
+        actual.backward(gradient=grad)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_matmul_matrix_with_left(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
+        test_left = torch.randn(*lazy_tensor.batch_shape, 6, lazy_tensor.size(-1), requires_grad=True)
+        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
+        test_left_copy = test_left.clone().detach().requires_grad_(True)
+        with gpytorch.settings.max_cg_iterations(100):
+            res = lazy_tensor.inv_matmul(test_vector, test_left)
+        actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
 
         grad = torch.randn_like(res)
@@ -316,6 +382,12 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
                 self.assertLess(
                     ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
                 )
+        self.assertLess(
+            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
 
     def test_diag(self):
         lazy_tensor = self.create_lazy_tensor()
diff --git a/test/lazy/test_cached_cg_lazy_tensor.py b/test/lazy/test_cached_cg_lazy_tensor.py
new file mode 100644
index 00000000..55d4cdee
--- /dev/null
+++ b/test/lazy/test_cached_cg_lazy_tensor.py
@@ -0,0 +1,307 @@
+#!/usr/bin/env python3
+
+import math
+import torch
+import gpytorch
+import unittest
+import warnings
+from gpytorch.lazy import CachedCGLazyTensor, NonLazyTensor
+from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
+
+
+class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        mat = torch.randn(5, 6)
+        mat = mat.matmul(mat.transpose(-1, -2))
+        mat.requires_grad_(True)
+
+        lazy_tensor = NonLazyTensor(mat)
+        eager_rhs = torch.randn(5, 10).detach()
+        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_log_det tests
+            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                lazy_tensor, eager_rhs.detach(), logdet_terms=False
+            )
+            eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+            solves = [solve.detach(), solve[..., -2:-1].detach()]
+
+        return CachedCGLazyTensor(
+            lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+        )
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        return lazy_tensor.base_lazy_tensor.tensor
+
+    def test_inv_matmul_vec(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = lazy_tensor.eager_rhss[1].squeeze(-1).clone().detach().requires_grad_(True)
+        test_vector_copy = lazy_tensor_copy.eager_rhss[1].squeeze(-1).clone().detach().requires_grad_(True)
+        # Make sure that we get no warning about CG
+        with gpytorch.settings.max_cg_iterations(200), warnings.catch_warnings(record=True) as w:
+            res = lazy_tensor.inv_matmul(test_vector)
+            actual = evaluated.inverse().matmul(test_vector_copy)
+            self.assertEqual(len(w), 0)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        # Make sure that we get a warning that CG was run
+        with warnings.catch_warnings(record=True) as w:
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            self.assertEqual(len(w), 1)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_matmul_vector_with_left(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = lazy_tensor.eager_rhss[0][..., -1].squeeze(-1).clone().detach().requires_grad_(True)
+        test_vector_copy = lazy_tensor_copy.eager_rhss[0][..., -1].squeeze(-1).clone().detach().requires_grad_(True)
+        test_left = lazy_tensor.eager_rhss[0][..., :-1].t().clone().detach().requires_grad_(True)
+        test_left_copy = lazy_tensor_copy.eager_rhss[0][..., :-1].t().clone().detach().requires_grad_(True)
+        # Make sure that we get no warning about CG
+        with gpytorch.settings.max_cg_iterations(200), warnings.catch_warnings(record=True) as w:
+            res = lazy_tensor.inv_matmul(test_vector, test_left)
+            actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
+            self.assertEqual(len(w), 0)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        # Make sure that we get no warning about CG
+        with warnings.catch_warnings(record=True) as w:
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            self.assertEqual(len(w), 0)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_matmul_matrix(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+        test_vector_copy = lazy_tensor_copy.eager_rhss[0].clone().detach().requires_grad_(True)
+        # Make sure that we get no warning about CG
+        with gpytorch.settings.max_cg_iterations(100), warnings.catch_warnings(record=True) as w:
+            res = lazy_tensor.inv_matmul(test_vector)
+            actual = evaluated.inverse().matmul(test_vector_copy)
+            self.assertEqual(len(w), 0)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        # Make sure that we get a warning that CG was run
+        with warnings.catch_warnings(record=True) as w:
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            self.assertEqual(len(w), 1)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_matmul_matrix_with_left(self):
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+
+        test_vector = lazy_tensor.eager_rhss[0][..., 2:].clone().detach().requires_grad_(True)
+        test_vector_copy = lazy_tensor_copy.eager_rhss[0][..., 2:].clone().detach().requires_grad_(True)
+        test_left = lazy_tensor.eager_rhss[0][..., :2].transpose(-1, -2).clone().detach().requires_grad_(True)
+        test_left_copy = lazy_tensor_copy.eager_rhss[0][..., :2].transpose(-1, -2).clone().detach().requires_grad_(True)
+        # Make sure that we get no warning about CG
+        with gpytorch.settings.max_cg_iterations(100), warnings.catch_warnings(record=True) as w:
+            res = lazy_tensor.inv_matmul(test_vector, test_left)
+            actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
+            self.assertEqual(len(w), 0)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        # Make sure that we get no warning about CG
+        with warnings.catch_warnings(record=True) as w:
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            self.assertEqual(len(w), 0)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+        self.assertLess(
+            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+        self.assertLess(
+            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+        )
+
+    def test_inv_quad_log_det(self):
+        pass
+
+    def test_inv_quad_log_det_no_reduce(self):
+        pass
+
+    def test_root_inv_decomposition(self):
+        lazy_tensor = self.create_lazy_tensor()
+        root_approx = lazy_tensor.root_inv_decomposition()
+
+        test_mat = lazy_tensor.eager_rhss[0].clone().detach()
+
+        res = root_approx.matmul(test_mat)
+        actual = lazy_tensor.inv_matmul(test_mat)
+        self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
+
+
+class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        mat = torch.randn(5, 6)
+        mat = mat.matmul(mat.transpose(-1, -2))
+        mat.requires_grad_(True)
+
+        lazy_tensor = NonLazyTensor(mat)
+        eager_rhs = torch.randn(5, 10).detach()
+        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_log_det tests
+            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                lazy_tensor, eager_rhs.detach()
+            )
+            eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+            solves = [solve.detach(), solve[..., -2:-1].detach()]
+
+        return CachedCGLazyTensor(
+            lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+        )
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        return lazy_tensor.base_lazy_tensor.tensor
+
+    def test_inv_quad_log_det(self):
+        # Forward
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
+
+        vecs = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+        vecs_copy = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+
+        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as w:
+            res_inv_quad, res_log_det = lazy_tensor.inv_quad_log_det(inv_quad_rhs=vecs, log_det=True)
+            self.assertEqual(len(w), 0)
+        res = res_inv_quad + res_log_det
+
+        actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
+        actual_log_det = torch.cat(
+            [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
+        ).view(lazy_tensor.batch_shape)
+        actual = actual_inv_quad + actual_log_det
+
+        diff = (res - actual).abs() / actual.abs().clamp(1, math.inf)
+        self.assertLess(diff.max().item(), 15e-2)
+
+    def test_inv_quad_log_det_no_reduce(self):
+        # Forward
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
+
+        vecs = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+        vecs_copy = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+
+        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as w:
+            res_inv_quad, res_log_det = lazy_tensor.inv_quad_log_det(
+                inv_quad_rhs=vecs, log_det=True, reduce_inv_quad=False
+            )
+            self.assertEqual(len(w), 0)
+        res = res_inv_quad.sum(-1) + res_log_det
+
+        actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
+        actual_log_det = torch.cat(
+            [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
+        ).view(lazy_tensor.batch_shape)
+        actual = actual_inv_quad + actual_log_det
+
+        diff = (res - actual).abs() / actual.abs().clamp(1, math.inf)
+        self.assertLess(diff.max().item(), 15e-2)
+
+
+class TestCachedCGLazyTensorNoLogdetBatch(TestCachedCGLazyTensorNoLogdet):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        mat = torch.randn(3, 5, 6)
+        mat = mat.matmul(mat.transpose(-1, -2))
+        mat.requires_grad_(True)
+
+        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_log_det tests
+            lazy_tensor = NonLazyTensor(mat)
+            eager_rhs = torch.randn(3, 5, 10).detach()
+            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                lazy_tensor, eager_rhs.detach(), logdet_terms=False
+            )
+
+        return CachedCGLazyTensor(
+            lazy_tensor, [eager_rhs], [solve], probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+        )
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        return lazy_tensor.base_lazy_tensor.tensor
+
+    def test_inv_matmul_vec(self):
+        pass
+
+    def test_inv_matmul_vector_with_left(self):
+        pass
+
+
+class TestCachedCGLazyTensorBatch(TestCachedCGLazyTensor):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        mat = torch.randn(3, 5, 6)
+        mat = mat.matmul(mat.transpose(-1, -2))
+        mat.requires_grad_(True)
+
+        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_log_det tests
+            lazy_tensor = NonLazyTensor(mat)
+            eager_rhs = torch.randn(3, 5, 10).detach()
+            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                lazy_tensor, eager_rhs.detach()
+            )
+
+        return CachedCGLazyTensor(
+            lazy_tensor, [eager_rhs], [solve], probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+        )
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        return lazy_tensor.base_lazy_tensor.tensor
+
+    def test_inv_matmul_vec(self):
+        pass
+
+    def test_inv_matmul_vector_with_left(self):
+        pass