Same Target, Different Origins: Aligning Latent Rollout for Consistent PDE Emulation

Xinying Li; Jing Xu; Wenjun Wang; Limei Han; Fenglei Cao; Ruoxi Jiang; Yuan Qi

cs.LGOct 15, 2026

Same Target, Different Origins: Aligning Latent Rollout for Consistent PDE Emulation

Authors: Xinying Li, Jing Xu, Wenjun Wang, Limei Han, Fenglei Cao, Ruoxi Jiang, Yuan Qi

Organizations: Fudan University · Shanghai Academy of AI for Science

Abstract

Neural PDE solvers commonly generate multi-step predictions through autoregressive rollout, recursively applying a learned evolution operator and feeding each prediction back as input for the next step. Approximation errors can therefore propagate through the rollout, causing predictions at a shared target time to depend on their temporal origins and rollout depths. For a well-posed deterministic PDE flow, however, states at different times along the same physical trajectory should, under exact evolution, reach the same target state. To this end, we propose Align Latent Rollout (ALR), a regularization approach that promotes consistency across rollout origins through the neural PDE solver's internal representations. A second-order Taylor analysis of latent discrepancy propagation motivates two objectives: zeroth-order latent state consistency, which aligns representations at a shared target time, and first-order latent dynamics consistency, which aligns local evolution responses to shared perturbations. ALR can be integrated into different neural PDE solvers by regularizing suitable internal representations, without modifying their backbone architectures. Experiments across multiple PDE benchmarks and neural-operator backbones show that ALR improves prediction accuracy while reducing shared-target discrepancies across rollout origins.

Figures & tables

Figure 1: Overview of Align Latent Rollout (ALR). Starting from temporal origins along the same physical trajectory, ALR aligns the latent states reached at a shared target time and further regularizes their local evolution through zeroth-order state consistency and first-order dynamics consistency.

Method	NS ( $\nu=10^{-4}$ )	NS ( $\nu=10^{-5}$ )	Shallow Water	Kolmogorov Flow
FNO	$0.0947\pm 0.0012$	$0.1004\pm 0.0015$	$0.0038\pm 0.0002$	$0.4095\pm 0.0164$
FNO+ALR (Ours)	$\mathbf{0.0670}\pm 0.0004$ $\downarrow\,29.25\%$	$\mathbf{0.0858}\pm 0.0003$ $\downarrow\,14.54\%$	$\mathbf{0.0027}\pm 0.0001$ $\downarrow\,28.95\%$	$\mathbf{0.2794}\pm 0.0089$ $\downarrow\,31.77\%$
FNO (Unroll)	$0.2184\pm 0.0035$	$0.1682\pm 0.0021$	$0.0032\pm 0.0002$	$0.3445\pm 0.0103$

Table 1: Results on four spatiotemporal PDE benchmarks. Unless otherwise specified, prediction errors hereafter are reported as mean

\pm

standard deviation of relative

\ell_{2}

error (

\downarrow

). FNO is the one-step baseline, FNO+ALR augments it with cross-orgin consistency regularization, and FNO (Unroll) is an unrolled-training reference. Percentages denote relative error reductions over the FNO baseline.

Table 3

Figure 2: Qualitative analysis of cross-path consistency on the Navier–Stokes benchmark ( $\nu=10^{-5}$ ). Each column corresponds to a rollout origin $t_{i}\in\{0,5,10,15,18\}$ , with all paths propagated to the shared target time $t=19$ . The top and bottom rows of each panel show the baseline FNO and FNO+ALR, respectively. (a) Physical predictions at $t=19$ . (b) Physical prediction errors relative to the ground truth. (c) Latent states at $t=19$ . (d) Latent-state discrepancies relative to the ground-truth-conditioned one-step path $18\rightarrow 19$ . (e) Local latent dynamics represented by $J_{G_{\theta}}(\bm{z}_{t_{i}\rightarrow 19})\bm{v}$ , using the same probing direction $\bm{v}$ for all paths. (f) Corresponding latent-dynamics discrepancies relative to the same reference path. The rightmost columns in (d) and (f) are zero by construction because the reference path is compared with itself.

Method	NS ( $\nu=10^{-4}$ )	SW		Method	NS ( $\nu=10^{-4}$ )	SW
F-FNO	0.2904	0.0073	$\|$	LSM	0.2812	0.0021
F-FNO+ALR	0.2126 ( $\downarrow 26.79\%$ )	0.0063 ( $\downarrow 13.70\%$ )	$\|$	LSM+ALR	0.2188 ( $\downarrow 22.19\%$ )	0.0019 ( $\downarrow 9.52\%$ )

Table 4: Generalization of ALR across neural operator backbones on Navier–Stokes (NS) and Shallow Water (SW), measured by relative

\ell_{2}

error (

\downarrow

). Percentages denote relative error reductions over the corresponding backbone.

