Guided Diffusion Sampling on Function Spaces with Applications to PDEs

TL;DR

This paper introduces FunDPS, a diffusion-based framework for solving PDE inverse problems by sampling in function spaces, achieving high accuracy with minimal data and no discretization dependence.

Contribution

It develops a discretization-agnostic neural operator diffusion model with a novel guidance mechanism and extends Tweedie's formula to infinite-dimensional spaces for PDE inverse problems.

Findings

32% accuracy improvement over baselines with 3% observations

Reduces sampling steps by 4x compared to existing methods

First discretization-independent diffusion framework for PDEs

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 Banach 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