Learned Coarse Models for Efficient Turbulence Simulation

TL;DR

This work demonstrates that fully learned turbulence simulators operating on coarse grids can outperform comparably coarse traditional solvers, particularly in preserving high-frequency content. A general-purpose Dilated ResNet with encode-process-decode structure achieves strong performance across diverse turbulent domains, including Athena++-generated data, and remains stable when stabilized by training noise and temporal downsampling. The approach offers significant speedups and flexibility, suggesting a viable path to distill expensive high-resolution solvers into efficient coarse-grid models. While generalization to out-of-distribution conditions remains imperfect, targeted regularization and training strategies improve robustness and point to practical applications in large-scale or real-time turbulence tasks.

Abstract

Turbulence simulation with classical numerical solvers requires high-resolution grids to accurately resolve dynamics. Here we train learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution. We show that our proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the comparably low resolutions across various scientifically relevant metrics. Our model is trained end-to-end from data and is capable of learning a range of challenging chaotic and turbulent dynamics at low resolution, including trajectories generated by the state-of-the-art Athena++ engine. We show that our simpler, general-purpose architecture outperforms various more specialized, turbulence-specific architectures from the learned turbulence simulation literature. In general, we see that learned simulators yield unstable trajectories; however, we show that tuning training noise and temporal downsampling solves this problem. We also find that while generalization beyond the training distribution is a challenge for learned models, training noise, added loss constraints, and dataset augmentation can help. Broadly, we conclude that our learned simulator outperforms traditional solvers run on coarser grids, and emphasize that simple design choices can offer stability and robust generalization.

Figures (16)

• Figure 1: (a) Learned simulator framework. The model is trained to predict the difference between the current and next state. Gaussian noise can be added to the inputs during training to produce models that are robust to small perturbations. Shown below is the schematic for the Dilated ResNet model (Dil-ResNet). (b-e) Frames predicted by rollouts from the learned simulator model compared to ground truth frames. (b) KS-1D: The model follows the ground truth closely for the first 150 steps (t $<$ 75) and remains plausible thereafter. (c) Incomp-2D: The model remains accurate after 119 model steps (91,392 solver steps). (d) CompDecay-3D: The model remains accurate after 31 model steps (1984 solver steps). (e) CoolMixLayer-3D. The model remains qualitatively accurate after 59 model steps (59,000 solver steps) of a box size of ($L=0.75$). Videos available at https://sites.google.com/view/learned-turbulence-simulators.
• Figure 2: Comparison between learned $32^3$ Dil-ResNet and same-resolution $32^3$Athena++, intermediate resolution $64^3$Athena++, and ground truth high resolution $128^3$Athena++ in CompDecay-3D. (a) Energy Field RMSE ($y$-axis) for Dil-ResNet ($32^3$ resolution (blue)) and Athena++ ($64^3$ (dark gray), and $32^3$ (light gray)) on a test trajectory the window of times seen during training. Dil-ResNet RMSE is lower than that of the comparably coarse Athena++$32^3$ but higher than Athena++$64^3$. (b) Energy Field RMSE ($y$-axis) as a function of rollout duration ($x$-axis) over the training window (white background) and generalizing beyond (pink background). The learned simulator's error grows faster starting around 3x the training window (at $t\approx 3000$). Each blue line corresponds to a different seed. Ground truth (red) error is at 0 and hidden by $x$-axis. (c) Energy Spectrum RMSE ($y$-axis). Dil-ResNet has lower error than both comparable and higher resolution $32^3$ and $64^3$Athena++ models. (d) Energy Spectrum RMSE ($y$-axis) over rollout. Dil-ResNet's rollouts maintain lower Energy Spectrum error for up to 4$\times$ the training duration. (e) Energy Spectrum for the learned and Athena++ simulators at rollout $t\approx1$ s (the end of the training window). The $x$-axis represents spatial frequency, and the $y$-axis represents spectral power. Compared to Dil-ResNet, the coarse Athena++ models lose power in the high frequency range. (f--i) Sample energy Fields at $t\approx1$ s for Dil-ResNet and the three Athena++ resolutions. States from coarse Athena++ resolutions lose high frequency detail with respect to the ground truth, which the coarse learned model captures.
• Figure 3: Effects of noise and temporal downsampling on rollout stability. (a) One step errors are larger for models trained with noise. Note the error spikes are very small and are not model-related general artifacts, but specific to particular frames of this test trajectory. (b) However, models trained without noise can yield unstable rollouts, especially when using very small time steps, which is not a problem for models trained with noise. (c, e) One-step model error rises monotonically with coarser temporal downsampling. (d, f) Rollout error has a U-shaped curve over temporal downsampling factors, for a trajectory of the same time duration, with minimum error around $\Delta t=0.032$.
• Figure 4: Comparison across learned models, contrasting noise and no-noise training conditions, across the three primary tasks (a) KS-1D, (b) Incomp-2D, and (c) CompDecay-3D. With a few exceptions, the various learned models had comparable performance, though the Dilated ResNets (Dil-ResNet, Con-Dil-ResNet) consistently have the lowest error. In KS-1D (a), the noise harmed performance; in Incomp-2D (b) the noise particularly benefits the FNO rollouts; in CompDecay-3D (c) the noise mainly stabilized rollouts.
• Figure 5: Generalization outside of the training distribution. (a) Generalization to different initial conditions. We vary the ratio of compressive and solenoidal components in the initial velocity field. Solid markers indicated the training region. We find that, compared to coarse Athena++, generalization to more compressive initial states is challenging for learned models and inconsistent across seeds, but that loss constraints (Con-Dil-ResNet) ameliorate this to some extent. (b) Cooling velocity generalization as function of the box size in the mixing layer. The black line indicates the ground truth cooling velocity averaged across time for 8 test trajectories with different box sizes, with one standard deviation represented as the shaded region. None of 3 seeds of Dil-ResNet trained on a single length of $L=0.75$ (blue) generalize outside the training range. Dil-ResNet trained on multiple lengths (orange) shows better generalization performance. However, generalization far beyond training box sizes remains a challenge.