Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows

Alberto Cabezas, Louis Sharrock, Christopher Nemeth

TL;DR

This paper proposes an adaptive Markov chain Monte Carlo algorithm, which combines a local Markov transition kernel with a non-local, flow-informed transition kernel, defined using a CNF, and adapted on-the-fly using samples from the Markov chain.

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.

Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows

TL;DR

This paper proposes an adaptive Markov chain Monte Carlo algorithm, which combines a local Markov transition kernel with a non-local, flow-informed transition kernel, defined using a CNF, and adapted on-the-fly using samples from the Markov chain.

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.
Paper Structure (30 sections, 1 theorem, 40 equations, 3 figures, 10 tables, 4 algorithms)

This paper contains 30 sections, 1 theorem, 40 equations, 3 figures, 10 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Our contributions
  3. Preliminaries
  4. Continuous Normalizing Flows
  5. Flow Matching
  6. Markovian Flow Matching
  7. MCMC with Flow Matching
  8. Training the CNF
  9. Adaptive MCMC with Flow Matching
  10. Overview
  11. Sampling
  12. Training
  13. Annealing
  14. Convergence
  15. Related work
  16. ...and 15 more sections

Key Result

Proposition 3.1

Assume that Assumptions assumption:step - assumption:score-function hold (see Appendix proof_prop1). Assume also that ${(\theta_K)_{K\in\mathbb{N}}}$ is a bounded sequence, which almost surely visits a compact subset of the domain of attraction of ${\theta^{*}}$ infinitely often. Then ${\theta_K \ri

Figures (3)

  • Figure 1: Comparison between MFM, FAB, DDS, and NF-MCMC. Samples from the target density for the 4-mode Gaussian mixture example.
  • Figure 2: Comparison between MFM, FAB, DDS, and NF-MCMC. Samples from the target density for the 16-mode Gaussian mixture example.
  • Figure 3: Comparison between MFM, FAB, DDS, and NF-MCMC. Representative samples from the target density for the Field system example.

Theorems & Definitions (1)

  • Proposition 3.1