PaperlayerView exact revision on arXiv

Source-linked AI summary

Decoupled Diffusion Sampling for Inverse Problems on Function Spaces

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

arXiv:2601.23280v3cs.LGmath.NA

TL;DR

Inverse PDE problems require posterior inference from sparse observations while paired coefficient-solution data is limited. DDIS decouples diffusion-based coefficient-prior learning from neural-operator modeling of the forward PDE and uses DAPS for physics-aware sampling. It achieves state-of-the-art performance, improving average ℓ2 error by 11% and spectral error by 54%, with a 40% ℓ2 advantage over joint models at 1% paired data.

  • Problem

    Inverse PDE solvers must infer coefficient fields from sparse observations, but joint-embedding approaches require scarce paired data and can lose effective guidance under data scarcity.

  • Method

    DDIS learns a diffusion prior over coefficients, models the forward PDE with a neural operator, and combines them during DAPS-based posterior sampling.

  • Results

    DDIS achieves state-of-the-art performance, improving average ℓ2 error by 11% and spectral error by 54%; under scarce paired data, its advantage reaches 40%.

  • Takeaways & Limitations

    DDIS remains stable down to 1% paired data, while joint-embedding methods degrade under scarce supervision.

  • Takeaways & Limitations

    Physics loss is not applicable for Navier-Stokes because of insufficient information.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Inverse PDE problems infer unknown coefficient fields from sparse or noisy observations, but paired coefficient-solution data is costly and joint-embedding models can lose effective cross-field guidance under scarcity. DDIS addresses this by separating coefficient-prior learning from explicit neural-operator physics guidance.

  • Motivation: Inverse PDE problems infer unknown coefficient fields from partial or noisy solution observations through an often ill-posed forward process.Sparse sensor coverage makes observations available only over a small fraction of the spatial domain.
  • Motivation: Paired coefficient-solution data is scarce because collecting it requires repeatedly solving the underlying PDE.This creates an imbalanced regime with abundant coefficients but limited paired samples.
  • Existing limitations: Joint-embedding methods learn p(a, u) from paired data and must recover physics through statistical cross-field correlations.They effectively reduce inverse PDE solving to inpainting in the joint space.
  • Existing limitations: Under data scarcity, joint models suffer guidance attenuation because effective coefficient guidance requires overlapping local support from multiple training components.The required overlap becomes difficult to satisfy in high-dimensional spaces.
  • DDIS: DDIS separates coefficient-prior learning from physics-induced likelihood evaluation, using a diffusion prior and a neural operator for forward-PDE guidance.This decoupling supports dense guidance and DAPS-based sampling under sparse observations.
  • Results: DDIS improves average ℓ2 error by 11% and spectral error by 54%, while its advantage reaches 40% under data scarcity.The reported results cover sparse supervision across challenging inverse PDE problems.

2 POSTERIOR SAMPLING FOR INVERSE PROBLEMS

Posterior sampling combines a learned coefficient prior with an observation likelihood, but direct posterior-score computation is intractable. DPS approximates the likelihood through a denoised estimate, whereas DAPS applies correction to the clean estimate before re-noising and can improve reconstruction quality.

  • Posterior formulation: The posterior satisfies p(a0 | uobs) ∝ p(a0) p(uobs | a0), combining a coefficient prior with an observation likelihood.Posterior sampling modifies the diffusion prior score using observation-dependent guidance.
  • Posterior formulation: Direct computation of the posterior score is intractable because it depends on the unknown conditional p(a0 | at).Practical samplers therefore use approximations or auxiliary latent transitions.
  • DPS: DPS approximates p(uobs | at) using E[a0 | at], but the resulting Jensen gap produces over-smoothed reconstructions.The approximation is needed because the likelihood requires marginalization over the unknown conditional.
  • DAPS: DAPS updates the noisy latent by first estimating the clean variable, applying likelihood correction there, and then re-noising it.This separates the likelihood-defined clean variable from the noisy diffusion latent.
  • DAPS: Under sparse observations and joint embeddings, DAPS likelihood gradients act only locally and provide ineffective updates.This is identified as a sparse-guidance failure of joint-embedding implementations.

3 DECOUPLED DIFFUSION INVERSE SOLVER (DDIS)

