Light-Weight Diffusion Multiplier and Uncertainty Quantification for Fourier Neural Operators

TL;DR

This work introduces DINOZAUR, a diffusion-based neural operator that replaces dense spectral multipliers with a diffusion multiplier inspired by the heat kernel, achieving dimensionality-independent parameter efficiency. By placing Bayesian priors on per-channel diffusion times and employing variational inference, DINOZAUR delivers spatially correlated uncertainty estimates within the operator itself, enabling calibrated UQ without expensive post hoc methods. Empirically, DINOZAUR attains competitive or superior deterministic performance across Darcy, Navier–Stokes, ShapeNet, and Ahmed bodies benchmarks while using far fewer parameters than traditional FNO/GINO baselines, and its Bayesian variant demonstrates improved NLL, MA, and IS metrics compared to MC dropout and Laplace methods. The approach advances neural operators toward scalable, uncertainty-aware surrogates for complex PDEs, with potential for transfer learning and anisotropic extensions to further improve modeling of physical systems.

Abstract

Operator learning is a powerful paradigm for solving partial differential equations, with Fourier Neural Operators serving as a widely adopted foundation. However, FNOs face significant scalability challenges due to overparameterization and offer no native uncertainty quantification -- a key requirement for reliable scientific and engineering applications. Instead, neural operators rely on post hoc UQ methods that ignore geometric inductive biases. In this work, we introduce DINOZAUR: a diffusion-based neural operator parametrization with uncertainty quantification. Inspired by the structure of the heat kernel, DINOZAUR replaces the dense tensor multiplier in FNOs with a dimensionality-independent diffusion multiplier that has a single learnable time parameter per channel, drastically reducing parameter count and memory footprint without compromising predictive performance. By defining priors over those time parameters, we cast DINOZAUR as a Bayesian neural operator to yield spatially correlated outputs and calibrated uncertainty estimates. Our method achieves competitive or superior performance across several PDE benchmarks while providing efficient uncertainty quantification.

Figures (7)

• Figure 1: Overview of DINOZAUR. NO block $\mathcal{N}_{\theta_i, \tau_i}$ is revised by updating the integral transform $\mathcal{I}_{\tau_i}$ with diffusion multiplier $\exp(-\lambda_k\tau_i)$ and including gradient features $\mathcal{G}_{\theta_i}$ to add anisotropy. Feed-forward network (FFN) is applied to mix two sets of features.
• Figure 2: (a) Probabilistic graph for DINOZAUR, showing how the latent diffusion-times, $\tau$, relate to the observations and non-mechanistic model parameters, $\theta, \sigma$. (b) We define priors for $\tau$ and sample multiple time parameters from posterior in each block at inference time, which translates into spatially correlated uncertainty at the output.
• Figure 3: Uncertainty predictions. Left: mean of pressure samples on a test mesh with a geodesic path traced along the top. Right, top: ground truth and 100 sampled predictions along the same path. Right, bottom: car profile with a qualitative visualization of the standard deviation of pressure samples along the outline. High uncertainty is near sharp field changes, e.g., between Points 0 and 10.
• Figure 4: Scalability and efficiency of DINOZAUR. (a) Test RL$_2$ error on ShapeNet car as width, depth, and modes are varied. (b) Model size under the same settings. (c) Average peak GPU memory usage across five train-predict cycles on synthetic data. (d) Average time under the same settings.
• Figure 5: Diffusion times gathered from FNO blocks of deterministic DINOZAUR trained on non-uniform datasets.

Theorems & Definitions (3)

• Proposition 1
• Proposition 1
• proof