Latent Evolution Operator (LEO): State-to-State Forecasting of Parametric PDEs

Khalid Rafiq; Wenjing Liao; Aditya G. Nair

cs.LGMay 14, 2025

Latent Evolution Operator (LEO): State-to-State Forecasting of Parametric PDEs

Authors: Khalid Rafiq, Wenjing Liao, Aditya G. Nair

Organizations: Department of Mechanical Engineering, University of Nevada, Reno, Reno, USA · School of Mathematics, Georgia Institute of Technology, Atlanta, USA

Abstract

Solving time-dependent parametric partial differential equations (PDEs) with high-fidelity solvers is expensive in multi-query settings. This motivates reduced-order models in latent space. Existing methods are typically either (i) recurrent latent dynamics models, which require step-by-step integration and accumulate error over long horizons, or (ii) neural operators, which enable direct prediction but are usually trained from initial conditions and do not explicitly enforce generalization from intermediate states. We introduce the Latent Evolution Operator (LEO), a parameterized map that takes a state at an arbitrary time, together with physical parameters and a forecast horizon, and predicts the future state in a single latent-space evaluation. Under standard manifold and regularity assumptions, finite-time evolution operators admit an encoder-propagator-decoder decomposition. We approximate this structure using neural networks trained on state pairs sampled across anchor times and horizons. We evaluate the method on the 1D Burgers' equation, a 2D advection-diffusion system, and the shallow-water equations. The model produces accurate predictions across parameters and time horizons, achieves significant speedups compared to Autoencoder-LSTM baselines, and maintains accuracy when initialized from intermediate states. These results indicate that learning finite-time evolution operators in latent space provides an effective alternative to recurrent latent dynamics models and operator learning methods trained only from initial conditions.

Figures & tables

Figure 1: The PDE evolution operator $\mathfrak{F}_{\tau}$ and its encoder–propagator–decoder factorization.

Figure 2: An illustration of the Latent Evolution Operator (LEO) framework. On the solution manifold $\mathcal{M}\subseteq\mathbb{R}^{n}$ , three solution trajectories correspond to the parameters $\zeta_{1},\zeta_{2},...\zeta_{n}$ . For a fixed parameter $\zeta_{1}$ , the evolution operator advances the state along the red trajectory from $\bm{u}(t,\zeta_{1})$ to $\bm{u}(t+\tau,\zeta_{1})$ . The encoder $\mathcal{E}_{\theta_{e}}$ maps $\bm{u}(t,\zeta_{1})$ to a latent state $\tilde{\bm{z}}(t,\zeta_{1})$ in the low–dimensional latent space. The latent propagator $\mathcal{P}_{\theta_{p}}$ , implemented by the concatenation network $\mathcal{F}_{\mathrm{concat}}$ inside the shaded block, takes the latent state $\tilde{\bm{z}}(t,\zeta_{1})$ together with the parameter vector and forecast horizon $[\bm{\zeta}_{1}\|\tau]$ and produces the propagated (parameter–informed) latent state $\hat{\bm{z}}(t+\tau,\zeta_{1})$ (magenta arrow). The decoder $\mathcal{D}_{\theta_{d}}$ then maps $\hat{\bm{z}}(t+\tau,\bm{\zeta}_{1})$ back to the forecasted state $\hat{\bm{u}}(t+\tau,\bm{\zeta}_{1})$ (maroon dashed curve).

