Linearization Turns Neural Operators into Function-Valued Gaussian Processes
This provides a practical uncertainty quantification method for neural operators in domains like PDE simulation, though it is incremental as it adapts existing Bayesian deep learning techniques.
The paper tackles the problem of quantifying predictive uncertainty in neural operators for high-stakes simulations by introducing LUNO, a framework that uses model linearization to propagate Gaussian weight-space uncertainty into predictions, resulting in a function-valued Gaussian process belief with minimal overhead and no retraining.
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.