Stochastic Latent Transformer: Efficient Modelling of Stochastically Forced Zonal Jets

Abstract

We present a novel probabilistic deep learning approach, the 'Stochastic Latent Transformer' (SLT), designed for the efficient reduced-order modelling of stochastic partial differential equations. Stochastically driven flow models are pertinent to a diverse range of natural phenomena, including jets on giant planets, ocean circulation, and the variability of midlatitude weather. However, much of the recent progress in deep learning has predominantly focused on deterministic systems. The SLT comprises a stochastically-forced transformer paired with a translation-equivariant autoencoder, trained towards the Continuous Ranked Probability Score. We showcase its effectiveness by applying it to a well-researched zonal jet system, where the interaction between stochastically forced eddies and the zonal mean flow results in a rich low-frequency variability. The SLT accurately reproduces system dynamics across various integration periods, validated through quantitative diagnostics that include spectral properties and the rate of transitions between distinct states. The SLT achieves a five-order-of-magnitude speedup in emulating the zonally-averaged flow compared to direct numerical simulations. This acceleration facilitates the cost-effective generation of large ensembles, enabling the exploration of statistical questions concerning the probabilities of spontaneous transition events.

arXiv comment25 pages, 11 figures

References in corpus (15)

• Fourier Neural Operator for Parametric Partial Differential Equations
• FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators
• Machine Learning for Stochastic Parameterization: Generative Adversarial Networks in the Lorenz '96 Model
• ClimaX: A foundation model for weather and climate
• Deep learning to discover and predict dynamics on an inertial manifold
• A transformer-based synthetic-inflow generator for spatially-developing turbulent boundary layers
• GenCast: Diffusion-based ensemble forecasting for medium-range weather
• Reduction of SO(2) symmetry for spatially extended dynamical systems
• Time2Vec: Learning a Vector Representation of Time
• Generative data-driven approaches for stochastic subgrid parameterizations in an idealized ocean model
• Periodic orbit analysis of a system with continuous symmetry - a tutorial
• Sinkformers: Transformers with Doubly Stochastic Attention
• Neural Operator with Regularity Structure for Modeling Dynamics Driven by SPDEs
• Convolutional Neural Networks Are Not Invariant to Translation, but They Can Learn to Be
• Statistics of extreme events in coarse-scale climate simulations via machine learning correction operators trained on nudged datasets

Cited by in corpus (3)

• Reconstructing unsteady flows from sparse, noisy measurements with a physics-constrained convolutional neural network
• Real-time forecasting of chaotic dynamics from sparse data and autoencoders
• Leveraging Scale Separation and Stochastic Closure for Data-Driven Prediction of Chaotic Dynamics