Autoregressive prediction of 2D MHD dynamics inferred from deep learning modeling

Abstract: We develop two deep learning surrogate autoregressive models for the prediction of the temporal evolution of two-dimensional ideal magnetohydrodynamic (MHD) Kelvin-Helmholtz instabilities across a range of magnetic field strengths. Using two neural network architectures, a Koopman-based Transformer model and a ConvLSTM-UNet, our approach enables simultaneous prediction of vorticity and current density directly from high-resolution simulations. The models are trained in an autoregressive manner and are able to reproduce key features of the multiscale dynamics over several instability growth and nonlinear saturation phases. Beyond accurate field reconstruction, the surrogates preserve essential physical structures of ideal MHD dynamics, including the conservation trends of global invariants and the propagation of Alfvénic fluctuations. Compared to direct numerical simulations, the proposed surrogates offer substantially reduced computational cost while maintaining good agreement with the reference dynamics. These results suggest that deep learning based surrogate models can provide a promising complementary tool for the efficient and physically consistent exploration of high-fidelity plasma and fluid simulations.

• The paper introduces autoregressive deep learning models (KT-MHD and ConvLSTM-UNet) for forecasting 2D MHD Kelvin-Helmholtz instability with high computational efficiency.
• The models integrate convolutional, transformer, and ConvLSTM architectures to capture complex spatio-temporal dynamics while preserving key MHD invariants.
• Quantitative evaluations show KT-MHD excels in current density prediction and ConvLSTM-UNet in local vorticity accuracy, demonstrating complementary strengths.

Autoregressive Deep Learning Models for 2D MHD Kelvin-Helmholtz Instability Forecasting

Introduction

The paper "Autoregressive prediction of 2D MHD dynamics inferred from deep learning modeling" (2604.18221) investigates the forecasting of two-dimensional incompressible ideal magnetohydrodynamic (MHD) Kelvin-Helmholtz instability via deep learning surrogates. The study introduces and benchmarks two neural architectures: a Koopman-based Transformer (KT-MHD) and a ConvLSTM-UNet, both trained in an autoregressive context for direct spatio-temporal prediction of coupled vorticity ($\omega$) and current density ($j$) fields. The models are designed to reproduce high-fidelity Direct Numerical Simulation (DNS) outputs with strong physical consistency and substantial computational efficiency gains.

Physical and Numerical Framework

The MHD system is formulated in terms of coupled nonlinear transport equations for vorticity and current density fields under periodic boundary conditions and in the ideal, incompressible limit. The canonical Kelvin-Helmholtz instability, which is strongly regulated by magnetic tension, is employed as the dynamical test case due to its paradigmatic role in plasma turbulence, mixing, and vortex roll-up. High-resolution training and test datasets are generated with a third-order accurate Characteristic Mapping Method (CMM), ensuring minimal numerical diffusion and enabling the capture of multiscale MHD structures, including thin current sheets.

Deep Learning Model Architectures

Koopman Transformer for MHD

The KT-MHD is a hybrid convolutional-transformer model integrating inductive biases from Koopman operator theory.

• A 3-layer convolutional encoder projects ($\omega$, $j$, $B_0$) into a spatio-temporal latent representation.
• Stacked temporal transformer blocks (multi-head self-attention) extract long-range temporal dependencies at fixed spatial positions.
• Latent temporal advancement is realized via a locally linear Koopman operator: a $3\times3$ convolution plus $1\times1$ mixing enforces spatially local linear evolution in latent space, as per the finite-dimensional Koopman operator surrogate requirement.
• Autoregressive rollout is performed in latent space, with independent convolutional decoders for vorticity and current density.

  Figure 1: KT-MHD architecture—convolutional encoders, temporal transformers, convolutional Koopman rollouts, variable-specific decoders.

ConvLSTM-UNet

This architecture fuses the multi-scale image-to-image regression capacity of the U-Net with the temporal memory of ConvLSTM layers.

• The encoder arm consists of convolutional downsampling blocks with integrated ConvLSTM operations for local temporal modeling at each spatial scale.
• The deepest layer (bottleneck) aggregates global temporal information.
• The decoder performs upsampling with skip connections, re-injecting high-resolution spatial features, and applies ConvLSTM at each stage to refine the temporal dynamics.

  Figure 2: ConvLSTM-UNet architecture—U-Net encoder-decoder with skip connections, ConvLSTM blocks at each spatial resolution.

