Wavelet Diffusion Neural Operator

• WDNO is a unified deep learning framework that performs diffusion in the wavelet domain to capture shocks and enables zero-shot super-resolution for time-dependent PDEs.
• It employs a multi-resolution training scheme with wavelet decompositions to ensure strong cross-resolution generalization in simulating complex physical systems.
• Empirical benchmarks demonstrate that WDNO achieves significantly lower error metrics on both simulation and control tasks compared to traditional diffusion and Fourier-based models.

Wavelet Diffusion Neural Operator (WDNO) is a unified deep-learning framework designed for simulation and control of physical systems modeled by time-dependent partial differential equations (PDEs). WDNO directly addresses two principal challenges for generative models in this domain: accurately representing abrupt, high-frequency state changes such as shocks, and generalizing across spatial and temporal resolutions for zero-shot super-resolution. Its central innovations are (i) conducting diffusion-based generative modeling entirely in the wavelet domain to leverage localized, multi-scale representations, and (ii) employing multi-resolution training schemes that exploit the approximate scale invariance of many PDEs for strong cross-resolution generalization (Hu et al., 2024).

1. Motivation and Conceptual Foundations

Classical PDE simulation and control tasks—spanning advection, Burgers' equations, compressible and incompressible Navier–Stokes systems, and weather prediction—feature sharp local variations demanding high-fidelity modeling. Standard diffusion generative models, which operate on fixed-resolution grids, have difficulty both with representing discontinuities and with providing predictions on finer, unseen grids. WDNO introduces a wavelet-domain approach and data-efficient multi-resolution training to address these foundational issues. By mapping the entire trajectory into a wavelet basis, WDNO encodes both low-frequency structures and high-frequency details, essential for capturing and propagating abrupt changes such as shock fronts (Hu et al., 2024).

2. Architectural Principles: Wavelet-Domain Diffusion

Wavelet Decomposition and Reconstruction

Let $u(t, x)$ denote the solution of the PDE over time $t \in [0,T]$ and space $x \in D$. WDNO represents $u$ in the discrete wavelet basis: $Wu = \{ c_L(m),\, d_L(m) \}_m$ where $c_L(m) = \langle u, \phi_{L, m} \rangle$ and $d_L(m) = \langle u, \psi_{L, m} \rangle$, with $\phi_{L, m}$ and $\psi_{L, m}$ as the respective scaling and mother wavelet functions at decomposition level $L$. The inverse transform $t \in [0,T]$0 reconstructs $t \in [0,T]$1 from these coefficients, ensuring numerical exactness for suitably chosen $t \in [0,T]$2 on an $t \in [0,T]$3-point grid.

Diffusion Process in Wavelet Coefficient Space

After transformation into wavelet space, WDNO applies a discrete diffusion process to the sequence of coefficients. Given coefficients $t \in [0,T]$4 of the state trajectory, the forward noising chain is: $t \in [0,T]$5 where $t \in [0,T]$6 is the noise schedule. The denoising reverse process, implemented via DDPM or DDIM, uses a learned score function $t \in [0,T]$7.

For simulation, inference recovers denoised $t \in [0,T]$8 through iterative updates and reconstructs the trajectory in the original domain. For control, an energy-objective guidance term augments this update, allowing trajectory-level control in the wavelet space. The score-matching loss for training in wavelet space is: $t \in [0,T]$9 where "cond" represents any wavelet-transformed conditioning variables, such as initial states or controls.

3. Multi-Resolution Training and Zero-Shot Super-Resolution

WDNO exploits approximate scale invariance in many PDEs to enable cross-resolution generalization. Training samples at a given resolution $x \in D$0 are downsampled to create coarser representations, and the resulting hierarchical learning process supports inference on unseen finer grids without retraining.

The hierarchical process consists of a Base-Resolution Model (BRM) trained on coarse data and a Super-Resolution Model (SRM), which is itself a diffusion denoiser trained to map (coarse $x \in D$1 fine) wavelet coefficient pairs. At inference, prediction proceeds by first generating a coarse output, then using SRM (conditioned on fine-resolution inputs) to lift to successively finer resolutions, supporting arbitrary multi-level super-resolution.

4. Empirical Performance

WDNO has been validated on five diverse benchmarks, demonstrating state-of-the-art performance in both simulation and control:

Task / System	Key Metric	WDNO vs. Baselines
1D Advection (Simulation, 80 steps)	MSE $x \in D$2	31% lower than DDPM, OFormer
1D Burgers' (Simulation & Control)	Simulation MSE $x \in D$3	Comparable to DDPM
	Control $x \in D$4	13–25% reduction
1D Compressible Navier–Stokes (Shock)	MSE $x \in D$5	15x lower than DDPM, 15% lower than FNO
2D Incompressible Fluid (Simulation)	MSE $x \in D$6	85% lower than DDPM, 43% over FNO
2D Fluid Control: Smoke Leakage (Control)	Leakage $x \in D$7	33.2% less than BPPO
ERA5 Weather (12→20hr, Simulation)	Rel. $x \in D$8 error $x \in D$9	5% lower than FNO

In settings with sharp state transitions (e.g., shock tubes), WDNO exceeds Fourier-based and standard diffusion models in both visual and quantitative accuracy. In indirect control (e.g., minimizing smoke leakage), WDNO reduces leakage by 78% relative to the next-best baseline (Hu et al., 2024).

5. Comparative Analysis and Ablation Insights

Empirical ablations show that replacing wavelets with Fourier transforms in the diffusion pipeline leads to significantly worse handling of discontinuities (e.g., MSE $u$0 for 1D compressible Navier–Stokes, vs. $u$1 for WDNO). The multi-resolution pipeline preserves prediction accuracy over several zero-shot super-resolution levels, outperforming linear interpolation and mesh-invariant operators, with 1D Burgers' super-res MSE dropping from $u$2 to $u$3 across three up-sampling stages. WDNO demonstrates slowest error accumulation in long time-horizon forecasting and robustness to noise, hyperparameter variations, and training-set size. For control, guidance strength $u$4 is empirically necessary for optimal performance (Hu et al., 2024).

6. Theoretical Underpinnings and Mechanistic Rationale

Wavelet transforms decompose PDE solutions into scale-localized structures, isolating abrupt features into detail coefficients and separating low- and high-frequency phenomena. This intrinsic sparsity and spatial–frequency locality allow the diffusion model to smooth high-frequency shocks and propagate low-frequency flows effectively. Multi-resolution training leverages PDEs' approximate scale invariance, ensuring the denoiser learns grid-independent transition dynamics, which the wavelet domain can faithfully “lift” to finer scales for zero-shot generalization. This multi-scale representation is critical to the empirical accuracy and super-resolution capacity of WDNO (Hu et al., 2024).

7. Limitations and Prospective Directions

WDNO operates under the assumption of uniform, structured spatial–temporal grids and utilizes fixed (e.g., biorthogonal) wavelet bases. Extension to irregular or adaptive meshes may necessitate geometric or graph-based wavelet constructions and graph-diffusion processes. The current framework does not directly incorporate explicit PDE constraints such as physics-informed losses, nor has it been validated in live, real-time experimental settings (e.g., turbulence or plasma control). Promising directions include augmenting WDNO with adaptive or graph-based wavelets, integrating physics constraints into loss functions, and extending to PDEs on unstructured domains (Hu et al., 2024).