Repository: cornellius-gp/gpytorch
Pull request #329 (2018-10-18 16:37:21)

[WIP] Default LazyTensor._quad_form_derivative implementation
Alright, this is going to be a fun one. 

This PR implements a default implementation of `LazyTensor._quad_form_derivative`. After playing around with it, this will have two huge advantages:

1. **Simplicity.** By far the most difficult method to implement when implementing a LazyTensor is the derivative portion. See for example the giant mess that is things like https://github.com/cornellius-gp/gpytorch/blob/9ebbf05e8ea70132dbcce25671bb06fa1721eb81/gpytorch/lazy/kronecker_product_lazy_tensor.py#L97

With this approach, implementing this from scratch becomes optional. Thus, we truly only need the end user to define a matmul if they want to implement a new LazyTensor.

2. **Speed.** I've been playing around with this for a bit, and initial results suggest that relying on `torch.autograd.grad` is crazy fast. We get a **2x** speed up in some of the initial things I've been testing. For example:

```python
c1 = torch.randn(100)
c2 = torch.randn(100)
# "Ensure" positive definite
c1[0] = 10
c2[0] = 10

tlv1 = ToeplitzLazyTensor(c1)
tlv2 = ToeplitzLazyTensor(c2)
klv = KroneckerProductLazyTensor(tlv1, tlv2)

p = torch.randn(10000, 2)
q = torch.randn(10000, 2)

# Wall time: 64 ms
%time klv._quad_form_derivative(p, q)

# Wall time: 31 ms
%time super(KroneckerProductLazyTensor, klv)._quad_form_derivative(p, q)
```

### Things to do
- Ensure compatibility with all of the various batch modes.
- Run more timing results to see whether we should make this default call the standard behavior. If we aren't comfortable entirely scrapping the custom derivative methods, one option here would be to introduce a new `gpytorch.settings` flag for whether to use the super call or the custom call.
- Let's add a unit test to the LazyTensorUnitTest that checks custom derivative calls against the super call.

Balandat: autograd ftw!
gpleiss: We should make sure that we’re doing careful profiling. I could imagine that our custom grads could be good when left_vecs/right_vecs is large, or when our LT is very simple - the extra matmul might be a bit of overhead.
jacobrgardner: Yeah for sure. Achieving simplicity is quite simple since once we verify correctness in batch mode, etc we could merge as-is and then at least quad_form_derivative is optional.

For speed, we'll definitely want to carefully analyze on a LazyTensor by LazyTensor basis whether using the default or a custom implementation makes sense. NonLazyTensor is the obvious example where we want custom behavior, but there could be others.

Patch:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 23114fc9..d37eb95c 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -105,9 +105,23 @@ class LazyTensor(object):
             :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
                                    this LazyTensor.
         """
-        raise NotImplementedError(
-            "The class {} requires a _quad_form_derivative function!".format(self.__class__.__name__)
-        )
+        args = self.representation()
+
+        toggled = [False] * len(args)
+
+        for i, arg in enumerate(args):
+            if not arg.requires_grad:
+                arg.requires_grad = True
+                toggled[i] = True
+
+        loss = (left_vecs * self._matmul(right_vecs)).sum()
+        grads = torch.autograd.grad(loss, args)
+
+        for i, arg in enumerate(args):
+            if toggled[i]:
+                arg.requires_grad = False
+
+        return grads
 
     def _size(self):
         """


Test patch:
diff --git a/test/examples/test_svgp_gp_classification.py b/test/examples/test_svgp_gp_classification.py
index 0636a04a..03c9ad49 100644
--- a/test/examples/test_svgp_gp_classification.py
+++ b/test/examples/test_svgp_gp_classification.py
@@ -25,10 +25,9 @@ train_y = torch.sign(torch.cos(train_x * (2 * pi))).squeeze()
 class GPClassificationModel(gpytorch.models.AbstractVariationalGP):
     def __init__(self):
         variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(16)
-        variational_strategy = gpytorch.variational.VariationalStrategy(self,
-                                                                        torch.randn(16, 1),
-                                                                        variational_distribution,
-                                                                        learn_inducing_locations=True)
+        variational_strategy = gpytorch.variational.VariationalStrategy(
+            self, torch.randn(16, 1), variational_distribution, learn_inducing_locations=True
+        )
 
         super(GPClassificationModel, self).__init__(variational_strategy)
         self.mean_module = ConstantMean(prior=SmoothedBoxPrior(-5, 5))
@@ -69,7 +68,7 @@ class TestSVGPClassification(unittest.TestCase):
 
         optimizer = optim.Adam(model.parameters(), lr=0.1)
         optimizer.n_iter = 0
-        for _ in range(50):
+        for _ in range(75):
             optimizer.zero_grad()
             output = model(train_x)
             loss = -mll(output, train_y)
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index b3218d4c..99ade3c4 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -113,6 +113,17 @@ class RectangularLazyTensorTestCase(object):
 class LazyTensorTestCase(RectangularLazyTensorTestCase):
     should_test_sample = False
 
