Accelerated surrogate dynamics for dynamical, stochastic system evolution

Marco Jochum; Ioannis Kouroudis; Gohar Ali Siddiqui; Taher Amine Hamzaoui; Manuel Gößwein; Alessio Gagliardi

cs.AISep 29, 2026

Accelerated surrogate dynamics for dynamical, stochastic system evolution

Authors: Marco Jochum, Ioannis Kouroudis, Gohar Ali Siddiqui, Taher Amine Hamzaoui, Manuel Gößwein, Alessio Gagliardi

Organizations: Chair of Simulation of Nanosystems for Energy Conversion, Department of Electrical Engineering, TUM School of Computation, Information and Technology, Atomistic Modeling Center (AMC), Munich Data Science Institute (MDSI), Technical University of Munich, Hans-Piloty-Straße 1, 85748 Garching, Germany

Abstract

Dynamic simulations are an entrenched way of gaining insight into the evolution of system dynamics. Their computational cost however is often prohibitively high, especially in cases of stochastic frameworks. Machine learning algorithms are especially suited as simulation surrogates. Nevertheless, they face some very distinct limitations. Firstly, the sheer dimensionality of these systems, however, precludes the use of traditional time series models who struggle with high dimensional feature spaces. Additionally, traditional time series focus exclusively on either long or short range effects, causing local or global drift given enough time. In this paper, we propose a framework that addresses those limitations. Our framework combines a Variational Autoencoder, with a convolutional or graph basis that reduces the dimensionality of the system. This latent vector is propagated in time using a Temporal Fusion Transformer model, which includes both long range and short range effect encoding, as well as static covariate support. We test our framework on three distinct cases, to prove its robustness and in all three we have achieved practically identical to the simulation results at a fraction of the time. Further, our framework is flexible enough to be adapted to any new system and provides an inbuilt uncertainty quantification for targeted experiment design.

Figures & tables

Figure 1 : Visualization of the entire surrogate model using a multi-horizon forecasting model as the propagator

\mathcal{P}

\mathcal{E}

and

\mathcal{D}

are the encoder/ decoder parts of the VAE,

c_{i}

are the input configurations and

\tilde{c_{i}}

are the predicted next configurations. This process is repeated to predict sequences that are longer than the prediction horizon of the multi-horizon model. All the models have a

S^{\mu}

branch that predicts the evolution of the dynamics. Stochastic models have an additional

S^{\sigma}

branch that predicts the stochastic behaviour of the evolution.

Figure 2 : Schematic representation of the encoding-decoding process. A high dimensional map $c_{t_{0}}(x,z)$ at a given time point $t_{0}$ is taken and encoded onto its respective latent space representation $z_{t_{0}}$ . It can then be decoded by propagating it through the decoder network $\mathcal{D}$ , recovering the high dimensional representation $\tilde{c}_{t_{0}}(x,z)$ , which should be as close to the original input as possible.

Figure 3 : Temporal Fusion Transformer model, adapted from [ 33 ] Static covariates S and time-varying inputs — past observations (

x_{t-k}...x_{t}

) and known future covariates (

x_{t+1}...x_{t+\tau_{max}}

) enter through dedicated Variable Selection Networks (orange), which learn instance-wise feature weights for interpretability. The Static Covariate Encoder (dark orange) produces four context vectors used throughout the network: one for variable selection, two to initialize the LSTM’s cell/hidden state, and one for static enrichment. A sequence-to-sequence LSTM Encoder–Decoder (blue = past, dark red = future) provides local temporal processing, replacing standard positional encoding. Every sub-block is wrapped in a Gate + Add and Norm unit (yellow/green), a GLU-based residual gate that lets the network skip unused components. LSTM outputs are enriched with static context via a Static Enrichment GRN (purple), then passed to Masked Interpretable Multi-Head Attention (teal), which lets each time step attend to relevant past and future points while preserving causal masking. This is followed by a Position-wise Feed-Forward GRN, another gated residual block, and a final FC layer producing quantile forecasts (

\tilde{z}_{t+\tau}

(0.1), (0.5), (0.9)) at each horizon step — enabling calibrated uncertainty estimates alongside point predictions.

