A Probabilistic Framework for Solving High-Frequency Helmholtz Equations via Diffusion Models

Deterministic neural operators perform well on many PDEs but can struggle with the approximation of high-frequency wave phenomena, where strong input-to-output sensitivity makes operator learning challenging, and spectral bias blurs oscillations. We argue for adopting a probabilistic approach for approximating waves in high-frequency regime, and develop our probabilistic framework using a score-based conditional diffusion operator. After demonstrating a stability analysis of the Helmholtz operator, we present our numerical experiments across a wide range of frequencies, benchmarked against other popular data-driven and machine learning approaches for waves. We show that our probabilistic neural operator consistently produces robust predictions with the lowest errors in $L^2$, $H^1$, and energy norms. Moreover, unlike all the other tested deterministic approaches, our framework remarkably captures uncertainties in the input sound speed map propagated to the solution field. We envision that our results position probabilistic operator learning as a principled and effective approach for solving complex PDEs such as Helmholtz in the challenging high-frequency regime.

💡 Research Summary

This paper tackles the long‑standing challenge of accurately solving the high‑frequency Helmholtz equation, a prototypical elliptic PDE that models time‑harmonic wave propagation in acoustics, elasticity, and scattering. Deterministic neural operators such as DeepONet, Fourier Neural Operators (FNO), and Helmholtz Neural Operators (HNO) have demonstrated impressive speed‑up for many PDEs, but they suffer from two fundamental drawbacks when applied to high‑frequency regimes: (i) spectral bias, whereby the networks preferentially learn low‑frequency components and oversmooth oscillatory features, leading to phase errors and loss of interference patterns; and (ii) extreme input‑to‑output sensitivity, as tiny perturbations in the sound‑speed field can cause order‑one changes in the solution because the operator norm grows roughly like k·ℓ (wave number times travel distance). The authors formalize this sensitivity using a WKB (geometric optics) analysis, showing that the relative error scales with the product of the wave number and the integral of the relative sound‑speed perturbation, which explains why deterministic single‑output surrogates collapse the phase information into a biased mean.

To overcome these issues, the authors propose a probabilistic operator learning framework based on conditional score‑based diffusion models. The forward diffusion process gradually corrupts the complex Helmholtz solution u₀ into a Gaussian field u_T using a fixed variance schedule {β_t}. The reverse process is parameterized by a time‑conditioned U‑Net that predicts the noise ε_θ given the noisy field u_t, the diffusion time t, and the conditioning variables z (sound‑speed map, source mask, and sinusoidal positional encodings). Training minimizes the conditional denoising score‑matching loss, aligning the network’s predicted score with the true Gaussian score. At inference, ancestral sampling from u_T∼N(0,I) through T≈1000 reverse steps yields a set of independent samples {u₀^{(s)}} that approximate the conditional distribution p_θ(u|z). This probabilistic formulation preserves phase variability across samples, quantifies epistemic uncertainty induced by input variability, and enables Bayes‑optimal estimation of stable functionals such as the total energy ∫Ω (|∇u|² + k²c⁻²|u|²) dx.

The experimental campaign is extensive. Six datasets are generated for frequencies ranging from 1.5 × 10⁵ Hz to 2.5 × 10⁶ Hz on a 256 × 256 Cartesian grid using the J‑Wave spectral solver with perfectly matched layers. Each dataset contains 10 000 pairs of random Gaussian‑field sound‑speed maps and the corresponding complex pressure fields; 8190 are used for training, 1020 for validation, and 500 for testing. Baselines include a vanilla FNO, the Helmholtz Neural Operator (HNO), and a U‑Net trained in a deterministic fashion with the same architecture as the diffusion score network.

Across all frequencies, the diffusion‑based operator achieves the lowest errors in L², H¹, and energy norms. The advantage becomes dramatic at the highest frequencies, where deterministic baselines suffer severe phase attenuation and large error spikes, while the diffusion model maintains stable accuracy. In a sensitivity study, the authors add Gaussian noise to the input sound‑speed maps and evaluate the propagated uncertainty. The diffusion model’s sample variance matches the true variability, providing well‑calibrated predictive intervals. Deterministic models, by contrast, are under‑dispersed and systematically underestimate uncertainty, leading to overconfident but inaccurate predictions.

The paper also discusses practical aspects such as the choice of diffusion schedule, the use of positional encodings to inject geometric information, and the computational cost of sampling (which remains modest compared to solving the Helmholtz equation directly). The authors argue that the probabilistic approach not only improves pointwise accuracy but also delivers valuable uncertainty quantification, which is essential for downstream decision‑making in applications like medical ultrasound imaging, seismic inversion, and acoustic design.

In conclusion, the work demonstrates that conditional diffusion models constitute a powerful and principled tool for operator learning in the high‑frequency Helmholtz regime. By explicitly modeling the conditional distribution of solutions, the method mitigates spectral bias, respects the intrinsic sensitivity of wave propagation, and yields calibrated uncertainty estimates. The theoretical analysis, thorough experiments, and clear comparisons position this probabilistic framework as a compelling alternative to deterministic neural operators, with promising extensions to multi‑physics, multi‑scale, and real‑time inference scenarios.