cc @jacobrgardner @gpleiss @Balandat

cc @saitcakmak who has done some investigations in the past into removing LinearOperator overhead. I'll let him give the details, but IIRC the high level tl;dr was that while we did see speedups those didn't amount to dramatic savings in the grand scheme of things for our typical use cases. But our use cases aren't necessarily standard so I think there could be meaningful value in allowing to short-circuit the linear operator overhead in the small data regime.

I think this is great! What I had done was to take out linear operator and any parts of GPyTorch that we didn't need for ExactGPs and create a bare-bone version of the library (~20% of GPyTorch and no linear operator). I had seen up to 2x faster execution for some model operations, but at the end of the day, it wasn't that significant when considered as part of the whole BO loop in Ax. I think introducing small (depending on where we measure from) improvements like this to core GPyTorch is great. I'd definitely use this for ExactGPs in BoTorch.

I had a question while implementing this PR, which I think worth some discussions. Note that CholLinearOperator does not implement a custom backward pass for the log determinant. Instead, it relies on PyTorch's default backward pass, which backprop through the Cholesky decomposition. Were there any special considerations there?

This PR does implement a custom backward pass for log determinant by d logdet(K) = K^{-1}, which involves two triangular solves with the cached Cholesky factor of K.

In terms of the running time, this custom backward pass for logdet is indeed faster than PyTorch's default backward pass. Numerically, computing the inverse by two triangular solves shouldn't be bad since the backward pass of Cholesky decomposition requires two triangular solves anyway.

(I am assuming PyTorch implements the Cholesky backward pass by something similar to Eq (9) in Iain Murray's note.)

I am not sure if there were any special considerations at the time - @gpleiss or @jacobrgardner might know?

I dug around and the only remotely related issue to cholesky derivatives I could find was this one pytorch/pytorch#18825.

I think we just assumed the pytorch default derivatives for these ops would be fast. I certainly have no problem with us merging a custom backward pass in the name of 10-40% speed ups.

@kayween One thing I'm noticing: the speedup is increasing with matrix size! That is unexpected: one would assume that the additional overhead of linear operator packaging becomes negligible as the matrix size increases.

Can we extend the benchmark out to larger N? If there isn't some point where we stop seeing increasing speed-ups, there's obviously some serious problem with the default pytorch ops.

I dug around and the only remotely related issue to Cholesky derivatives I could find was this one pytorch/pytorch#18825.

I dug along this line a bit. @jacobrgardner and anyone who is curious, this is how the current PyTorch implements the Cholesky backward pass, which involves two triangular solves and a matmul.

https://github.com/pytorch/pytorch/blob/a3cc252e03572835c15afde54b81fc5e8616ad27/torch/csrc/autograd/FunctionsManual.cpp#L1984-L2010

I haven't dug into cholesky_inverse too deep since it is a part of LAPACK. But I assume it's based on two triangular solves. Thus, the custom backward pass of logdet in this PR saves (at least) a matmul compared to backproping through Cholesky decomposition.

@jacobrgardner Here are more benchmark results for larger N. The code generating the results is available in this gist. As with before, all benchmark runs use Cholesky decomposition for GP training.

Running time (in seconds) on synthetic datasets

The time reduction doesn't seem to saturate as we increase the size of training data! For example, we're seeing about 61% time reduction on CPU when n = 10000.

If there isn't some point where we stop seeing increasing speed-ups, there's obviously some serious problem with the default pytorch ops.

I don't think the issue is necessarily with PyTorch ops. There are probably two explanations why the speed-up does not stop with larger dataset sizes.

1. My hunch is that linear operators' representation tree implementation is not very efficient compared to PyTorch ops, which is where the overhead comes from. For example, InvQuad needs to reconstruct those linear operators in the forward and backward passes.
2. This PR has a custom backward pass for logdet, which is not available in linear operators.

cc @gpleiss @Balandat @saitcakmak

My hunch is that linear operators' representation tree implementation is not very efficient compared to PyTorch ops, which is where the overhead comes from. For example, InvQuad needs to reconstruct those linear operators in the forward and backward passes.

@kayween this doesn't make sense to me. The representation tree at the end of the day should just be carrying around the underlying pytorch tensors as pointers and reconstructing the thin python wrapper.

Like, in a vanilla GP, we have an AddedDiagLinearOperator(DenseLinearOperator(K), DiagLinearOperator(\sigma^2)). The representation tree for this has, at its roots, tensor(K) and tensor(\sigma^2). The pytorch Function takes those tensors as inputs, and just reconstructs AddedDiagLinearOperator(...) in the forward pass of the Function.

All that is to say: there's no way that python overhead, which should essentially be constant, dominates the linear algebra being done by n = 10,000. Or, more importantly, that overhead certainly shouldn't be growing with n!

This PR has a custom backward pass for logdet, which is not available in linear operators.

Sure, this makes sense to me.

@jacobrgardner This following code snippet shows that linear operators still have about 30% overhead compared to PyTorch even at the scale of n = 10000. This overhead may not come from the representation tree, but it seems like it has to be somewhere inside the inv_quad function call.

import linear_operator
import torch

from gpytorch.kernels import RBFKernel

torch.manual_seed(42)

n = 10_000

train_x = torch.rand(n, 5)
train_y = torch.rand(n)


def compute_inv_quad(use_linop: bool):
    covar_module = RBFKernel()
    covar = covar_module(train_x).evaluate_kernel().add_jitter(1e-3)

    if use_linop:
        with linear_operator.settings.fast_computations(False, False, False):
            return covar.inv_quad(inv_quad_rhs=train_y.unsqueeze(-1))
    else:
        covar = covar.to_dense()
        chol = torch.linalg.cholesky(covar)

        inv_chol_rhs = torch.linalg.solve_triangular(chol, train_y.unsqueeze(-1), upper=False).squeeze(-1)
        return (inv_chol_rhs**2).sum(-1)


%timeit compute_inv_quad(use_linop=True)
%timeit compute_inv_quad(use_linop=False)

# outputs:
# 1.31 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 912 ms ± 2.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)