Bidirectional Autoregressive Latent Diffusion for Forward and Inverse Magnetohydrodynamics

Abstract: This work presents a new bidirectional autoregressive latent diffusion approach for predicting the evolution of multiple fields (mass density, pressure, velocity, and magnetic field components) for magnetohydrodynamics. We show that this bidirectional flow can be used as a self-supervised consistency metric for uncertainty and error estimation, which enables the model to estimate test-time uncertainty and error without access to ground truth, by comparing how closely flowing forwards and backwards in time returns to the same predicted fields. We also demonstrate this methods's potential to serve as a non-invasive plasma diagnostic, and show how adaptive feedback can be used to make the model more robust based on sparse diagnostics or limited views/measurements.

• The paper introduces a bidirectional autoregressive latent diffusion model that accurately predicts forward MHD evolution and recovers past states.
• It employs a VAE-based encoder-decoder and conditional diffusion framework with a self-supervised cycle-consistency metric to quantify prediction error without ground truth.
• The model leverages adaptive feedback using bounded extremum seeking to enhance prediction fidelity under sparse diagnostics, making it viable for real-time plasma control.

Bidirectional Autoregressive Latent Diffusion for Forward and Inverse Magnetohydrodynamics

Introduction

"Bidirectional Autoregressive Latent Diffusion for Forward and Inverse Magnetohydrodynamics" (2606.29620) presents a comprehensive framework employing a latent-space conditional diffusion model for accurate, efficient, and flexible surrogate modeling of magnetohydrodynamics (MHD). The approach supports both forward-time and backward-time rollouts, providing solutions to both direct evolution and inverse recovery tasks across multiple physical fields—including mass density, pressure, velocity, and magnetic components—in time-dependent 2D MHD. The authors introduce a self-supervised cycle-consistency metric for error estimation at test-time in the absence of ground truth and demonstrate adaptive feedback to improve predictions under partial observation, positioning the model as a viable practical tool for real-time plasma diagnostics and control.

Figure 1: Conditional diffusion overview showing per-field encoding of each MHD field into a lower-dimensional latent representation.

Model Architecture

The core pipeline is built around an encoder-decoder variational autoencoder (VAE) that compresses each $512 \times 512$ MHD field into a $\mathbf{z} \in \mathbb{R}^{16\times 16 \times 4}$ latent. This delivers a 256$\times$ reduction in dimensionality, balancing reconstruction fidelity and tractable computational cost.

The VAE architecture uses a low KL-divergence penalty to maintain expressive latent structure without sacrificing generative flexibility. The entire dataset is encoded into this latent manifold, on which a denoising diffusion probabilistic model (DDPM) is trained to learn conditional transitions between adjacent time steps.

The DDPM predicts latent states either forward (as a surrogate for time evolution) or backward (for inferring past states from partial observations). The model is trained to jointly learn both directions, with a direction indicator $c_d \in \{+1, -1\}$ input to the transition network, leading to parameter sharing and enforcing a form of implicit consistency.

Figure 2: A: Forward rollouts in latent space predict MHD evolution; B: Backward rollouts solve the inverse problem by inferring prior states.

Self-Supervised Consistency Metric

A central contribution is the design of a self-supervised cycle-consistency mechanism to estimate prediction error in the absence of ground truth. The procedure involves rolling forward $i$ steps from an initial latent seed, then rolling backward $i$ steps from the forward-generated terminal latent. The discrepancy between the original seed and the roundtrip-recovered latent is the consistency error $\mathcal{C}_i$, which empirically correlates with true rollout error. This diagnostic is entirely self-supervised and can be averaged over stochastic samples or performed deterministically.

Figure 3: Consistency cycling for self-supervised error estimation—round-trip consistency deteriorates with cycle length, furnishing a test-time error surrogate.

Adaptive Prediction with Sparse Diagnostics

The framework includes a feedback adaptation loop leveraging limited observations (e.g., measuring density at final time only) to iteratively adjust the predicted latent states via bounded extremum seeking (ES), an algorithm which does not require analytic gradients. The approach incrementally minimizes the $L^1$ error between the measured and predicted fields, enhancing robustness and prediction fidelity when faced with partial or sparse diagnostics. This adaptive loop positions the latent diffusion model as a flexible digital twin for plasma diagnostics and control.

Empirical Evaluation

Model accuracy is established first by evaluating VAE reconstruction error in image space and in spectral domains. The reconstructions maintain fine-scale details as indicated by close agreement in predicted versus true spectral energy distributions up to high wavenumber, though the method is ultimately bottlenecked by decoder fidelity.

Figure 4: VAE reconstructions: true fields, latent embeddings, predicted fields, and per-image errors for $t=0,1$.

Figure 5: Predicted versus true energy spectra for randomly sampled validation examples.

Spatial and temporal evolution fidelity is examined via RGB-stacked field visualizations and autoregressive sampling, confirming the preservation of multi-field coupling and coherence over extended horizons.

Figure 6: True and predicted field composites at several timesteps for qualitative assessment.

The iterative nature of DDPM generation and time-stepwise predictions is illustrated, with error accumulation shown for extended rollouts in both time directions, revealing the compounding of transition model error—quantified precisely in mean squared error (MSE) metrics.

Figure 7: Denoising generative process visualized for multiple diffusion steps per time target in latent space.

Figure 8: Forward rollout latent prediction sequences, visualized for all six MHD fields.

Figure 9: Backward rollout latent prediction sequences, demonstrating recovery of earlier MHD states.

Statistical evaluation over the test set confirms symmetric error buildup in both forward and backward directions, supporting the validity of the bidirectional model.

Figure 10: Mean per-step error with standard deviation for both forward and backward rollouts across the test dataset.

Figure 11: Long-horizon reconstructions: model accuracy maintained over dozens of steps backward and forward; degradation patterns visible at deep horizons.

Implications and Theoretical Perspective

The approach provides a tractable surrogate for computationally expensive MHD simulations, supporting both predictive modeling and inverse analysis with the same set of parameters. The ability to roll out backwards opens new avenues for real-time operational inverse problems, such as state recovery from limited diagnostics.

The self-supervised error estimation establishes a paradigm for trustable error quantification in scientific generative models, overcoming the common challenge of ground-truth absence in deployment. The adaptive feedback using bounded ES demonstrates the extensibility to real-world scenarios with limited, noisy, and partial observations, a key requirement for autonomous plasma control systems.

From a theoretical perspective, the tight integration of latent diffusion models, cycle-consistency diagnostics, and adaptive optimization provides a blueprint for future AI-accelerated physical modeling, particularly for high-dimensional, time-dependent, multi-field systems where ground truth is expensive or unavailable.

Future Directions

The study outlines several future enhancements, notably the integration of hard physics constraints using Physics-Constrained Neural Networks (PCNNs)—e.g., enforcing $\nabla\cdot\mathbf{B}=0$—to regularize the generative process and guarantee compliance with MHD and Maxwellian invariants. Other relevant directions include hybridization with neural operators, further incorporating out-of-distribution detection for robust on-the-fly adaptation, and pushing the framework to 3D or real-world experimental geometries.

Conclusion

This work advances the state-of-the-art in deep learning-based surrogate modeling for MHD by providing a bidirectional latent diffusion methodology, equipped with rigorous self-supervised error metrics and adaptive feedback mechanisms. The model is capable of efficient forward and inverse predictions, real-time error estimation, and online adaptation, making it directly relevant to operational plasma diagnosis and control. The flexible design and extensible evaluation suggest strong promise for practical application and further integration with physics-based priors and feedback control systems.