+    def test_quad_form_derivative(self):
+        lazy_tensor = self.create_lazy_tensor()
+        left_vecs = torch.randn(lazy_tensor.size(-1), 2)
+        right_vecs = torch.randn(lazy_tensor.size(-1), 2)
+
+        deriv_custom = lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
+        deriv_auto = gpytorch.lazy.LazyTensor._quad_form_derivative(lazy_tensor, left_vecs, right_vecs)
+
+        for dc, da in zip(deriv_custom, deriv_auto):
+            self.assertLess(torch.norm(dc - da), 1e-1)
+
     def test_add_diag(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
@@ -311,6 +322,17 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
     def evaluate_lazy_tensor(self):
         raise NotImplementedError()
 
+    def test_quad_form_derivative(self):
+        lazy_tensor = self.create_lazy_tensor()
+        left_vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 2)
+        right_vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 2)
+
+        deriv_custom = lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
+        deriv_auto = gpytorch.lazy.LazyTensor._quad_form_derivative(lazy_tensor, left_vecs, right_vecs)
+
+        for dc, da in zip(deriv_custom, deriv_auto):
+            self.assertLess(torch.norm(dc - da), 1e-1)
+
     def setUp(self):
         if hasattr(self.__class__, "seed"):
             seed = self.__class__.seed
diff --git a/test/lazy/test_interpolated_lazy_tensor.py b/test/lazy/test_interpolated_lazy_tensor.py
index b7876b5b..dece6893 100644
--- a/test/lazy/test_interpolated_lazy_tensor.py
+++ b/test/lazy/test_interpolated_lazy_tensor.py
@@ -13,6 +13,11 @@ class TestInterpolatedLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 1
     should_test_sample = True
 
+    def test_quad_form_derivative(self):
+        # InterpolatedLazyTensor's representation includes int variables (the interp. indices),
+        # so the default derivative doesn't apply
+        pass
+
     def create_lazy_tensor(self):
         left_interp_indices = torch.LongTensor([[0, 1], [2, 3], [3, 4], [4, 5]])
         left_interp_values = torch.tensor([[0.1, 0.9], [1, 2], [0.5, 1], [1, 3]], dtype=torch.float)
@@ -45,6 +50,11 @@ class TestInterpolatedLazyTensorBatch(BatchLazyTensorTestCase, unittest.TestCase
     seed = 0
     should_test_sample = True
 
+    def test_quad_form_derivative(self):
+        # InterpolatedLazyTensor's representation includes int variables (the interp. indices),
+        # so the default derivative doesn't apply
+        pass
+
     def create_lazy_tensor(self):
         left_interp_indices = torch.LongTensor([[0, 1], [2, 3], [3, 4], [4, 5]]).repeat(5, 1, 1)
         left_interp_values = torch.tensor([[0.1, 0.9], [1, 2], [0.5, 1], [1, 3]], dtype=torch.float).repeat(5, 1, 1)
diff --git a/test/lazy/test_matmul_lazy_tensor.py b/test/lazy/test_matmul_lazy_tensor.py
index 6f1ead0a..ba1b1520 100644
--- a/test/lazy/test_matmul_lazy_tensor.py
+++ b/test/lazy/test_matmul_lazy_tensor.py
@@ -15,7 +15,7 @@ class TestMatmulLazyTensor(LazyTensorTestCase, unittest.TestCase):
 
     def create_lazy_tensor(self):
         lhs = torch.randn(5, 6, requires_grad=True)
-        rhs = lhs.transpose(-1, -2)
+        rhs = lhs.clone().detach().transpose(-1, -2)
         covar = MatmulLazyTensor(lhs, rhs)
         return covar
 
@@ -28,7 +28,7 @@ class TestMatmulLazyTensorBatch(BatchLazyTensorTestCase, unittest.TestCase):
 
     def create_lazy_tensor(self):
         lhs = torch.randn(5, 5, 6, requires_grad=True)
-        rhs = lhs.transpose(-1, -2)
+        rhs = lhs.clone().detach().transpose(-1, -2)
         covar = MatmulLazyTensor(lhs, rhs)
         return covar
 
diff --git a/test/lazy/test_mul_lazy_tensor.py b/test/lazy/test_mul_lazy_tensor.py
index f0d7028e..3e0cca36 100644
--- a/test/lazy/test_mul_lazy_tensor.py
+++ b/test/lazy/test_mul_lazy_tensor.py
@@ -42,6 +42,11 @@ class TestMulLazyTensor(LazyTensorTestCase, unittest.TestCase):
 class TestMulLazyTensorMulti(LazyTensorTestCase, unittest.TestCase):
     seed = 10
 
+    def test_quad_form_derivative(self):
+        # MulLazyTensor creates non-leaf variables, so the default derivative
+        # doesn't apply
+        pass
+
     def create_lazy_tensor(self):
         mat1 = make_random_mat(30, 3)
         mat2 = make_random_mat(30, 3)
@@ -85,6 +90,11 @@ class TestMulLazyTensorBatch(BatchLazyTensorTestCase, unittest.TestCase):
 class TestMulLazyTensorMultiBatch(BatchLazyTensorTestCase, unittest.TestCase):
     seed = 3
 
+    def test_quad_form_derivative(self):
+        # MulLazyTensor creates non-leaf variables, so the default derivative
+        # doesn't apply
+        pass
+
     def create_lazy_tensor(self):
         mat1 = make_random_mat(40, rank=5, batch_size=2)
         mat2 = make_random_mat(40, rank=5, batch_size=2)
@@ -109,6 +119,11 @@ class TestMulLazyTensorMultiBatch(BatchLazyTensorTestCase, unittest.TestCase):
 class TestMulLazyTensorWithConstantMul(BatchLazyTensorTestCase, unittest.TestCase):
     seed = 2
 
+    def test_quad_form_derivative(self):
+        # MulLazyTensor creates non-leaf variables, so the default derivative
+        # doesn't apply
+        pass
+
     def create_lazy_tensor(self):
         mat1 = make_random_mat(20, rank=5, batch_size=2)
         mat2 = make_random_mat(20, rank=5, batch_size=2)