Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Guided Diffusion Sampling on Function Spaces with Applications to PDEs

Jiachen Yao, Abbas Mammadov, Julius Berner, Gavin Kerrigan, Jong Chul Ye, Kamyar Azizzadenesheli, Anima Anandkumar

TL;DR

FunDPS introduces a function-space diffusion posterior sampling framework for PDE-based inverse problems, enabling discretization-agnostic, posterior-consistent reconstruction from highly sparse or noisy measurements. The method couples an unconditional diffusion prior on function spaces (trained via score matching with Gaussian random-field noise and neural-operator denoisers) with test-time, plug-and-play likelihood guidance derived from an infinite-dimensional Tweedie’s formula, yielding a practical gradient-based conditioning mechanism. A key theoretical contribution is the generalization of Tweedie’s formula to Banach spaces, linking conditional expectations to the Fréchet score and enabling efficient posterior sampling in infinite dimensions. The approach is complemented by a multi-resolution training/inference pipeline (including ReNoise) that reduces computation while preserving accuracy, and is validated across five PDE tasks where it surpasses fixed-resolution diffusion baselines by about 32% on average and achieves substantial speedups, including up to 25× faster wall-clock times in certain settings. This work provides a versatile, discretization-agnostic framework for forward and inverse PDE problems with potential broad impact on scientific computing and uncertainty-aware simulations.

Abstract

We propose a general framework for conditional sampling in PDE-based inverse problems, targeting the recovery of whole solutions from extremely sparse or noisy measurements. This is accomplished by a function-space diffusion model and plug-and-play guidance for conditioning. Our method first trains an unconditional discretization-agnostic denoising model using neural operator architectures. At inference, we refine the samples to satisfy sparse observation data via a gradient-based guidance mechanism. Through rigorous mathematical analysis, we extend Tweedie's formula to infinite-dimensional Hilbert spaces, providing the theoretical foundation for our posterior sampling approach. Our method (FunDPS) accurately captures posterior distributions in function spaces under minimal supervision and severe data scarcity. Across five PDE tasks with only 3% observation, our method achieves an average 32% accuracy improvement over state-of-the-art fixed-resolution diffusion baselines while reducing sampling steps by 4x. Furthermore, multi-resolution fine-tuning ensures strong cross-resolution generalizability. To the best of our knowledge, this is the first diffusion-based framework to operate independently of discretization, offering a practical and flexible solution for forward and inverse problems in the context of PDEs. Code is available at https://github.com/neuraloperator/FunDPS

Guided Diffusion Sampling on Function Spaces with Applications to PDEs

TL;DR

FunDPS introduces a function-space diffusion posterior sampling framework for PDE-based inverse problems, enabling discretization-agnostic, posterior-consistent reconstruction from highly sparse or noisy measurements. The method couples an unconditional diffusion prior on function spaces (trained via score matching with Gaussian random-field noise and neural-operator denoisers) with test-time, plug-and-play likelihood guidance derived from an infinite-dimensional Tweedie’s formula, yielding a practical gradient-based conditioning mechanism. A key theoretical contribution is the generalization of Tweedie’s formula to Banach spaces, linking conditional expectations to the Fréchet score and enabling efficient posterior sampling in infinite dimensions. The approach is complemented by a multi-resolution training/inference pipeline (including ReNoise) that reduces computation while preserving accuracy, and is validated across five PDE tasks where it surpasses fixed-resolution diffusion baselines by about 32% on average and achieves substantial speedups, including up to 25× faster wall-clock times in certain settings. This work provides a versatile, discretization-agnostic framework for forward and inverse PDE problems with potential broad impact on scientific computing and uncertainty-aware simulations.

Abstract

We propose a general framework for conditional sampling in PDE-based inverse problems, targeting the recovery of whole solutions from extremely sparse or noisy measurements. This is accomplished by a function-space diffusion model and plug-and-play guidance for conditioning. Our method first trains an unconditional discretization-agnostic denoising model using neural operator architectures. At inference, we refine the samples to satisfy sparse observation data via a gradient-based guidance mechanism. Through rigorous mathematical analysis, we extend Tweedie's formula to infinite-dimensional Hilbert spaces, providing the theoretical foundation for our posterior sampling approach. Our method (FunDPS) accurately captures posterior distributions in function spaces under minimal supervision and severe data scarcity. Across five PDE tasks with only 3% observation, our method achieves an average 32% accuracy improvement over state-of-the-art fixed-resolution diffusion baselines while reducing sampling steps by 4x. Furthermore, multi-resolution fine-tuning ensures strong cross-resolution generalizability. To the best of our knowledge, this is the first diffusion-based framework to operate independently of discretization, offering a practical and flexible solution for forward and inverse problems in the context of PDEs. Code is available at https://github.com/neuraloperator/FunDPS

Paper Structure

This paper contains 71 sections, 6 theorems, 57 equations, 12 figures, 13 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Bayesian approach to conditional sampling
  3. Diffusion Models in Function Spaces
  4. Contributions
  5. Related Works
  6. Neural PDE solvers
  7. Inverse PDE problems
  8. Diffusion-based posterior sampling
  9. Generative neural operators
  10. Method
  11. Problem Settings
  12. PDE-based inverse problems
  13. Bayesian perspective
  14. Diffusion prior
  15. Conditional Score via Likelihood and Prior Score
  16. ...and 56 more sections

Key Result

Theorem 3.1

Let $B$ be a separable Banach space. Assume that $\mu(H(\gamma)) = 1$, the score of $\nu$ is Fréchet differentiable along $H(\gamma)$, and that the Fréchet derivatives of $\, \mathrm{d} \gamma^x / \, \mathrm{d} \gamma$ are $\mu$-almost surely bounded by a $\mu$-integrable function. Then, for $\nu$-a

Figures (12)

  • Figure 1: Comparison between our method (FunDPS) with the state-of-the-art diffusion baseline, DiffusionPDE, on the Darcy Flow and Helmholtz problems. The left column shows the sparse observation measurements (3% of total points), while the other two columns display the absolute reconstruction error of our method and DiffusionPDE, respectively. FunDPS achieves superior accuracy with an order of magnitude fewer sampling steps.
  • Figure 2: The sampling and training pipelines of FunDPS. During inference, we utilize a standard reverse diffusion approach with additional FunDPS guidance to drag the samples to the posterior. During training, the model is based on a U-shaped neural operator, and Gaussian random fields are used as the noise sampler to ensure consistency within function spaces. Notations are detailed in blue box in the bottom-left corner, where function ${\boldsymbol a}$ jointly represents the PDE parameters and solution.
  • Figure 3: (a) Comparison of FunDPS and DiffusionPDE in terms of accuracy and inference time with varying step sizes; (b) Demonstration of multi-resolution inference pipeline. $\sigma$ is in log-scale.
  • Figure 4: Comparison of the accuracy of FunDPS with ReNoise under different ratios of low-resolution inference steps.
  • Figure 5: Reconstruction of coefficient functions from partially-observed Darcy Flow problems.
  • ...and 7 more figures

Theorems & Definitions (11)

  • Theorem 3.1: Tweedie's formula in infinite-dimensional Banach spaces
  • proof
  • Theorem C.1: Cameron-Martin
  • Proposition C.2
  • proof
  • Lemma C.3
  • proof
  • Proposition C.4
  • proof
  • Proposition C.5
  • ...and 1 more