Flow matching Operators for Residual-Augmented Probabilistic Learning of Partial Differential Equations

TL;DR

This paper introduces a novel flow matching approach in function space for probabilistic learning of PDE solutions, enabling resolution-invariant inference and effective correction from low- to high-fidelity data.

Contribution

It develops a residual-augmented neural operator framework with a new training strategy for efficient, resolution-invariant PDE surrogate modeling from limited data.

Findings

Accurately models PDE solution operators across resolutions.

Provides meaningful uncertainty estimates with limited high-fidelity data.

Demonstrates effectiveness on 1D and 2D PDE benchmarks.

Abstract

Learning probabilistic surrogates for partial differential equations remains challenging in data-scarce regimes: neural operators require large amounts of high-fidelity data, while generative approaches typically sacrifice resolution invariance. We formulate flow matching in an infinite-dimensional function space to learn a probabilistic transport that maps low-fidelity approximations to the manifold of high-fidelity PDE solutions via learned residual corrections. We develop a conditional neural operator architecture based on feature-wise linear modulation for flow matching vector fields directly in function space, enabling inference at arbitrary spatial resolutions without retraining. To improve stability and representational control of the induced neural ODE, we parameterize the flow vector field as a sum of a linear operator and a nonlinear operator, combining lightweight linear components with a conditioned Fourier neural operator for expressive, input-dependent dynamics. We then formulate a residual-augmented learning strategy where the flow model learns probabilistic corrections from inexpensive low-fidelity surrogates to high-fidelity solutions, rather than learning the full solution mapping from scratch. Finally, we derive tractable training objectives that extend conditional flow matching to the operator setting with input-function-dependent couplings. To demonstrate the effectiveness of our approach, we present numerical experiments on a range of PDEs, including the 1D advection and Burgers' equation, and a 2D Darcy flow problem for flow through a porous medium. We show that the proposed method can accurately learn solution operators across different resolutions and fidelities and produces uncertainty estimates that appropriately reflect model confidence, even when trained on limited high-fidelity data.