Figure 3: Latent Evolution–Operator (LEO) variants. Schematic comparison of the four architectures considered in this work. In all cases an encoder $\mathcal{E}_{\theta_{e}}$ maps the input state $\bm{u}(t)$ to a latent code $\tilde{\bm{z}}(t)$ . The propagator $\mathcal{P}_{\theta_{p}}$ advances the latent conditioned on parameter inputs $[\bm{\zeta},\tau]$ to $\hat{\bm{z}}(t+\tau)$ , and a decoder $\mathcal{D}_{\theta_{d}}$ reconstructs states. Decoder is shared for both reconstruction $\tilde{\bm{u}}(t)$ and forecasts $\hat{\bm{u}}(t+\tau)$ . Direct LEO AE backbone (D–LEO–AE) uses only reconstruction and direct forecasting with loss $L=L_{\mathrm{rec}}+\lambda_{\mathrm{dir}}L_{\mathrm{dir}}$ . Consistent LEO AE backbone (C–LEO–AE) adds a latent consistency term enforcing approximate commutativity $\mathcal{P}\circ\mathcal{E}\approx\mathcal{E}\circ\mathfrak{F}_{\tau}$ , with loss $L=L_{\mathrm{rec}}+\lambda_{\mathrm{con}}L_{\mathrm{con}}$ . Hybrid LEO AE backbone (H–LEO–AE) combines both supervision signals, $L=L_{\mathrm{rec}}+\lambda_{\mathrm{dir}}L_{\mathrm{dir}}+\lambda_{\mathrm{con}}L_{\mathrm{con}}$ . Direct LEO VAE backbone (D–LEO–VAE) replaces the deterministic encoder with a stochastic one, $\tilde{\bm{z}}(t)=\mu(t)+\epsilon\sigma(t)$ with $\epsilon\sim\mathcal{N}(0,I)$ , and augments the direct loss with a KL regularizer, $L=L_{\mathrm{rec}}+\lambda_{\mathrm{dir}}L_{\mathrm{dir}}+\beta L_{\mathrm{KL}}$ , encouraging a well–structured latent space.

Figure 4: Train–validation partitioning in a representative $(\zeta_{1},\tau)$ plane. Hatched cells denote training pairs; gray cells denote held-out validation pairs, stratified into nine zones (LE/IN/RE along each axis).

	Symbol	Meaning
PDE & evolution	$\bm{u}(\bm{x},t,\bm{\zeta})$	State vector at space $\bm{x}$ , time $t$ , parameter $\bm{\zeta}$ .
	$\bm{u}(t,\bm{\zeta})\in\mathbb{R}^{n}$	Spatially discretized state vector at time $t$ .
	$\bm{F}(\cdot)$	PDE right-hand side.
	$\mathfrak{F}_{\tau}(\bm{u}(t,\bm{\zeta}),\bm{\zeta})$	Time- $\tau$ evolution: $\bm{u}(t+\tau,\bm{\zeta})$ .
	$\bm{\zeta}$	Parameter vector (e.g., Reynolds number).
	$\tau$	Forecast horizon (time offset).

Table 1: Summary of notation used in the LEO formulation and experiments.

Figure 5: Direct forecasting performance for Burgers’ equation. (A) Zone-wise mean-squared error (MSE, log scale) across the $3\times 3$ grid of $(\mathrm{Re},\tau)$ validation zones combining left/interpolation/right ranges in each variable. (B) Mean MSE averaged over all nine zones for each LEO variant. (C) High-dimensional forecasts for the two best models evaluated in their worst-performing low- $\mathrm{Re}$ extrapolation configuration: the black curve is the anchor state $u(t_{a})$ with $t_{a}=0.15$ , solid curves are exact solutions at two future horizons $\tau_{\mathrm{LE}}$ and $\tau_{\mathrm{RE}}$ , and scatter markers are the corresponding model predictions. (D) Anchor agnostic prediction: the grey curve is the initial state $u(0)$ , the green curve is an anchor state at $t_{a}=0.80$ , and colored solid curves/markers are exact solutions and predictions at the same two horizons, illustrating that LEO can start from arbitrary anchor states along the trajectory.

