cs.LGSep 24, 2026
SaveBeyond Compression: Training Latent Representations for Stable Long-Horizon Rollout in Neural Surrogate Solvers
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
Abstract
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 error for the vanilla LDM, the UNet, and the proposed LDM, respectively, at rollout steps ; 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, . (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 step rollout (averaged over the entire dataset). (d) Model comparison of allocated VRAM (Gigabytes - GB) versus RelMSE at 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, , respectively, at rollout steps ; each model is given 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, and . (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, . Row 2 depicts corresponding samples of one of the 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 sample depicted in Row 1. Variance was calculated over samples generated by injecting standard deviation 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 | |||||||||
|---|---|---|---|---|---|---|---|---|---|
| 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 .
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 ). (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 error of the predicted -percentile FIP at each cycle order of magnitude. (b) A reinterpretation of (a): 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 ( 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 standard deviation across 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 |
|---|---|---|
| 175 | MPa | |
| 25 | MPa | |
| 200 | MPa | |
| 7000 | MPa | |
| 100 | MPa | |
| 4 | – |
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 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 | ||||
|---|---|---|---|---|
| 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 |
Table 3: Average ReLMSE for predictions for the spinodal decomposition benchmark for various models and rollouts .
| Model | ||||
|---|---|---|---|---|
| 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 |
Table 4: Average VRMSE for predictions for the spinodal decomposition benchmark for various models and rollouts .
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 -step rollout error for the same hyperparameter sweep. Note the lack of correlation between trends in each.
| Model | Channels: | Channels: | ||
|---|---|---|---|---|
| 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 |
Table 5: Sensitivity of cumulative intervention gains to latent channel count. Average VRMSE at rollout on the active matter benchmark, comparing an -channel and an -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 -step VRMSE error for its corresponding field variable. Average error over the validation set is reported for 3 starting times: Time Steps , , and . At each time step, the first steps are provided to the model as context to begin rollout.
| Model | |||||||||
|---|---|---|---|---|---|---|---|---|---|
| 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 |
Table 6: Cumulative ablation of training-level interventions. Error by model intervention and rollout length 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 error of the predicted -percentile FIP at each cycle order of magnitude. (b) A reinterpretation of (a): 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 ( 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.