feat: Efficiently Sampling Functions from Gaussian Process Posteriors

[adam-hartshorne]

Wilson and Terenin introduced "Efficiently Sampling Functions from Gaussian Process Posteriors" which is now widely used for, as the title states, efficient sampling (and won an honourable mention for best paper at ICML 2020).

The technique is clearly described in this blog post

https://sml-group.cc/blog/2020-gp-sampling/

and code for use with GPFlow (built on Tensorflow) is provided here.

https://github.com/j-wilson/GPflowSampling

I think it would be an extremely good addition to an GP library for JAX.

[thomaspinder]

Thanks for raising this @adam-hartshorne, this is very much on our radar. We are hoping to sketch out an initial idea of an implementation next week. If you have any thoughts or ideas about how you'd like to see this in GPJax, then please do raise them here, even if it's as simple as some pseudocode.

[aterenin]

Just a note that I have an ultra-lightweight research prototype on what pathwise conditioning support in a JAX-based package might look like here: https://github.com/aterenin/SparseGPAX.

I'd be delighted to help however I can for this to be implemented! Thank you very much for the interest in our work! 😊

[adam-hartshorne]

Obviously, you will have lots of different features to work on, but wondering if there has been any progress on this feature?

[thomaspinder]

Hey @adam-hartshorne. In short, things are moving a little slower than I would like as I'm in the process of finishing up my thesis. I'd hope to have this feature merged within the next 4-6 weeks though.

We first plan to merge a branch that enables spectral kernels to computed using a set of finite basis functions e.g., random Fourier features. Once this is merged, we'll finish up the pathwise sampling approach.

[adam-hartshorne]

Thank you for the update. Good luck with finishing up your thesis.