1 Introduction↩︎

Magnetohydrodynamics (MHD) describes plasma dynamics modeled as an electrically conducting fluid under the influence of its own electromagnetic fields [1]. MHD describes systems that include plasma turbulence [2], [3], quantum plasmas [4], Tokamaks [5]–[9], plasma detachment in magnetic nozzles [10], liquid metals [11], [12], and ion beams in particle accelerators [13].

In this work we study the use of bidirectional autoregressive latent diffusion models that can flow forward or backback in time to predict the evolution of multiple fields (mass density, pressure, velocity, and magnetic field components) for magnetohydrodynamics. We start by creating and training with arbitrary mixture-of-Gaussian synthetic MHD densities and then test those models on the Orszag–Tang vortex, which is a widely used test case for 2D MHD simulations [14].

We show that this bidirectional flow can be used as a self-supervised consistency check for our generative model for test-time estimation of prediction errors by comparing whether cycling forwards and backwards in time returns to the same predicted fields. We also demonstrate this method’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/measurements.

Such generative models have the potential to provide accurate surrogates of MHD simulations which run much faster than traditional approaches, which could be useful for complex plasma processes that would benefit from real-time data processing [15] and real-time data analysis [16]. Although not as accurate as physics models, these methods have two main benefits in being very fast for real-time applications such as finding approximate solutions, and the ability to roll backward in time is useful for solving inverse problems for operational systems with limited direct observations by providing a good initial guess of the plasma’s state. Overall, such an approach can potentially contribute to provide virtual plasma diagnostics which can be used for real-time feedback control for complex plasma physics systems. This would extend to plasmas work that is already underway for using adaptively tuned advanced physics-constrained latent diffusion-based generative models as real-time diagnostics for time-varying beams in charged particle accelerators based on limited views/measurements [17].

1.1 Summary of Main Contributions↩︎

In this work, we present a new autoregressive latent diffusion model for predicting the evolution of multiple fields (mass density, pressure, velocity, and magnetic field components) for magnetohydrodynamics. Our main contributions are summarized as:

• Ability to roll out predictions autoregressively both forward and backward in time, thus serving as both a surrogate model and an inverse problem solver that can recover initial and intermediate plasma states from final measured output conditions.

• Self-supervised consistency metric that allows the model to estimate test-time uncertainty and error without access to ground truth.

• Demonstration of how adaptive feedback coupled with this approach can be used to make the model more robust as a virtual plasma diagnostic based on sparse physical diagnostics or limited views/measurements.

2 2D MHD Model↩︎

For modeling 2D MHD, we utilize the tool and approach of P. Mocz [18]–[21], which considers the conservation law dynamics \[\frac{\partial \mathbf{U}}{\partial t} + \nabla \cdot \mathbf{F}(\mathbf{U}) = 0,\] \[\mathbf{U} = \begin{pmatrix} \rho \\ \rho\mathbf{v} \\ \rho e \\ \mathbf{B} \end{pmatrix}, \quad \mathbf{F}(\mathbf{U}) = \begin{pmatrix} \rho \mathbf{v} \\ \rho\mathbf{v}\mathbf{v}^T + p - \mathbf{B}\mathbf{B}^T\\ \rho e \mathbf{v} + p\mathbf{v} - \mathbf{B}(\mathbf{v}\cdot \mathbf{B}) \\ \mathbf{B} \mathbf{v}^T - \mathbf{v}\mathbf{B}^T, \end{pmatrix}\] where \(p=p_\mathrm{gas} + \frac{1}{2}\mathbf{B}^2\) is the total gass pressure, \(e=u+\frac{1}{2}\mathbf{v}^2+\frac{1}{2\rho}\mathbf{B}^2\) is the total energy per unit mass, and \(u\) is the thermal energy per unit mass.

For this work we generate synthetic 2D fields with randomized initial density, pressure, and velocity profiles and the magnetic fields are generated by first creating a random vector potential which is a mixture of 100 Gaussians, from which the initial magnetic field components are then calculated. All of the simulations are run at a \(512\times 512\)-pixel resolution, with each simulation saving 100 time-steps over a 1-second run. We generate 500 such random MHD trajectories for training and 50 additional trajectories as test data.

3 Latent Diffusion↩︎

Although generative diffusion directly in image-space gives the highest accuracy results, working with such large \((512\times 512)\) images can make the iterative diffusion flow-based generative process very slow, both for training and for inference, because to generate data, the models learn to flow from an initial analytically known noise distribution \(\mathcal{N}(\mathbf{0},I_{d\times d\times c})\) to their estimate \(\hat{\rho}_\mathrm{data}\) of the probability distribution describing the provided data samples \(\rho_\mathrm{data}\). Here \(d\times d\) represents image size and \(c\) is the number of image channels, which, for example, for RGB images is typically 3, or 4 if they include opacity. This flow takes the form of a finite-difference approximation of a stochastic diffusion differential equation, which requires tens to hundreds of steps for the generative process. For very large images, using a very large denoising diffusion model, this becomes very computationally expensive.