Autoregressive Strategy and Training Procedure

Both architectures employ an autoregressive block-prediction strategy: given a temporal window, the model predicts multiple frames ahead; predictions are then recycled as input for further forecasting. Hyperparameters for window size, horizon, and rollouts are optimized via Optuna. Model training is performed using MSE loss aggregated over all autoregressive steps.

Figure 3: Autoregressive prediction mechanism—predicted frames (stars) fed back recursively for long-horizon rollout.

Quantitative Performance and Physical Consistency

Short- and Long-Term Accuracy

On test sequences at unseen magnetic field strengths ($B_0=0.08$ and $0.12$), both surrogates robustly reproduce the nonlinear roll-up and turbulent features of the Kelvin-Helmholtz instability. ConvLSTM-UNet yields superior scores for vorticity in terms of MSE, SSIM, and PSNR, indicating higher-fidelity image reconstruction and better preservation of localized gradients. The KT-MHD exhibits superior performance on current density—especially at longer horizons—retaining sharper magnetic sheets and demonstrating reduced smoothing compared to ConvLSTM-UNet.

Figure 4: Comparison of true and predicted vorticity fields—KT-MHD, DNS, ConvLSTM-UNet, progression over rollout time.

Figure 5: Comparison of true and predicted current density fields—KT-MHD, DNS, ConvLSTM-UNet.

Error accumulation is inherent with increasing forecast horizon; ConvLSTM-UNet's local history mechanisms slow error propagation in $\omega$, while KT-MHD's global latent coherence stabilizes $j$0 (Tables 1–4).

Spectral and Invariant Diagnostics

Energy spectra are analyzed via compensated plots, revealing that both models replicate the correct $j$1 scaling in the inertial range and accurately capture energy transfer at large and intermediate scales. Small scales are over- or under-dissipated depending on the model, with ConvLSTM-UNet more prone to fine-scale diffusion and KT-MHD better preserving multiscale energy.

Figure 6: Compensated total-energy spectra $j$2 showing inertial range scaling—ground truth, KT-MHD, ConvLSTM-UNet.

Both models approximately conserve global MHD invariants, but ConvLSTM-UNet demonstrates significantly reduced drift in total energy and cross helicity over long rollouts (1.24% versus 8.10% mean relative error for energy), as measured by invariant-based error metrics. Nevertheless, KT-MHD sometimes outperforms on certain sequences, reflecting greater variance and regime dependence across initial conditions (Figure 7, Table of invariants).

Figure 7: Evolution of total energy $j$3 for DNS, KT-MHD, and ConvLSTM-UNet.

Generalization and Regime Interpolation

The models are explicitly tested for interpolation ($j$4; within training regime) and extrapolation ($j$5; outside training regime). Both demonstrate stable prediction and robust generalization, but performance metrics (MSE, SSIM, PSNR) consistently improve under interpolation, with a relative benefit exceeding 35–50% in MSE for both variables (see Appendix tables).

Computational Efficiency

Once trained, both models predict high-resolution field time-series several orders of magnitude faster than DNS (speedup of $j$6 per frame). KT-MHD achieves this with just 2.5M parameters and $j$71.5 hours of GPU training versus 4.8M parameters and 7+ hours for ConvLSTM-UNet.

Implications and Future Directions

This work confirms that modern deep time-series models can act as physically consistent, data-driven surrogates for ideal 2D MHD dynamics even with limited training data and minimal parameter complexity. The KT-MHD's capacity for spectral and current sheet consistency complements the ConvLSTM-UNet's advantage in global invariant conservation and local structure retention. Hybridizations integrating global attention, latent linearity, and local convolutional memory may further enhance stability, spectral accuracy, and physical constraints preservation.

Potential future research avenues include:

• Incorporation of explicit physical constraints (divergence-free projections, invariant penalties) during training.
• Extension to more challenging MHD regimes (e.g., resistive-MHD, forced turbulence, or Hasegawa--Wakatani systems).
• Increased dataset diversity for improved generalization and uncertainty quantification.

Conclusion

Autoregressive deep-learning models, when carefully constructed and tuned, provide efficient, physically faithful surrogates for high-dimensional MHD turbulence. They reliably reproduce multiscale dynamics, conserve global invariants, and deliver strong generalization, enabling rapid exploration and prototyping for plasma physics and related disciplines.