Single-shot prediction of parametric partial differential equations

TL;DR

Flexi-VAE is position as a scalable and interpretable surrogate modeling tool for accelerating high-fidelity simulations in computational fluid dynamics (CFD) and other parametric PDE-driven applications, with extensibility to higher-dimensional and more complex systems.

Abstract

We introduce Flexi-VAE, a data-driven framework for efficient single-shot forecasting of nonlinear parametric partial differential equations (PDEs), eliminating the need for iterative time-stepping while maintaining high accuracy and stability. Flexi-VAE incorporates a neural propagator that advances latent representations forward in time, aligning latent evolution with physical state reconstruction in a variational autoencoder setting. We evaluate two propagation strategies, the Direct Concatenation Propagator (DCP) and the Positional Encoding Propagator (PEP), and demonstrate, through representation-theoretic analysis, that DCP offers superior long-term generalization by fostering disentangled and physically meaningful latent spaces. Geometric diagnostics, including Jacobian spectral analysis, reveal that propagated latent states reside in regions of lower decoder sensitivity and more stable local geometry than those derived via direct encoding, enhancing robustness for long-horizon predictions. We validate Flexi-VAE on canonical PDE benchmarks, the 1D viscous Burgers equation and the 2D advection-diffusion equation, achieving accurate forecasts across wide parametric ranges. The model delivers over 50x CPU and 90x GPU speedups compared to autoencoder-LSTM baselines for large temporal shifts. These results position Flexi-VAE as a scalable and interpretable surrogate modeling tool for accelerating high-fidelity simulations in computational fluid dynamics (CFD) and other parametric PDE-driven applications, with extensibility to higher-dimensional and more complex systems.

Figures (17)

• Figure 1: Conceptual overview of the Flexi-VAE framework for single-shot prediction. The high-dimensional solution manifold $\mathcal{M} \subseteq \mathbb{R}^n$ is shown on the left, while the corresponding latent space $\boldsymbol{z}\subseteq \mathbb{R}^m$ is depicted on the right (illustrated here as a 2D manifold). Given an initial high-dimensional state $\boldsymbol{u}(\boldsymbol{x}, t_{kj}, \boldsymbol{\zeta_k})$, the encoder $\mathcal{E}_{\theta_e}$ maps it to a low-dimensional latent representation $\tilde{\boldsymbol{z}}(t_{kj}, \boldsymbol{\zeta_k})$. Here, $k$ refers to the parameter configuration and $j$ refers to the initial condition. The latent state is then evolved forward in time by the propagator $\mathcal{P}_{\theta_p}$ (magenta arrow), conditioned on the forecasting horizon $\tau_{kji}$ and system parameters $\zeta_k$, producing the propagated latent vector $\hat{\mathbf{z}}(t_{kj} + \tau_{kji}, \boldsymbol{\zeta_k})$. Here, $i$ refers to the predictive time horizon. Finally, the shared decoder $\mathcal{D}_{\theta_d}$ forecasts the high-dimensional prediction $\hat{u}(\mathbf{x}, t_{kj} + \tau_{kji}, \boldsymbol{\zeta}_k)$ from the propagated latent.
• Figure 2: Schematic representation of the Flexi-VAE architecture. The encoder compresses the high-dimensional input into a low-dimensional latent space. The parametric neural propagator advances the latent state across multiple time steps, incorporating both temporal and parametric information. Finally, the shared decoder is used both to reconstruct the high-dimensional state from the latent vector at the current time step and to forecast the future high-dimensional state from the propagated latent representation.
• Figure 3: Comparison of two parameter-informed latent propagation strategies in the Flexi-VAE framework. Top: The Positional Encoding Propagator (PEP) maps the latent state $\tilde{\boldsymbol{z}}(t,\boldsymbol{\zeta})$ to a higher-dimensional space via $\mathcal{F}_{\text{up}}$, where it is enriched with sinusoidal embeddings of the parameters $\boldsymbol{\zeta}$ and time offset $\tau$, followed by summation and down-projection through $\mathcal{F}_{\text{down}}$ to yield the propagated latent $\hat{\boldsymbol{z}}(t+\tau,\boldsymbol{\zeta})$. Bottom: The Direct Concatenation Propagator (DCP) concatenates $\tilde{\boldsymbol{z}}(t,\boldsymbol{\zeta})$ with $\boldsymbol{\zeta}$ and $\tau$, and passes the augmented vector through a fully connected network $\mathcal{F}_{\text{concat}}$. In both approaches, the propagated latent is decoded by the shared decoder $\mathcal{D}_{\theta_d}$ to produce the forecasted physical state $\hat{\boldsymbol{u}}(\boldsymbol{x}, t+\tau, \boldsymbol{\zeta})$.
• Figure 4: Overview of the Flexi-VAE architecture. The encoder $\mathcal{E}_{\theta_e}$ transforms the input state $\boldsymbol{u}(\boldsymbol{x}, t, \boldsymbol{\zeta})$ into a latent representation $\tilde{\boldsymbol{z}}$. The propagator $\mathcal{P}_{\theta_p}$ evolves the latent state, conditioned on parameters $\boldsymbol{\zeta}$ and temporal offset $\tau$. The decoder $\mathcal{D}_{\theta_d}$ reconstructs the instantaneous $\tilde{\boldsymbol{u}}(\boldsymbol{x}, t, \boldsymbol{\zeta})$ and the propagated state $\hat{\boldsymbol{u}}(\boldsymbol{x}, t + \tau, \boldsymbol{\zeta})$.
• Figure 5: Forecasting performance of the DCP based Flexi-VAE model on the 1D viscous Burgers' equation across interpolation and extrapolation regimes. Bottom center: The $(\text{Re},\, t+\tau)$ parameter space. Axes 1 and 3 (marked in red) indicate representative slices through low and high Reynolds numbers, respectively. Top left and bottom right: Predicted and exact solutions for Axis 1 (low-Re, left extrapolation) and Axis 3 (high-Re, right extrapolation) are shown at three future time offsets: $\tau_{\mathrm{LE}}$, $\tau_{\mathrm{IN}}$, and $\tau_{\mathrm{RE}}$.

Theorems & Definitions (12)

• Theorem 1
• Definition 1: Chart
• Definition 2: $C^k$ Atlas
• Definition 3: Smooth Manifold
• Definition 4: $C^k$ Functions on $\mathcal{M}$
• Definition 5: Partition of Unity
• Proposition 1: Existence of a $C^\infty$ partition of unity
• Definition 6: Reach, Definition 2.1 in aamari2019estimating
• Definition 7: Lipschitz function
• Lemma 1: Proposition 3 in yarotsky2017error