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

May 20, 2026cs.LG

Deep Learning Surrogates for Emulating Stochastic Climate Tipping Dynamics

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.

Adeline Hillier, Jennifer Sleeman, Jay Brett +3

Johns Hopkins Applied Physics Laboratory Hopkins University · Johns

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

Ask PDF

Sep 24, 2026cs.LG

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

Latent neural surrogate solvers, or latent dynamics models, accelerate simulations of time-dependent physical systems by evolving a compressed latent space rather than resolving full-resolution fields directly. In principle this reduces computational cost and simplifies learning, but in practice errors often accumulate rapidly during long autoregressive rollouts, limiting predictive utility. We show that this instability does not stem from the latent representation itself, but arises when it is trained solely for reconstruction, producing representations poorly suited to long-horizon forecasting. We systematically evaluate training-level interventions that align latent representations with long-horizon rollout: Koopman operator learning and Hamming noise injection during autoencoder training to improve compression, together with noise injection and multi-step rollout fine-tuning to improve dynamics. Interventions that improve long-horizon rollout stability often degrade conventional training metrics, including reconstruction and one-step prediction accuracy. Collectively, these interventions reduce long-rollout error by approximately 40% and match or exceed the accuracy of full-resolution models on two physics benchmarks, while requiring 2 orders of magnitude fewer floating point operations and half the GPU memory. Applied to mesoscale crystal-plasticity simulations of high-cycle fatigue, the resulting surrogate achieves stable extrapolation over horizons orders of magnitude beyond those observed during training. More broadly, these results show that neural compression should be designed not merely to reduce dimensionality, but to restructure the solution space for stable dynamical evolution, a key requirement for reliable, efficient neural surrogates in scientific applications.

Andreas E. Robertson, Ashley T. Lenau, John D. Shimanek +5

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

Latent Dynamics Future Latent Representations Surrogate Models Data Compression Methods

Ask PDF

May 7, 2026cs.LG

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

Fast surrogate modeling for high-dimensional physical dynamics requires more than low short-term error: useful models must roll out efficiently while preserving the statistical structure of long trajectories. Neural operators provide inexpensive autoregressive forecasts but can drift in turbulent regimes, whereas rolling diffusion and latent generative surrogates can represent stochastic transitions at the cost of multi-step denoising, noise-schedule design, or auxiliary compression models. We propose MeanFlow Long-term Invariant Spatiotemporal Consistency Autoregressive Models (MeLISA), a latent-free autoregressive generative surrogate built on pixel-space MeanFlow. MeLISA defines a blockwise stochastic transition kernel that generates each forecast block with a single model evaluation, avoiding latent encoders and iterative diffusion solvers at inference time. To stabilize long-horizon rollouts, MeLISA combines a Window-Consistency MeanFlow objective that learns conditional spatiotemporal generation from partially observed temporal windows with a Time Increment Consistency loss that constrains multi-lag finite increments and targets temporal-correlation structure. We evaluate MeLISA with compact UNet and scalable DiT backbones on two high-resolution benchmarks, extended 2D Kolmogorov flow at

256 \times 256

and turbulent channel-flow slice at

192 \times 192

. MeLISA outperforms neural-operator baselines on short-term forecasting accuracy and long-horizon statistical metrics, including energy spectra, turbulent kinetic energy, and mixing-rate-related dynamics, while achieving inference speeds comparable to, and in some cases faster than, neural operators. Compact 3.7-5.7M-parameter variants already deliver strong parameter efficiency, and DiT variants provide a scalable path up to 150M parameters. Overall, MeLISA benefits both rollout efficiency and long-horizon statistical accuracy.

Tianyue Yang, Xiao Xue

The Center for Computational Science University College London

Long-Horizon Forecasting Autoregressive Model Generative Models Spatiotemporal Representations Nonlinear Dynamics

Ask PDF