The latent diffusion approach is to first compress the data to a much smaller latent representation by using the encoder of a variational autoencoder (VAE), \(\mathcal{E}_\theta\), which maps images \(X\in\mathbb{R}^{d\times d \times c}\) down to latent representations \(\mathcal{E}_\theta(X)=\mathbf{z}\in\mathbb{R}^{L_d \times L_d \times L_c}\) with \(L_d \times L_d \times L_c \ll d\times d\times c\). The overall setup is shown in Figure 1.

3.1 Variational Autoencoder↩︎

In our approach, for any single field, \(d\times d = 512\times 512 = 262144\), while \(L_d \times L_d \times L_c = 16 \times 16 \times 4 = 1024\), resulting in a \(256\times\) factor of data compression. The job of the VAE decoder, \(\mathcal{D}_\theta\), is to then reconstruct the input as accurately as possible based only on the low-dimensional latent representation, resulting in an overall data flow of \[X \quad \underbrace{\xrightarrow{}}_{\mathcal{E}_\theta}\quad \mathbf{z} \quad \underbrace{\xrightarrow{}}_{\mathcal{D}_\theta} \quad \hat{X}.\] Decreasing latent space dimension increases compression amount, making the VAE’s job more difficult and decreases the accuracy of reconstructions \(\hat{X}\). Optimal latent dimension is unique to each data set, it depends on the image resolution, the intrinsic dimensionality of the manifold on which the data lives, and on the number of weights \(\theta\) and the architecture of the VAE. In practice this needs to be tested for each new dataset. In this approach we found that going down to \(\mathbf{z}\in\mathbb{R}^{8\times 8 \times 4}\) made it too difficult for our VAE to recover fine details of the fields, while keeping a larger \(\mathbf{z}\in\mathbb{R}^{32\times 32 \times 4}\) only showed minor improvement relative to \(\mathbf{z}\in\mathbb{R}^{8\times 8 \times 4}\).

The VAE’s architecture is actually a little bit different from that of a standard autoencoder, rather than simply generating \(\mathbf{z}=\mathcal{E}_\theta(X)\) directly, the encoder’s output is actually split into two parts which are then used to define a normal distribution from which the latent representation is sampled before being passed to the decoder \[\begin{align} && X \;\underbrace{\xrightarrow{}}_{\mathcal{E}_\theta} \;[\Sigma_\theta, \mu_\theta]=\mathcal{E}_\theta(X),\\ && \mathbf{z} = \mu_\theta + \Sigma_\theta \epsilon, \quad \epsilon \sim \mathcal{N}(\mathbf{0},I_{L_d \times L_d \times L_c}), \\ && \mathbf{z} \;\underbrace{\xrightarrow{}}_{\mathcal{D}_\theta} \;\hat{X} = \mathcal{D}_\theta(\mathbf{z}). \end{align}\] Note that compared to a standard autoencoder (AE), whose generative process takes the form \[\begin{align} && X \;\underbrace{\xrightarrow{}}_{\mathcal{E}_\theta} \;\mu_\theta=\mathcal{E}_\theta(X),\\ && \mathbf{z} = \mu_\theta, \\ && \mathbf{z} \;\underbrace{\xrightarrow{}}_{\mathcal{D}_\theta} \;\hat{X} = \mathcal{D}_\theta(\mathbf{z}) = \mathcal{D}_\theta(\mathcal{E}_\theta(X)), \end{align}\] the VAE can be thought of as an AE in which even for a fixed input \(X\), random noise is added to slightly perturb the latent representations \(\mathbf{z}=\mathcal{E}_\theta(X)\), which makes the decoder’s job more difficult as it must try and reconstruct the same \(X\) from slightly perturbed latent representations.

