Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

Guided Diffusion Sampling on Function Spaces with Applications to PDEs

About

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

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

Related benchmarks

TaskDatasetResultRank
Inverse ProblemPoisson
PDE Residual429.9
21
Inverse PDE solvingHelmholtz full observations
Relative Error0.1388
14
PDE Forward ProblemDarcy Flow Noisy
L2 Relative Error8.5
12
PDE Forward ProblemPoisson Noisy
L2 Relative Error10.63
12
PDE Inverse ProblemDarcy Flow Noisy
Error Rate31.34
12
Spatio-temporal PDE reconstructionHelmholtz
PDE Residual4.60e+3
12
Shape InpaintingMNIST contours (test)
RMSE0.64
12
PDE ReconstructionDarcy Flow 3% uniformly random measurements
Coefficient Rel L2 Error5.18
10
PDE ReconstructionHelmholtz 3% uniformly random measurements
Coefficient Relative L2 Error17.16
10
PDE ReconstructionNavier-Stokes 3% uniformly random measurements
Coefficient Relative L2 Error8.48
10
Showing 10 of 64 rows

Other info

Follow for update