📊 statistics

Flow Annealing Posterior Sampling for Function-Space Regression and Inverse Problems

This paper introduces Flow Annealing Posterior Sampling (FAPS), a novel framework that unifies stochastic-process regression and PDE inverse problems by leveraging pretrained function-space flow-matching priors to generate coherent posterior samples with accurate uncertainty quantification, while outperforming existing baselines in performance and reducing test-time sampling costs.

Original authors: Yaozhong Shi, Zachary E. Ross, Yisong Yue

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yaozhong Shi, Zachary E. Ross, Yisong Yue

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to reconstruct a beautiful, complex painting (like a weather map, a fluid flow, or a black hole image) based only on a few scattered, blurry dots of color that someone left on a canvas. You know the general style of the artist (the "prior"), and you have a few clues about what the final picture should look like (the "observations"), but the clues are noisy and incomplete.

This is the core challenge the paper addresses: How do we fill in the missing parts of a picture in a way that is both accurate and honest about our uncertainty?

The authors introduce a new method called FAPS (Flow Annealing Posterior Sampling). Here is how it works, broken down into simple concepts and analogies:

1. The Problem: The "Guessing Game"

In science, we often deal with "stochastic processes"—basically, rules that describe how things like weather, fluids, or magnetic fields behave.

The Old Way: Traditional methods (like Gaussian Processes) are great for simple, smooth things but struggle with complex, chaotic systems (like a storm or a black hole).
The New Way (Before FAPS): Recent AI methods learned to generate these complex pictures from scratch. However, when you give them a few clues, they often get stuck or produce blurry, unrealistic guesses. They struggle to say, "Here is one possibility, but here are five other equally likely possibilities."

2. The Solution: FAPS (The "Guided Sculptor")

FAPS is a new framework that combines two powerful ideas: a pre-trained artist and a guided sculpting process.

Step A: The Pre-Trained Artist (The "Flow-Matching Prior")

Imagine an artist who has spent years studying thousands of weather maps or fluid flows. They have learned the "grammar" of these systems. They know that if there is a swirl here, there is likely a vortex there.

In FAPS, this artist is a pre-trained AI model (called an Operator Flow Matching prior).
This model doesn't just guess; it knows how to smoothly morph a blank canvas (random noise) into a realistic-looking physical field. It has already learned the "shape" of the universe it's working in.

Step B: The Guided Sculpting (The "Annealing" Process)

Now, you give this artist a few noisy clues (the data). The artist needs to adjust their masterpiece to match these clues without ruining the realistic style they learned.

The Journey: FAPS doesn't try to fix the picture in one giant leap. Instead, it takes a slow, step-by-step journey (called "annealing").
The Process:
1. Transport: The artist starts with a random, messy sketch and uses their training to turn it into a clean, realistic-looking image.
2. Correction: The artist looks at the messy clues you gave them. They gently nudge the image to match the clues, but they do it carefully so they don't break the realistic style.
3. Re-bridging: They step back, look at the whole picture, and smooth out any weird jagged edges caused by the nudging.
4. Repeat: They do this over and over, getting closer and closer to the perfect balance between "what the data says" and "what physics says is possible."

3. The Secret Sauce: The "Low-Rank Map"

One of the paper's biggest innovations is how the artist handles the "nudges" (corrections).

The Problem: If you have a huge grid of data (like a 128x128 image) but only 128 clues, the clues are very sparse. If the artist tries to fix the image based on just one clue, they might accidentally mess up the whole picture.
The FAPS Fix: The artist uses a Low-Rank Covariance Preconditioner. Think of this as a special map that shows the artist which parts of the painting are connected.
- Analogy: If you touch the left side of a drum, the whole drum vibrates. The map tells the artist, "If you change this spot, you must also change these other spots because they are physically connected."
- This allows the method to spread the information from your few sparse clues across the entire image in a way that respects the physics, rather than just guessing randomly.

4. Why It's Better (The Results)

The paper tested FAPS on three types of challenges:

Simple Math (Gaussian Processes): It matched the "perfect" mathematical answer better than other AI methods.
Complex Science (Non-Gaussian): It successfully reconstructed complex fields like Navier-Stokes fluid flows, Black Hole images, and Global Climate patterns from very few data points. Other methods tended to make the images too blurry or lose the fine details; FAPS kept the sharp, realistic structures.
Inverse Problems (Solving Equations Backwards): In science, we often know the result (e.g., the temperature of a room) and want to find the cause (e.g., where the heater is). FAPS solved these "reverse engineering" problems faster and more accurately than previous diffusion-based methods.