Another major difference is that for a VAE, the training loss is not only reconstruction error, but there is an additional Kullback-Leibler (KL)-divergence term which tries to force the latent space to look like \(p(\mathbf{z})=\mathcal{N}(\mathbf{0},I_{L_d \times L_d \times L_c})\). VAE training loss is defined as \[\mathcal{L} = -\mathbb{E}_{q_\theta(\mathbf{z}|X)}\left[\log p_\theta(X|\mathbf{z})\right] + \beta\, D_{\mathrm{KL}}\!\left(q_\theta(\mathbf{z}|X)\,\|\,p(\mathbf{z})\right),\] where \(q_\theta(\mathbf{z}|X)\) is the probability distribution of the encoder and \(p(\mathbf{z})\) is the prior that we want to match. The KL divergence term \(D_{KL}\) is defined in general as \[D_{\mathrm{KL}}\!\left(q_\theta(\mathbf{z}|X)\,\|\,p(\mathbf{z})\right) = \int q_\theta(\mathbf{z}|X)\,\log\frac{q_\theta(\mathbf{z}|X)}{p(\mathbf{z})}\,d\mathbf{z},\] and for our VAE setup with Gaussian \(p(\mathbf{z})\) it simplifies to \[\frac{1}{2}\sum_{i=1}^{L_d}\sum_{j=1}^{L_d}\sum_{c=1}^{L_c}\left([\mu_\theta]_{ijc}^2 + [\Sigma_\theta]_{ijc}^2 - 1 - \log[\Sigma_\theta]_{ijc}^2\right).\] The idea is that if the KL loss term \(\beta>0\) is weighted heavily and the latent space truly converges to \(p(\mathbf{z})=\mathcal{N}(\mathbf{0},I_{L_d \times L_d \times L_c})\), then the VAE can be used in a purely generative form by simply sampling from \(\mathcal{N}(\mathbf{0},I_{L_d \times L_d \times L_c})\) and passing those samples through the decoder to generate random samples \(\hat{X}\). If we instead use a smaller \(\beta \ll 1\), then the VAE has more freedom to arrange the latent distribution in an optimal way while developing interesting latent structure. Because our approach will be to do conditional diffusion-based sampling, where we guide the diffusion generative process, we do not need to strictly force convergence to the latent prior and therefore we use \(\beta=10^{-3}\) to allow more flexibility in the learned latent space while also gently pushing that latent representations to stay small as a regularization. It is worth noting that very similar results are achieved with a vanilla AE by simply adding an \(L_2\) norm-based regularization on the latent encoding.

3.2 Denoising Diffusion Probabilistic Models↩︎

To get an idea of how generative diffusion models work, we consider an analytically intractable data distribution \(p_{\mathrm{data}}\) of high-dimensional objects \(\mathbf{x}\) such as our latent embeddings of the 2D MHD projections with \(\mathbf{x}\in\mathbb{R}^{d} = \mathbb{R}^{16\times 16\times 4}\), so \(d=1024\). The goal of generative diffusion models is to generate new samples of \(p_{\mathrm{data}}\) by transporting samples from an easy-to-sample analytical distribution \(p_{\mathrm{target}}\). To do this we consider the diffusion SDE: \[\mathbf{x}(0) \sim p_{\mathrm{data}}, \quad \, \text{d}\mathbf{x}= \boldsymbol{\mu}(\mathbf{x},t)\, \text{d}t + \boldsymbol{\sigma}(\mathbf{x},t)\, \text{d}w, \label{diffSDE}\tag{1}\] where \(w(t)\) is a Wiener process, \(\boldsymbol{\mu}(\mathbf{x},t)\) is known as the drift term and \(\boldsymbol{\sigma}(\mathbf{x},t)\) defines a diffusion tensor \(\mathbf{D}(\mathbf{x},t)=\boldsymbol{\sigma}(\mathbf{x},t)\,\boldsymbol{\sigma}^T(\mathbf{x},t)/2\). Consider the probability path \(p(\mathbf{x},t)\) driven by \(\boldsymbol{\mu}\) and \(\boldsymbol{\sigma}\) such that \(p(\mathbf{x},0)=p_{\mathrm{data}}(\mathbf{x})\) and \(p(\mathbf{x},1)=p_{\mathrm{target}}(\mathbf{x})\), then SDE (1 ) transports samples \(\mathbf{x}(0) \sim p_{\mathrm{data}}\) to samples \(\mathbf{x}(1) \sim p_{\mathrm{target}}\).

As proven in [22], there exists a reverse-time SDE (1 ): \[\begin{align} \text{d}\mathbf{x}&=& \left [ \boldsymbol{\mu}(\mathbf{x},t) - \nabla_\mathbf{x}\cdot \left(\boldsymbol{\sigma}(\mathbf{x},t)\boldsymbol{\sigma}^T(\mathbf{x},t)\right) \right . \nonumber \\ && \left . - \boldsymbol{\sigma}(\mathbf{x},t)\,\boldsymbol{\sigma}^T(\mathbf{x},t) \nabla_\mathbf{x}\log p\left(\mathbf{x}, t\right) \right ]\, \text{d}t \nonumber \\ && + \boldsymbol{\sigma}(\mathbf{x},t) \,\text{d} w, \label{diffSDEreverse} \end{align}\tag{2}\] which is used to transport samples \(\mathbf{x}(1)\sim p_{\mathrm{target}}\) for generating novel samples \(\mathbf{x}(0)\sim p_{\mathrm{data}}\). A generative diffusion model utilizes a UNet or Diffusion Transformer \(\mathbf{M}^\theta(\mathbf{x},t)\) to model the reverse SDE (2 ) as: \[\mathbf{x}(1)\sim p_{\mathrm{target}}, \quad \,\text{d}\mathbf{x}= \mathbf{M}^\theta(\mathbf{x},t)\,\text{d}t + \boldsymbol{\sigma}(\mathbf{x},t)\,\text{d}w \label{diffSDEreverse95NN}\tag{3}\] where model weights \(\theta\) are parameterized by forward sample paths of Eq. (1 ) from data samples.

