Linearization Turns Neural Operators into Function-Valued Gaussian Processes

TL;DR

This paper introduces LUNO, a framework that applies linearization and Bayesian methods to neural operators, enabling reliable uncertainty quantification in function space predictions for PDE solutions.

Contribution

LUNO is the first to interpret neural operator uncertainty as a function-valued Gaussian process using model linearization, facilitating practical Bayesian uncertainty quantification.

Findings

LUNO provides accurate uncertainty estimates with minimal overhead.

The method is resolution-agnostic and applicable post-hoc.

Scalability to large models and datasets is demonstrated.

Abstract

Neural operators generalize neural networks to learn mappings between function spaces from data. They are commonly used to learn solution operators of parametric partial differential equations (PDEs) or propagators of time-dependent PDEs. However, to make them useful in high-stakes simulation scenarios, their inherent predictive error must be quantified reliably. We introduce LUNO, a novel framework for approximate Bayesian uncertainty quantification in trained neural operators. Our approach leverages model linearization to push (Gaussian) weight-space uncertainty forward to the neural operator's predictions. We show that this can be interpreted as a probabilistic version of the concept of currying from functional programming, yielding a function-valued (Gaussian) random process belief. Our framework provides a practical yet theoretically sound way to apply existing Bayesian deep learning methods such as the linearized Laplace approximation to neural operators. Just as the underlying neural operator, our approach is resolution-agnostic by design. The method adds minimal prediction overhead, can be applied post-hoc without retraining the network, and scales to large models and datasets. We evaluate these aspects in a case study on Fourier neural operators.