5. The "Plug-and-Play" Advantage

Perhaps the most practical feature is that FAPS is modular.

You can train the "Artist" (the prior) once on a specific type of data (e.g., weather).
Then, you can use that same trained artist to solve different problems (e.g., predicting wind speed, or solving a heat equation) just by changing the "clues" you give it.
It doesn't need to be retrained for every new experiment. It's like having a master chef who can cook any dish if you just give them the ingredients, rather than hiring a new chef for every meal.

Summary

FAPS is a smart, efficient way to fill in the blanks of complex scientific pictures. It uses a pre-trained AI that knows the "rules of physics" to generate realistic possibilities, then carefully guides those possibilities to match your noisy, sparse data. It does this faster and with more accurate uncertainty estimates than previous methods, making it a powerful tool for scientists trying to understand the world from limited observations.

Technical Summary: Flow Annealing Posterior Sampling (FAPS)

Problem Statement

The paper addresses the challenge of principled regression for stochastic processes, a task central to scientific inverse problems, data assimilation, and uncertainty quantification. In these settings, the goal is to infer a posterior distribution over an unknown function or physical field given sparse and noisy observations. While classical Gaussian Process (GP) regression offers a Bayesian solution for Gaussian priors and simple observation models, many scientific processes are non-Gaussian, high-dimensional, and observed through indirect physical measurements (e.g., solutions to Partial Differential Equations, or PDEs).

Existing approaches face limitations:

Neural Processes learn amortized conditional predictors but often lack rigorous posterior consistency across varying query discretizations and do not explicitly separate reusable prior learning from test-time inference.
Neural Operators provide efficient PDE solution maps but are typically deterministic and do not yield posterior distributions from sparse data.
Diffusion-based Posterior Samplers (e.g., DAPS, FunDPS) combine learned priors with observation guidance but are often diffusion-based, computationally expensive, and primarily focused on specific PDE settings rather than unifying regression and inverse problems.

There is a need for a general posterior sampling framework that leverages pretrained function-space priors to perform likelihood-guided inference for both direct stochastic-process regression and indirect PDE inverse problems without explicit prior-density evaluation.

Methodology: Flow Annealing Posterior Sampling (FAPS)

FAPS is a function-space posterior sampling framework built on pretrained Operator Flow Matching (OFM) priors. It unifies stochastic-process regression and PDE inverse problems through a decoupled annealing procedure.

1. Function-Space Flow-Matching Prior

The method utilizes a pretrained OFM prior, which learns a probability flow transporting a reference Gaussian measure $\gamma = \mathcal{N}(0, C_0)$ to a target data distribution $\mu_1$ . This is achieved via a neural operator $v_\theta$ trained to approximate the conditional velocity field $v_t(u_t | u_0, u_1) = u_1 - u_0$ along a noisy straight interpolation path. The resulting flow map $\Phi_{s \to t}$ allows transporting samples from the reference measure to the target process.

2. Unified Posterior Formulation

FAPS treats both functional regression and PDE inverse problems under a single observation model:
$y = A(u_1) + \epsilon$
where $u_1$ is the unknown function, $A$ is the observation operator, and $\epsilon$ is noise.

For functional regression, $A$ is a direct point-evaluation or masking operator.
For PDE inverse problems, $A$ is a composition of a sparse observation operator and a pretrained neural-operator surrogate $G_\phi$ representing the forward PDE solution map.

3. The Annealing and Re-bridging Procedure

FAPS samples from the posterior $\pi_1^y$ by alternating between three steps across an annealing schedule $0 = t_0 < t_1 < \dots < t_K = 1$ :

