Repository: cornellius-gp/gpytorch

Issue:
[Feature Request] Allow a diagonalize operation analogous to symeig + root decomposition
# 🚀 Feature Request

Create a `lazy_tensor.diagonalize()` operation analogous to `torch.symeig` that caches the "eigenvectors" and "eigenvalues" similar to how `root_decomposition()` caches the root of the matrix. 

Diagonalize is slightly more general than `root_decomposition()` and `root_inv_decomposition()` as root decompositions can be straightforwardly created from diagonalized lazy tensors.

This will also enable setting different methods to perform root_inv_decompositions, rather than just Lanczos.

## Motivation

**Is your feature request related to a problem? Please describe.**

a) It's currently tricky to use approximate eigen-decompositions for Kronecker solves of the form `(K + \sigma^2 I)` while also ensuring that `K` has the correct root decomposition, e.g. for Lanczos. This pops up a bit in sampling routines for larger Kronecker structured models, e.g. [botorch PR #631](https://github.com/pytorch/botorch/pull/631).

b) The current Kronecker plus diagonal solve is only correct for `K + c I` and I'd like to implement `K + D` for general diagonal, but Kronecker structured diagonals.

d) Implementing Jordan normal form decompositions for general lazy tensors.

## Pitch

**Describe the solution you'd like**
```python
cached(name="diagonalization")
def diagonalization(self, method: Optional[str] = None): 
    if method == "lanczos":
          qmat, tmat = lanczos(self.matmul, .....)
          evals, eigenvectors = lanczos_tridiag_to_diag(tmat)
          return evals, qmat @ eigenvectors
    if method == "symeig":
         evals, eigenvectors = torch.symeig(self.evaluate(), eigenvectors = True)
         return evals, eigenvectors

def root_decomposition(self, method=None):
    if hasattr("diagonalization"):
        evals, eigenvectors = self.diagonalization(method=method)
        return eigenvectors.matmul(evals.clamp(0.0).sqrt())
```
    

**Describe alternatives you've considered**
It's possible to do all of these manually but the book keeping is pretty intense.

**Are you willing to open a pull request?** (We LOVE contributions!!!)

Yes, I can take a look at it and setup a base PR.

## Additional context

Will be quite useful for LinearOperator.



Discussion:

Resolving patch:
diff --git a/gpytorch/functions/_diagonalization.py b/gpytorch/functions/_diagonalization.py
new file mode 100644
--- /dev/null
+++ b/gpytorch/functions/_diagonalization.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+
+import torch
+from torch.autograd import Function
+
+from .. import settings
+from ..utils import lanczos
+
+
+class Diagonalization(Function):
+    @staticmethod
+    def forward(ctx, representation_tree, device, dtype, matrix_shape, max_iter, batch_shape, *matrix_args):
+        r"""
+        :param list matrix_args: The arguments representing the symmetric matrix A (or batch of PSD matrices A)
+
+        :rtype: (torch.Tensor, torch.Tensor)
+        :return: :attr:`Q`, :attr: `S` such that :math:`Q S Q^T \approx A`
+        """
+
+        ctx.representation_tree = representation_tree
+        ctx.device = device
+        ctx.dtype = dtype
+        ctx.matrix_shape = matrix_shape
+        ctx.max_iter = max_iter
+        ctx.batch_shape = batch_shape
+
+        # Get closure for matmul
+        lazy_tsr = ctx.representation_tree(*matrix_args)
+        matmul_closure = lazy_tsr._matmul
+        # Do lanczos
+        q_mat, t_mat = lanczos.lanczos_tridiag(
+            matmul_closure,
+            ctx.max_iter,
+            dtype=ctx.dtype,
+            device=ctx.device,
+            matrix_shape=ctx.matrix_shape,
+            batch_shape=ctx.batch_shape,
+        )
+
+        if ctx.batch_shape is None:
+            q_mat = q_mat.unsqueeze(-3)
+            t_mat = t_mat.unsqueeze(-3)
+        if t_mat.ndimension() == 3:  # If we only used one probe vector
+            q_mat = q_mat.unsqueeze(0)
+            t_mat = t_mat.unsqueeze(0)
+
+        mins = torch.diagonal(t_mat, dim1=-1, dim2=-2).min(dim=-1, keepdim=True)[0]
+        jitter_val = settings.tridiagonal_jitter.value()
+        jitter_mat = torch.diag_embed(jitter_val * mins).expand_as(t_mat)
+        eigenvalues, eigenvectors = lanczos.lanczos_tridiag_to_diag(t_mat + jitter_mat)
+
+        # Get orthogonal matrix and eigenvalues
+        q_mat = q_mat.matmul(eigenvectors)
+
+        if settings.memory_efficient.off():
+            ctx._lazy_tsr = lazy_tsr
+
+        if ctx.batch_shape is None:
+            q_mat = q_mat.squeeze(1)
+        q_mat = q_mat.squeeze(0)
+        eigenvalues = eigenvalues.squeeze(0)
+
+        to_save = list(matrix_args) + [q_mat, eigenvalues]
+        ctx.save_for_backward(*to_save)
+        return eigenvalues, q_mat
+
+    @staticmethod
+    def backward(ctx, evals_grad_output, evecs_grad_output):
+        # backwards pass uses explicit gradients from
+        # Matrix Backpropagation for Deep Networks with Structured Lazyers,
+        # Ionescu, et al CVPR, 2015. https://arxiv.org/pdf/1509.07838.pdf
+        # TODO: check matrix friendly backpropagation
+
+        q_mat = ctx.saved_tensors[-2]
+        eigenvalues = ctx.saved_tensors[-1]
+
+        # (\tilde K)_{ij} = 1_{i\neq j} (\sigma_i - \sigma_j)^{-1}
+        # add a small amount of jitter to ensure that no zeros are produced
+        kmat = (eigenvalues.unsqueeze(-1) - eigenvalues.unsqueeze(-2) + 1e-10).reciprocal()
+        torch.diagonal(kmat, dim1=-1, dim2=-2).zero_()
+
+        # dU = U(\tilde K^T \hadamard (U^T dL/dU)U^T
+        inner_term = kmat.transpose(-1, -2) * q_mat.transpose(-1, -2).matmul(evecs_grad_output)
+        term1 = q_mat.matmul(inner_term).matmul(q_mat.transpose(-1, -2))
+
+        # d\Sigma = U dL/d\Sigma U^T
+        term2 = q_mat.matmul(torch.diag_embed(evals_grad_output)).matmul(q_mat.transpose(-1, -2))
+
+        # finally sum the two
+        dL_dM = term1 + term2
+        output = tuple([None] * 6 + [dL_dM])
+
+        return output
diff --git a/gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py b/gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py
--- a/gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py
@@ -19,7 +19,7 @@ def _constant_kpadlt_constructor(lt, dlt):
     # eigendecomposition of K_i by its inverse to get an eigendecomposition of K_i D_i^{-1}.
     sub_evals, sub_evecs = [], []
     for lt_, dlt_ in zip(lt.lazy_tensors, dlt.lazy_tensors):
