Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Bayesian Inverse Problems Meet Flow Matching: Efficient and Flexible Inference via Transformers

Daniil Sherki, Ivan Oseledets, Ekaterina Muravleva

TL;DR

This work tackles the computational bottleneck in Bayesian inverse problems by introducing CFM-Tr, a transformer-based Conditional Flow Matching framework that learns conditional probability trajectories to sample from the posterior $\,\pi(m|d,e)\$ without iterative forward solves. By generating synthetic triplets $(m,d,e)$ and training a velocity field on interpolation paths, the method supports arbitrary numbers of observations and achieves fast, deterministic posterior sampling. Across a simple nonlinear model, SEIR dynamics, and a 2D Darcy flow inversion, it attains relative parameter errors as low as about $1.5\%$ to $2.8\%$ while delivering speedups up to $2000\times$ on CPU versus MCMC. The approach enables real-time Bayesian inference and lays groundwork for Bayesian experimental design by efficiently learning transport maps from data.

Abstract

The efficient resolution of Bayesian inverse problems remains challenging due to the high computational cost of traditional sampling methods. In this paper, we propose a novel framework that integrates Conditional Flow Matching (CFM) with a transformer-based architecture to enable fast and flexible sampling from complex posterior distributions. The proposed methodology involves the direct learning of conditional probability trajectories from the data, leveraging CFM's ability to bypass iterative simulation and transformers' capacity to process arbitrary numbers of observations. The efficacy of the proposed framework is demonstrated through its application to three problems: a simple nonlinear model, a disease dynamics framework, and a two-dimensional Darcy flow Partial Differential Equation. The primary outcomes demonstrate that the relative errors in parameters recovery are as low as 1.5%, and that the inference time is reduced by up to 2000 times on CPU in comparison with the Monte Carlo Markov Chain. This framework facilitates the expeditious resolution of Bayesian problems through the utilisation of sampling from the learned conditional distribution.

Bayesian Inverse Problems Meet Flow Matching: Efficient and Flexible Inference via Transformers

TL;DR

This work tackles the computational bottleneck in Bayesian inverse problems by introducing CFM-Tr, a transformer-based Conditional Flow Matching framework that learns conditional probability trajectories to sample from the posterior without iterative forward solves. By generating synthetic triplets and training a velocity field on interpolation paths, the method supports arbitrary numbers of observations and achieves fast, deterministic posterior sampling. Across a simple nonlinear model, SEIR dynamics, and a 2D Darcy flow inversion, it attains relative parameter errors as low as about to while delivering speedups up to on CPU versus MCMC. The approach enables real-time Bayesian inference and lays groundwork for Bayesian experimental design by efficiently learning transport maps from data.

Abstract

The efficient resolution of Bayesian inverse problems remains challenging due to the high computational cost of traditional sampling methods. In this paper, we propose a novel framework that integrates Conditional Flow Matching (CFM) with a transformer-based architecture to enable fast and flexible sampling from complex posterior distributions. The proposed methodology involves the direct learning of conditional probability trajectories from the data, leveraging CFM's ability to bypass iterative simulation and transformers' capacity to process arbitrary numbers of observations. The efficacy of the proposed framework is demonstrated through its application to three problems: a simple nonlinear model, a disease dynamics framework, and a two-dimensional Darcy flow Partial Differential Equation. The primary outcomes demonstrate that the relative errors in parameters recovery are as low as 1.5%, and that the inference time is reduced by up to 2000 times on CPU in comparison with the Monte Carlo Markov Chain. This framework facilitates the expeditious resolution of Bayesian problems through the utilisation of sampling from the learned conditional distribution.

Paper Structure

This paper contains 28 sections, 13 equations, 6 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Our contribution
  3. Background and Related Work
  4. Classical Approaches to Solving Bayesian Inverse Problems
  5. Deep Learning Models for Solving Inverse Problems
  6. Generative Adversarial Networks and Variational Autoencoders
  7. Flow Matching
  8. Methodology
  9. Dataset
  10. Training
  11. Architecture
  12. Handling variable number of observations
  13. Numerical Experiments
  14. Simple nonlinear model
  15. SEIR disease model
  16. ...and 13 more sections

Figures (6)

  • Figure 1: Solving the inverse problem using flow-matching scheme
  • Figure 2: Transformer architecture
  • Figure 3: Probabilistic solutions to the inverse problem for $\mathbf{m}_{\textbf{true}} = [0.4, 0.3, 0.3, 0.1, 0.15, 0.6]$
  • Figure 4: PDE coefficient and solution: true (left) and reconstructed using Flow Matching (right)
  • Figure 5: Generation paths of variable $m$ conditional on different $d$, $e$ from prior uniform distribution
  • ...and 1 more figures