Figure 6: Latent consistency across LEO variants. Panel A: anchor state $\bm{u}(t_{a})$ and one representative future state $\bm{u}(t_{a}+\tau)$ , shown for a fixed Burgers’ configuration with $t_{a}=0.1$ and $\mathrm{Re}=2400$ . For the latent trajectory analysis in Panels B–E, forecast horizons are varied over $\tau\in[0.7,1.7]$ in steps of $0.04$ and shows latent trajectories for C-LEO-AE, H-LEO-AE, D-LEO-AE, and D-LEO-VAE, respectively. In each panel, the green triangle marks the encoded anchor state, diamonds denote direct encodings of $\bm{u}(t_{a}+\tau)$ , and circles denote propagated latents $\hat{\bm{z}}(t_{a}+\tau,\bm{\zeta})$ for the same horizons; colors encode $\tau$ . Panel F: mean Euclidean distance in latent space between encoded and propagated futures, averaged over randomly sampled $(\mathrm{Re},t_{a},\tau)$ triples from the Burgers’ validation set that decode to the same high-dimensional forecasted state.

Figure 7: Scaling and latent structure of D-LEO-VAE on Burgers’. Left: Mean squared error (MSE, log scale) in the four joint extrapolation zones obtained by combining the left/right extrapolation ranges of $\mathrm{Re}$ and $\tau$ , averaged over randomly sampled triples $(\mathrm{Re},\tau,t_{a})$ in each zone, as a function of dataset size (number of state–to–state training tuples, log scale). Between roughly $7\text{k}$ and $80\text{k}$ tuples the error follows an approximate power law with slope $\approx-0.54$ , indicating data-efficient generalization in the most extrapolative regimes. Right: Latent-space structure for three training sizes (rows 1–3). Each point $(\hat{z}_{1},\hat{z}_{2})$ in a uniform grid is decoded to a high-dimensional field, from which we extract sharpness ( $\max_{x}|\partial u/\partial x|$ ; left column) and peak position (right column).

$\mathrm{Re}$	L-DeepONet	D-LEO-VAE
LE	$\mathbf{1.51{\times}10^{-5}}\ \pm\ 1.82{\times}10^{-5}$	$2.09{\times}10^{-5}\ \pm\ 9.60{\times}10^{-5}$
IN	$6.30{\times}10^{-7}\ \pm\ 8.14{\times}10^{-7}$	$\mathbf{4.97{\times}10^{-7}}\ \pm\ 1.06{\times}10^{-7}$
RE	$2.38{\times}10^{-6}\ \pm\ 2.80{\times}10^{-6}$	$\mathbf{1.48{\times}10^{-6}}\ \pm\ 9.74{\times}10^{-7}$
Par.	$540{,}606$	$486{,}198$

Table 2: Fixed-IC (

t_{a}=0

) generalization over approximately

n{=}1000

random

(\mathrm{Re},t)

pairs per region. Entries report mean

\pm

standard deviation of the MSE (lower is better) over the spatial domain. The last row lists the total number of trainable parameters in each model.

Figure 8: Comparison with latent neural operators and anchor-time generalization. A : Fixed-IC ( $t_{a}=0$ ) performance across Reynolds-number regions for L-DeepONet and D-LEO-VAE; bars show mean MSE on a logarithmic $y$ -axis. B – D : Temporal generalization across varying anchor times. Mean MSE versus anchor time $t_{a}$ for three Reynolds regimes, comparing L-DeepONet, D-LEO-VAE trained in the fixed-IC regime, and D-LEO-VAE trained on varying anchor states. For each $t_{a}$ , we draw $500$ random pairs $(\mathrm{Re},\tau)$ from the indicated regime and report the mean test MSE, enforcing $t_{a}+\tau\leq 2$ to remain within the simulation time window. L-DeepONet and fixed-IC D-LEO-VAE both experience anchor-time distribution shift for $t_{a}>0$ , whereas the varying-anchor D-LEO-VAE remains accurate across the full range. B : left-extrapolation $\mathrm{Re}\in[400,600]$ ; C : interpolation $\mathrm{Re}\in[1006,1206]$ ; D : right-extrapolation $\mathrm{Re}\in[2200,2400]$ .

