[Paper] Decoupled Diffusion Sampling for Inverse Problems on Function Spaces

Source: arXiv - 2601.23280v1

Overview

A new generative framework called Decoupled Diffusion Inverse Solver (DDIS) tackles inverse problems for partial differential equations (PDEs) directly in function space. By separating the learning of the unknown coefficient distribution from the physics of the forward PDE, DDIS dramatically cuts the amount of paired training data needed while delivering sharper reconstructions than prior diffusion‑based solvers.

Key Contributions

• Decoupled architecture – an unconditional diffusion model learns a prior over the PDE coefficient, while a neural operator (e.g., Fourier Neural Operator) encodes the forward PDE for physics‑aware guidance.
• Data‑efficient learning – provably avoids the “guidance attenuation” problem that plagues joint coefficient‑solution diffusion models when training data are scarce.
• Decoupled Annealing Posterior Sampling (DAPS) – a novel sampling schedule that prevents the over‑smoothing typical of standard Diffusion Posterior Sampling (DPS).
• State‑of‑the‑art performance – on benchmark inverse PDE tasks, DDIS reduces ℓ₂ reconstruction error by ~11 % and spectral error by ~54 % on average; with only 1 % of paired data it still outperforms joint models by ~40 % in ℓ₂ error.
• Theoretical guarantees – the paper includes a proof that the decoupled design yields a non‑vanishing guidance term even under extreme data sparsity.

Methodology

1. Unconditional diffusion prior – Train a standard diffusion model on a large collection of coefficient fields (e.g., conductivity maps) without any observation information. This learns a rich, high‑dimensional prior in function space.
2. Neural operator forward model – Independently train a neural operator (such as a Fourier Neural Operator or DeepONet) to map a coefficient field to its PDE solution. This model is physics‑aware but does not need paired coefficient‑solution data; it can be trained on synthetic solves.
3. Guided sampling – During inference, start from a noisy coefficient sample drawn from the diffusion prior. At each diffusion step, the neural operator evaluates the forward PDE and compares the simulated observation to the real measurement. The resulting residual is fed back as a guidance term that nudges the diffusion trajectory toward coefficient fields consistent with the observed data.
4. Decoupled Annealing Posterior Sampling (DAPS) – Instead of using a fixed noise schedule (as in DPS), DAPS gradually reduces the strength of the guidance term, allowing the sampler to first explore the prior manifold and then fine‑tune to the data. This mitigates the “over‑smoothing” where the posterior collapses to a blurry average.

The overall pipeline requires only a modest set of real paired observations (coefficient + measurement) for the guidance step; the heavy lifting is done by the unconditional diffusion prior and the neural operator, both of which can be trained on abundant synthetic data.

Results & Findings

• Robustness to noise – DDIS maintains high fidelity even when observations are corrupted, thanks to the physics‑guided correction at each diffusion step.
• Generalization – Because the diffusion prior is learned unconditionally, DDIS can be re‑used across multiple inverse tasks (different observation operators) without retraining the prior.
• Ablation – Removing DAPS leads to noticeable over‑smoothing; removing the neural operator guidance collapses performance to that of a pure prior sampler.

Practical Implications

• Reduced data collection costs – Engineers can train the heavy‑weight diffusion prior on cheap simulated coefficient fields, needing only a handful of real measurements to calibrate the solver.
• Plug‑and‑play physics integration – The neural operator can be swapped out for any differentiable PDE solver (e.g., finite‑element, spectral), making the approach adaptable to fluid dynamics, electromagnetics, material design, etc.
• Fast prototyping – Once the prior and operator are trained, inference is a few hundred diffusion steps, comparable to existing diffusion‑based generative models and easily parallelizable on GPUs.
• Improved design loops – Inverse design pipelines (e.g., topology optimization, metamaterial synthesis) can benefit from sharper reconstructions, leading to fewer design iterations and tighter performance margins.

Limitations & Future Work

• Scalability to very high‑dimensional domains – While the method works well on 2‑D and modest 3‑D grids, diffusion models in extremely high‑resolution function spaces can become memory‑intensive.
• Operator accuracy dependence – The quality of the neural operator directly impacts guidance; errors in the forward model may bias the posterior.
• Limited to PDEs with known forward operator – The approach assumes an analytically or numerically tractable forward map; extending to black‑box simulators remains an open challenge.
• Future directions – The authors suggest exploring hierarchical diffusion priors for multi‑scale PDEs, integrating uncertainty quantification into the guidance term, and applying DDIS to time‑dependent inverse problems (e.g., seismic imaging).

Authors

• Thomas Y. L. Lin
• Jiachen Yao
• Lufang Chiang
• Julius Berner
• Anima Anandkumar

Paper Information

• arXiv ID: 2601.23280v1
• Categories: cs.LG, math.NA
• Published: January 30, 2026
• PDF: Download PDF