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

2025/09/11 by Ira J. S. Shokar, Shokar, Ira J. S., Rich R. Kerswell +3
Computer Science · Decision Sciences · Physics and Astronomy · #Atmospheric and Oceanic Physics (physics.ao-ph) #Chaotic Dynamics (nlin.CD) #Dynamical Systems (math.DS) #FOS: Computer and information sciences #FOS: Mathematics #FOS: Physical sciences #Gaussian Processes and Bayesian Inference #Machine Learning (cs.LG) #Model Reduction and Neural Networks #Probabilistic and Robust Engineering Design

paper · pdf · doi:10.48550/arxiv.2509.09599

openalex publication_date 2025/09/11 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• Deep Learning of the Evolution Operator Enables Forecasting of Out-of-Training Dynamics in Chaotic Systems
• GenCast: Diffusion-based ensemble forecasting for medium-range weather
• Multiple Physics Pretraining for Physical Surrogate Models
• Learning Neural PDE Solvers with Parameter-Guided Channel Attention
• GraphCast: Learning skillful medium-range global weather forecasting
• Transfer Learning with Physics-Informed Neural Networks for Efficient Simulation of Branched Flows
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators
• Fourier Neural Operator for Parametric Partial Differential Equations
• Denoising Diffusion Probabilistic Models
• Stand-Alone Self-Attention in Vision Models
• FiLM: Visual Reasoning with a General Conditioning Layer
• Attention Is All You Need
• Layer Normalization
• Gaussian Error Linear Units (GELUs)
• Julia: A Fresh Approach to Numerical Computing
• Julia: A Fresh Approach to Numerical Computing
• A Survey on Transfer Learning
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

Related

• Deep Learning of the Evolution Operator Enables Forecasting of Out-of-Training Dynamics in Chaotic Systems
• Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics
• Dynamics-Encoded Deep Learning for Robust System Identification and Parameter Estimation
• Data-driven prediction of large-scale spatiotemporal chaos with distributed low-dimensional models
• AFD-STA: Adaptive Filtering Denoising with Spatiotemporal Attention for Chaotic System Prediction