Repository: cornellius-gp/gpytorch
Pull request #519 (2019-02-19 00:32:31)

Improve speed, stability of woodbury for preconditioner
This PR does several things:

- Use lower-level BLAS calls for the woodbury formula. This improves the speed slightly
- Use lower-level BLAS calls for AddedDiagLazyTensor._matmul (the matmul function we use the most out of any LazyTensor)
- Move all woodbury stuff into its own file (we'll probably be reusing it for some future preconditioners and stuff)

The big thing that this PR does:

- Scale the diagonal component of the woodbury formula by its' smallest number. When the diagonal component of woodbury is very small, it's possible that dividing by the diagonal could be numerically unstable. This PR should improve stability by scaling before applying the formula and undoing the scale after.


Patch:
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index bc1b3d93..05d39d1e 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -2,12 +2,10 @@
 
 import torch
 import warnings
-from .non_lazy_tensor import lazify
 from .sum_lazy_tensor import SumLazyTensor
 from .diag_lazy_tensor import DiagLazyTensor
-from ..utils import pivoted_cholesky
+from ..utils import broadcasting, pivoted_cholesky, woodbury
 from .. import settings
-from ..utils import broadcasting
 
 
 class AddedDiagLazyTensor(SumLazyTensor):
@@ -35,6 +33,13 @@ class AddedDiagLazyTensor(SumLazyTensor):
         else:
             raise RuntimeError("One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!")
 
+    def _matmul(self, rhs):
+        return torch.addcmul(
+            self._lazy_tensor._matmul(rhs),
+            self._diag_tensor._diag.unsqueeze(-1),
+            rhs
+        )
+
     def add_diag(self, added_diag):
         return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))
 
@@ -58,26 +63,18 @@ class AddedDiagLazyTensor(SumLazyTensor):
                     "NaNs encountered in preconditioner computation. Attempting to continue without " "preconditioning."
                 )
                 return None, None
-            self._woodbury_cache = pivoted_cholesky.woodbury_factor(self._piv_chol_self, self._diag_tensor.diag())
+            self._woodbury_cache, self._inv_scale, self._precond_logdet_cache = woodbury.woodbury_factor(
+                self._piv_chol_self, self._piv_chol_self, self._diag_tensor.diag(), logdet=True
+            )
+            self._scaled_inv_diag = self._inv_scale / self._diag_tensor.diag().unsqueeze(-1)
+            self._scaled_inv_diag_piv_chol_self = self._piv_chol_self * self._scaled_inv_diag
 
         # preconditioner
         def precondition_closure(tensor):