This is a very general class of generative models based on the diffusion SDEs, they provide a unifying framework that connects a broad spectrum of modern generative models, including normalizing flows [23], flow matching [24], consistency models [25], posterior mean matching [26], energy models [27], and optimal transport [28]. For example, continuous-time normalizing flow (Neural ODEs) can be interpreted as diffusion processes with vanishing diffusion and learnable drift and consistency models can be viewed as alternative parametrizations of diffusion-based score learning.

Once our VAE is trained, we freeze the encoder and decoder and we generate latent representations of all of our data, which gives us our latent training dataset \(D_\mathbf{z}=\{ \mathbf{z}_1, \mathbf{z}_2, \dots, \mathbf{z}_N\}\) with \(\mathbf{z} \sim \rho_{\theta}(\mathbf{z})\) which is the VAE-based encoding our our entire training data set \(D_X=\{X_1,X_2,\dots,X_N\}\) with \(X_i \sim \rho_{\mathrm{data}}(X)\). The goal of our denoising diffusion probabilistic model (DDPM) is to learn how to generate estimates \(\hat{\mathbf{z}}\) of samples \(\mathbf{z} \sim \rho_{\theta}(\mathbf{z})\) in two particular forward time flow and backward time flow conditioned auto-regressive ways such that we recover samples according to \[\begin{align} \hat{\mathbf{z}}_{t+1} &\sim& p_\theta(\mathbf{z}_{t+1} | \mathbf{z}_t, \mathbf{z}_{t-1},c_d=+1), \\ \hat{\mathbf{z}}_{t-2} &\sim& p_\theta(\mathbf{z}_{t-2} | \mathbf{z}_t, \mathbf{z}_{t-1},c_d=-1). \end{align}\]

Encoding a full MHD trajectory \(\{X_t\}_{t=1}^{T}\) with the frozen VAE yields a latent trajectory \(\{\mathbf{z}_t\}_{t=1}^{T}\) distributed according to a joint latent density \(\rho_\theta(\mathbf{z}_1,\dots,\mathbf{z}_T)\) that inherits the temporal correlations of the underlying plasma dynamics. Rather than model this joint density directly, the diffusion model learns its local conditional factors. Because the dynamics may be unrolled in either temporal direction, we condition on an ordered pair of adjacent latent states together with a direction flag \(c_d\in\{+1,-1\}\) and learn the single bidirectional transition density \[\hat{\mathbf{z}}_{t+c_d}\;\sim\; p_\theta\!\left(\mathbf{z}_{t+c_d}\;\middle|\;\mathbf{z}_t,\;\mathbf{z}_{t-c_d},\;c_d\right),\] in which a single set of weights \(\theta\) represents both temporal directions. Taking \(c_d=+1\) gives the forward transition \(p_\theta(\mathbf{z}_{t+1}\mid\mathbf{z}_t,\mathbf{z}_{t-1},+1)\), which advances the system one step into the future, while \(c_d=-1\) gives the backward transition \(p_\theta(\mathbf{z}_{t-1}\mid\mathbf{z}_t,\mathbf{z}_{t+1},-1)\), which reconstructs the preceding state and lets the same model serve as an inverse solver. Under the second-order Markov approximation implied by conditioning on two states, these are exactly the two-sided conditional factors of \(\rho_\theta(\mathbf{z}_1,\dots,\mathbf{z}_T)\) obtained by applying the chain rule in increasing or decreasing time order.