Figure 3: Qualitative rollout prediction results on Shallow Water under the rollout-length- $60$ training setting. Ground truth, FNO, and FNO+ALR are shown from top to bottom at selected rollout steps $T\in\{1,10,20,30,40,50,60\}$ .

Rollout Length	20	40	60
FNO	0.0038	0.0130	0.0517
FNO+ALR	0.0027	0.0035	0.0085

Table 5: Autoregressive rollout performance at increasing horizons on the Shallow Water benchmark. Results are reported using relative

\ell_{2}

error (

\downarrow

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Definition
$\bm{x}$	Spatial coordinate, $\bm{x}\in\Omega\subset\mathbb{R}^{d}$ .
$\tau$	Continuous physical time, $\tau\in[0,T]$ .
$t,t_{i},t_{j}$	Discrete temporal indices; $t_{i}$ denotes a rollout origin and $t_{j}$ a target time.
$t_{1},t_{2},t$	Two temporal origins and a shared target time satisfying $t_{1}<t_{2}<t$ .
$N$	Final discrete temporal index of a trajectory, $\{\bm{u}_{t}\}_{t=0}^{N}$ .
$\Delta\tau$	Temporal discretization interval, with $\bm{u}_{t}=\bm{u}(\cdot,t\Delta\tau)$ .

Appendix

Table 8

$K$	Relative $\ell_{2}$ $\downarrow$	Training Time ( $\times$ ) $\downarrow$	Inference Time ( $\times$ )
–	0.1004	1.00	1.00
4	0.0893	2.04	1.00
6	0.0858	2.66	1.00
8	0.0851	3.42	1.00
10	0.0856	4.04	1.00

Appendix

Table 7: Computational efficiency under different consistency-window lengths

K

on Navier–Stokes (

\nu=10^{-5}

). Training and inference times are normalized by the FNO baseline. Lower relative

\ell_{2}

error and runtime are better.

Figure 4: Qualitative comparison of the complete 29-step autoregressive rollout on Navier–Stokes ( $\nu=10^{-4}$ ). The initial physical state is shown at $T=0$ . For subsequent temporal steps, ground truth, FNO baseline, and FNO+ALR predictions are shown from top to bottom. All models start from the same initial state and are rolled out autoregressively over the full temporal horizon.

Neural PDE Solvers

Neural PDE solvers use trainable models to approximate solutions or solution operators for partial differential equations, often with physical constraints.

See all 269 papers

ReAl: Refine neural operators via aligning dynamical flow
Oct 14, 2026Jing Xu, Zenglin Xu, Li Yi +2PDE Surrogate Modeling PDE Operator Learning
Shanghai Academy of AI for Science · Fudan University
Posterior-First Neural PDE Simulation: Inferring What Downstream Tasks Need from a Single Field
May 5, 2026Wenshuo Wang, Fan ZhangNeural PDE Solvers PDE Surrogate Modeling
School of Future Technology, South China University of Technology, China · State Key Laboratory of Ocean Sensing & Ocean College, Zhejiang University, China · Kavli Institute for Astrophysics and Space Research, Massachusetts Institute of Technology, USA
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

Oct 14, 2026J/K move · Enter open · S save

ReAl: Refine neural operators via aligning dynamical flow

Jing Xu, Zenglin Xu, Li Yi, Ruoxi Jiang+1

Shanghai Academy of AI for Science · Fudan University

PDE Surrogate Modeling PDE Operator Learning Long-Term Time Series Forecasting

Open paper AskPDF arXiv

Neural partial differential equation (PDE) surrogates trained with mean squared error (MSE) can capture dominant dynamics but often lose high-frequency information and fine-scale structures. Generative models (e.g., Flow Matching) can capture the distribution of fine-scale details, yet stochastic errors may accumulate during autoregressive rollouts and compromise long-term stability. To address these limitations, we introduce Refine and Align(ReAl), a model-agnostic plug-in that augments pretrained PDE surrogates to improve the reconstruction of high-frequency structures associated with complex physical phenomena. In the refinement stage, a conditional flow-matching model learns the residual distribution around the deterministic base prediction, recovering missing fine-scale structures while retaining the dominant dynamics. In the alignment stage, to mitigate the generative discrepancies amplified in the autoregressive rollout, particularly in chaotic systems, we generate multiple trajectories from the same input and reinforce the refiner with trajectory-level rewards, thereby achieving long-horizon dynamical alignment. Experiments across diverse surrogate architectures and PDE tasks demonstrate that ReAl consistently improves high-frequency fidelity, forecasting accuracy, and long-horizon stability.

Abstract

Neural PDE Solvers

Explore similar work

ReAl: Refine neural operators via aligning dynamical flow

Posterior-First Neural PDE Simulation: Inferring What Downstream Tasks Need from a Single Field

Autoregression-Free Neural Operators for Time-Dependent PDEs