# Conditioning on PDE Parameters to Generalise Deep Learning Emulation of Stochastic and Chaotic Dynamics

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 43 references that are in the map.
- First posted on arXiv 11 Sep 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [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)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [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)
- [Stand-Alone Self-Attention in Vision Models](https://mlsanity.com/paper/1906.05909/) (2019)
- [GraphCast: Learning skillful medium-range global weather forecasting](https://mlsanity.com/paper/2212.12794/) (2022)
- [Neural General Circulation Models for Weather and Climate](https://mlsanity.com/paper/2311.07222/) (2023)
- [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://mlsanity.com/paper/2312.15796/) (2023)
- [Enhancing Computational Fluid Dynamics with Machine Learning](https://mlsanity.com/paper/2110.02085/) (2021)
- [Modeling the Dynamics of PDE Systems with Physics-Constrained Deep Auto-Regressive Networks](https://mlsanity.com/paper/1906.05747/) (2019)
- [Multiple Physics Pretraining for Physical Surrogate Models](https://mlsanity.com/paper/2310.02994/) (2023)
- [Deep transfer operator learning for partial differential equations under conditional shift](https://mlsanity.com/paper/2204.09810/) (2022)
- [Deep learning to discover and predict dynamics on an inertial manifold](https://mlsanity.com/paper/2001.04263/) (2019)
- [Learning Neural PDE Solvers with Parameter-Guided Channel Attention](https://mlsanity.com/paper/2304.14118/) (2023)
- [Knowledge-Based Learning of Nonlinear Dynamics and Chaos](https://mlsanity.com/paper/2010.03415/) (2020)
- [Transfer Learning with Physics-Informed Neural Networks for Efficient Simulation of Branched Flows](https://mlsanity.com/paper/2211.00214/) (2022)
- [Stochastic Latent Transformer: Efficient Modelling of Stochastically Forced Zonal Jets](https://mlsanity.com/paper/2310.16741/) (2023)
- [Deep Learning of the Evolution Operator Enables Forecasting of Out-of-Training Dynamics in Chaotic Systems](https://mlsanity.com/paper/2502.20603/) (2025)

## Abstract (arXiv)

We present a deep learning emulator for stochastic and chaotic spatio-temporal systems, explicitly conditioned on the parameter values of the underlying partial differential equations (PDEs). Our approach involves pre-training the model on a single parameter domain, followed by fine-tuning on a smaller, yet diverse dataset, enabling generalisation across a broad range of parameter values. By incorporating local attention mechanisms, the network is capable of handling varying domain sizes and resolutions. This enables computationally efficient pre-training on smaller domains while requiring only a small additional dataset to learn how to generalise to larger domain sizes. We demonstrate the model's capabilities on the chaotic Kuramoto-Sivashinsky equation and stochastically-forced beta-plane turbulence, showcasing its ability to capture phenomena at interpolated parameter values. The emulator provides significant computational speed-ups over conventional numerical integration, facilitating efficient exploration of parameter space, while a probabilistic variant of the emulator provides uncertainty quantification, allowing for the statistical study of rare events.