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

TL;DR

DINOZAUR introduces a diffusion-based neural operator that reduces parameters and memory usage in Fourier Neural Operators while providing reliable uncertainty quantification, enhancing scalability and trustworthiness in PDE solving.

Contribution

The paper presents DINOZAUR, a novel diffusion-based neural operator that replaces dense multipliers with a lightweight diffusion multiplier, enabling scalable PDE solutions with integrated uncertainty quantification.

Findings

Achieves competitive or superior performance on PDE benchmarks.

Reduces parameter count and memory footprint significantly.

Provides calibrated uncertainty estimates with spatial correlation.

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.