Uncertainty Quantification for Fourier Neural Operators

Published: 03 Mar 2024, Last Modified: 05 May 2024AI4DiffEqtnsInSci @ ICLR 2024 PosterEveryoneRevisionsBibTeXCC BY 4.0
Keywords: Partial Differential Equations, Fourier Neural Operator, Uncertainty Quantification, Ensemble Predictions, Bayesian Deep Learning, Laplace Approximation
TL;DR: We formulate weight-space uncertainty for Fourier Neural Operators via Laplace approximation and compare it in the context of ensemble predictions to alternative approaches from the field of Deep Learning for weather data on a toy dataset.
Abstract: In medium-term weather forecasting, deep learning techniques have emerged as a strong alternative to classical numerical solvers for partial differential equations that describe the underlying physical system. While well-established deep learning models such as Fourier Neural Operators are effective at predicting future states of the system, extending these methods to provide ensemble predictions still poses a challenge. However, it is known that ensemble predictions are crucial in real-world applications such as weather, where local dynamics are not necessarily accounted for due to the coarse data resolution. In this paper, we explore different methods for generating ensemble predictions with Fourier Neural Operators trained on a simple one-dimensional PDE dataset: input perturbations and training for multiple outputs via a statistical loss function. Moreover, we formulate a new Laplace approximation for Fourier layers and show that it exhibits better uncertainty quantification for short training runs.
Submission Number: 75
Loading