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

TL;DR

This work tackles high-frequency Helmholtz equations by shifting from deterministic surrogates to a probabilistic operator learning framework. It develops a conditional diffusion model that maps inputs (sound-speed maps, source masks, and coordinates) to a conditional distribution over Helmholtz solutions, enabling sampling and uncertainty quantification. Across extensive 2D experiments and preliminary 3D results, the diffusion-based operator achieves lower $L^2$, $H^1$, and energy errors than strong baselines and preserves oscillatory content and far-field structure that deterministic models tend to smooth out; it also yields calibrated uncertainty that reflects input sensitivity. The approach promises principled, uncertainty-aware surrogates for complex PDEs in challenging regimes and points toward efficiency-enhanced diffusion or transformer-based backbones for scalable 3D applications.

Abstract

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.

Figures (22)

• Figure 1: Conditional diffusion for Helmholtz.Forward diffusion (top): the wavefield $u_0$ is progressively noised by a fixed schedule to a Gaussian field $u_T$. Reverse denoising (bottom): sampling starts from pure Gaussian $u_T$ and is conditioned on the inputs $z$ comprised of sound-speed map $c$, as well as source mask and positional encodings that are not shown. A time-indexed U-Net (“Trained Score Net”) predicts noise and removes it iteratively over $T\!\approx\!1000$ steps to produce samples $u_0^{(s)}$ that approximate the conditional distribution of solutions.
• Figure 2: Qualitative comparisons across selected frequencies. Each row shows, left-to-right, Sound Map (c), Ground Truth (GT), U-Net, FNO, HNO, Diffusion, followed by the comparison to GT.
• Figure 3: 2D spectral power comparison at $f=2.5\times10^6$ Hz. Example Fourier power spectra $|F(k_x,k_y)|^2$ for GT and each model on a single test sample.
• Figure 4: Illustration of coefficient-function space and interpolation paths. The red point denotes the reference (baseline) medium $c_0$. The surrounding black points $\{c^{(d)}\}_{d=1}^{D}$ indicate distinct GRF directions/realizations in coefficient space. For each direction $d$, we define a linear homotopy (dashed ray) $c^{(d)}(s)=(1-s)c_0+s\,c^{(d)}$, which traces a path from the center $c_0$ toward the endpoint $c^{(d)}$. The concentric contour circles visualize fixed interpolation levels $s\in\{0.3,0.5,0.7,1.0\}$: moving outward corresponds to increasing deviation from $c_0$, with $s=1.0$ reaching the full GRF realization.
• Figure 5: Sampling along a coefficient path $\mathbf{d{=}1}$ (all 4 near vs. all 4 far). For each pair of media, we linearly interpolate the sound speed $c_s=(1-s)c_0+s\,c_1$ with $s\in[0,1]$ and track the wavefield amplitude at fixed probe pixels. Left column: near-source probes. Right column: near-boundary probes.