# Stochastic Latent Transformer: Efficient Modelling of Stochastically Forced Zonal Jets

arXiv: https://arxiv.org/abs/2310.16741 · HTML: https://mlsanity.com/paper/2310.16741/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 69 references that are in the map.
- First posted on arXiv 25 Oct 2023.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [On the Spectral Bias of Neural Networks](https://mlsanity.com/paper/1806.08734/) (2018)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Deep learning for universal linear embeddings of nonlinear dynamics](https://mlsanity.com/paper/1712.09707/) (2017)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [Skillful Precipitation Nowcasting using Deep Generative Models of Radar](https://mlsanity.com/paper/2104.00954/) (2021)
- [Time2Vec: Learning a Vector Representation of Time](https://mlsanity.com/paper/1907.05321/) (2019)
- [Enhancing Computational Fluid Dynamics with Machine Learning](https://mlsanity.com/paper/2110.02085/) (2021)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Multiscale Simulations of Complex Systems by Learning their Effective Dynamics](https://mlsanity.com/paper/2006.13431/) (2020)
- [Deep transfer operator learning for partial differential equations under conditional shift](https://mlsanity.com/paper/2204.09810/) (2022)
- [β-Variational autoencoders and transformers for reduced-order modelling of fluid flows](https://mlsanity.com/paper/2304.03571/) (2023)
- [Machine Learning for Stochastic Parameterization: Generative Adversarial Networks in the Lorenz '96 Model](https://mlsanity.com/paper/1909.04711/) (2019)
- [Data-driven discovery of intrinsic dynamics](https://mlsanity.com/paper/2108.05928/) (2021)
- [Deep learning to discover and predict dynamics on an inertial manifold](https://mlsanity.com/paper/2001.04263/) (2019)
- [Probabilistic Forecasting with Generative Networks via Scoring Rule Minimization](https://mlsanity.com/paper/2112.08217/) (2021)
- [Stochastic Transformer Networks with Linear Competing Units: Application to end-to-end SL Translation](https://mlsanity.com/paper/2109.13318/) (2021)
- [Discovering and forecasting extreme events via active learning in neural operators](https://mlsanity.com/paper/2204.02488/) (2022)
- [Machine learning emulation of a local-scale UK climate model](https://mlsanity.com/paper/2211.16116/) (2022)
- [Convolutional Neural Networks Are Not Invariant to Translation, but They Can Learn to Be](https://mlsanity.com/paper/2110.05861/) (2021)
- [Neural Operator with Regularity Structure for Modeling Dynamics Driven by SPDEs](https://mlsanity.com/paper/2204.06255/) (2022)
- [The Monte Carlo Transformer: a stochastic self-attention model for sequence prediction](https://mlsanity.com/paper/2007.08620/) (2020)

## Built on by (most cited first, 8 of 8)

- [Demystifying Data-Driven Probabilistic Medium-Range Weather Forecasting](https://mlsanity.com/paper/2601.18111/) (2026)
- [A multi-scale loss formulation for learning a probabilistic model with proper score optimisation](https://mlsanity.com/paper/2506.10868/) (2025)
- [Causal Climate Emulation with Bayesian Filtering](https://mlsanity.com/paper/2506.09891/) (2025)
- [Hierarchy of extreme-event predictability in turbulence revealed by machine learning](https://mlsanity.com/paper/2603.13789/) (2026)
- [UP-dROM : Uncertainty-Aware and Parametrised dynamic Reduced-Order Model, application to unsteady flows](https://mlsanity.com/paper/2503.23236/) (2025)
- [AttentionDrop: A Novel Regularization Method for Transformer Models](https://mlsanity.com/paper/2504.12088/) (2025)
- [Conditioning on PDE Parameters to Generalise Deep Learning Emulation of Stochastic and Chaotic Dynamics](https://mlsanity.com/paper/2509.09599/) (2025)
- [Efficient Adaptation of ROMs for Unsteady Flows Using Data Assimilation](https://mlsanity.com/paper/2602.23188/) (2026)

## Abstract (arXiv)

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.