DDIS combines a function-space diffusion prior, a neural-operator surrogate for the forward PDE, and DAPS-based posterior sampling. The neural operator supplies dense physics-aware guidance from sparse observations while preserving the prior’s role in coefficient space.

  • Architecture: DDIS learns the coefficient prior, forward physics, and posterior sampler as three decoupled components.The components are diffusion prior learning, neural-operator learning, and DAPS-based physics-aware sampling.
  • Diffusion prior: The diffusion model learns p(a) over PDE coefficients without paired supervision.Training uses a score-based diffusion model on noisy coefficient samples.
  • Neural operator: A neural operator Lϕ learns the forward map L : A → U from scarce paired samples, with optional PDE-residual regularization.The surrogate is trained by supervised regression and can incorporate a physics regularizer.
  • Neural operator: Operating in function space makes the neural operator resolution-invariant and compatible with low- or multi-resolution paired data and high-resolution inference.Physics-residual regularization can further reduce reliance on paired supervision.
  • Sampling: At each annealing step, DDIS alternates reverse diffusion, neural-operator-guided Langevin dynamics, and forward diffusion re-noising.This iterative process refines samples toward p(a0 | uobs).
  • Sampling: The neural operator’s global receptive field propagates sparse observation errors across coefficient space, providing dense guidance without retraining for new observation patterns.This addresses the sparse-guidance failure of prior joint-embedding models.

4 THEORETICAL ANALYSIS

The analysis identifies guidance attenuation in joint-embedding diffusion models under data scarcity and sparse observations, then contrasts it with DDIS's operator-based guidance. It also analyzes sparse-guidance failure in joint-model DAPS and the resulting covariance collapse.

  • Guidance attenuation in joint-embedding models: Joint-embedding guidance depends on learned cross-field coupling between coefficients and solutions, so observations influence coefficients only through cross-partial score terms.The coefficient update depends on the cross-partial ∂a t sθ,u, while single isotropic Gaussian components cannot couple a and u.
  • Guidance attenuation in joint-embedding models: Under data scarcity, joint-model guidance attenuates when the diffusion state is dominated by one mixture component or lies outside mixture overlap.Non-vanishing coefficient guidance requires an overlap region with at least two mixture components having comparable responsibility.
  • Guidance robustness in DDIS: The DDIS formulation applies likelihood corrections to clean coefficient estimates using a neural-operator surrogate before re-noising for posterior sampling.This decoupled construction supports DAPS while using operator-mediated guidance rather than learned joint correlations.
  • Guidance robustness in DDIS: DDIS avoids this data-dependent attenuation mechanism because its likelihood guidance is governed by the explicit Jacobian of a neural forward operator.The operator is trained by regression independently of joint data density, so scarce data affects approximation accuracy but does not induce guidance attenuation.
  • Joint embeddings and DAPS: Sparse observations in joint-model DAPS produce discontinuous updates that alter covariance structure and can push samples off the diffusion model's learned manifold.The constrained location becomes nearly independent of neighboring locations, degrading reconstruction quality.
  • Empirical evaluation: Table 3 compares relative ℓ2 error across 100%, 5%, and 1% paired-data supervision, while Table 2 reports standard-supervision accuracy and runtime under 3% sparse observations.These evaluations connect the theoretical guidance analysis to scarce-data and sparse-observation settings.

5 EXPERIMENTS

Experiments evaluate DDIS on inverse Helmholtz, Poisson, and Navier–Stokes problems under sparse observations, varying supervision, resolution, inference time, and sampling strategies. DDIS provides strong accuracy–efficiency trade-offs, retains accuracy with scarce paired data, and remains robust across resolutions.

  • Experimental setup: DDIS is evaluated on inverse Helmholtz, Poisson, and Navier–Stokes problems using sparse observations and relative ℓ2 and spectral energy errors.Each instance uses 500 randomly sampled solution points, approximately 3% of the domain.
  • Experimental setup: DDIS is compared with DiffusionPDE, FunDPS, ECI-sampling, OFM, FunDAPS, and DecoupledDPS across accuracy, runtime, supervision, and sampling settings.The ablations isolate the effects of posterior sampling and the decoupled design.
  • Standard supervision: DDIS remains near the accuracy–runtime Pareto frontier and achieves spectral error up to 3.2× lower than FunDPS across tasks.The advantage persists on Navier–Stokes, where ℓ2 differences are small, because DDIS better preserves high-frequency features.
  • Scarce paired-data supervision: At 1% paired-data supervision, DDIS maintains stable reconstruction accuracy while FunDPS degrades sharply.DecoupledDPS improves over FunDPS, but DDIS achieves larger gains through the additional benefit of DAPS enabled by decoupling.
  • Low-/multi-resolution supervision: Training on low- or mixed-resolution data causes only modest degradation relative to full-resolution training.The reported robustness is attributed to the resolution-invariant neural operator.

6 CONCLUSION AND DISCUSSION

