Beyond Loss Guidance: Using PDE Residuals as Spectral Attention in Diffusion Neural Operators

TL;DR

PRISMA tackles slow and unstable inference in diffusion-based PDE solvers by embedding PDE residuals directly into the model architecture. It introduces Spectral Residual Attention (SRA) within a conditional U-shaped diffusion operator to provide physics-guided attention in the spectral domain, enabling gradient-descent free inference for both forward and inverse PDE problems. The approach yields 15x–250x faster inference with competitive or superior accuracy, especially under noisy observations, and remains robust across full, sparse, and noisy data configurations. This unified, residual-informed diffusion framework has practical impact for fast, reliable PDE solving in scientific computing with imperfect data.

Abstract

Diffusion-based solvers for partial differential equations (PDEs) are often bottle-necked by slow gradient-based test-time optimization routines that use PDE residuals for loss guidance. They additionally suffer from optimization instabilities and are unable to dynamically adapt their inference scheme in the presence of noisy PDE residuals. To address these limitations, we introduce PRISMA (PDE Residual Informed Spectral Modulation with Attention), a conditional diffusion neural operator that embeds PDE residuals directly into the model's architecture via attention mechanisms in the spectral domain, enabling gradient-descent free inference. In contrast to previous methods that use PDE loss solely as external optimization targets, PRISMA integrates PDE residuals as integral architectural features, making it inherently fast, robust, accurate, and free from sensitive hyperparameter tuning. We show that PRISMA has competitive accuracy, at substantially lower inference costs, compared to previous methods across five benchmark PDEs, especially with noisy observations, while using 10x to 100x fewer denoising steps, leading to 15x to 250x faster inference.

Figures (17)

• Figure 1: The PRISMA inference pipeline. We use a U-shaped Diffusion Neural Operator (UNO) to iteratively refine estimates of $\mathbf{a}$ and $\mathbf{u}$ starting from noise $(\mathbf{a}, \mathbf{u})_T$ to a clean solution $(\mathbf{a}, \mathbf{u})_0$. At each denoising step, PDE residuals are architecturally injected via a novel SRA block at every layer of the UNO, enabling fast, gradient-descent free inference for both forward and inverse problems.
• Figure 2: Overview of the PRISMA model architecture. (Left) Computation of the observation-informed PDE residual: The model's current estimate ($\mathbf{x}_\sigma$) is mixed with the known observations ($\mathbf{x}_{\text{obs}}$) using the mask ($\mathbf{M_x}$) to produce a residual field. (Right)Spectral Residual Attention (SRA) block: The SRA block operates in the Fourier domain to compute a physics-informed attention mask between the network state ($\mathbf{x}^l_t$) and the residual. This mask is applied in a gated skip-connection, controlled by a learned guidance strength ($g_{\text{res}}$), to produce a modulated state ($\mathbf{x}^{t}_{\text{SRA}}$) that is then passed to the UNO block.
• Figure 3: Qualitative results for PRISMA and baseline models on four PDE benchmarks. We evaluate performance under three distinct conditions: (a, b) noisy observations (corrupted by $\mathcal{N}(0,1)$ Gaussian noise), (c) full, clean observations, and (d) sparse observations (3% of data known). The relative $\ell_2$ error is reported below each prediction (pixel-wise error rate for the Darcy inverse case).
• Figure 4: Inference time vs. accuracy trade-off on the Darcy Flow forward problem under full observations. Point size corresponds to the number of denoising steps.
• Figure 5: Skewness of the PDE residual field plotted against inference iterations for the Poisson equation. (Left: Forward problem, Right: Inverse problem)