Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows

TL;DR

This paper introduces Markovian Flow Matching, a novel method combining continuous normalizing flows with adaptive MCMC to improve sampling efficiency and discover multiple modes in complex distributions.

Contribution

It repurposes flow matching for probabilistic inference, integrating CNFs into MCMC with on-the-fly adaptation and tempering, enabling efficient, multi-modal sampling.

Findings

Achieves comparable performance to state-of-the-art methods

Reduces computational cost significantly

Successfully discovers multiple modes in target distributions

Abstract

Continuous normalizing flows (CNFs) learn the probability path between a reference distribution and a target distribution by modeling the vector field generating said path using neural networks. Recently, Lipman et al. (2022) introduced a simple and inexpensive method for training CNFs in generative modeling, termed flow matching (FM). In this paper, we repurpose this method for probabilistic inference by incorporating Markovian sampling methods in evaluating the FM objective, and using the learned CNF to improve Monte Carlo sampling. Specifically, we propose an adaptive Markov chain Monte Carlo (MCMC) algorithm, which combines a local Markov transition kernel with a non-local, flow-informed transition kernel, defined using a CNF. This CNF is adapted on-the-fly using samples from the Markov chain, which are used to specify the probability path for the FM objective. Our method also includes an adaptive tempering mechanism that allows the discovery of multiple modes in the target distribution. Under mild assumptions, we establish convergence of our method to a local optimum of the FM objective. We then benchmark our approach on several synthetic and real-world examples, achieving similar performance to other state-of-the-art methods, but often at a significantly lower computational cost.