-        evals_, evecs_ = lt_.symeig(eigenvectors=True)
+        evals_, evecs_ = lt_.diagonalization()
         sub_evals.append(DiagLazyTensor(evals_ / dlt_.diag_values))
         sub_evecs.append(evecs_)
     evals = KroneckerProductDiagLazyTensor(*sub_evals)
@@ -36,7 +36,7 @@ def _symmetrize_kpadlt_constructor(lt, dlt):
     symm_prod = KroneckerProductLazyTensor(
         *[d.matmul(k).matmul(d) for k, d in zip(lt.lazy_tensors, dlt_inv_root.lazy_tensors)]
     )
-    evals, evecs = symm_prod.symeig(eigenvectors=True)
+    evals, evecs = symm_prod.diagonalization()
     evals_plus_i = DiagLazyTensor(evals + 1.0)
 
     return dlt_inv_root, evals_plus_i, evecs
@@ -123,7 +123,7 @@ def _solve(self, rhs, preconditioner=None, num_tridiag=0):
         if self._diag_is_constant:
             # we perform the solve in double for numerical stability issues
             # TODO: Use fp64 registry once #1213 is addressed
-            evals, q_matrix = self.lazy_tensor.to(torch.double).symeig(eigenvectors=True)
+            evals, q_matrix = self.lazy_tensor.to(torch.double).diagonalization()
             evals_plus_diagonal = evals + self.diag_tensor.diag().double()
             evals_root = evals_plus_diagonal.pow(0.5)
             inv_mat_sqrt = DiagLazyTensor(evals_root.reciprocal())
@@ -198,7 +198,7 @@ def _solve(self, rhs, preconditioner=None, num_tridiag=0):
 
     def _root_decomposition(self):
         if self._diag_is_constant:
-            evals, q_matrix = self.lazy_tensor.symeig(eigenvectors=True)
+            evals, q_matrix = self.lazy_tensor.diagonalization()
             updated_evals = DiagLazyTensor((evals + self.diag_tensor.diag()).pow(0.5))
             return MatmulLazyTensor(q_matrix, updated_evals)
 
@@ -229,7 +229,7 @@ def _root_decomposition(self):
 
     def _root_inv_decomposition(self, initial_vectors=None):
         if self._diag_is_constant:
