FM-DeepRV: Deep Learning for Bayesian Inference with Flow Matching

Bayesian inference with Gaussian processes(GPs) is intractable at scale, because drawing a GP sample requires doing the Cholesky factorization of the kernel matrix which is an $\mathcal{O}(N^3)$ operation. DeepRV is a neural surrogate approach that trains a decoder network to approximate this factorization, replacing the Cholesky with a single forward pass. As diffusion models are state of the art for image generation, we propose replacing the Mean Square Error trained decoder in DeepRV with a conditional flow matching (CFM) network that learns the optimal transport vector field between a Gaussian and the GP sample distribution. Having trained this neural surrogate, at inference (i.e.~for a new dataset) we place the learned ODE into a Bayesian probabilistic programming framework and use MCMC to infer the posterior distribution. Because the learned ODE surrogate is differentiable, we can use Hamiltonian Monte Carlo, which is very efficient in this setting. We call the resulting method FM-DeepRV. We evaluate on spatial GP benchmarks with Matérn-$\tfrac{1}{2}$ and Matérn-$\tfrac{3}{2}$ kernels across two lengthscales and two grid sizes, and on an image inpainting task with MNIST digits. We explore settings in which our method improves performance over DeepRV.