Parameter	Description
J	Exchange interaction energy between spins	0.1, 0.5, 1, 2, 3, 5, 6, 7, 8
$K_{BT}$	Thermal energy (Boltzmann constant multiplied by temperature)	0.1, 0.5, 0.8, 1, 2, 2.5, 3, 5.
Ratio	The probability $P(X=1)$	0.1, 0.25, 0.4, 0.5, 0.75, 0.8

Table 1 : Simulation Parameters, as this is an abstract case the units are irrelevant.

Panel of Figure 4 — Figure 4 : Snapshots generated through the physical simulation (up) and snapshots generated through our framework (down) with their energies and corresponding timesteps

Figure 5 : Predictions of the surrogate model for the mean field of the three different parameter configurations a) - c . The combination of parameters strongly affects the final concentration profile as well as the dynamic process that precedes the steady-state concentration distribution. The time stamps correspond to the following list of time steps: [10, 20, 50, 100, 200, 300]. The represetation shows the evolution of the concentration summed up along the Z axis and averaged over multiple simulation runs.

Figure 6 : Prediction for the standard deviation fields of the kMC simulation. The represetation shows the evolution of the concentration summed up along the Z axis and averaged over multiple simulation runs. Parameter configurations are the same as in Figure 5 .

Figure 7 : The system used for surrogate of molecular dynamics. The hydrogens are removed from the MD configurations. The two ramachandran angle $\phi$ and $\psi$ are indicated.

Panel of Figure 8 — Figure 8 : Comparison of Free Energy Surface (FES) of the dihedral angles $\phi$ and $\psi$ between true and surrogate dynamics.

Panel of Figure 9 — Figure 9 : Comparison of Transition probabilities between simulation and surrogate dyanamics. States labeled in fig. 8

Figure 10 : Time required for the execution of the simulation, the training of the model and the inference.

Explore similar work

Deep Learning Surrogates for Emulating Stochastic Climate Tipping Dynamics
May 20, 2026Adeline Hillier, Jennifer Sleeman, Jay Brett +3Surrogate Models Global Climate Models
Johns Hopkins Applied Physics Laboratory Hopkins University · Johns
Beyond Compression: Training Latent Representations for Stable Long-Horizon Rollout in Neural Surrogate Solvers
Sep 24, 2026Andreas E. Robertson, Ashley T. Lenau, John D. Shimanek +5Latent Dynamics Future Latent Representations
Center for Integrated Nanotechnologies, Sandia National Laboratories, New Mexico, USA · Center for Computing Research, Sandia National Laboratories, Albuquerque, NM, USA · School of Engineering, Brown University, Providence, RI, USA +2
Towards Scalable One-Step Generative Modeling for Autoregressive Dynamical System Forecasting
May 7, 2026Tianyue Yang, Xiao XueLong-Horizon Forecasting Autoregressive Model
The Center for Computational Science University College London

May 20, 2026, cs.LGJ/K move · Enter open · S save

Deep Learning Surrogates for Emulating Stochastic Climate Tipping Dynamics

Adeline Hillier, Jennifer Sleeman, Jay Brett, Caroline Tang+2

Johns Hopkins Applied Physics Laboratory Hopkins University · Johns

Surrogate Models Global Climate Models Multi-Stream Temporal Fusion Surrogates Deep Learning Maritime Navigation

Open paper AskPDF arXiv

This work explores a dynamics-informed Temporal Fusion Transformer (TFT) as a data-driven surrogate for computationally intensive Earth system simulations. Focusing on multivariate time series describing global ocean transport, we demonstrate the surrogate's ability to forecast tip events across thousands of time steps. The data involve up to 21 non-stationary time series in addition to static covariates describing free parameters and initial conditions. Modifications to the architecture and objective function yield a surrogate that anticipates the timing of Atlantic and Pacific collapses to high fidelity and captures the stochastic uncertainty in transition timing across ensemble predictions. The learned surrogate achieves a 465x computational speedup over the numerical simulator while maintaining differentiability with respect to parameters and initial conditions.

Abstract

Explore similar work

Deep Learning Surrogates for Emulating Stochastic Climate Tipping Dynamics

Beyond Compression: Training Latent Representations for Stable Long-Horizon Rollout in Neural Surrogate Solvers

Towards Scalable One-Step Generative Modeling for Autoregressive Dynamical System Forecasting