Figure 9: 2D advection–diffusion: latent evolution performance. Here, Axis 1, Axis 2, and Axis 3 denote representative slices through the validation zones at fixed Reynolds-number bands ( $Re_{\mathrm{LE}}$ , $Re_{\mathrm{IN}}$ , and $Re_{\mathrm{RE}}$ ), with $\tau$ varied across $\tau_{\mathrm{LE}}$ , $\tau_{\mathrm{IN}}$ , and $\tau_{\mathrm{RE}}$ . A : Anchor state and representative forecasted states along Axis 1 (low- $Re$ slice, i.e., $Re_{\mathrm{LE}}$ ); top row: reference, bottom row: D-LEO-VAE prediction. B : Anchor state and representative forecasted states along Axis 3 (high- $Re$ slice, i.e., $Re_{\mathrm{RE}}$ ); top row: reference, bottom row: D-LEO-VAE prediction. C : Mean MSE (log scale) across all $(Re,\tau)$ validation zones, computed from $3000$ random $(t_{a},Re,\tau)$ triplets per region, with $t_{a}+\tau\leq 2$ and $\tau$ sampled from the corresponding $\tau$ interval. D : Mean mass ratio $M_{\mathrm{pred}}/M_{\mathrm{exact}}$ in each $(Re,\tau)$ validation zone, showing that the learned evolution operator preserves total scalar mass to within a few percent across all regions.

Re	L-DeepONet	D-LEO-VAE
LE	$9.96\times 10^{-4}\pm 6.74\times 10^{-4}$	$3.19\times 10^{-5}\pm 3.69\times 10^{-5}$
IN	$9.51\times 10^{-4}\pm 6.56\times 10^{-4}$	$3.63\times 10^{-6}\pm 4.94\times 10^{-6}$
RE	$3.04\times 10^{-3}\pm 2.03\times 10^{-3}$	$2.70\times 10^{-5}\pm 1.30\times 10^{-4}$
Par.	1,011,236	962,314

Table 3: Zone-aggregated MSE (mean

\pm

std) for L-DeepONet and D-LEO-VAE over temporal generalization trials. For each regime we aggregate all sampled triplets

(t_{a},Re,\tau)

used in the

t_{a}

-sweep evaluation.

Figure 10: 2D advection–diffusion: temporal generalization. Mean MSE versus anchor time $t_{a}$ for L-DeepONet and D-LEO-VAE across three Reynolds regimes: left extrapolation, interpolation, and right extrapolation. For each $t_{a}$ , we sample $1000$ random pairs $(Re,\tau)$ per regime and average over feasible horizons with $t_{a}+\tau<2$ .

Figure 11: Single–shot forecasts on the shallow–water benchmark. One randomly selected test trajectory. The left panel shows the observed zonal velocity field at $t=0$ . The three columns show ground truth (top) and LEO one–shot forecasts (bottom) for small, medium, and large lead times (approximately $3.4$ , $10.1$ , and $19.5$ days).

Figure 12: POD-based comparison of SWE dynamics. (A) Energy fraction of the leading POD modes and (B) Top three spatial POD modes (left) and their temporal coefficients for truth (solid black) and LEO forecasts projected onto the same basis (dashed blue). For each mode we report the dominant frequencies of the true and model coefficients ( $f_{T}$ , $f_{M}$ ) and their correlation, quantifying how well the learned evolution operator reproduces the large-scale dynamics.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

	D–LEO–AE	C–LEO–AE	H–LEO–AE	D–LEO–VAE
$B$	64	64	128	64
$\alpha$	$9.88\times 10^{-4}$	$6.93\times 10^{-4}$	$9.41\times 10^{-4}$	$7.0\times 10^{-4}$
$N_{e}$	200	75	150	100
$\lambda_{\mathrm{dir}}$	0.514	–	1.179	1.70
$\lambda_{\mathrm{con,max}}$	–	3.135	0.671	–
$f_{w}$	–	0.3768	0.313	–

Appendix

Table 4: Optimal training hyperparameters for Burgers’ LEO variants (Batch size

B

, Learning rate

\alpha

, Number of epochs

N_{e}

, Direct loss

\lambda_{\mathrm{dir}}

, Consistency max

\lambda_{\mathrm{con,max}}

, Warm-up fraction

f_{w}

, KL weight

\beta

)