A single denoiser network realizes this density by learning to reverse a Gaussian noising of the target latent \(\mathbf{z}_{t+c_d}\) while conditioned on the context \((\mathbf{z}_t,\mathbf{z}_{t-c_d},c_d)\); drawing both signs of \(c_d\) from the training trajectories forces the shared weights to denoise consistently toward the future and toward the past. Longer trajectories are then generated autoregressively: starting from a seed pair of adjacent latent states, each newly sampled state is appended to the context and used to predict the next in the chosen direction, so that the joint density of a generated segment factorizes as a product of learned transitions, \[\begin{align} && p_\theta\!\left(\mathbf{z}_{t_0+c_d},\,\mathbf{z}_{t_0+2c_d},\dots\;\middle|\; \mathbf{z}_{t_0},\,\mathbf{z}_{t_0-c_d},\,c_d\right) \nonumber \\ && =\prod_{m\ge 1} p_\theta\!\left(\mathbf{z}_{t_0+mc_d}\;\middle|\; \mathbf{z}_{t_0+(m-1)c_d},\,\mathbf{z}_{t_0+(m-2)c_d},\,c_d\right). \end{align}\] The associated physical fields are recovered at each step by decoding through the frozen VAE, \(\hat{X}_{t+c_d}=\mathcal{D}_\theta(\hat{\mathbf{z}}_{t+c_d})\). Because the same density drives both directions, a trajectory can be integrated forward from an initial condition as a surrogate solver, or backward from a final measured state to infer the plasma history.

We denote the diffusion (denoising) step by a superscript \(k=0,\dots,K\) to keep it distinct from the physical time index \(t\), with \(\mathbf{z}_{t+c_d}^{(0)}\equiv \mathbf{z}_{t+c_d}\) the clean target and \(\mathbf{z}_{t+c_d}^{(K)}\sim \mathcal{N}(\mathbf{0},I)\) pure noise. The target is corrupted by a variance schedule \(\{\beta_k\}_{k=1}^{K}\), which in closed form reads \[\begin{align} \mathbf{z}_{t+c_d}^{(k)}&=&\sqrt{\bar\alpha_k}\,\mathbf{z}_{t+c_d} +\sqrt{1-\bar\alpha_k}\,\boldsymbol{\epsilon},\qquad \boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},I),\nonumber \\ \alpha_k&=&1-\beta_k,\;\;\bar\alpha_k=\prod_{j=1}^{k}\alpha_j . \end{align}\] Generation runs the context-conditioned reverse process \[p_\theta\!\left(\mathbf{z}_{t+c_d}^{(k-1)}\;\middle|\; \mathbf{z}_{t+c_d}^{(k)},\,\mathbf{z}_t,\,\mathbf{z}_{t-c_d},\,c_d\right) =\mathcal{N}\!\left(\mathbf{z}_{t+c_d}^{(k-1)};\;\boldsymbol{\mu}_\theta,\;\sigma_k^2 I\right),\] where \[\boldsymbol{\mu}_\theta=\frac{1}{\sqrt{\alpha_k}}\!\left(\mathbf{z}_{t+c_d}^{(k)} -\frac{\beta_k}{\sqrt{1-\bar\alpha_k}}\,\boldsymbol{\epsilon}_\theta\right),\] where the denoiser \(\boldsymbol{\epsilon}_\theta=\boldsymbol{\epsilon}_\theta( \mathbf{z}_{t+c_d}^{(k)},k,\mathbf{z}_t,\mathbf{z}_{t-c_d},c_d)\) is the network that is actually trained, by the conditional noise-prediction loss \(\mathcal{L}(\theta)\) defined as \[\mathcal{L}(\theta) = \mathbb{E}_{t,c_d,k,\boldsymbol{\epsilon}} \left[\left\|\,\boldsymbol{\epsilon} -\boldsymbol{\epsilon}_\theta \right\|_2^2\right],\] where \[\boldsymbol{\epsilon}_\theta = \boldsymbol{\epsilon}_\theta\!\left( \sqrt{\bar\alpha_k}\,\mathbf{z}_{t+c_d}+\sqrt{1-\bar\alpha_k}\,\boldsymbol{\epsilon},\; k,\;\mathbf{z}_t,\,\mathbf{z}_{t-c_d},\,c_d\right),\] with the physical-time index \(t\) and direction \(c_d\) drawn uniformly over the training trajectories and \(k\sim\mathcal{U}\{1,\dots,K\}\). Marginalizing the reverse chain over the intermediate noised states \(\mathbf{z}_{t+c_d}^{(1:K)}\) returns the transition density quoted above. The forward vs backward autoregressive rollout setup is shown in Figure 2.

4 Self-supervised Consistency Error Detection↩︎

When the model is deployed, it is rolled out autoregressively from a known seed pair into a regime where no ground-truth latents are available. In this setting the quantity we actually care about, the forward rollout error at depth \(i\), \[\mathcal{E}_i \;=\; \mathrm{MSE}\!\left(\mathbf{z}_{t+i},\,\hat{\mathbf{z}}_{t+i}\right), \; \mathrm{MSE}(\mathbf{a},\mathbf{b})=\frac{1}{L_d^2 L_c}\sum_{p,q,c}\bigl(a_{pqc}-b_{pqc}\bigr)^2,\] cannot be measured, because \(\mathbf{z}_{t+i}\) is precisely the unknown state the model is predicting. Since each learned transition is only approximate and these approximations compound from step to step, \(\mathcal{E}_i\) grows with rollout depth, and we would like a per-trajectory estimate of how far we can trust a given rollout before committing to it. The bidirectionality of the model supplies such an estimate at no additional training cost.