-            return pivoted_cholesky.woodbury_solve(
-                tensor, self._piv_chol_self, self._woodbury_cache, self._diag_tensor.diag()
-            )
-
-        # logdet correction
-        if not hasattr(self, "_precond_logdet_cache"):
-            lr_flipped = self._piv_chol_self.matmul(
-                self._piv_chol_self.transpose(-2, -1).div(self._diag_tensor.diag().unsqueeze(-1))
+            res = woodbury.woodbury_solve(
+                tensor, self._scaled_inv_diag_piv_chol_self, self._woodbury_cache,
+                self._scaled_inv_diag, self._inv_scale
             )
-            lr_flipped = lr_flipped + torch.eye(n=lr_flipped.size(-2), dtype=lr_flipped.dtype, device=lr_flipped.device)
-            if lr_flipped.ndimension() == 3:
-                ld_one = (lazify(torch.cholesky(lr_flipped, upper=True)).diag().log().sum(-1)) * 2
-                ld_two = self._diag_tensor.diag().log().sum(-1)
-            else:
-                ld_one = lr_flipped.cholesky(upper=True).diag().log().sum() * 2
-                ld_two = self._diag_tensor.diag().log().sum().item()
-            self._precond_logdet_cache = ld_one + ld_two
+            return res
 
         return precondition_closure, self._precond_logdet_cache
diff --git a/gpytorch/utils/pivoted_cholesky.py b/gpytorch/utils/pivoted_cholesky.py
index 9f2b6205..bdcc4b7e 100644
--- a/gpytorch/utils/pivoted_cholesky.py
+++ b/gpytorch/utils/pivoted_cholesky.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import torch
-from .cholesky import cholesky_solve
 
 
 def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
@@ -78,44 +77,4 @@ def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
             errors = torch.norm(matrix_diag.gather(-1, pi_i), 1, dim=-1) / orig_error
         m = m + 1
 
-    return L[..., :m, :].contiguous()
-
-
-def woodbury_factor(low_rank_mat, shift):
-    r"""
-    Given a low rank (k x n) matrix V and a shift, returns the
-    matrix R so that
-
-    .. math::
-
-        \begin{equation*}
-            R = (I_k + 1/shift VV')^{-1}V
-        \end{equation*}
-
-    to be used in solves with (V'V + shift I) via the Woodbury formula
-    """
-    k = low_rank_mat.size(-2)
-    shifted_mat = low_rank_mat.matmul(low_rank_mat.transpose(-1, -2) / shift.unsqueeze(-1))
-
-    shifted_mat = shifted_mat + torch.eye(k, dtype=shifted_mat.dtype, device=shifted_mat.device)
-
-    R = cholesky_solve(low_rank_mat, torch.cholesky(shifted_mat))
-    return R
-
-
-def woodbury_solve(vector, low_rank_mat, woodbury_factor, shift):
-    """
-    Solves the system of equations: :math:`(sigma*I + VV')x = b`
-    Using the Woodbury formula.
-
-    Input:
-        - vector (size n) - right hand side vector b to solve with.
-        - woodbury_factor (k x n) - The result of calling woodbury_factor on V
-          and the shift, \sigma
-        - shift (vector) - shift value sigma
-    """
-    if vector.ndimension() > 1:
-        shift = shift.unsqueeze(-1)
-
-    right = low_rank_mat.transpose(-1, -2).matmul(woodbury_factor.matmul(vector / shift))
-    return (vector - right) / shift
+    return L[..., :m, :].transpose(-1, -2).contiguous()
diff --git a/gpytorch/utils/woodbury.py b/gpytorch/utils/woodbury.py
new file mode 100644
index 00000000..0afc69ef
--- /dev/null
+++ b/gpytorch/utils/woodbury.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+
+import torch
+from .cholesky import cholesky_solve
+from .. import settings
+
+
+def woodbury_factor(umat, vmat, diag, logdet=False):
+    r"""
+    Given a matrix factorized as :math:`(D + U V^T)`, where
+    :math:`U`, :math:`V` are (n x k) and :math:`D` is a (n x n) diagonal matrix,
+    returns the matrix :math:`R` so that
+
+    .. math::
+        \begin{equation*}
+            R = (I_k + V^T D^-1 U)^{-1} V^T
+        \end{equation*}
+
+    to be used in solves with (D + U V^T) via the Woodbury formula.
+    Can also be used in batch mode, where U, V, and D are batches of matrices
+
+    If the :attr:`logdet=True` flag is set, this function will also return the
+    log determinant of :math:`(D + U V^T)`.
+
+    Args:
+        :attr:`umat` (Tensor n x k):
+            The left matrix factor
+        :attr:`vmat` (Tensor n x k):
+            The right matrix factor
+        :attr:`diag` (Tensor n):
+            The diagonal of D
+        :attr:`logdet` (bool):
+            Whether or not to return the log determinant
+
+    Returns:
+        Tensor (k x n) (if `logdet=False`)
+        Tensor (k x n), Tensor () (if `logdet=True`)
+    """
+    if settings.debug.on():
+        if umat.shape != vmat.shape:
+            raise ValueError("umat ({}) and vmat ({}) must have the same shape.".format(umat.shape, vmat.shape))
+        if umat.shape[:-1] != diag.shape:
+            raise ValueError("Incompatible shape for diag ({}) given umat shape ({}).".format(diag.shape, umat.shape))
+
+    # Sizes
+    *batch_shape, n, k = umat.shape
+
+    # These reshapes make it easier to use faster blas calls
+    umat = umat.view(-1, n, k)
+    vmat = vmat.view(-1, n, k)
+    diag = diag.view(-1, n, 1)
+
+    # Scale the diagonal
+    # s = scale = max |1 / diag|
+    inv_scale = diag.abs().min()
+    scaled_inv_diag = inv_scale / diag
+
+    # Compute (1/s (I_k + V^T D^-1 U)), where s is a scale factor
+    inner_mat = torch.baddbmm(
+        inv_scale,
+        torch.eye(k, dtype=scaled_inv_diag.dtype, device=scaled_inv_diag.device),
+        1,
+        vmat.transpose(-1, -2),
+        umat * scaled_inv_diag
+    )
+
+    # Compute the cholesky factor of (1/s (I_k + V^T D^-1 U))
+    chol = torch.cholesky(inner_mat)
+
+    # Compute s (I_k + V^T D^-1 U)^-1 V^T
+    R = cholesky_solve(vmat.transpose(-1, -2), chol).view(*batch_shape, k, n)
+
+    # Maybe compute the log determinant
+    if logdet:
+        # Using the matrix determinant lemma here
+        logdet = chol.diagonal(dim1=-1, dim2=-2).log().sum(-1).mul(2) - scaled_inv_diag.log().sum([-1, -2])
+        # Undo the effect of scaling on D^{-1}
+        logdet = logdet + inv_scale.log().mul(n - k)
+        # Reshape
+        logdet = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
+        return R, inv_scale, logdet
+    else:
+        return R, inv_scale
+
+
+def woodbury_solve(rhs, scaled_inv_diag_umat, woodbury_factor, scaled_inv_diag, inv_scale):
+    """
+    Solves the system of equations: :math:`(D + U V^T)x = b` using the Woodbury formula,
+    where x is the right-hand-side (size n), U, V are (n x k), and D is a (n x n) diagonal matrix.
+    Can also be used in batch mode, where U, V, and D are batches of matrices and rhs is a batch of right-hand-side.
+
+    This should be used after calling woodbury_factor.
+
+    Args:
+        :attr:`rhs` (size n x t)
+            Right hand side vector b to solve with.
+        :attr:`scaled_inv_diag_umat` (n x k)
+            The `D^{-1} U` matrix
+        :attr:`woodbury_factor` (n x k)
+            The result of calling woodbury_factor on U, V, and D.
+        :attr:`diag` (vector)
+            The diagonal of D
+    """
+    # (D + UV^T)x = D^-1 x - D^-1 U ((I + V^T D^-1 U)^-1 V^T) D^-1 x
+    #             = D^-1 x - D^-1 U (1/s (woodbury_factor)) D^-1 x
+    #             = s( E^-1 x - E^-1 U (woodbury_factor) E^-1 x )
+    scaled_inv_diag_rhs = rhs * scaled_inv_diag
+    res = torch.add(scaled_inv_diag_rhs, -1, scaled_inv_diag_umat @ (woodbury_factor @ scaled_inv_diag_rhs))
+    res = res.div_(inv_scale)
+
+    # Reshape the result to be the correct shape
+    return res


Test patch:
diff --git a/test/utils/test_pivoted_cholesky.py b/test/utils/test_pivoted_cholesky.py
index b07f9ab8..d2a69b12 100644
--- a/test/utils/test_pivoted_cholesky.py
+++ b/test/utils/test_pivoted_cholesky.py
@@ -3,11 +3,11 @@
 import os
 import random
 import unittest
+from gpytorch.utils import pivoted_cholesky, woodbury
 from test._utils import approx_equal
 
 import torch
 from gpytorch.kernels import RBFKernel
-from gpytorch.utils import pivoted_cholesky
 
 
 class TestPivotedCholesky(unittest.TestCase):
@@ -28,34 +28,24 @@ class TestPivotedCholesky(unittest.TestCase):
         train_x = torch.linspace(0, 1, size)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        covar_approx = piv_chol.t().matmul(piv_chol)
+        covar_approx = piv_chol @ piv_chol.transpose(-1, -2)
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
-    def test_solve_vector(self):
-        size = 100
-        train_x = torch.linspace(0, 1, size)
-        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
-        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        woodbury_factor = pivoted_cholesky.woodbury_factor(piv_chol, torch.ones(100))
-
-        rhs_vector = torch.randn(100)
-        shifted_covar_matrix = covar_matrix + torch.eye(size)
-        real_solve = shifted_covar_matrix.inverse().matmul(rhs_vector)
-        approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(100))
-
-        self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
-
     def test_solve(self):
         size = 100
         train_x = torch.linspace(0, 1, size)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        woodbury_factor = pivoted_cholesky.woodbury_factor(piv_chol, torch.ones(100))
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(piv_chol, piv_chol, torch.ones(100), logdet=True)
+        self.assertTrue(approx_equal(logdet, (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).logdet(), 2e-4))
 
         rhs_vector = torch.randn(100, 50)
         shifted_covar_matrix = covar_matrix + torch.eye(size)
         real_solve = shifted_covar_matrix.inverse().matmul(rhs_vector)