Figure 13: Slice–wise Burgers’ forecasts across the $(Re,\tau)$ grid. The central panel shows the $3\times 3$ validation layout in $(Re,t+\tau)$ with grey rectangles indicating left/interpolation/right ranges in each variable. Two vertical slices are highlighted: Axis 1 (low Reynolds) and Axis 3 (high Reynolds). The top and right panels plot, for these slices, the initial condition $u(\cdot,t_{0})$ (black), the exact solutions at three forecast horizons $\tau_{\mathrm{LE}}$ , $\tau_{\mathrm{IN}}$ , $\tau_{\mathrm{RE}}$ (dashed, dotted lines), and the corresponding Direct LEO (VAE) predictions (coloured markers).

Figure 14: Positional–encoding propagator (PEP) for Burgers’. (A) Schematic of the positional–encoding propagator: the parameter pair $(\zeta,\tau)$ is mapped to a sinusoidal embedding and mixed with the latent code $\tilde{z}(t,\zeta)$ via up/down networks $\mathcal{F}_{\text{up}}$ and $\mathcal{F}_{\text{down}}$ . (B) Zone–wise mean MSE (log scale) for D–LEO–VAE with concatenation (blue) versus positional encoding (orange) across the nine $(\mathrm{Re},\tau)$ validation zones. The concatenation–based propagator outperforms PEP in all zones. (C) Sharpness ( $\max_{x}|\partial u/\partial x|$ , top row) and peak position (bottom row) of decoded fields over a grid of latent points $(\hat{z}_{1},\hat{z}_{2})$ for training set sizes of 10k, 20k, 40k, and 80k tuples. Compared with the smooth, disentangled latent maps obtained with concatenation (Fig. 7 ), the PEP latent structure is more irregular and entangled, especially at smaller data sizes.

Figure 15: Decoder sensitivity along the reconstructed profile. Reconstructions of

u(x,t,\zeta)

from the encoded latent

\tilde{z}

(left) and propagated latent

\hat{z}

(right), colour–coded by

|\partial\mathcal{D}_{\theta_{d}}(z)/\partial z_{1}|

. Sensitivity is more concentrated and anisotropic near the shock at

\tilde{z}

than at

\hat{z}

Figure 16: Effect of small latent perturbations. Decoder outputs for $\tilde{z}$ and $\hat{z}$ (blue) and for their perturbed versions $z+\varepsilon e_{j}$ along each latent axis (red/green). Insets zoom into the shock region, showing stronger sensitivity at $\tilde{z}$ and more stable behaviour at $\hat{z}$ .

Figure 17: Local geometry metrics of the decoder. Singular values $\Sigma(J(z))$ (left), Frobenius norms $\|J(z)\|_{F}$ (middle), and $\log\det(J(z)^{\top}J(z))$ (right) at the encoded latent $\tilde{z}$ and propagated latent $\hat{z}$ . The smaller values at $\hat{z}$ indicate a better–conditioned, less expansive region of latent space.

Figure 18: Prediction time comparison for D-LEO-VAE and AE–LSTM across temporal offsets $\tau$ on CPU (A) and GPU (B). AE–LSTM runtime increases approximately linearly with $\tau$ due to its autoregressive nature ( $\mathcal{O}(n)$ ), whereas D-LEO-VAE achieves essentially constant cost ( $\mathcal{O}(1)$ ). For large offsets, D-LEO-VAE attains up to $50\times$ speedup on CPU and $90\times$ on GPU.

Component	D-LEO-VAE
Latent dimension ( $m$ )	3
Batch size ( $B$ )	64
Learning rate ( $\alpha$ )	$7.0\times 10^{-4}$
Number of epochs	75
KL-divergence weight ( $\beta$ )	$1.152\times 10^{-5}$
Propagation loss weight ( $\lambda_{dir}$ )	1.15

Appendix

Table 5: Model configuration for the 2D advection–diffusion benchmark.