The key observation is that, because a single network represents both temporal directions, any forward rollout can be reversed by a backward rollout of equal length, and an error-free model must return exactly to where it started. Concretely, we roll forward \(i\) steps from the true seed pair, \[\hat{\mathbf{z}}_{t+m}\;\sim\;p_\theta\!\left(\,\cdot\;\middle|\; \hat{\mathbf{z}}_{t+m-1},\,\hat{\mathbf{z}}_{t+m-2},\,c_d=+1\right), \qquad m=1,\dots,i,\] with the seed convention \(\hat{\mathbf{z}}_t\equiv\mathbf{z}_t\), \(\hat{\mathbf{z}}_{t-1}\equiv\mathbf{z}_{t-1}\). We then flip the direction flag and, seeding the return pass with the terminal pair \((\hat{\mathbf{z}}_{t+i-1},\hat{\mathbf{z}}_{t+i})\), roll backward the same \(i\) steps, \[\tilde{\mathbf{z}}^{(i)}_{t+i-1-m}\;\sim\;p_\theta\!\left(\,\cdot\;\middle|\; \tilde{\mathbf{z}}^{(i)}_{t+i-m},\,\tilde{\mathbf{z}}^{(i)}_{t+i-m+1},\,c_d=-1\right), \;m=1,\dots,i,\] where the superscript \((i)\) records the depth at which the trajectory was turned around. The composition of the two passes is a closed cycle that maps the seed pair back onto itself, \[\begin{align} && (\mathbf{z}_{t-1},\mathbf{z}_t) \;\xrightarrow{\;\Phi^{i}_{+}\;}\; (\hat{\mathbf{z}}_{t+i-1},\hat{\mathbf{z}}_{t+i}) \;\xrightarrow{\;\Phi^{i}_{-}\;}\; (\tilde{\mathbf{z}}^{(i)}_{t-1},\,\tilde{\mathbf{z}}^{(i)}_{t}), \\ && \Phi^{i}_{-}\!\circ\Phi^{i}_{+}\stackrel{\text{ideal}}{=}\mathrm{Id}, \end{align}\] and we define the self-supervised consistency error at depth \(i\) as the round-trip discrepancy on the two recovered seed latents, \[\mathcal{C}_i \;=\; \tfrac{1}{2}\Bigl[\, \mathrm{MSE}\!\left(\mathbf{z}_{t-1},\,\tilde{\mathbf{z}}^{(i)}_{t-1}\right) +\mathrm{MSE}\!\left(\mathbf{z}_{t},\,\tilde{\mathbf{z}}^{(i)}_{t}\right)\Bigr].\] Crucially, every quantity entering \(\mathcal{C}_i\) is available at test time: the anchor pair \((\mathbf{z}_{t-1},\mathbf{z}_t)\) is the encoding of measured data, and \((\tilde{\mathbf{z}}^{(i)}_{t-1},\tilde{\mathbf{z}}^{(i)}_{t})\) is produced entirely by the model. No ground truth for any of the intermediate rolled-out states \(\hat{\mathbf{z}}_{t+1},\dots,\hat{\mathbf{z}}_{t+i}\) is ever required.

The diagnostic value of \(\mathcal{C}_i\) follows from a simple consistency argument. If the forward and backward maps were mutually exact, the cycle would be the identity and \(\mathcal{C}_i=0\) for every \(i\); any departure from zero therefore certifies that error has entered somewhere along the \(2i\)-step round trip. Cycle consistency is thus a necessary condition for an accurate rollout, and while it is not strictly sufficient, in that compensating forward and backward errors could in principle cancel, we find empirically that \(\mathcal{C}_i\) increases monotonically with depth in lockstep with the unobservable \(\mathcal{E}_i\), so that the measurable quantity \(\mathcal{C}_i\) serves as a faithful surrogate for the error we cannot see. Because each diffusion sample is stochastic, we make the cycle well-defined by either using the deterministic probability-flow sampler or by averaging over \(S\) independent rollouts, \(\bar{\mathcal{C}}_i=\tfrac{1}{S}\sum_{s}\mathcal{C}_i^{(s)}\), in which case the spread across samples provides an additional, complementary uncertainty estimate. In practice this yields a direct trust criterion: given a tolerance \(\tau\), a rollout is trusted out to the largest depth for which \(\mathcal{C}_i\le\tau\), and \(\mathcal{C}_i\) can equally be used to weight, branch, or terminate trajectories. The same signal drives the adaptive feedback scheme of the next section. The anchor need not be the original seed either; the turnaround can be performed at any point along a long rollout and checked against any earlier latent pair that is trusted, and if a diagnostic in physical units is preferred the discrepancy can be evaluated after decoding, \(\mathrm{MSE}(\mathcal{D}_\theta(\mathbf{z}_t), \mathcal{D}_\theta(\tilde{\mathbf{z}}^{(i)}_{t}))\), at the cost of one extra decode. A high level overview of this consistency cylcing approach is shown in Figure 3.

