Correcting Neural Operator Spectral Bias via Diffusion Posterior Sampling with Sparse Observations

FreqNO-DPS: Correcting Neural Operator Spectral Bias via Diffusion Posterior Sampling

AI models that predict physics, such as fluid dynamics or wave propagation, often suffer from a fundamental flaw. They smooth out fine details like high-frequency oscillations. This phenomenon, known as spectral bias, means the model might capture the general shape of a wave but misses the sharp, critical peaks that engineers rely on. A new paper introduces FreqNO-DPS. This method uses a "diffusion" model—a type of generative AI that iteratively removes noise—to fill in these missing details. It does this by combining a fast but blurry AI prediction with a few real, high-fidelity sensor measurements.

The Problem

Neural operator (NO) surrogates are designed to approximate the solutions to Partial Differential Equations (PDEs) orders of magnitude faster than traditional numerical solvers. However, they are plagued by spectral bias. They tend to reproduce low-frequency content faithfully. Conversely, they systematically attenuate high-frequency content. In fields like seismology, this isn't just a mathematical curiosity. It leads to underestimating peak amplitudes and missing fine-scale geological features.

Current remedies typically involve changing the architecture of the surrogate itself. These include retaining more Fourier modes or using multi-scale training objectives. These approaches are expensive because they require retraining the model. In many production environments, you are handed a frozen, pretrained surrogate that you cannot change. Furthermore, while sparse sensor networks can provide pointwise ground truth to anchor a reconstruction, they cover a tiny fraction of the domain. Relying on sensors alone is an ill-posed problem (one where many solutions fit the data). Relying on the surrogate alone reintroduces the blur. As shown in, both the deterministic surrogate and sensor-only diffusion sampling (DPS) suffer from significant high-frequency attenuation.

How It Works

The core idea of FreqNO-DPS is to treat the frozen neural operator's prediction as an auxiliary observation within a diffusion posterior sampling framework. Instead of trying to retrain the surrogate, the authors incorporate its output into the generative process.

The method operates in three main stages, as outlined in :

1. Calibration: Before inference, the authors use paired data (the surrogate's prediction vs. high-fidelity ground truth) to estimate three spectral statistics. These are the transfer function $H(k)$, the residual variance $\sigma^2_{NO}(k)$, and the signal power $P_u(k)$. This effectively builds a "reliability map" of the surrogate in the Fourier domain.
2. Prior Scoring: An unconditional diffusion model is trained on high-fidelity simulations to serve as a prior. This prior captures the valid manifold of possible physical fields.
3. Spectrally Shaped Guidance: During the reverse diffusion process, the model is guided by two separate signals. First, the sparse sensor measurements provide spatial anchoring. Second, the neural operator provides global structural information.

Crucially, the authors do not simply add the surrogate's prediction as a standard observation. Doing so (referred to as "isotropic guidance") causes the surrogate's spectral bias to leak into the final result .

Instead, they derive a closed-form, spectrally shaped guidance score using a Linear Minimum Mean Squared Error (LMMSE) estimator. This score weights the surrogate's contribution according to its frequency-dependent accuracy. At low frequencies where the surrogate is reliable, the guidance is strong. At high frequencies where the surrogate is known to be inaccurate, the guidance is suppressed. This allows the diffusion prior to generate the necessary high-frequency detail.

Numbers

The authors evaluate the framework on 3D elastic wavefield prediction using the HEMEWS-3D dataset. The results are striking when compared to the baselines.

The most important metric for assessing the fix is the banded relative FFT bias ($rFFT$). This measures how much the model underestimates specific frequency bands. The paper reports that at 5% sensor coverage, FreqNO-DPS achieves a near-zero high-frequency bias ($rFFT_{high}$ of $+0.002$). In contrast, the isotropic guidance baseline fails significantly with a bias of $-0.210$. Even in ultra-sparse regimes with only 2% sensor coverage, the method maintains high spectral fidelity ($rFFT_{high}$ of $+0.047$). Meanwhile, the sensor-only DPS baseline collapses to a massive $-0.557$ [Table 1].

Regarding pointwise accuracy, the authors measure the relative Mean Absolute Error (rMAE). At 5% coverage, FreqNO-DPS reports an rMAE of $0.100$. This is statistically indistinguishable from the isotropic guidance baseline ($0.099$). This is a critical result. It means you can fix the spectral inaccuracies without sacrificing the spatial precision of the model.

What's Missing

While the results are mathematically elegant, there are several gaps that a practitioner should note:

• Distribution Shift Sensitivity: The spectral calibration relies on the assumption that the surrogate's error profile remains consistent. If you move from the training geology to a significantly different environment, the "reliability map" might become invalid. This could lead to incorrect guidance. The authors acknowledge this in Section 6.2 but do not provide a formal way to detect it at runtime.
• Computational Overhead: This is not a "free" lunch. While the spectral guidance itself is cheap (requiring only two FFTs), the overall process requires iterative diffusion sampling. This is significantly slower than a single forward pass of a neural operator. For real-time applications, the latency penalty of 64+ denoising steps will be a major bottleneck.
• Assumption of Stationarity: The method assumes the surrogate's residual error is Wide-Sense Stationary (WSS). This means its statistical properties are translation-invariant. While the authors provide a coherence diagnostic to verify this, it is an extra step that must be performed for every new surrogate architecture you deploy.

Should You Prototype This

Yes, but only if your primary pain point is spectral accuracy.

If you are currently deploying neural operators and finding that they are "too smooth," this is a highly principled way to fix it. Because the method uses a frozen surrogate, you don't need to spend weeks retraining your existing models. You just need a small calibration set of paired data.

However, if your application requires millisecond-level latency, the iterative nature of diffusion sampling makes this a difficult sell. You should prototype this if you have a high-fidelity "gold standard" dataset available for calibration and a requirement for high-frequency structural integrity.

Code is reportedly available; see paper for the canonical link: https://github.com/niccoloperrone/FreqNO-DPS

Figures from the paper