Beyond Compression: Training Latent Representations for Stable Long-Horizon Rollout in Neural Surrogate Solvers
Authors: Andreas E. Robertson, Ashley T. Lenau, John D. Shimanek, Benjamin A. Jasperson, Vivek Oommen, David L. Damm, Krishna Garikipati, Remi Dingreville
Organizations: 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 · Sandia National Laboratories, New Mexico, USA · Sandia National Laboratories
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.
Figures & tables
Figure 1: Stabilizing latent dynamics models. Comparison between conventional sequential training (where compression and dynamics and trained separately and in sequence) and sequential training with proposed interventions. (a) The autoencoder compression model is trained for reconstruction (black compute path), producing misalignment for long-horizon dynamics predictions (red compute path). (b) Training-level interventions are applied to the autoencoder training leading to better alignment with the dynamics task and stronger long-horizon rollout predictions. Models are still trained sequentially.
Figure 2: Comparison between the pixel-space model dynamic model (UNet) and the proposed LDM incorporating the interventions of Section 3.2 (LDM) on the spinodal decomposition benchmark. Rows 1–4 show the ground-truth concentration field and the relative L1 error for the vanilla LDM, the UNet, and the proposed LDM, respectively, at rollout steps Δt=1,3,6,11 ; each model is given five initial time steps as context before predictions are generated autoregressively. (a) RelMSE as a function of rollout step, averaged across the test dataset, starting at time, ti=25 . (b) Distribution of mean absolute latent-space velocity magnitude for the vanilla LDM and the proposed LDM. The UNET is not applicable to this analysis. (c) Model comparison of (Giga) Floating Point Operations (GFLOP) versus RelMSE at 40 step rollout (averaged over the entire dataset). (d) Model comparison of allocated VRAM (Gigabytes - GB) versus RelMSE at 40 step rollout (averaged over the entire dataset).
Figure 3: Comparison between the vanilla LDM and proposed LDM incorporating the interventions of Section 3.2 (LDM) on the active matter benchmark. Rows 1–3 show the ground-truth of the first component of the second-order moment tensor, D0,0 , respectively, at rollout steps Δt=1,11,21,31 ; each model is given 9 time steps of historical context before predictions are generated autoregressively. (a,b) VRMSE as a function of rollout step, averaged across the test dataset, from two starting times, t0=9 and t0=29 . (c,d) Distribution of mean and median absolute latent-space velocity magnitude for the vanilla LDM and the proposed model.
Figure 4: Overview of the impact of the spatially structured latent encodings, Koopman-inspired dynamics constraints, and Hamming noise injection interventions. Row 1 depicts real space samples of the first component of the second order moment tensor, D0,0 . Row 2 depicts corresponding samples of one of the 8 latent encodings from the LiteVAE-type autoencoder. Row 3 similarly depicts a single latent field from the Koopman-trained LiteVAE autoencoder. (a) Comparison of the decoder sensitivity, Section 3.3 , between the basic LiteVAE autoencoder, koopman training, and koopman + hamming noise injection. (b-d) variance of the reconstruction for the t=6 sample depicted in Row 1. Variance was calculated over 30 samples generated by injecting standard deviation 0.01 noise onto the latent encoding. (e-g) hyperparameter optimization for the hyperparameters involved in proposed autoencoder interventions. Error is reported for the second order moment tensor. See Appendix D.2.3 for optimization over remaining parameters.
Model
Δt=1
Δt=5
Δt=15
Error
Relative
Absolute
Error
Relative
Absolute
Error
Relative
Absolute
Baseline
0.281
–
–
0.593
–
–
0.982
–
–
Spatial
0.049
0.827
0.827
0.269
0.546
0.546
1.020
-0.039
-0.039
Koopman
0.033
0.327
0.884
0.232
0.140
0.609
0.772
0.243
0.213
KL
0.036
-0.100
0.872
0.241
-0.038
0.595
0.906
-0.174
0.077
Hamming Noise
0.032
0.106
0.886
0.228
0.050
0.615
0.768
0.153
0.218
Table 1: Cumulative effect of training-level interventions on rollout error at different rollout time Δt .
Figure 5: Summary of the cumulative improvement caused by each addition. Methods are included cumulatively on the x-axis (e.g., the ‘Hamming’ model is trained using ‘Spatial’, ‘Koopman’, ‘KL’, and ‘Hamming’ methods). (a) Plot of the absolute percent improvement versus the baseline (the baseline would be 0 ). (b) The relative improvement between each addition. (c) The correlation between the 1-step and 15-step absolute improvements.
Figure 6: Overview of proposed LDM performance on the Cyclic Fatigue Benchmark. Rows 1,2 depict an example stress and strain field prediction. Rows 3-5 contrast predictions on an example Fatigue Indicator Parameter (FIP) field. Row 4 depicts the prediction from the LDM trained to predict all fields, Row 5 is the prediction from an LDM specialized to FIP predictions. All examples are derived from one of the expensive long horizon calculations performed. As a result, they depict time steps both within the training range and outside of it. Rows 6 and 7 report rollout errors over the training time distribution. Rollout errors for the two long-time examples are superimposed for context.
Figure 7: Summary of cyclic fatigue based failure prediction on long rollout Example 1 using the LDM (Fatigue Indicator Parameter (FIP)) model. In all plots, lines are colored by the number of simulation context steps provided before switching to the neural solver. (a) the relative L1 error of the predicted 95 -percentile FIP at each cycle order of magnitude. (b) A reinterpretation of (a): L1 error in the number of cycles to failure if failure occurred at each cycle order of magnitude. (c) predicted empirical FIP CDF at the temporal boundary of training data. (d) predicted Empirical FIP CDF at maximum ground truth cycle count ( 215 cycles).
Figure 8: Effect of foundation-model (Walrus) pretraining on validation error across benchmarks. Validation error versus training epoch for Walrus initialized from its pretrained weights compared to a randomly initialized model of identical architecture, on the active matter (Sec. 4.2 ) and high-cycle fatigue (Sec. 4.3 ) case studies. Shaded regions denote ±1 standard deviation across 10 random seeds. The pretrained model reaches lower validation error consistently across both datasets and throughout training, despite two compounding distribution shifts: Walrus was pretrained on real-space physical fields, not the compressed latent encodings used here, and, in the fatigue case, on a dataset it never saw during pretraining at all.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Calibration of the crystal plasticity model against experimental cyclic loading data. Simulated stress–strain response using the parameters in Table 2 , compared against experimental data for aluminum [ 49 ] . Simulated curves are shaded from dark to light with increasing cycle number, tracking the model’s evolving hardening response across cycles.
Parameter
Value
Unit
τ0
175
MPa
τ1
25
MPa
θ0
200
MPa
h1
7000
MPa
r1
100
MPa
d1
4
–
Appendix
Table 2: Crystal-plasticity parameters. Calibrated parameter values for fully reversed loading using an example microstructure from the present work and mechanical data from Ref. [ 49 ] .
Figure 10: Overview of the microstructure preparation pipeline. Step 1: individual 2D slices are extracted from the 3D DCT-measured aluminum microstructure along each principal direction. Step 2: each slice is periodically extended into a 128×128 patch to accommodate the CPFFT simulation code, using a periodic-extension conditioning process, using the PolyMicro Foundation Model [ 50 ] . Step 3: Step 3: pixel values are transformed from PolyMicro’s ROGSH basis back to Euler angle space.
Model
Δt=1
Δt=5
Δt=15
Δt=40
UNET
0.00003
0.00024
0.00167
0.00598
LDM
0.00016
0.00051
0.00203
0.00586
Vanilla LDM
0.00261
0.02099
0.08853
0.18081
Appendix
Table 3: Average ReLMSE for predictions for the spinodal decomposition benchmark for various models and rollouts Δt .
Model
Δt=1
Δt=5
Δt=15
Δt=40
UNET
0.01274
0.02715
0.17213
0.13762
LDM
0.04619
0.04988
0.18177
0.14900
Vanilla LDM
0.90402
1.11307
1.27472
1.28404
Appendix
Table 4: Average VRMSE for predictions for the spinodal decomposition benchmark for various models and rollouts Δt .
Figure 11: Optimization of autoencoder hyperparameters for reconstruction weakens rollout performance. (a-d) The reconstruction and total loss of the autoencoder swept over different latent space resolutions and channel numbers. (e) the corresponding 15 -step rollout error for the same hyperparameter sweep. Note the lack of correlation between trends in each.
Model
Channels: 8
Channels: 11
Error
Absolute
Error
Absolute
Baseline
0.98166
–
0.98166
–
Hamming noise
0.76793
0.21773
0.73952
0.24667
Dynamic noise
0.60120
0.38757
0.57556
0.41368
Roll 1
0.62439
0.36394
0.55968
0.42987
Roll 2
0.57895
0.41024
0.51364
0.47677
Appendix
Table 5: Sensitivity of cumulative intervention gains to latent channel count. Average VRMSE at rollout Δt=15 on the active matter benchmark, comparing an 8 -channel and an 11 -channel latent space. Each row adds one training-level intervention cumulatively on top of all rows above it, in the order introduced in Section 3.2 . Error reports the raw VRMSE for that row’s model; Absolute reports its improvement over the baseline (row 1).
Figure 12: Hyperparameter optimization for dynamics training. (a) optimization of the standard deviation of the injected noise into the dynamics model’s input. (b,c) optimization of the number of fine-tuning rollout steps. (b) contrasts noise injection based pretraining of the dynamics model against pretraining without noise. (c) reports performance with a autoencoder trained with hamming noise injection. Rollout fine tuning monotonically improves all three.
Figure 13: Summary of hyperparameter optimization. Each figure reports 15 -step VRMSE error for its corresponding field variable. Average error over the validation set is reported for 3 starting times: Time Steps 0 , 10 , and 20 . At each time step, the first 9 steps are provided to the model as context to begin rollout.
Model
Δt=1
Δt=5
Δt=15
Error
Relative
Absolute
Error
Relative
Absolute
Error
Relative
Absolute
Baseline
0.07137
–
–
0.40241
–
–
0.95216
–
–
Spatial
0.00152
0.97871
0.97871
0.06127
0.84775
0.84775
0.85100
0.10625
0.10625
Koopman
0.00130
0.14157
0.98172
0.08094
-0.32116
0.79885
0.59237
0.30391
0.37786
KL
0.00137
-0.05239
0.98076
0.07923
0.02119
0.80311
0.58977
0.00440
0.38060
Hamming noise
0.00125
0.09187
0.98253
0.08019
-0.01212
0.80073
0.61107
-0.03612
0.35823
Appendix
Table 6: Cumulative ablation of training-level interventions. Error by model intervention and rollout length Δt on the active matter benchmark; each row adds one intervention cumulatively on top of all rows above it, in the order introduced in Section 3.2 . Relative reports the percent change from the row directly above; Absolute reports the percent change from the baseline (row 1).
Figure 14: Comparison of the MASSIF Ground Truth FIP 2 field against the prediction from the LDM (FIP) model for long rollout example 1. Row 3 reports relative absolute error where each pixel’s value is used to define the relative baseline.
Figure 15: Summary of cyclic fatigue based failure prediction on long rollout Example 2 using the LDM (FIP) model. In all plots, lines are colored by the number of simulation context steps provided before switching to the neural solver. (a) the relative L1 error of the predicted 95 -percentile FIP at each cycle order of magnitude. (b) A reinterpretation of (a): L1 error in the number of cycles to failure if failure occurred at each cycle order of magnitude. (c) predicted Emperical FIP CDF at the temporal boundary of training data. (d) predicted Empirical FIP CDF at maximum ground truth cycle count ( 215 cycles).
Figure 16: Comparison of the MASSIF Ground Truth FIP 2 field against the prediction from the LDM (FIP) model for long rollout example 2. Row 3 reports relative absolute error where each pixel’s value is used to define the relative baseline.
Figure 17: Hyperparameter optimization of the AViT Dynamics model for the Cyclic Fatigue Case Study. (a) comparison of standard training against training where the loss is weighted towards larger times. (b) hyperparameter optimization of the number of heads in the attention mechanism. (c) hyperparameter optimization of the dimensionality of the AViT embedding space. (d) hyperparameter optimization of the number of attention blocks.
Temporal surrogate models are effective for predicting chaotic dynamical systems where computational cost can be prohibitive. Several deep neural network architectures can be used for such purposes. In this work, a few commonly used architectures are compared using a common training protocol. The objective is to fairly assess the impact of model architectures for long-horizon prediction stability. Experiments are carried out for three problems, the double pendulum, the Kuramoto-Sivashinsky equations, and the Kolmogorov flow. The experiments are carried out with matching model capacity. Analysis is also carried out for a scenario where each model is individually optimized. It is observed that in both scenarios, the models exhibit categorical differences in long-horizon rollouts. For a concrete quantification, stepwise error injections and perturbation amplifications are analyzed using metrics such as local jacobian, relative one-step bias, and finite-time Lyapunov growth. Additionally, an attractor analysis is also conducted to assess how well the learned models replicate the underlying system geometry. An ablation study to isolate the impact of each component of a continuous-update architecture is also carried out. It is concluded that models that having integrator-like updates show lower bias and perturbation amplification yielding stable long-horizon rollout and more accurate predictions.
Neural PDE solvers provide efficient surrogates for time-dependent physical systems, but autoregressive prediction over long horizons remains challenging because local errors can induce distribution shift and accumulate under recursive deployment. We develop a variational approach to this problem by introducing latent Markov dynamics in which physical states are represented by latent distributions and evolved through probabilistic transitions. The framework is formulated directly on function spaces and specialized to functional Gaussian models, where structured latent perturbations induce a spectral geometry and variational transition alignment regularizes the learned dynamics. We further analyze how these mechanisms affect autoregressive error propagation, providing a theoretical connection between variational training and long-horizon prediction. We instantiate the framework as the Variational Autoencoding Markov Operator (VAMO), which combines spatially resolved latent fields, structured Gaussian perturbations, and a neural-operator transition. Empirically, we demonstrate the effectiveness of VAMO on several fluid-dynamics benchmarks with prediction horizons extending substantially beyond those represented during training, where it consistently reduces error accumulation and improves rollout stability over several deterministic and noise-injection baselines. Overall, these results highlight variational modeling as a complementary approach to robust long-horizon neural PDE dynamics.
Junyi Liao, Johann Guilleminot, Vahid Tarokh
†Department of Electrical and Computer Engineering, Duke University · ‡Department of Mechanical Engineering and Materials Science, Duke University
Autoregressive neural simulators now match classical solvers on short-horizon prediction of physical systems, yet their accuracy degrades rapidly when rolled out over long horizons. In this work, we identify transient amplification of perturbations around rollout trajectories as a structural mechanism driving rollout error. Using a linearization analysis we show that when the Jacobians along an autoregressive trajectory are non-normal and non-commuting, the model amplifies errors transiently, resulting in model rollout drift even when the overall system is asymptotically stable. Building on the analysis, we propose commutativity regularization: a combination of two penalties designed to reduce the normality defect of individual Jacobians and the commutator norm of Jacobians across steps. The penalties are estimated with Jacobian-vector products and have no inference-time cost. We show a propagator bound that quantifies rollout error under approximate commutativity and normality. We evaluate UNet and FNO variants with commutativity regularization on 1D and 2D spatio-temporal data in synthetic and real settings, showing successful long-horizon rollouts over thousands of steps. Further, we show that the method improves FourCastNet climate forecasts on ERA5 without using any new data. The gain is most pronounced out-of-distribution: trained on trajectories of a few hundred steps, regularized models remain in-distribution for thousands of rollout steps on initial conditions where baselines diverge.
Adeel Pervez, Francesco Locatello
Institute of Science and Technology Austria Klosterneuburg, Austria