5 Adaptive Predictions From Limited Views↩︎

We consider a case in which we have rolled-out predictions \[\mathbf{z}_0, \;\mathbf{z}_1, \;\hat{\mathbf{z}}_2, \dots \hat{\mathbf{z}}_{T},\] and we have access to a measurement of a limited number of fields at \(t=T\). For example maybe we are able to measure just the density field \(\rho(x,y,T)\), and we would like to use that measurement to improve our predictions of the fields which we do not have access to: \[v_x(x,y,T), \;v_y(x,y,T), \;P(x,y,T), \;B_x(x,y,T), \;B_y(x,y,T).\] For this case we consider the application of a model-independent adaptive feedback approach known as bounded extremum seeking (ES). Bounded ES is applicable to dynamic systems of the form \[\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x},t) + \mathbf{g}(\mathbf{x},t)\mathbf{u}(\mathbf{x},t),\] where \(\mathbf{f}\) and \(\mathbf{g}\) are unknown and time-varying functions, and the goal is to minimize an analytically unknown cost function \(C(\mathbf{x},t)\). Bounded ES was invented in [29], it provides guaranteed bounds on control efforts and parameter update rates despite acting on noisy, analytically unknown time-varying systems, and has been studied for a wide range of systems including ones with non-differentiable and discontinuous controllers [30]. The bounded ES method has also been combined with generative multi-modal diffusion models to make them more robust for time-varying particle accelerator applications [31], it has been applied in-hardware in particle accelerators for optimizing the rise-time of high-voltage converter modulators [32] and for RF cavity tuning [33], and it has been studied for Tokamak stabilization [34]. Bounded ES has been combined with deep reinforcement learning (DRL) for time-varying systems and showed that bounded ES complements the learned RL policy by maintaining robustness as the system departs from the training regime, while DRL supplies rapid control when conditions remain near those seen in training [35], which is especially useful for robotic control in time-varying scenarios [36]. A method has also been created for automatic Mahalanobis distance-guided switching beetween ES and the DRL controller when an out of distribution state is detected [37].

In this case, our cost function is defined as \[\begin{align} C(\hat{\mathbf{z}}_T,T) &=& \left \| \rho(x,y,T) - \hat{\rho}(x,y,T) \right \|^2_1 \nonumber \\ &=& \iint_{x,y} \left | \rho(x,y,T) - \hat{\rho}(x,y,T) \right | dx dy \nonumber \\ &=& \iint_{x,y} \left | \rho(x,y,T) - D_\theta(\hat{\mathbf{z}}_{T,\rho}, \mathbf{c}_\rho ) \right | dx dy. \end{align}\] which we can minimize by using the model-independent bounded extremum seeking approach of [29]–[31]. We consider a class of systems of the form \[\dot{\mathbf{p}} = \mathbf{f}(\mathbf{p},t) + \mathbf{g}(\mathbf{p},t)\mathbf{u}(C(\mathbf{p},t)), \label{dpES}\tag{4}\] where \(\mathbf{p}\in \mathbb{R}^n\) is a set of adjustable parameters, \(\mathbf{f}:\mathbb{R}^n\times \mathbb{R} \rightarrow \mathbb{R}^n\) and \(\mathbf{g}:\mathbb{R}^n \times \mathbb{R} \rightarrow \mathbb{R}^{n\times n}\) are analytically unknown time-varying functions, \(C(\mathbf{p},t):\mathbb{R}^n\times \mathbb{R} \rightarrow \mathbb{R}\) is an analytically unknown time-varying function which we hope to minimize or maximize, and \(\mathbf{u}:\mathbb{R}\times \mathbb{R}\rightarrow \mathbb{R}^n\) is our control vector. Then based on the results of [29]–[31], if we choose our controller \(\mathbf{u}\) such that \[u_i = \sqrt{\alpha\omega_i}\cos(\omega_i t + k C(\mathbf{p},t)),\] where \(\omega_i = \omega r_i\) with \(r_i \neq r_j\) for all \(i \neq j\), then in the limit \(\omega \rightarrow \infty\), the dynamics (4 ), will on average be given by \[\dot{\bar{\mathbf{p}}} = \mathbf{f}(\bar{\mathbf{p}},t) -\frac{k\alpha}{2} \mathbf{g}(\mathbf{p},t)^T\mathbf{g}(\mathbf{p},t)\nabla_{\mathbf{p}}C(\bar{\mathbf{p}},t), \label{dpES95ave}\tag{5}\] where \(\mathbf{g}(\mathbf{p},t)^T\mathbf{g}(\mathbf{p},t) \geq 0\) ensures that we can overpower \(\mathbf{f}(\mathbf{p},t)\) and drive \(\mathbf{p}(t)\) towards a minimizer of \(C(\mathbf{p},t)\) by choosing sufficiently large \(k\alpha>0\).

