PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.21370v1cs.LG

TL;DR

Autoregressive neural PDE solvers can accumulate errors, leaving predictions at a shared target dependent on their temporal origins. ALR uses Taylor-motivated latent state and dynamics consistency objectives; across benchmarks and backbones, experiments report improved accuracy and reduced cross-origin discrepancies.

  • Problem

    Autoregressive errors can accumulate across rollout steps, and standard prediction losses do not explicitly couple internal representations reached from different origins at a shared target.

    Source: §1 · pp. 1–2

  • Method

    ALR uses second-order Taylor analysis to motivate regularizing shared-target latent states and their local evolution responses across rollout origins.

    Source: §1 · p. 2, §3.2 · p. 5, §3.3 · p. 5

  • Results

    Across multiple PDE benchmarks and neural-operator backbones, ALR improves prediction accuracy and reduces cross-origin discrepancies; FNO+ALR reduces relative ℓ2 error versus FNO across four benchmarks.

    Source: §4.2 · p. 7, §5 · p. 10

  • Takeaways & Limitations

    ALR regularizes suitable internal representations and can be integrated into different neural PDE solvers without modifying their backbone architectures.

    Source: §5 · pp. 9–10

    The consistency formulation assumes a fixed governing configuration and deterministic, well-posed PDE evolution.

    Source: §3.1 · p. 3, §3.2 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

Autoregressive neural PDE solvers can accumulate prediction errors, and their predictions at a shared target time may depend on rollout origin. ALR regularizes internal representations across origins, with experiments reporting improved accuracy and reduced discrepancies.

1 INTRODUCTION

  • Recursive autoregressive predictions feed prior outputs back as inputs, allowing local approximation errors to propagate and accumulate.This can degrade prediction accuracy.Source: §1 · p. 1
  • Existing mitigation methods use multistep training, solver interactions, iterative refinement, or structural constraints, but generally do not directly characterize origin-dependent discrepancy propagation through learned evolution.Standard prediction losses also do not explicitly couple internal representations reached from different origins at a shared target time.Source: §1 · pp. 1–2
  • ALR uses second-order Taylor analysis to motivate state and dynamics consistency objectives, and evaluations report improved accuracy and reduced latent discrepancies across PDE benchmarks and neural-operator backbones.The objectives align shared-target representations and their local responses to perturbations.Source: §1 · p. 2

2 RELATED WORK

Related work spans neural operators for time-dependent PDEs, methods to mitigate autoregressive rollout error, and latent-dynamics regularization. ALR is framed as aligning both shared-target latent states and their local evolution responses.

Neural Operators for Time-Dependent PDEs

  • Neural operators learn mappings between function spaces, with spectral, multiscale, geometry-aware, convolutional, attention-based, and latent-space architectures used for PDE modeling.The cited approaches address global interactions, efficiency, resolution handling, irregular domains, or heterogeneous interactions.Source: §2 · p. 2

Autoregressive Prediction and Rollout Error Mitigation

  • Rollout-error mitigation includes data-centric training, solver-in-the-loop methods, controls on unstable operations or spectral evolution, iterative refinement, and trajectory-level alternatives.These approaches generally leave propagation of discrepancies across rollout histories largely implicit.Source: §2 · p. 2

Latent Dynamics Regularization

  • Latent-dynamics work includes compact latent evolution and structure-preserving or consistency-based methods; ALR considers shared-target states and first-order responses without assuming a particular latent evolution structure.Figure 1 depicts aligning representations reached from different temporal origins at a shared target time, plus their local evolution.Source: §2 · p. 3, Figure 1 · p. 3

3 METHOD

The method formulates cross-origin consistency for deterministic PDE evolution and regularizes representations reached at a shared target by different rollout histories. A Taylor analysis motivates state and local-dynamics objectives for the induced latent evolution.

