Repository: cornellius-gp/gpytorch

Issue:
Improving memory requirements for latent function inference (low-level Pyro-GPyTorch integration)
I have been working on a program performing latent function inference via the low-level Pyro-GPyTorch integration. The code is really similar to the example available in the [documentation](https://docs.gpytorch.ai/en/v1.5.1/examples/07_Pyro_Integration/Pyro_GPyTorch_Low_Level.html) but it is supposed to train on hundreds of thousands of samples.

During training, the memory requirements increase quadratically with the number of training data used. This has been pinned down to the `forward` method of the `VariationalStrategy` class. As pointed out in #1826, the `add_jitter` method of the `LazyEvaluatedKernelTensor` class evaluates the covariance matrix on the training samples and thus allocates the corresponding memory.

It seems to me that the full matrix does not need to be evaluated. Indeed, only the `mean` and the `stddev` of the variational distribution are used:
https://github.com/cornellius-gp/gpytorch/blob/aa0ff73cbd065a3a29b893c89bdab5497da63298/gpytorch/models/pyro/_pyro_mixin.py#L16-L17

As the output of the variational strategy is a lazy distribution, the latter involves only the diagonal of the covariance matrix:

https://github.com/cornellius-gp/gpytorch/blob/aa0ff73cbd065a3a29b893c89bdab5497da63298/gpytorch/distributions/multivariate_normal.py#L249

- Modules that will be modified:

In [`gpytorch.lazy.lazy_evaluated_kernel_tensor.py`](https://github.com/cornellius-gp/gpytorch/blob/c58dc3a77c7be7ee4434484e359a0bbc1c2d27f6/gpytorch/lazy/lazy_evaluated_kernel_tensor.py#L236-L237) use the lazy implementation of `add_jitter` from the `LazyTensor` class:

```diff
    def add_jitter(self, jitter_val=1e-3):
-        return self.evaluate_kernel().add_jitter(jitter_val)
+        return super(LazyEvaluatedKernelTensor, self).add_jitter(jitter_val)
```

- Impact on speed:

The spatial complexity of the algorithm drops from being quadratic to linear w.r.t. the
number of training data. Benefits are visible also in terms of throughput. Using the [example in the documentation](https://docs.gpytorch.ai/en/v1.5.1/examples/07_Pyro_Integration/Pyro_GPyTorch_Low_Level.html) as reference:

![time_mem](https://user-images.githubusercontent.com/32060225/159301194-e8801ed7-4a99-42a5-b224-49e1f5790924.png)

The above plots report the time and memory required to perform 10 training steps of a  GP model on an increasing number of training data (`NSamp`).  Details and code available [here](https://github.com/pierocor/gpytorch_test).


Discussion:

Resolving patch:
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -352,6 +352,58 @@ def covar_dist(
 
         return res
 
+    def expand_batch(self, *sizes: Union[torch.Size, Tuple[int, ...]]) -> Kernel:
+        r"""
+        Constructs a new kernel where the lengthscale (and other kernel parameters)
+        are expanded to match the batch dimension determined by `sizes`.
+
+        :param sizes: The batch shape of the new tensor
+        """
+        # Type checking
+        if len(sizes) == 1 and hasattr(sizes, "__iter__"):
+            new_batch_shape = torch.Size(sizes[0])
+        elif all(isinstance(size, int) for size in sizes):
+            new_batch_shape = torch.Size(sizes)
+        else:
+            raise RuntimeError("Invalid arguments {} to expand_batch.".format(sizes))
+
+        # Check for easy case:
+        orig_batch_shape = self.batch_shape
+        if new_batch_shape == orig_batch_shape:
+            return self
+
+        # Ensure that the expansion size is compatible with the given batch shape
+        try:
+            torch.broadcast_shapes(new_batch_shape, orig_batch_shape)
+        except RuntimeError:
+            raise RuntimeError(
+                f"Cannot expand a kernel with batch shape {self.batch_shape} to new shape {new_batch_shape}"
+            )
+
+        # Create a new kernel with updated batch shape
+        new_kernel = deepcopy(self)
+        new_kernel._batch_shape = new_batch_shape
+
+        # Reshape the parameters of the kernel
+        for param_name, param in self.named_parameters(recurse=False):
+            # For a given parameter, get the number of dimensions that do not correspond to the batch shape
+            non_batch_shape = param.shape[len(orig_batch_shape) :]
+            new_param_shape = torch.Size([*new_batch_shape, *non_batch_shape])
+            new_kernel.__getattr__(param_name).data = param.expand(new_param_shape)
+
+        # Reshape the buffers of the kernel
+        for buffr_name, buffr in self.named_buffers(recurse=False):
+            # For a given buffer, get the number of dimensions that do not correspond to the batch shape
+            non_batch_shape = buffr.shape[len(orig_batch_shape) :]
+            new_buffer_shape = torch.Size([*new_batch_shape, *non_batch_shape])
+            new_kernel.__getattr__(buffr_name).data = buffr.expand(new_buffer_shape)
+
+        # Recurse, if necessary
+        for sub_module_name, sub_module in self.named_sub_kernels():
+            new_kernel.__setattr__(sub_module_name, sub_module.expand_batch(new_batch_shape))
+
+        return new_kernel
+
     def named_sub_kernels(self) -> Iterable[Tuple[str, Kernel]]:
         """
         For compositional Kernel classes (e.g. :class:`~gpytorch.kernels.AdditiveKernel`
@@ -510,14 +562,23 @@ def __getitem__(self, index) -> Kernel:
         # Process the index
         index = index if isinstance(index, tuple) else (index,)
 
-        for param_name, param in self._parameters.items():
-            new_kernel._parameters[param_name].data = param.__getitem__(index)
-            ndim_removed = len(param.shape) - len(new_kernel._parameters[param_name].shape)
+        for param_name, param in self.named_parameters(recurse=False):
+            new_param = new_kernel.__getattr__(param_name)
+            new_param.data = new_param.__getitem__(index)
+            ndim_removed = len(param.shape) - len(new_param.shape)
+            new_batch_shape_len = len(self.batch_shape) - ndim_removed
+            new_kernel.batch_shape = new_param.shape[:new_batch_shape_len]
+
+        for buffr_name, buffr in self.named_buffers(recurse=False):
+            # For a given buffer, get the number of dimensions that do not correspond to the batch shape
+            new_buffr = new_kernel.__getattr__(buffr_name)
+            new_buffr.data = new_buffr.__getitem__(index)
+            ndim_removed = len(buffr.shape) - len(new_buffr.shape)
             new_batch_shape_len = len(self.batch_shape) - ndim_removed
-            new_kernel.batch_shape = new_kernel._parameters[param_name].shape[:new_batch_shape_len]
+            new_kernel.batch_shape = new_buffr.shape[:new_batch_shape_len]
 
         for sub_module_name, sub_module in self.named_sub_kernels():
-            new_kernel._modules[sub_module_name] = sub_module.__getitem__(index)
+            new_kernel.__setattr__(sub_module_name, sub_module.__getitem__(index))
 
         return new_kernel
 
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -124,6 +124,18 @@ def _getitem(self, row_index, col_index, *batch_indices):
             num_outs_per_in_rows = num_outs_per_in
             num_outs_per_in_cols = num_outs_per_in
 
+        # We will be running the __getitem__ command on x1, x2, and the kernel parameters
+        # Since kernels can broadcast, x1, x2, and the kernel parameters may not have all of the batch dimensions
+        # that are being indexed by the __getitem__ operation
+        # Therefore, we begin by figuring out the broadcasted shape, and expanding all of these objects to that shape
+        try:
+            batch_shape = torch.broadcast_shapes(x1.shape[:-2], x2.shape[:-2], self.kernel.batch_shape)
+        except RuntimeError:
+            raise RuntimeError(
+                f"The kernel inputs (sizes {x1.shape} and {x2.shape}) are incompatible with the kernel "
+                f"(batch size {self.kernel.batch_shape}). This is likely a bug in GPyTorch."
+            )
+
         # The row index and col index should exactly correspond to which entries of x1 and x2 we need
         # So we'll basically call x1[*batch_indices, row_index, :], x2[*batch_indices, col_index, :]
 
@@ -174,38 +186,33 @@ def _getitem(self, row_index, col_index, *batch_indices):
         try:
             x1 = x1[(*batch_indices, row_index, dim_index)]
         # We're going to handle multi-batch indexing with a try-catch loop
-        # This way - in the default case, we can avoid doing expansions of x1 which can be timely
+        # This way - in the default case, we can avoid doing expansions of x1 which can be
+        # costly in terms of time
         except IndexError:
-            if any(not isinstance(bi, slice) for bi in batch_indices):
-                raise RuntimeError(
-                    "Attempting to tensor index a non-batch matrix's batch dimensions. "
-                    f"Got batch index {batch_indices} but my shape was {self.shape}"
-                )
-            x1 = x1.expand(*([1] * (len(batch_indices) - self.x1.dim() + 2)), *self.x1.shape)
+            x1 = x1.expand(*batch_shape, *x1.shape[-2:])
             x1 = x1[(*batch_indices, row_index, dim_index)]
 
         # Call x2[*batch_indices, col_index, :]
         try:
             x2 = x2[(*batch_indices, col_index, dim_index)]
         # We're going to handle multi-batch indexing with a try-catch loop
-        # This way - in the default case, we can avoid doing expansions of x1 which can be timely
+        # This way - in the default case, we can avoid doing expansions of x2 which can be
+        # costly in terms of time
         except IndexError:
-            if any(not isinstance(bi, slice) for bi in batch_indices):
-                raise RuntimeError(
-                    "Attempting to tensor index a non-batch matrix's batch dimensions. "
-                    f"Got batch index {batch_indices} but my shape was {self.shape}"
-                )
-            x2 = x2.expand(*([1] * (len(batch_indices) - self.x2.dim() + 2)), *self.x2.shape)
+            x2 = x2.expand(*batch_shape, *x2.shape[-2:])
             x2 = x2[(*batch_indices, col_index, dim_index)]
 
         if len(batch_indices) == 0 or all(ind == slice(None, None, None) for ind in batch_indices):
             new_kernel = self.kernel  # Avoid unnecessary copying when we aren't explicitly indexing batch dims
         else:
-            # Deal with the fact that the batch shape could be different from the kernel's batch shape
-            # (e.g. this can happen due to broadcasting)
-            while len(batch_indices) > len(self.kernel._batch_shape) and batch_indices[0] == _noop_index:
-                batch_indices = batch_indices[1:]
-            new_kernel = self.kernel.__getitem__(batch_indices)
+            try:
+                new_kernel = self.kernel.__getitem__(batch_indices)
+            # We're going to handle multi-batch indexing with a try-catch loop
+            # This way - in the default case, we can avoid doing expansions of self.kernel which can be
+            # costly in terms of time
+            except IndexError:
+                expanded_kernel = self.kernel.expand_batch(batch_shape)
+                new_kernel = expanded_kernel.__getitem__(batch_indices)
 
         # Now construct a kernel with those indices
         return self.__class__(
@@ -300,9 +307,6 @@ def _transpose_nonbatch(self):
             **self.params,
         )
 
-    def add_jitter(self, jitter_val=1e-3):
-        return self.evaluate_kernel().add_jitter(jitter_val)
-
     def _unsqueeze_batch(self, dim):
         x1 = self.x1.unsqueeze(dim)
         x2 = self.x2.unsqueeze(dim)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -40,7 +40,7 @@ def find_version(*file_paths):
 torch_min = "1.11"
 install_requires = [
     "scikit-learn",
-    "linear_operator>=0.1.1",
+    "linear_operator>=0.2.0",
 ]
 # if recent dev version of PyTorch is installed, no need to install stable
 try:


Test patch:
diff --git a/gpytorch/test/base_kernel_test_case.py b/gpytorch/test/base_kernel_test_case.py
--- a/gpytorch/test/base_kernel_test_case.py
+++ b/gpytorch/test/base_kernel_test_case.py
@@ -162,6 +162,20 @@ def test_kernel_getitem_double_batch(self):
 
         self.assertAllClose(res1, res2, rtol=1e-3, atol=1e-5)
 
+    def test_kernel_getitem_broadcast(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([2]))
+        x = self.create_data_double_batch()
+        kernel = kernel.expand_batch(torch.broadcast_shapes(kernel.batch_shape, x.shape[:-2]))
+
+        idx1 = torch.LongTensor([0, 2, 1]).unsqueeze(-1)
+        idx2 = torch.LongTensor([1, 0, 0]).unsqueeze(-2)
+        res1 = kernel(x).to_dense()[idx1, idx2]  # Result of first kernel on first batch of data
+
+        new_kernel = kernel[idx1, idx2]
+        res2 = new_kernel(x[idx1, idx2]).to_dense()  # Should also be result of first kernel on first batch of data.
+
+        self.assertAllClose(res1, res2, rtol=1e-3, atol=1e-5)
+
     def test_kernel_pickle_unpickle(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
         pickle.loads(pickle.dumps(kernel))  # Should be able to pickle and unpickle a kernel