6 VAE Reconstructions↩︎

As a first step, we need to train the VAE and confirm that it can compress and accurately reconstruct the MHD fields because the accuracy of the overall final predictions, regardless of how the latent space dynamics are modeled, will be bottle-necked by the decoder’s ability to create accurate reconstructions. One example of MHD fields along with their latent embeddings and reconstructions, as well as image-wise errors for one simulation at time steps \(t=0\) and \(t=1\) is shown in Figure 4.

By eye, the reconstructions look good and have very low mean squared error, but the more important quantitative test is to look at the power spectrum of the predictions versus the true fields to quantify how well high frequency low-level details are captured. In order to calculate the energy spectrum we first compute the Fourier transforms \[\hat{f}(\mathbf{k}) = \sum_{\mathbf{x}} f(\mathbf{x})\, e^{-i 2\pi\, \mathbf{k}\cdot\mathbf{x}/N}, \quad f \in \{v_x, v_y, B_x, B_y\},\] the energy spectral density is then calculated as \[E_{2\mathrm{D}}(\mathbf{k}) = \frac{1}{2N^2}\left( |\hat{v}_x(\mathbf{k})|^2 + |\hat{v}_y(\mathbf{k})|^2 + |\hat{B}_x(\mathbf{k})|^2 + |\hat{B}_y(\mathbf{k})|^2 \right),\] where \[\sum_{\mathbf{k}} E_{2\mathrm{D}}(\mathbf{k}) = \frac{1}{2}\sum_{\mathbf{x}}\left( v_x^2 + v_y^2 + B_x^2 + B_y^2 \right).\] We then perform radial binning \[E(k) = \sum_{\,k \,\le\, |\mathbf{k}| \,<\, k+1} E_{2\mathrm{D}}(\mathbf{k}), \quad k = 0, 1, \dots, \tfrac{N}{2}-1,\] where \[k =\|\mathbf{k}\|^2=\sqrt{k_x^2 + k_y^2},\] and plot the results up to \(k_\mathrm{max}=N/2\), truncating to the Nyquist wavenumber cutoff. Figure 5 shows true vs predicted power spectra for 6 randomly chosen time steps from the validation data set, which are seen to have good agreement to high wavenumber.

Another way to visualize the true vs predicted data, and to see that only very fine-grained details are missed, is to stack several 2D MHD fields together as the RGB channels of single color images, as shown in Figure 6.

7 Generative Diffusion↩︎

At any one time-step, we build up the latent images by our iterative denoising generative process which is applying our diffusion model as an iterative denoising diffusion probabilistic model (DDPM), as shown in Figure 7. For the diffusion model, we utilize both a UNet-based approach as well as vision transformer (ViT)-based diffusion models. We find that the UNet trains faster, on less data, but that the ViT reaches the same and sometimes slightly better accuracy as training goes on longer and larger datasets are used.

We then test the roll-out ability of our overall setup both forward and backward in time. Figures 8 and 9 show the first 15 time-steps of both forward and backward latent space roll outs for all 6 MHD fields compared to their true values, respectively.

In either direction, we expect to see that error is accumulated as we autoregressively produce the MHD fields with only the first two time-steps using ground truth latent field representations as inputs and subsequent predictions having to use the diffusion model’s imperfect predictions as their own inputs. This is quantified in Figure 10 where we see the mean error values and their standard deviations, for both forward and backward roll outs, averaged over the entire test data set.

Figure 11 shows predicted reconstructed and true fields for autoregressively predicted roll outs of length 2, 20, 40, 50, and 98 time-steps, going backward in time (left) and forward in time (right) for a test data set.

8 Conclusions↩︎

In future work, we plan on extending these results to include hard physics constraints on various aspects of the MHD generative process. One approach is the use of physics constrained neural networks (PCNNs), which utilize a novel neural operator approach for building hard physics constraints into convolutional neural networks by simulating potentials (such as the vector potential \(\mathbf{A}\)) from which electromagnetic fields are then defined according to Maxwell’s equations (\(\mathbf{B}=\nabla\times\mathbf{A}\)), thus guaranteeing hard constraints such as \(\nabla\cdot\mathbf{B}=0\), as developed in [38].

References↩︎