Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Process Learning via Operator Flow Matching

Yaozhong Shi, Zachary E. Ross, Domniki Asimaki, Kamyar Azizzadenesheli

TL;DR

OFM introduces a flow-based operator framework to learn priors over stochastic processes on function spaces and to deliver tractable densities for arbitrary point sets, enabling Bayesian functional regression (UFR). By extending flow matching to infinite dimensions via neural operators and marginal optimal transport, OFM achieves exact likelihoods and calibrated posterior samples for both GP and non-GP data. The approach yields state-of-the-art performance across diverse tasks, including Navier–Stokes and black-hole simulations, and provides a principled mechanism to sample from posterior function values at arbitrary query sets. This work bridges operator learning, optimal transport, and Bayesian regression, with potential to generalize stochastic-process priors across scientific domains.

Abstract

Expanding on neural operators, we propose a novel framework for stochastic process learning across arbitrary domains. In particular, we develop operator flow matching (OFM) for learning stochastic process priors on function spaces. OFM provides the probability density of the values of any collection of points and enables mathematically tractable functional regression at new points with mean and density estimation. Our method outperforms state-of-the-art models in stochastic process learning, functional regression, and prior learning.

Stochastic Process Learning via Operator Flow Matching

TL;DR

OFM introduces a flow-based operator framework to learn priors over stochastic processes on function spaces and to deliver tractable densities for arbitrary point sets, enabling Bayesian functional regression (UFR). By extending flow matching to infinite dimensions via neural operators and marginal optimal transport, OFM achieves exact likelihoods and calibrated posterior samples for both GP and non-GP data. The approach yields state-of-the-art performance across diverse tasks, including Navier–Stokes and black-hole simulations, and provides a principled mechanism to sample from posterior function values at arbitrary query sets. This work bridges operator learning, optimal transport, and Bayesian regression, with potential to generalize stochastic-process priors across scientific domains.

Abstract

Expanding on neural operators, we propose a novel framework for stochastic process learning across arbitrary domains. In particular, we develop operator flow matching (OFM) for learning stochastic process priors on function spaces. OFM provides the probability density of the values of any collection of points and enables mathematically tractable functional regression at new points with mean and density estimation. Our method outperforms state-of-the-art models in stochastic process learning, functional regression, and prior learning.
Paper Structure (24 sections, 1 theorem, 55 equations, 25 figures, 9 tables, 2 algorithms)

This paper contains 24 sections, 1 theorem, 55 equations, 25 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Operator Flow Matching
  4. Generalizing Flow Matching to Stochastic Processes
  5. Likelihood Estimation and Bayesian Universal Functional Regression
  6. Experiments
  7. Conclusion
  8. Potential questions and answers
  9. Background: Flow Matching, Gaussian Measures on Function Spaces, and the Cameron–Martin Theorem
  10. Stochastic process learning
  11. Model stochastic process with infinite-dimensional flow matching via Kolmogorov Extension Theorem
  12. Universal Functional Regression
  13. Marginal (dynamic) optimal-transport flow matching in function space via optimal coupling and dynamic Kantorovich formulation
  14. Derivation of Eq. \ref{['eq:expand']}
  15. Derivation of Eq. \ref{['eq:generalize_cond_u']}
  16. ...and 9 more sections

Key Result

Proposition 3.1

Given noisy observations $\lbrace \widehat{u}(x_i) \rbrace_{i=1}^n$, the posterior distribution is Where the constant $C = - \frac{n}{2}\log(2\pi\sigma^2) -\log \mathbb P \left( \lbrace \widehat{u}(x_i) \rbrace_{i=1}^n \right)$.

Figures (25)

  • Figure 1: Two‑phase strategy for prior learning and posterior sampling. In the prior‑learning phase, OFM leverages the marginal optimal‑transport path to learn a bijective mapping between a predefined GP and the unknown stochastic process that generates the training data. In the posterior‑sampling phase, the learned prior is frozen; given noisy, partial observations, the exact posterior is obtained via Bayes’ theorem. SGLD, aided by the Hutchinson trace estimator, then enables efficient and robust sampling.
  • Figure 2: Operator Flow Matching (OFM) regression on Navier-Stokes functional data with resolution $64\times64$. (a) 32 random observations (only $0.7\%$). (b) Ground truth sample. (c) Predicted mean from OFM. (d) One posterior sample from OFM. (e) One posterior sample from the best fitted GP. (f) Predicted mean from ConvCNP. (g) One posterior sample from ConvCNP.
  • Figure 3: OFM regression on GP data. (a) Ground truth GP regression with observed data and predicted samples. (b) OFM regression with observed data and predicted samples. (c) Standard deviation comparison between true GP and OFM predictions.
  • Figure 4: OFM regression on TGP data.
  • Figure 5: OFM regression on black hole data with resolution $64\times64$. (a) 32 random observations. (b) Ground truth sample. (c) Predicted mean from OFM. (d) One posterior sample from OFM. (e) One posterior sample from best fitted GP. (f) Predicted mean from ConvCNP. (g) One posterior sample from ConvCNP.
  • ...and 20 more figures

Theorems & Definitions (3)

  • Proposition 3.1
  • proof
  • proof