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

Makkunda Sharma, Jhonathan Navott, Daniel Jenson, Elizaveta Semenova, Seth Flaxman

Keywords: Bayesian Inference, Conditional Flow Matching, MCMC, Deep Learning

TL;DR: Using Flow Matching Surrogates to accelarate MCMC inference in spatial datasets

Abstract:

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 127