The discussion presents DDIS as a decoupled framework that separates coefficient-prior learning from physics-based likelihood evaluation. It links this design to more reliable guidance and improved behavior under sparse observations and limited paired data, while contrasting it with limitations of alternative approaches.

  • DDIS design: DDIS separates coefficient-space prior modeling from physics-based likelihood evaluation through a neural operator.The framework combines diffusion prior learning, forward-physics operator learning, and DAPS-based posterior sampling.
  • Theoretical and empirical findings: Joint-embedding models lose effective cross-field guidance under data scarcity or sparse sensor layouts, whereas DDIS provides reliable neural-operator guidance.The analysis characterizes guidance attenuation geometrically and identifies sparse-observation failures with DAPS-based sampling.
  • Theoretical and empirical findings: DDIS creates accuracy–runtime Pareto frontiers, lowers spectral error, remains stable to 1% paired data, and is robust to low- and mixed-resolution supervision.These outcomes are reported across the paper’s empirical evaluations.
  • Alternative approaches: Conditional diffusion models require retraining when the observation mask or forward operator changes and require large paired datasets covering measurement configurations.DDIS instead learns the prior unconditionally from unpaired data and can generalize to observation masks and operators at inference time without retraining.
  • Alternative approaches: Direct numerical-solver guidance is constrained by unstable backpropagation, solver-stability violations, and ill-conditioning under sparse observations.DDIS uses learned neural operators to provide differentiable, stable guidance and rapid likelihood evaluation.
  • Alternative approaches: DAPS avoids the approximation issue in DPS, but its sparse-observation use with joint embeddings remains ineffective because likelihood gradients act only locally.DDIS uses DAPS after decoupling physics from prior modeling.

A.5 CURRENT CHALLENGES IN FLOW-BASED POSTERIOR SAMPLING

The paper identifies structural and computational limitations in flow-based and other posterior-sampling alternatives under sparse observations and nonlinear PDEs. Across figures, DDIS is associated with stronger spectral fidelity and better accuracy–speed trade-offs.

  • Flow-based posterior sampling: Flow-based joint-embedding extensions such as ECI-sampling and OFM exhibit reconstruction or computational limitations under sparse observations.The paper attributes these limitations to structural problems in joint-embedding inference.
  • Flow-based posterior sampling: ECI-sampling can drift from the joint manifold because heuristic observation replacement lacks coefficient-channel likelihood guidance and manifold projection.The paper also reports limited robustness to discontinuities.
  • Flow-based posterior sampling: OFM regression requires many Langevin iterations and backpropagation through the full trajectory, creating substantial computational and memory overhead.The paper characterizes this cost as impractical for the setting considered.
  • Comparative results: DDIS achieves superior accuracy–speed trade-offs over joint-embedding baselines across three PDEs and reduces Navier–Stokes spectral error by 60% relative to FunDPS.The ℓ2 advantage on Navier–Stokes is modest, but the spectral improvement yields higher-quality reconstructions.
  • Comparative results: Power-spectrum comparisons assess predicted spectral density against ground truth across wave number for Poisson, Helmholtz, and Navier–Stokes problems.The Navier–Stokes comparison highlights FunDPS’s high-wave-number gap and DDIS’s ability to capture those components.

B.3 MIXED-RESOLUTION TRAINING

This section situates mixed-resolution and decoupled posterior sampling within diffusion-based PDE inverse problems. It emphasizes resolution flexibility, sparse-observation challenges, and the distinction between decoupled and joint approaches.

  • Mixed-resolution training uses full 642 data plus an additional 10% of 1282 data while maintaining accuracy and reducing training costs.
  • Under 3% sparse observations, DDIS (Mixed) uses coarse-to-fine sampling and achieves comparable accuracy with 30% lower runtime.
  • Sparse observations challenge posterior samplers because solver-based updates can violate PDE stability conditions and become spatially discontinuous.
  • Joint-embedding diffusion methods learn coefficient-solution correlations from paired data and effectively recast inverse PDE solving as joint-space inpainting.
  • DecoupledDPS improves over joint embeddings but remains inferior to DDIS under both 100% and 1% paired-data settings.

D.3.1 ASYMPTOTIC GUARANTEE FOR DAPS

The analysis establishes that ideal DAPS preserves observation-conditioned time marginals through its two-phase update and converges asymptotically to the target posterior. DDIS transports likelihood information from solution space to coefficient space using a surrogate forward operator.

  • DDIS uses a neural operator Lϕ to transport likelihood information from solution space into coefficient-space sampling and provide dense guidance.
  • Ideal DAPS first samples a clean coefficient from the posterior conditional and then re-noises it to the preceding diffusion level.
  • If at follows the observation-conditioned marginal, the ideal DAPS update yields the corresponding marginal at the next noise level.
  • As t approaches 0, the DAPS state converges in distribution to the target posterior p(a0 | uobs).
  • When Lϕ matches the true forward map L, the coefficient-space DAPS transition is recovered and the asymptotic guarantee continues to hold.