3.1 PROBLEM SETTING

  • A rollout repeatedly applies the learned one-step operator, so the same target can be reached from different origins and depths with potentially different predictions.The setup uses fixed governing conditions and assumes deterministic, well-posed PDE evolution.Source: §3.1 · pp. 3–4

3.2 CROSS-ORIGIN CONSISTENCY IN DETERMINISTIC PDE EVOLUTION

  • Exact states along one deterministic trajectory reach the same target state whether evolution proceeds directly or through an intermediate time.The paper calls this shared-target property cross-origin consistency.Source: §3.2 · p. 4
  • Learned rollouts may produce different predictions and latent representations at a shared target, since standard prediction loss does not directly couple representations from different origins.The method selects an intermediate representation of the shared operator as the latent state, without requiring an explicit encoder–decoder architecture.Source: §3.2 · pp. 4–5

3.3 ALIGN LATENT ROLLOUT

  • ALR models progression between latent representations with an induced evolution map and analyzes how a shared-target latent discrepancy propagates to the next step.The analysis assumes this map is locally three times continuously differentiable for the second-order Taylor expansion.Source: §3.3 · p. 5
  • The Taylor expansion separates local linearized discrepancy propagation from a curvature correction, motivating state alignment and alignment of local linearized evolution.The Jacobian term propagates the existing mismatch; variation in the local Jacobian contributes the leading correction.Source: §3.3 · pp. 5–6
  • Zeroth-order state consistency reduces shared-target latent mismatch, while first-order dynamics consistency aligns local evolution responses to perturbations.The dynamics objective addresses local evolution differences that state alignment alone does not constrain.Source: §3.3 · p. 6
  • The dynamics discrepancy can be estimated with Jacobian–vector products using a shared stochastic direction, avoiding explicit construction of full Jacobians.The same probe makes the JVP difference the action of the Jacobian difference on a common perturbation.Source: §3.3 · p. 6
  • Training retains the original prediction loss and augments it with the state and dynamics consistency objectives.Source: §3.3 · p. 6

4 EXPERIMENT

Across four PDE benchmarks, FNO+ALR reduces relative ℓ2 error versus FNO; ablations show benefits from both consistency objectives and a moderate consistency window. ALR also improves results across backbones and rollout lengths, with inference unchanged from FNO.

4.1 EXPERIMENTAL SETUP

  • The FNO setup uses the second Fourier block as the latent representation and an auxiliary local autoregressive rollout to construct cross-origin pairs, with K = 6 by default.The original ground-truth-conditioned prediction objective is retained; the auxiliary rollout is used only for consistency regularization.Source: §4.1 · pp. 6–7

4.2 MAIN RESULTS

  • FNO+ALR reduces relative ℓ2 error versus FNO by 29.25%, 14.54%, 28.95%, and 31.77% across four PDE benchmarks.The benchmarks are Navier–Stokes at ν = 10−4 and ν = 10−5, Shallow Water, and Kolmogorov Flow.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 ABLATION STUDY

  • Combining state and dynamics consistency yields the lowest errors: 0.0858 on Navier–Stokes and 0.0027 on Shallow Water.Dynamics consistency alone outperforms state consistency alone on both benchmarks.Source: §4.3 · p. 7, Table 2 · p. 7
  • Across origins propagated to t = 19, FNO+ALR maintains more consistent physical predictions and reduces shared-target latent-state and local-dynamics discrepancies.The comparisons use the ground-truth-conditioned one-step path 18 →19 as the reference for latent discrepancies.Source: §4.3 · p. 8, Figure 2 · p. 8
  • Increasing the consistency window K from 4 to 6 reduces relative ℓ2 error to 0.0858 on Navier–Stokes and 0.0027 on Shallow Water; larger windows add only marginal gains.The default K = 6 captures most of the improvement at moderate rollout depth; larger K increases training cost, while inference remains unchanged from FNO.Source: Table 3 · p. 7, §4.3 · p. 8