-        approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(100))
+        scaled_inv_diag = (inv_scale / torch.ones(100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
 
         self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
 
@@ -80,7 +70,7 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
         ).unsqueeze(-1)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        covar_approx = piv_chol.transpose(1, 2).matmul(piv_chol)
+        covar_approx = piv_chol @ piv_chol.transpose(-1, -2)
 
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
@@ -91,7 +81,13 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
         ).unsqueeze(-1)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        woodbury_factor = pivoted_cholesky.woodbury_factor(piv_chol, torch.ones(2, 100))
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(
+            piv_chol, piv_chol, torch.ones(2, 100), logdet=True
+        )
+        actual_logdet = torch.stack([
+            mat.logdet() for mat in (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).view(-1, 100, 100)
+        ], 0).view(2)
+        self.assertTrue(approx_equal(logdet, actual_logdet, 2e-4))
 
         rhs_vector = torch.randn(2, 100, 5)
         shifted_covar_matrix = covar_matrix + torch.eye(size)
@@ -102,7 +98,10 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
             ],
             0,
         )
-        approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(2, 100))
+        scaled_inv_diag = (inv_scale / torch.ones(2, 100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
 
         self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
 
@@ -141,7 +140,7 @@ class TestPivotedCholeskyMultiBatch(unittest.TestCase):
         ).unsqueeze(-1)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(2, 2, 3, size, size)
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        covar_approx = piv_chol.transpose(-1, -2).matmul(piv_chol)
+        covar_approx = piv_chol @ piv_chol.transpose(-1, -2)
 
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
@@ -166,7 +165,13 @@ class TestPivotedCholeskyMultiBatch(unittest.TestCase):
         ).unsqueeze(-1)
         covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(2, 2, 3, size, size)
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
-        woodbury_factor = pivoted_cholesky.woodbury_factor(piv_chol, torch.ones(2, 2, 3, 100))
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(
+            piv_chol, piv_chol, torch.ones(2, 2, 3, 100), logdet=True
+        )
+        actual_logdet = torch.stack([
+            mat.logdet() for mat in (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).view(-1, 100, 100)
+        ], 0).view(2, 2, 3)
+        self.assertTrue(approx_equal(logdet, actual_logdet, 2e-4))
 
         rhs_vector = torch.randn(2, 2, 3, 100, 5)
         shifted_covar_matrix = covar_matrix + torch.eye(size)
@@ -187,7 +192,10 @@ class TestPivotedCholeskyMultiBatch(unittest.TestCase):
             ],
             0,
         ).view_as(rhs_vector)
-        approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(2, 3, 100))
+        scaled_inv_diag = (inv_scale / torch.ones(2, 3, 100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
 
         self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))