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

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

1. Problem Statement

The paper addresses the challenge of solving the Helmholtz equation in high-frequency regimes within heterogeneous media. While deterministic neural operators (such as DeepONet and Fourier Neural Operators) have shown success in learning solution maps for various PDEs, they struggle significantly with high-frequency wave phenomena. The authors identify two primary limitations of deterministic approaches in this context:

• Spectral Bias: Deterministic models tend to preferentially capture low-frequency content, leading to the oversmoothing of oscillatory structures. This degrades phase accuracy and interference patterns, which are critical in acoustics and seismics.
• Input Sensitivity and Multi-modality: In high-frequency regimes, the solution map becomes extremely sensitive to small perturbations in input coefficients (e.g., sound speed). Tiny variations can induce large phase shifts or multi-modal responses. Deterministic point-estimate predictors fail to represent this sensitivity, often collapsing multi-modal distributions into a single mean, which results in amplitude attenuation and blurred interference.

The authors argue that a probabilistic approach is necessary to model the conditional distribution of solutions, thereby capturing both the fine-scale oscillatory features and the uncertainty propagated from input coefficients.

2. Methodology

The proposed framework utilizes conditional score-based diffusion models to learn a probabilistic operator that maps problem inputs to a distribution over wavefields.

• Theoretical Motivation: The authors provide a stability analysis of the Helmholtz operator, demonstrating that the Lipschitz constant of the solution map scales linearly with the wavenumber $k$ and the travel distance. Using a WKB (geometric optics) approximation, they show that small coefficient errors accumulate into significant phase errors. Consequently, the Mean Squared Error (MSE) optimal deterministic predictor (the conditional mean) effectively averages over phase variability, causing destructive interference and loss of high-frequency contrast. A probabilistic formulation avoids this collapse by modeling the full conditional distribution $p(u|z)$.
• Model Architecture:
  • Inputs: The conditioning vector $z$ concatenates the sound-speed map, a binary source mask, and sinusoidal positional encodings.
  • Diffusion Process: The framework employs a forward noising chain that progressively corrupts the target complex wavefield $u_0$ into Gaussian noise. The reverse process is learned via a time-indexed U-Net (the "score network") that predicts the noise component $\epsilon_\theta$.
  • Conditioning Mechanism: The input $z$ is not noised. It is incorporated into the network via concatenation with the input image stack and through FiLM (Feature-wise Linear Modulation) layers, where a context vector derived from $z$ scales and shifts the feature maps.
  • Training Objective: The model is trained using conditional denoising score matching to minimize the difference between the true noise and the predicted noise.
• 3D Extension: For three-dimensional problems, the authors replace the standard U-Net backbone with a U-shaped Vision Transformer (UViT). This architecture is chosen to better capture long-range spatial dependencies and global interference patterns inherent in 3D wave propagation.

3. Key Contributions

• Probabilistic Operator Learning: The paper introduces a conditional diffusion operator specifically designed for high-frequency wave propagation, moving beyond single-output surrogates to learn a distribution over solutions.
• Stability Analysis: It provides a theoretical justification for why deterministic operators fail at high frequencies, linking the failure to the accumulation of phase errors and the resulting "one-to-many" nature of the mapping under unresolved variability.
• Uncertainty Quantification: Unlike deterministic baselines, the framework explicitly captures and propagates uncertainties from the input sound-speed map to the solution field, providing calibrated uncertainty estimates.
• Architectural Innovation: The successful adaptation of diffusion models to steady-state wave problems (without temporal stacking) and the extension to 3D using UViT backbones.

4. Experimental Results

The framework was evaluated across a range of frequencies (150 kHz to 2.5 MHz) on 2D and preliminary 3D datasets, benchmarked against deterministic baselines including FNO, HNO (Helmholtz Neural Operator), and a U-Net.

• Error Analysis:
  • The diffusion model consistently achieved the lowest errors in $L^2$, $H^1$, and energy norms across all tested frequencies.
  • At the highest frequency (2.5 MHz), the diffusion model's relative $L^2$ error was 0.095, whereas the best deterministic baseline (FNO) had an error of 0.412, and the U-Net/HNO exceeded 0.76.
  • Spectral analysis confirmed that deterministic models exhibited significant spectral bias (oversmoothing high wavenumbers), while the diffusion model closely tracked the ground-truth power spectrum across the full frequency range.
• Sensitivity and Uncertainty:
  • In sensitivity analyses involving controlled perturbations of the sound-speed field, the diffusion model successfully reproduced the rapid, non-monotone oscillations and turning points of the solution trajectory, particularly in the far-field.
  • Deterministic models tended to drift or flatten, failing to capture the variability.
  • The diffusion model produced calibrated uncertainty distributions that covered the ground-truth variability, whereas deterministic models were systematically under-dispersed.
• Energy Preservation: The framework demonstrated superior ability to predict stable functionals, such as energy, even when pointwise phase predictions were complex.
• 3D Performance: In preliminary 3D experiments on $64^3$ grids, the Diffusion-UViT variant significantly outperformed both the Diffusion-U-Net and the deterministic U-Net, achieving the lowest errors and providing the sharpest boundary reconstructions.

5. Significance and Claims

The authors claim that their results position probabilistic operator learning as a principled and effective approach for solving complex PDEs like the Helmholtz equation in the challenging high-frequency regime.

• Robustness: The framework offers robustness where deterministic surrogates fail, specifically by decoupling fragile quantities (pointwise phase) from stable ones (energy, band-limited summaries).
• Handling Indeterminacy: While the underlying forward map is deterministic, the authors argue that practical indeterminacy at high frequencies (due to discretization mismatch, modeling errors, or unresolved input variability) makes single-output surrogates brittle. The probabilistic formulation learns the conditional distribution, using its mean for accuracy and its dispersion to encode epistemic sensitivity.
• Future Outlook: The paper envisions that this approach can be extended to more complex geometries and multi-source configurations. It acknowledges that inference time is currently higher than deterministic operators due to the iterative nature of diffusion but suggests that future work could address this via flow matching, improved solvers, or latent diffusion techniques.

The paper concludes that by internalizing phase ambiguity rather than collapsing it, probabilistic operators provide a more faithful surrogate for high-frequency wave propagation in heterogeneous media.