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

2025-08-01

AI Summary

Fourier Neural Operators (FNOs) suffer from poor scalability due to over-parameterization and lack intrinsic uncertainty quantification (UQ), while existing posterior UQ methods compromise their geometric inductive bias. To address these issues, we propose DINOZAUR: the first FNO variant that embeds a heat-kernel diffusion process into the spectral multiplier design, replacing high-dimensional tensor parameters with a single time-varying scalar—enabling lightweight modeling. Concurrently, we introduce a Bayesian prior directly in the frequency domain, enabling geometrically consistent, calibration-aware UQ. DINOZAUR thus achieves both efficiency and reliability: it attains state-of-the-art or competitive accuracy across multiple PDE benchmarks, with significantly reduced parameter count and memory footprint, while producing spatially correlated, statistically calibrated uncertainty estimates.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Scalability challenges in Fourier Neural Operators due to overparameterization
• Lack of native uncertainty quantification in neural operators
• Need for efficient uncertainty quantification in PDE solutions

Innovation

• Diffusion-based neural operator parametrization
• Dimensionality-independent diffusion multiplier
• Bayesian neural operator with uncertainty quantification