[Question] Best way to create a bidiagonal LazyTensor

Hi everyone,

I am looking for some advice on how to create a specific LazyTensor such that matrix products/inversions are done efficiently.

Goal

My goal is to create the following LazyTensor, of size nx(n+1):

[[1, -g,  0, 0, ...,  0,  0],
 [0,  1, -g, 0, ...,  0,  0],
...
 [0,  0,  0, 0, ..., -g,  0],
 [0,  0,  0, 0, ...,  1, -g]]

Current solution and issues

The current way I do this is:

from gpytorch.lazy import ZeroLazyTensor, DiagLazyTensor, CatLazyTensor
import torch

g = 0.5
n = 5

eye_n = DiagLazyTensor(torch.tensor([1.] * n))
g_n = DiagLazyTensor(torch.tensor([-g] * n))
# Create two different zero_col to make sure gradients are
# correctly backpropagated
zero_col_1 = ZeroLazyTensor(n, 1, dtype=torch.float)
zero_col_2 = ZeroLazyTensor(n, 1, dtype=torch.float)

diag_part = CatLazyTensor(eye_n, zero_col_1, dim=-1)
superdiag_part = CatLazyTensor(zero_col_2, g_n, dim=-1)

out = diag_part + superdiag_part
print(out.evaluate())

This creates the correct tensor, but it scales very poorly: since the tensor has a complicated structure, it cannot be inverted or multiplied efficiently. Consequently, when I use this tensor in a custom kernel (with n the number of points), I can only use the GP on a few thousand samples, after which my GPU’s memory is not large enough to do the operations.

Is there a more efficient way to define this LazyTensor?

[PFMassiani]

Wonderful, thank you. I will keep this thread open for now just in case, and I will close it when I manage to implement the custom LazyTensor.

[jacobrgardner]

I’d write your own lazy tensor.

All a LazyTensor is is a way to define how to do operations efficiently on a matrix whose entries are specified in a potentially unusual way.

In your case, you could imagine a LazyTensor that takes as input g and defines at a minimum the three abstract methods: _matmul, _size, and _transpose_nonbatch, as well as _solve since bidiagonal matrices can’t use CG for solves.

Something like this might be the minimum viable implementation:

class BidiagonalLazyTensor(LazyTensor):
    def __init__(self, off_diag, upper=True):
        super().__init__(off_diag, upper=upper)
        self._off_diag = off_diag
        self.upper = upper

    def _matmul(self, rhs):
        # TODO: Implement (ideally batch aware) bidiagonal matmul routine

    def _size(self):
        n = self._off_diag.size(-1) + 1   # Size of matrix is size of superdiagonal + 1
        return self._off_diag.shape[:-1] + torch.Size([n, n])  # Include batch shape
    
    def _transpose_nonbatch(self):
        return BidiagonalLazyTensor(self._off_diag, upper=not self.upper)

    def _solve(self):
        # Default solve method won't work with bidiagonal matrices -- see below.

Now, in your specific case I wouldn’t stop there. By just defining matmul, the default behavior of _solve, the method that computes A^{-1}b with the lazy tensor will be to use CG, which obviously isn’t applicable to bidiagonal matrices. With a bidiagonal matrix with constant diagonals we’d need to write our own _solve method, which we can do because we know that the inverse of such a matrix is upper triangular with a specific structure.

In particular, if C = A^{-1} with A having your specified structure, then C[i, j] = 0 whenever i > j (i.e., C is upper triangular), and C[i, j] = (g)^j whenever i <= j (here assuming 0 indexing – in other words, the diagonal of C is all ones). In other words, the jth superdiagonal of C is constant with value(g)^j, where the “0th” superdiagonal is taken to be the diagonal.

Thus, you could override _solve to exploit this structure. If you wanted to keep things linear space, you’d need to code up your own version of backward substitution specific to this C, which shouldn’t be too hard.

A good place to start would be to look at TriangularLazyTensor here, since your new LT has very similar properties in that it isn’t positive definite, so many of the methods overwritten in TriangularLazyTensor will also be overwritten in BidiagonalLazyTensor.