-            evals, q_matrix = self.lazy_tensor.symeig(eigenvectors=True)
+            evals, q_matrix = self.lazy_tensor.diagonalization()
             inv_sqrt_evals = DiagLazyTensor((evals + self.diag_tensor.diag()).pow(-0.5))
             return MatmulLazyTensor(q_matrix, inv_sqrt_evals)
 
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -11,6 +11,7 @@
 import gpytorch
 
 from .. import settings, utils
+from ..functions._diagonalization import Diagonalization
 from ..functions._inv_matmul import InvMatmul
 from ..functions._inv_quad import InvQuad
 from ..functions._inv_quad_log_det import InvQuadLogDet
@@ -22,7 +23,7 @@
 from ..utils.deprecation import _deprecate_renamed_methods
 from ..utils.errors import CachingError
 from ..utils.getitem import _compute_getitem_size, _convert_indices_to_tensors, _is_noop_index, _noop_index
-from ..utils.memoize import add_to_cache, cached, pop_from_cache
+from ..utils.memoize import add_to_cache, cached, get_from_cache, pop_from_cache
 from ..utils.pivoted_cholesky import pivoted_cholesky
 from ..utils.warnings import NumericalWarning
 from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
@@ -1329,6 +1330,45 @@ def requires_grad_(self, val):
         self.requires_grad = val
         return self
 