4.4 GENERALIZATION ACROSS BACKBONES

  • ALR reduces relative ℓ2 error for F-FNO and LSM on both Navier–Stokes and Shallow Water, using an appropriate intermediate representation for each backbone.The results support applicability across distinct backbone architectures and latent-state realizations.Source: §4.4 · p. 9, Table 4 · p. 9

4.5 PERFORMANCE ACROSS ROLLOUT LENGTHS

  • At 60 Shallow Water rollout steps, FNO error is 0.0517 versus 0.0085 for FNO+ALR; relative to 20 steps, errors increase 13.6× and 3.1×, respectively.Each rollout-length setting was trained and evaluated at its corresponding horizon; Figure 3 shows predictions at selected steps through 60.Source: §4.5 · p. 9, Figure 3 · p. 9

5 CONCLUSION

ALR promotes cross-origin consistency in autoregressive neural PDE solvers by regularizing internal representations, with two objectives motivated by second-order analysis. It is evaluated across benchmarks and backbones, and the notation defines the rollout and latent quantities used to describe the method.

5 CONCLUSION

  • ALR regularizes internal representations to promote cross-origin consistency in autoregressive neural PDE solvers.For deterministic, well-posed PDE flow, states from different origins on the same trajectory should reach the same target state.Source: §5 · p. 9
  • Second-order analysis motivates latent state and dynamics objectives that align shared-target representations and their local linearized evolution.The method can be integrated into different neural PDE solvers without modifying their backbone architectures, and experiments report improved accuracy, reduced discrepancies, and greater stability under longer rollouts.Source: §5 · p. 10

A NOTATION

  • The notation distinguishes rollout origin ti from target time tj and uses ti→tj to mark origin-dependent quantities.It also defines zti→tj as the latent representation at the target along a rollout from ti, and Lstate and Ldynamic as the state and dynamics consistency losses.Source: A NOTATION · p. 14

B DATASET

The dataset appendix describes the governing equations and settings for Navier–Stokes, shallow-water, and Kolmogorov-flow benchmarks. It also specifies the FNO, F-FNO, and LSM architectures and the intermediate representations selected for ALR.

B.1 NAVIER–STOKES EQUATION

  • The Navier–Stokes benchmark describes two-dimensional incompressible flow in vorticity form, with viscosity ν, external forcing, and periodic boundary conditions.The velocity field u and vorticity w are related by w = ∇×u, and incompressibility is imposed by ∇·u = 0.Source: §B.1 · pp. 14–15

B.2 SHALLOW-WATER EQUATION

  • The shallow-water benchmark describes free-surface flow in conservative form, with water depth h and velocity components u and v.The momentum components are hu and hv; bathymetry is denoted by b(x), and g is gravitational acceleration.Source: §B.2 · p. 15

B.3 KOLMOGOROV FLOW

  • Kolmogorov Flow uses two-dimensional incompressible Navier–Stokes dynamics with forcing; its setup specifies n = 8 and a linear drag term −0.1ω(x, t).The system uses a periodic domain and initial vorticity sampled from a prescribed Gaussian random field.Source: §B.3 · pp. 15–16

C BASELINE NEURAL SOLVERS

  • ALR uses the activated second Fourier-layer output as FNO's latent state, the feature after F-FNO's second residual spectral block, and LSM's fourth-downsampling-stage output.The remaining backbone components retain their original structures.Source: C BASELINE NEURAL SOLVERS · p. 16

D.1 DATASET CONFIGURATION

The experiments use Navier–Stokes, Shallow Water, and Kolmogorov Flow datasets with specified trajectory splits, spatial resolutions, and rollout settings. Shallow Water additionally includes longer rollout horizons for a horizon analysis.