Flow Transport (To Endpoint): Given a "noisy" state $u_{t_k}$ at annealing level $t_k$ , the method integrates the pretrained flow to the endpoint to obtain a "clean" anchor $\hat{u}_k = \Phi_{t_k \to 1}(u_{t_k})$ . This anchor serves as a prior-consistent estimate of the clean function.
Likelihood-Guided Langevin Correction: Around the anchor $\hat{u}_k$ , FAPS performs $L$ steps of Langevin dynamics to correct the sample based on the observation likelihood $p(y|u_1)$ . The update rule is:
$u_1^{(\ell+1)} = u_1^{(\ell)} + \eta \left[ -\frac{u_1^{(\ell)} - \hat{u}_k}{\lambda_k^2} + C_X \nabla_{u_1} \log p(y|u_1^{(\ell)}) \right] + \sqrt{2\eta} \zeta^{(\ell)}$
Crucially, the gradient is preconditioned by a low-rank covariance matrix $C_X$ , estimated offline from the target process. This preconditioner propagates sparse observation information along dominant function-space correlations, significantly improving updates in unobserved regions.
Re-bridging: The corrected clean sample is re-bridged to the next annealing level $t_{k+1}$ using the bridge kernel induced by the flow path, effectively adding noise back to the state to prepare for the next iteration.

This process repeats until $t_K = 1$ , yielding posterior samples without ever explicitly evaluating the prior density.

Key Contributions

First Unified Function-Space Flow-Matching Posterior Sampler: FAPS is presented as the first framework to unify stochastic-process regression and PDE inverse problems using pretrained function-space flow-matching priors. It provides a common likelihood-guided sampler for both direct and indirect observations.
Low-Rank Covariance-Preconditioned Langevin Correction: The method introduces a low-rank covariance preconditioner for the Langevin correction step. This exploits dominant correlations in the function space, significantly improving posterior updates under sparse observations compared to unpreconditioned methods.
Flexible and Computationally Efficient Inference: FAPS is "plug-and-play," leveraging pretrained priors and forward surrogates. It supports variable query discretizations and noisy observations without retraining. It avoids explicit prior-density evaluation, making it substantially more memory- and time-efficient than likelihood-based sampling with flow priors or diffusion-based samplers.
Strong Empirical Performance: Extensive experiments demonstrate that FAPS achieves state-of-the-art performance in functional regression (outperforming Neural Process variants) and competitive or superior performance in PDE inverse problems compared to diffusion-based baselines, with lower test-time sampling costs.

Experimental Results

The paper evaluates FAPS across three settings:

Gaussian Process Regression: On 1D Matérn and Gibbs-kernel benchmarks, FAPS closely matches reference posterior means and uncertainties across varying query resolutions (128, 512, 1024). It outperforms Neural Processes (CNP, ANP, NDP, FlowNP) and direct OFM sampling in Sliced Wasserstein Distance (SWD) and Maximum Mean Discrepancy (MMD), particularly at higher resolutions where baselines degrade.
Non-Gaussian Functional Regression: On Navier-Stokes, Black Hole imaging, and Global Climate datasets, FAPS reconstructs dominant spatial structures and produces coherent posterior samples. It achieves the best or near-best calibration (SSR) and reconstruction metrics (CRPS, PSNR, SSIM) compared to baselines, which often over-smooth or underestimate uncertainty.
PDE Inverse Problems: On Darcy flow, Poisson, Helmholtz, and Navier-Stokes inverse benchmarks (128x128 resolution, 128 noisy observations), FAPS (both FNO and UNet variants) achieves competitive or best results in CRPS, PSNR, and SSIM compared to DiffusionPDE, FunDPS, and DDIS.
- Efficiency: FAPS is significantly faster than diffusion-based baselines (e.g., 1.73x faster than DiffusionPDE, 2.37x faster than DDIS) and requires substantially less GPU memory (avoiding the high memory overhead of explicit likelihood evaluation in flow priors or large diffusion step counts).
- Zero-Shot Super-Resolution: FAPS demonstrates the ability to perform inverse inference on query meshes (e.g., 160x160) not seen during prior training, a capability enabled by the function-space nature of the OFM prior.

Significance and Claims

The authors claim that FAPS establishes a new paradigm for uncertainty-aware inference in scientific machine learning. By converting pretrained function-space flow-matching priors into likelihood-guided posterior samplers, FAPS offers a reusable, flexible, and computationally efficient alternative to existing methods.

The paper modestly acknowledges that the practical sampler relies on an approximate local Gaussian correction and that a full non-asymptotic convergence theory for the practical sampler remains future work. However, the authors provide empirical evidence and error propagation analysis suggesting that local transition errors accumulate additively and that the method converges to the target posterior as the annealing mesh and sampling parameters are refined.

Ultimately, the work suggests that function-space flow-matching priors can serve as a robust foundation for scalable Bayesian inference in complex scientific domains, unifying regression and inverse modeling under a single, efficient framework.