+    @cached(name="diagonalization")
+    def diagonalization(self, method: Optional[str] = None):
+        """
+        Returns a (usually partial) diagonalization of a symmetric PSD matrix.
+        Options are either "lanczos" or "symeig". "lanczos" runs Lanczos while
+        "symeig" runs LazyTensor.symeig.
+        """
+        if not self.is_square:
+            raise RuntimeError(
+                "diagonalization only operates on (batches of) square (symmetric) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if method is None:
+            if self.size(-1) <= settings.max_cholesky_size.value():
+                method = "symeig"
+            else:
+                method = "lanczos"
+
+        if method == "lanczos":
+            from ..lazy import lazify
+
+            func = Diagonalization()
+            evals, evecs = func.apply(
+                self.representation_tree(),
+                self.device,
+                self.dtype,
+                self.matrix_shape,
+                self._root_decomposition_size(),
+                self.batch_shape,
+                *self.representation(),
+            )
+            evecs = lazify(evecs)
+
+        elif method == "symeig":
+            evals, evecs = self.symeig(eigenvectors=True)
+
+        return evals, evecs
+
     @cached(name="root_decomposition")
     def root_decomposition(self, method: Optional[str] = None):
         """
@@ -1381,6 +1421,15 @@ def root_decomposition(self, method: Optional[str] = None):
             return RootLazyTensor(F)
 
         if method == "lanczos":
+            # check to see if we have already run lanczos for a diagonalization
+            try:
+                evals, evecs = get_from_cache(self, "diagonalization")
+                F = evecs * evals.clamp(1e-7).sqrt().unsqueeze(-2)
+                return RootLazyTensor(F)
+            except CachingError:
+                pass
+
+            # if not run standard lanczos
             return RootLazyTensor(self._root_decomposition())
 
         raise RuntimeError(f"Unknown method '{method}'")
@@ -1422,6 +1471,15 @@ def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
                 "Got a {} of size {}.".format(self.__class__.__name__, self.size())
             )
 
+        # check to see if we have already run lanczos for a diagonalization
+        try:
+            evals, evecs = get_from_cache(self, "diagonalization")
+            F = evecs * evals.clamp(1e-7).sqrt().reciprocal().unsqueeze(-2)
+            return RootLazyTensor(F)
+        except CachingError:
+            pass
+
+        # if not run standard lanczos
         if initial_vectors is not None:
             if self.dim() == 2 and initial_vectors.dim() == 1:
                 if self.shape[-1] != initial_vectors.numel():


Test patch:
diff --git a/gpytorch/test/lazy_tensor_test_case.py b/gpytorch/test/lazy_tensor_test_case.py
--- a/gpytorch/test/lazy_tensor_test_case.py
+++ b/gpytorch/test/lazy_tensor_test_case.py
@@ -278,6 +278,7 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
     skip_slq_tests = False
     should_call_cg = True
     should_call_lanczos = True
+    should_call_lanczos_diagonalization = True
     tolerances = {
         **RectangularLazyTensorTestCase.tolerances,
         "cholesky": {"rtol": 1e-3, "atol": 1e-5},
@@ -520,6 +521,28 @@ def test_root_decomposition(self, cholesky=False):
             else:
                 self.assertFalse(lanczos_mock.called)
 
+    def test_diagonalization(self, symeig=False):
+        _wrapped_lanczos = MagicMock(wraps=gpytorch.utils.lanczos.lanczos_tridiag)
+        with patch("gpytorch.utils.lanczos.lanczos_tridiag", new=_wrapped_lanczos) as lanczos_mock:
+            lazy_tensor = self.create_lazy_tensor()
+            test_mat = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+            with gpytorch.settings.max_cholesky_size(math.inf if symeig else 0):
+                evals, evecs = lazy_tensor.diagonalization()
+                evecs = evecs.evaluate()
+                approx = evecs.matmul(torch.diag_embed(evals)).matmul(evecs.transpose(-2, -1))
+                res = approx.matmul(test_mat)
+                actual = lazy_tensor.matmul(test_mat)
+                self.assertAllClose(res, actual, rtol=0.05)
+
+            # Make sure that we're calling the correct function
+            if not symeig and self.__class__.should_call_lanczos_diagonalization:
+                self.assertTrue(lanczos_mock.called)
+            else:
+                self.assertFalse(lanczos_mock.called)
+
+    def test_diagonalization_symeig(self):
+        return self.test_diagonalization(symeig=True)
+
     def test_root_decomposition_cholesky(self):
         return self.test_root_decomposition(cholesky=True)
 
diff --git a/test/functions/test_diagonalization.py b/test/functions/test_diagonalization.py
new file mode 100644
--- /dev/null
+++ b/test/functions/test_diagonalization.py
@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+
+from gpytorch.lazy import NonLazyTensor
+from gpytorch.test.base_test_case import BaseTestCase
+
+
+class TestDiagonalization(BaseTestCase, unittest.TestCase):
+    seed = 0
+
+    def _create_mat(self):
+        mat = torch.randn(4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4))
+        return mat
+
+    def test_diagonalization(self):
+        mat = self._create_mat().detach().requires_grad_(True)
+        mat_clone = mat.detach().clone().requires_grad_(True)
+
+        for method in ["symeig", "lanczos"]:
+            # Forward
+            evals, evecs = NonLazyTensor(mat).diagonalization(method=method)
+            evecs = evecs.evaluate()
+            res = evecs.matmul(torch.diag_embed(evals)).matmul(evecs.transpose(-1, -2))
+            self.assertAllClose(res, mat)
+
+            # Backward
+            sum([mat.trace() for mat in res.view(-1, mat.size(-2), mat.size(-1))]).backward()
+            sum([mat.trace() for mat in mat_clone.view(-1, mat.size(-2), mat.size(-1))]).backward()
+            self.assertAllClose(mat.grad, mat_clone.grad)
+
+
+class TestDiagonalizationBatch(TestDiagonalization):
+    seed = 0
+
+    def _create_mat(self):
+        mat = torch.randn(3, 4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        return mat
+
+
+class TestDiagonalizationMultiBatch(TestDiagonalization):
+    seed = 0
+
+    def _create_mat(self):
+        mat = torch.randn(2, 3, 4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        return mat
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/lazy/test_chol_lazy_tensor.py b/test/lazy/test_chol_lazy_tensor.py
--- a/test/lazy/test_chol_lazy_tensor.py
+++ b/test/lazy/test_chol_lazy_tensor.py
@@ -13,6 +13,7 @@ class TestCholLazyTensor(LazyTensorTestCase, unittest.TestCase):
     should_test_sample = True
     should_call_cg = False
     should_call_lanczos = False
+    should_call_lanczos_diagonalization = True
 
     def create_lazy_tensor(self):
         chol = torch.tensor(
diff --git a/test/lazy/test_kronecker_product_added_diag_lazy_tensor.py b/test/lazy/test_kronecker_product_added_diag_lazy_tensor.py
--- a/test/lazy/test_kronecker_product_added_diag_lazy_tensor.py
+++ b/test/lazy/test_kronecker_product_added_diag_lazy_tensor.py
@@ -50,7 +50,6 @@ class TestKroneckerProductAddedKroneckerDiagLazyTensor(TestKroneckerProductAdded
     # this lazy tensor has an explicit inverse so we don't need to run these
     skip_slq_tests = True
     should_call_cg = False
-    should_call_lanczos = False
 
     def create_lazy_tensor(self):
         a = torch.tensor([[4, 0, 2], [0, 3, -1], [2, -1, 3]], dtype=torch.float)
@@ -96,7 +95,6 @@ def create_lazy_tensor(self):
 
 
 class TestKroneckerProductAddedConstDiagLazyTensor(TestKroneckerProductAddedDiagLazyTensor):
-    should_call_lanczos = False
     should_call_cg = False
 
     def create_lazy_tensor(self):
diff --git a/test/lazy/test_root_lazy_tensor.py b/test/lazy/test_root_lazy_tensor.py
--- a/test/lazy/test_root_lazy_tensor.py
+++ b/test/lazy/test_root_lazy_tensor.py
@@ -12,6 +12,7 @@ class TestRootLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
     should_call_lanczos = False
+    should_call_lanczos_diagonalization = True
 
     def create_lazy_tensor(self):
         root = torch.randn(3, 5, requires_grad=True)