Neural SPDE solver for uncertainty quantification in high-dimensional space-time dynamics

TL;DR

This paper introduces a neural SPDE solver that leverages Gaussian Processes and deep learning to improve uncertainty quantification and data assimilation in high-dimensional space-time geophysical datasets, outperforming traditional methods.

Contribution

It develops a neural architecture that learns both state and SPDE parameters, providing a stochastic prior model with analytical sampling form for enhanced uncertainty quantification.

Findings

Improves over Optimal Interpolation baseline.

Enables flexible posterior estimation with large ensembles.

Demonstrates effectiveness on Sea Surface Height data.

Abstract

Historically, the interpolation of large geophysical datasets has been tackled using methods like Optimal Interpolation (OI) or model-based data assimilation schemes. However, the recent connection between Stochastic Partial Differential Equations (SPDE) and Gaussian Markov Random Fields (GMRF) introduced a novel approach to handle large datasets making use of sparse precision matrices in OI. Recent advancements in deep learning also addressed this issue by incorporating data assimilation into neural architectures: it treats the reconstruction task as a joint learning problem involving both prior model and solver as neural networks. Though, it requires further developments to quantify the associated uncertainties. In our work, we leverage SPDEbased Gaussian Processes to estimate complex prior models capable of handling nonstationary covariances in space and time. We develop a specific architecture able to learn both state and SPDE parameters as a neural SPDE solver, while providing the precisionbased analytical form of the SPDE sampling. The latter is used as a surrogate model along the data assimilation window. Because the prior is stochastic, we can easily draw samples from it and condition the members by our neural solver, allowing flexible estimation of the posterior distribution based on large ensemble. We demonstrate this framework on realistic Sea Surface Height datasets. Our solution improves the OI baseline, aligns with neural prior while enabling uncertainty quantification and online parameter estimation.