JaxGaussianProcesses / GPJax

Gaussian processes in JAX.
https://docs.jaxgaussianprocesses.com/
Apache License 2.0
459 stars 53 forks source link

feat: Efficiently Sampling Functions from Gaussian Process Posteriors #218

Closed adam-hartshorne closed 1 year ago

adam-hartshorne commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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

thomaspinder commented 2 years ago

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 commented 2 years ago

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