D.1 DATASET CONFIGURATION

  • The ν = 10−5 Navier–Stokes dataset has 1,200 trajectories of 20 states on a 64 × 64 grid, split 800/200/200 for training, validation, and testing.The solver starts from the initial state and advances for 19 steps without additional downsampling or normalization.Source: §D.1 · p. 16
  • The ν = 10−4 Navier–Stokes dataset has 10,000 trajectories of 30 states on a 64 × 64 grid, with 800/100/100 trajectories assigned to training, validation, and testing.The rollout advances 29 steps, retaining the original resolution and applying no data normalization.Source: §D.1 · p. 17
  • Shallow Water uses 1,000 radial dam-break trajectories, downsampled to 64 × 64, with an 800/100/100 split and a standard 20-step rollout.The horizon analysis also tests 40- and 60-step settings; water depth is standardized using training-split statistics.Source: §D.1 · p. 17
  • Kolmogorov Flow uses 120 trajectories of 320 states, subsampled to 128 × 128, with 100/10/10 trajectories for training, validation, and testing.The standardized scalar field is evaluated over a 20-step autoregressive rollout.Source: §D.1 · p. 17

D.2 TRAINING CONFIGURATION

Training uses a ground-truth-conditioned one-step prediction objective, while auxiliary rollouts construct ALR consistency pairs. State and dynamics consistency terms use equal weights, with warm-up schedules varying by benchmark.

D.2 TRAINING CONFIGURATION

  • Models train for 500 epochs with AdamW, a 5 × 10−4 learning rate, 10−5 weight decay, and a OneCycle schedule, using relative ℓ2 prediction loss.For ALR, an origin is sampled per trajectory and an auxiliary local rollout constructs consistency pairs without adding a prediction loss; K = 6 and both consistency weights are 0.1 unless specified otherwise. Warm-up epochs are 50/100 for Navier–Stokes and 100/200 for Shallow Water and Kolmogorov Flow, for state/dynamics terms respectively.Source: §D.2 · p. 17

D.3 EVALUATION METRIC

The evaluation uses relative ℓ2 error, with lower values indicating more accurate predictions.

D.3 EVALUATION METRIC

  • Relative ℓ2 error evaluates prediction accuracy, with lower values indicating more accurate predictions.The metric compares the ground-truth and predicted solutions.Source: §D.3 · pp. 17–18

E COMPUTATIONAL EFFICIENCY

ALR adds training computation but leaves inference time unchanged. K = 6 is the default, balancing the reported accuracy gain against training efficiency.

E COMPUTATIONAL EFFICIENCY

  • Training cost rises from 2.04× at K = 4 to 4.04× at K = 10, while normalized inference time remains 1.00× for every setting.Runtime is normalized to the original FNO baseline.Source: E COMPUTATIONAL EFFICIENCY · p. 18, Table 7 · p. 18
  • Relative ℓ2 error falls from 0.0893 at K = 4 to 0.0858 at K = 6; K = 6 is the default trade-off between accuracy and training efficiency.Further increases in K produce only marginal accuracy changes despite noticeably higher training cost.Source: Table 7 · p. 18, E COMPUTATIONAL EFFICIENCY · p. 18

F QUALITATIVE VISUALIZATION OF AUTOREGRESSIVE ROLLOUT

The qualitative comparison follows FNO and FNO+ALR for 29 autoregressive steps on Navier–Stokes from the same initial state. FNO+ALR remains more aligned with ground truth at intermediate and late rollout stages.

F QUALITATIVE VISUALIZATION OF AUTOREGRESSIVE ROLLOUT

  • Across the 29-step Navier–Stokes rollout, FNO+ALR stays more closely aligned with ground truth than FNO at intermediate and late stages.The baseline's visible deviations begin around T = 9–10 and become clearer at T ≥ 16.Source: F QUALITATIVE VISUALIZATION OF… · p. 18, Figure 4 · p. 19
  • Both methods remain close to ground truth at T = 1–8, while FNO later shows phase shifts and boundary distortions near the central vorticity transition.The differences emerge around T = 9–10 and become more pronounced around T = 13–15.Source: F QUALITATIVE VISUALIZATION OF… · p. 18

Related papers

Loading 2610.21370v1…