How difficult are sparse linear operators?

[Turakar]

In the BMMM paper and on the docs of the linear-operators package you already point out that your method is extensible to sparse matrices, as it only requires MVM-operations. I would like to investigate extending linear-operators with a new SparseLinearOperator which just contains a sparse torch matrix and delegates all features to existing torch implementations. What would be the complexity of implementing this?

• Is it the correct approach to create a new linear operator?
• Which functions would one need to implement?
• Can one maybe just put a sparse matrix into DenseLinearOperator?

[Balandat]

Yeah you could do this, and just "putting a sparse matrix into DenseLinearOperator" would probably be the easiest thing to try. I'm not super familiar with the current limitations / pitfalls of the torch sparse tensors though, so there may be more to this.

In an ideal world, since linear_operator implements a bunch of standard torch functions via the __torch_function__ interface, we'd be able to just use both LinearOperator objects and all torch tensors interchangeably. We made a good amount of progress to enable this for (non-sparse) tensors (still could have more coverage of course), but again I'm not sure what the torch limitations on the sparse tensor side are.

[Turakar]

What do you expect to be the benefit in this case by allowing LinearOperator and Tensor to be used interchangeably? Do you intend to just return a sparse tensor instead of a linear operator in "an ideal world"?

[Balandat]

The high level idea is to get to a place where anyone can just replace their torch.Tensor objects in their existing code to reap the benefit of structure-exploiting linear algebra. I talked about this in a poster for the pytorch conference last year: https://www.youtube.com/watch?v=HE85_4AQqzQ&list=PL_lsbAsL_o2CzovdREM1BpwfuHl-tNwPl&index=2&ab_channel=PyTorch

[Turakar]

Thank you, that's interesting :)

[Turakar]

Thank you for your answers. I tried some initial things, and e.g. passing a sparse CSR tensor into DenseLinearOperator allowed me to implement an ExactGP with a locally periodic kernel over a large dataset by setting elements beyond a cutoff distance to 0. I think this might be an interesting approach, so I will proceed with this and report back.

I already had to circumvent some limitations of CSR tensors:

• CSR tensors require all elements of a batch to have the same number of observed values (of course, observed may mean 0). Luckily, this is no problem in my case, an LMC.
• The CSR indexing functionality does not support strides on the row and column indices, but it does not detect no-ops like slice(None, None, None), so I have to filter them manually.
• Batched matmul seems to be weird: torch.matmul with batched CSR matrix pytorch/pytorch#98675

[Turakar]

Ok, so I ended up concluding the following:

• I was able to implement a batch sparse linear operator using above tricks (batch is implemented just as a list over CSR tensors).
• It does allow me to scale to relatively large matrices, but I quickly hit the limits of my GPU memory (16 GB). This sparse approach is not faster than dense matrices in the vast majority of cases.
• KeOps seemed to be a better solution for this case. KeOps does not store the covariance matrix in memory, but computes its entries on the fly.
• Sparsity on the KeOps operator did reduce computing time, but the support is kinda limited and I also hit some bugs there (Do not include exact ranges in cache argument, but only their presence getkeops/keops#309, Genred sparse reduction fails silently for int64 indices getkeops/keops#307, Question: Extending ranges to 3D+ tensors getkeops/keops#120).
• According to Nsight Compute, the KeOps implementation was also quite inefficient with a low occupancy and compute/memory throughput.

In the end, I ended up coding a CUDA kernel that follows the KeOps approach. This is quite fast, but also highly specialized, so I think it is of limited use to other people. However, the general approach of writing a custom CUDA Matmul kernel is quite appealing if performance is a primary concern and the covariance matrix becomes large.