Table of Contents
Fetching ...

Decoupled Diffusion Sampling for Inverse Problems on Function Spaces

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

TL;DR

This paper tackles inverse PDE problems in function spaces under sparse observations by decoupling prior modeling from physics-based likelihoods. It introduces Decoupled Diffusion Inverse Solver (DDIS), which learns a coefficient-space diffusion prior $p(a)$ with unpaired data and trains a neural operator $L_{\phi}$ to surrogate the forward map $L(a)$ from limited paired data. Inference leverages Decoupled Annealing Posterior Sampling (DAPS) to apply physics-guided corrections via the operator surrogate, addressing attenuation and sparse-guidance failures seen in joint-embedding diffusion models. Theoretical analysis shows that joint embeddings suffer guidance attenuation under data scarcity, while DDIS maintains robust guidance and achieves state-of-the-art performance on inverse Poisson, Helmholtz, and Navier–Stokes problems, particularly with scarce data and under resolution mismatches. The work has practical implications for scalable, physics-informed inverse problems where paired data are expensive and sensor layouts are sparse.

Abstract

We propose a data-efficient, physics-aware generative framework in function space for inverse PDE problems. Existing plug-and-play diffusion posterior samplers represent physics implicitly through joint coefficient-solution modeling, requiring substantial paired supervision. In contrast, our Decoupled Diffusion Inverse Solver (DDIS) employs a decoupled design: an unconditional diffusion learns the coefficient prior, while a neural operator explicitly models the forward PDE for guidance. This decoupling enables superior data efficiency and effective physics-informed learning, while naturally supporting Decoupled Annealing Posterior Sampling (DAPS) to avoid over-smoothing in Diffusion Posterior Sampling (DPS). Theoretically, we prove that DDIS avoids the guidance attenuation failure of joint models when training data is scarce. Empirically, DDIS achieves state-of-the-art performance under sparse observation, improving $l_2$ error by 11% and spectral error by 54% on average; when data is limited to 1%, DDIS maintains accuracy with 40% advantage in $l_2$ error compared to joint models.

Decoupled Diffusion Sampling for Inverse Problems on Function Spaces

TL;DR

This paper tackles inverse PDE problems in function spaces under sparse observations by decoupling prior modeling from physics-based likelihoods. It introduces Decoupled Diffusion Inverse Solver (DDIS), which learns a coefficient-space diffusion prior with unpaired data and trains a neural operator to surrogate the forward map from limited paired data. Inference leverages Decoupled Annealing Posterior Sampling (DAPS) to apply physics-guided corrections via the operator surrogate, addressing attenuation and sparse-guidance failures seen in joint-embedding diffusion models. Theoretical analysis shows that joint embeddings suffer guidance attenuation under data scarcity, while DDIS maintains robust guidance and achieves state-of-the-art performance on inverse Poisson, Helmholtz, and Navier–Stokes problems, particularly with scarce data and under resolution mismatches. The work has practical implications for scalable, physics-informed inverse problems where paired data are expensive and sensor layouts are sparse.

Abstract

We propose a data-efficient, physics-aware generative framework in function space for inverse PDE problems. Existing plug-and-play diffusion posterior samplers represent physics implicitly through joint coefficient-solution modeling, requiring substantial paired supervision. In contrast, our Decoupled Diffusion Inverse Solver (DDIS) employs a decoupled design: an unconditional diffusion learns the coefficient prior, while a neural operator explicitly models the forward PDE for guidance. This decoupling enables superior data efficiency and effective physics-informed learning, while naturally supporting Decoupled Annealing Posterior Sampling (DAPS) to avoid over-smoothing in Diffusion Posterior Sampling (DPS). Theoretically, we prove that DDIS avoids the guidance attenuation failure of joint models when training data is scarce. Empirically, DDIS achieves state-of-the-art performance under sparse observation, improving error by 11% and spectral error by 54% on average; when data is limited to 1%, DDIS maintains accuracy with 40% advantage in error compared to joint models.
Paper Structure (78 sections, 29 theorems, 117 equations, 20 figures, 11 tables, 1 algorithm)

This paper contains 78 sections, 29 theorems, 117 equations, 20 figures, 11 tables, 1 algorithm.

Key Result

Theorem 4.1

Fix $(x,t)$. If there exists $k\in[N]$ such that then all responsibility gradients vanish:

Figures (20)

  • Figure 1: Gradient guidance by joint-embedding models vanishes under scarce paired data. Blue blobs visualize regions supported by the learned joint model around individual training examples (black dots), and the circle marker denotes diffusion state $x_t$. We consider three cases: when $x_t$ is far from all blobs (i), near a single training sample (ii), or close to multiple training samples (iii). In all cases, the update on $u$ (dashed arrows) remains valid; however, the update on $a$ (solid arrows) vanishes in (i) and (ii), and is nonzero only in (iii). In high dimensions with limited data, case (iii) is rare, rendering coefficient-space guidance ineffective.
  • Figure 2: Training comparison. Left (Gray): Joint-embedding methods rely on paired data $(a, u)$ to learn the joint distribution $p(a,u)$. Right (Cyan): DDIS decouples the architecture: the diffusion prior learns $p(a)$ utilizing abundant unpaired coefficients, while the neural operator takes paired data $(a, u)$ to directly learn the forward physics map.
  • Figure 3: Comparison of FunDPS and our DDIS on inverse Poisson problem reconstruction under 1% paired data scarcity. FunDPS suffers from over-smoothing due to Jensen’s gap while DDIS achieves sharp and dense guidance with improved accuracy.
  • Figure 4: DDIS sampling process. Each annealing step alternates between ① reverse diffusion, which estimates an unconditional sample $\widehat{a}_0(a_t)$ by diffusion model; ② Langevin dynamics guided by the neural operator $L_\phi$ to enforce physics consistency; and ③ forward diffusion, which re-injects noise for the next annealing level. The process iteratively refines the posterior sample $a_0 \mid u_{\mathrm{obs}}$.
  • Figure 5: Pareto frontiers for the inverse Poisson problem. (Left) Relative $\ell_2$ error versus inference time. (Right) Relative $\ell_2$ error versus training set size. DDIS dominates baselines across both computational and data efficiency dimensions, with the physics-informed variant (DDIS+Phys) further enhances sample efficiency.
  • ...and 15 more figures

Theorems & Definitions (51)

  • Definition 2.1: Target Posterior
  • Theorem 4.1: Local dominance $\Rightarrow$ vanishing responsibility gradients (informal \ref{['thm:local-dominance-gradw']})
  • Corollary 4.1.1: Locality implies attenuation (informal \ref{['cor:local-dominance-attenuation']})
  • Theorem 4.2: Non-vanishing responsibility gradients require overlap (informal \ref{['thm:dw-implies-overlap']})
  • Corollary 4.2.1: Non-vanishing guidance requires overlap (informal \ref{['cor:nonvanishing-requires-overlap']})
  • Proposition 4.1: Structural robustness of DDIS guidance
  • Theorem 4.3: Sparse constraint induces correlation shrinkage (informal \ref{['thm:constrained_terms']})
  • Corollary 4.3.1: Strong sparse guidance collapses correlations
  • Remark 4.1: Jensen-gap Approximation of DPS.
  • Remark 5.1: Sparse-Guidance Failure of DAPS
  • ...and 41 more