E DETAILED DERIVATION OF GUIDANCE ATTENUATION IN JOINT-EMBEDDING MODELS

The derivation identifies cross-field coupling as the mechanism that transfers observations into coefficient updates in joint-embedding models. Under local dominance or non-overlap, this coupling vanishes, causing guidance attenuation.

  • Joint-embedding models define observations directly on the clean joint variable x0=(a,u), unlike DDIS’s forward-operator likelihood.
  • The coefficient component of joint-model guidance depends on the cross-partial that transmits information from the observation component into coefficient updates.
  • Effective coefficient guidance requires a non-factorized joint distribution pt(at,ut) ≠ pt(at)pt(ut).
  • A single isotropic Gaussian component induces no coefficient-solution coupling; nonzero coupling arises through mixture responsibilities.
  • When one mixture component locally dominates, responsibility gradients and the coefficient guidance attenuate toward zero.
  • Non-vanishing coefficient guidance requires overlap, meaning at least two mixture components have comparable mass near the diffusion state.
  • Under data scarcity, overlap may be absent, so coefficient-space guidance becomes ineffective.

F.2 GENERALIZATION BOUND

The generalization analysis compares DDIS with joint-embedding methods through hypothesis-class complexity and sample usage. It concludes that DDIS has a more favorable bound when unpaired data greatly exceeds paired data.

  • The analysis uses Rademacher complexity to derive generalization bounds for DDIS and joint-embedding methods.
  • DDIS’s estimation error is bounded using ν unpaired samples and np paired samples.
  • Joint-embedding methods are analyzed using np paired samples בלבד.
  • When ν ≫ np, the DDIS sample-complexity bound is more favorable than the joint-embedding bound.

G PROOFS OF THE FAILURE MODES OF JOINT EMBEDDINGS + DAPS

The appendix formalizes covariance shrinkage caused by sparse constraints in preconditioned Langevin dynamics and describes the PDE settings and training procedures used for evaluation.

  • Covariance-collapse proof: The proof derives the stationary covariance through a Lyapunov equation and applies Sherman–Morrison inversion to the rank-1 constrained operator.The derivation introduces B = C−1 + 1/σ2s δxi ⊗δxi and verifies the covariance equation.
  • Covariance-collapse proof: Theorem G.1 analyzes how a sparse constraint changes the stationary covariance of preconditioned Langevin dynamics.The setup uses a prior covariance operator, noise covariance operator, constraint strength, and drift element.
  • Experimental settings: The experiments cover inverse Poisson, Helmholtz, and Navier–Stokes problems with sparse observations and specified coefficient-field or initial-vorticity priors.The Helmholtz setting is oscillatory and multiscale, while Navier–Stokes uses terminal-time observations to recover the initial vorticity.
  • Implementation constraints: The appendix notes that PDE-residual guidance is unavailable when information is lost, and reports experiments on a single NVIDIA RTX 4090 GPU.It also describes padding and cropping to mitigate boundary artifacts from FFT-based FNO layers.
  • Training procedure: Physics-informed training combines limited paired supervision with PDE-residual regularization and boundary-condition penalties for the neural operator.The reported weights are λ_PDE = 0.1 and λ_BC = 10.

H.3.3 QUANTITATIVE RESULTS AND ANALYSIS

The flow-based comparison finds OFM more competitive than ECI-sampling, but identifies computational and numerical bottlenecks that constrain practical inference in complex PDE settings.

  • Quantitative comparison: OFM is more competitive than ECI-sampling, but both are evaluated across PDE tasks using forward and inverse relative ℓ2 errors.Table 12 reports performance and normalized time per sample across tasks.
  • Computational efficiency: 20,000 Langevin steps would exceed 20 hours per sample in the PDE setting, making the official OFM convergence recommendation computationally prohibitive.Adaptive ODE solvers and joint (a, u) backpropagation further increase VRAM consumption and make standard inference budgets impractical.
  • Numerical stability: Steep data-fidelity gradients can drive OFM samples off the learned prior manifold, causing NaN, Inf, or out-of-memory failures under standard GPU budgets.Performance is therefore averaged across 10 independent runs under practical constraints.
  • Baselines: The comparison includes joint-embedding diffusion baselines DiffusionPDE and FunDPS alongside flow-based methods.FunDPS is identified as the prior state-of-the-art for sparse-observation inverse PDE problems.
Loading 2601.23280v3…