RBF kernel with gradient and extending with keops

[JGittles]

Hi Everyone,

Lately I've been interested in better understanding the RBF kernel with grad information (https://docs.gpytorch.ai/en/stable/_modules/gpytorch/kernels/rbf_kernel_grad.html#RBFKernelGrad) with the intention of building out a custom kernel utilizing keops to copy this functionality. I've been working through the linked source code but was still a bit confused on how I could best convert the operations within to work with keops lazy tensors as well as what is already "lazy-fied" so to speak.

From what I've learned so far, each output kernel matrix is some n1d by n2d matrix (assuming no batch size) that consists of the typical rbf kernel as well as two gradient blocks on the upper and leftmost sides with a hessian block taking up the remaining space. Naively, I am a bit lost on where the limits exist on how many operations I can "subject" a lazy tensor to for these sections. with that, I noticed the hessian block itself is computed in part with a kroenecker product lazy tensor so off the bat it seems like that portion of the code may not need to be changed to achieve improved memory/compute efficiency?

Has anyone here had experience doing these sorts of operations? I am admittedly a novice with tensor operations and computational efficiency so any input/resource would be greatly appreciated!

[gpleiss]

Lately I've been interested in better understanding the RBF kernel with grad information (https://docs.gpytorch.ai/en/stable/_modules/gpytorch/kernels/rbf_kernel_grad.html#RBFKernelGrad) with the intention of building out a custom kernel utilizing keops to copy this functionality.

You should follow the keeps RBF kernel as a guide for developing this - and for a first pass avoid anything that involves kronecker structure. Incorporating kronecker structure will make everything more efficient, but you will already get a massive performance boost by just using keeps.