Figure 19: Single–shot SWE forecasts from multiple observed states. Top row: ground–truth zonal velocity snapshots at $\tau=0,248h,337h,469h$ time steps. Rows 2–4: three observed states along the same trajectory (green dashed boxes) and LEO forecasts at later times. Dashed arrows are labeled with the corresponding temporal jump $\tau$ and colored by jump length (small, medium, long). In all cases the forecasts closely match the future snapshots qualitatively, illustrating that the latent evolution operator generalizes across anchor times and supports genuine state–to–state forecasting.

PDE Surrogate Modeling

PDE surrogate modeling builds fast approximations of partial differential equation solutions for simulation, prediction, and design.

See all 325 papers

AE-ViT: Stable Long-Horizon Parametric Partial Differential Equations Modeling
Apr 7, 2026Iva Mikuš, Boris Muha, Domagoj VlahPDE Surrogate Modeling Reduced-Order Modeling
University of Zagreb Faculty of Electrical Engineering and Computer Science · University of Zagreb Faculty of Science Department of Mathematics
Autoregression-Free Neural Operators for Time-Dependent PDEs
May 25, 2026Jiaquan Zhang, Caiyan Qin, Haoyu Bian +7Latent Dynamics Modeling PDE Operator Learning
School of Computer Science and Engineering, University of Electronic Science and Technology of China, Chengdu, Sichuan 611731, China · School of Robotics and Advanced Manufacture, Harbin Institute of Technology, Shenzhen, Guangdong 518055, China · School of Mathematical Sciences, Capital Normal University, Beijing 100048, China +3
Late Fusion Neural Operators for Extrapolation Across Parameter Space in Partial Differential Equations
Apr 17, 2026Eva van Tegelen, Taniya Kapoor, George A. K. van Voorn +2Surrogate Modeling PDE Surrogate Modeling
Artificial Intelligence Group, Wageningen University and Research, Netherlands · Mathematical and Statistical Methods, Wageningen University and Research, Netherlands

Apr 7, 2026J/K move · Enter open · S save

AE-ViT: Stable Long-Horizon Parametric Partial Differential Equations Modeling

Iva Mikuš, Boris Muha, Domagoj Vlah

University of Zagreb Faculty of Electrical Engineering and Computer Science · University of Zagreb Faculty of Science Department of Mathematics

PDE Surrogate Modeling Reduced-Order Modeling

Open paper AskPDF arXiv

Deep Learning Reduced Order Models (ROMs) are becoming increasingly popular as surrogate models for parametric partial differential equations (PDEs) due to their ability to handle high-dimensional data, approximate highly nonlinear mappings, and utilize GPUs. Existing approaches typically learn evolution either on the full solution field, which requires capturing long-range spatial interactions at high computational cost, or on compressed latent representations obtained from autoencoders, which reduces the cost but often yields latent vectors that are difficult to evolve, since they primarily encode spatial information. Moreover, in parametric PDEs, the initial condition alone is not sufficient to determine the trajectory, and most current approaches are not evaluated on jointly predicting multiple solution components with differing magnitudes and parameter sensitivities. To address these challenges, we propose a joint model consisting of a convolutional encoder, a transformer operating on latent representations, and a decoder for reconstruction. The main novelties are joint training with multi-stage parameter injection and coordinate channel injection. Parameters are injected at multiple stages to improve conditioning. Physical coordinates are encoded to provide spatial information. This allows the model to dynamically adapt its computations to the specific PDE parameters governing each system, rather than learning a single fixed response. Experiments on the Advection-Diffusion-Reaction equation and Navier-Stokes flow around the cylinder wake demonstrate that our approach combines the efficiency of latent evolution with the fidelity of full-field models, outperforming DL-ROMs, latent transformers, and plain ViTs in multi-field prediction, reducing the relative rollout error by approximately

5

times.

Abstract

PDE Surrogate Modeling

Explore similar work

AE-ViT: Stable Long-Horizon Parametric Partial Differential Equations Modeling

Autoregression-Free Neural Operators for Time-Dependent PDEs

Late Fusion Neural Operators for Extrapolation Across Parameter Space in Partial Differential Equations