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

Abstract: 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.

• The paper introduces AE-ViT, a jointly trained convolutional encoder-decoder and latent vision transformer that injects physical parameters through FiLM modulation, attention conditioning, and parameter tokens.
• AE-ViT reduces long-horizon rollout errors by roughly fivefold versus leading baselines, achieving 0.0059 on advection-diffusion-reaction and 0.0276 for cylinder-wake velocity prediction.
• The method maintains accurate rollouts hundreds of steps beyond its four-step training window, but remains limited by rectangular grids, quadratic attention cost, and untested out-of-distribution parameter extrapolation.

Overview

The paper introduces AE-ViT, an autoregressive surrogate architecture for parametric time-dependent PDEs that combines a fully convolutional encoder-decoder with a vision transformer (ViT) processor operating on latent tensor representations. The work targets the mapping $(\lambda, t) \mapsto \phi(\cdot, t; \lambda)$ for parameter-dependent evolution problems, where $\lambda$ may include material coefficients, forcing parameters, and geometric quantities. Two design concerns motivate the architecture: first, full-field models must capture long-range spatial interactions at high computational cost, while autoencoder-based latent evolution models typically produce latent vectors that encode spatial structure but are difficult to evolve; second, in parametric PDEs the initial condition alone does not determine the trajectory, so explicit parameter conditioning is essential. The authors' central claim is that joint training of encoder, processor, and decoder—with parameter information injected at multiple stages—yields latent dynamics that are both cheap to evolve and accurate over long horizons. On their benchmarks, AE-ViT reduces relative rollout error by roughly a factor of five relative to DL-ROMs, latent transformers, and plain ViTs.

Architecture and method

AE-ViT consists of three components trained end-to-end. A fully convolutional encoder with ResNet blocks and Group Normalization compresses the solution snapshot to a spatially structured latent tensor ($16 \times 16 \times 256$ on a $32\times32$ grid), avoiding the lossy projection to a one-dimensional latent vector used in most latent-evolution ROMs. A ViT encoder then processes this tensor: patches are extracted via strided convolution, enriched with learnable positional encodings, and passed through transformer layers. Finally, a mirrored convolutional decoder reconstructs the next snapshot at original resolution.

Parameter conditioning is applied throughout, which is the paper's principal novelty:

• FiLM modulation in the encoder and decoder: each hidden state $h$ is transformed as $$h \leftarrow \alpha(\lambda) \odot h + \beta(\lambda)$$, acting as channel-wise reweighting.
• Parametric attention: an MLP maps $\lambda$ into per-head, per-channel scaling and shifting of the query, key, and value matrices, e.g., $$V_t \leftarrow V_t \odot (1 + \eta\,\alpha_V(\lambda)) + \eta\,\beta_V(\lambda)$$, where $\eta$ is a learnable per-layer strength bounded by a fixed cap. Layer-normalization modulations are initialized near zero ($\theta = 10^{-3}$), so training begins close to an unconditioned transformer—a stabilization strategy borrowed from DiT-style conditioning and LoRA-style identity initialization.
• Parameter token: $\lambda$0 can be appended as an additional attention token.

Coordinate channels are concatenated to the input using Fourier features at frequencies $\lambda$1, providing explicit spatial awareness beyond what patch positional encodings supply. Notably, time is deliberately excluded from all injection modules because it leaves the training range under time extrapolation; the authors report that including time aids short-term fit but degrades extrapolation.

Training uses scheduled sampling with inverse-sigmoid decay over a fixed window of four steps, minimizing average MSE across the unrolled window. Despite this short-horizon supervision, rollouts remain accurate over hundreds of steps, which the authors interpret as evidence that the model learns a stable approximation of the solution operator rather than merely minimizing one-step error. They also offer a kernel-regression interpretation: the softmax attention weights $\lambda$2 form a learned, state- and parameter-dependent effective interaction kernel on latent tokens—an appropriate view for nonlinear problems such as Navier-Stokes, though the authors are careful not to claim formal derivation from Fourier/Galerkin attention theory.

Positioning relative to prior work

The paper distinguishes its approach from operator-learning methods such as DeepONet and FNO, which treat physical parameters as part of the input function rather than explicitly conditioned features; when trajectories share an initial condition, such models cannot differentiate parameter instances. It also contrasts with separately trained autoencoder-plus-dynamics pipelines (DL-ROM, Neural ODE latent dynamics, latent transformers). The authors hypothesize that separate training produces latents optimized for reconstruction rather than predictive robustness, and they report a concrete failure supporting this: training a ViT on a frozen autoencoder's latent representations did not converge at all, whereas joint training succeeds. Compared to FactFormer's factorized axial attention and VCNeF's conditional neural fields, AE-ViT emphasizes joint parameter-aware learning; VCNeF avoids autoregressive rollout but is memory-heavy and evaluated only on short horizons without assessing compounding error, while FactFormer lacks parameter conditioning.

Results

Advection-Diffusion-Reaction. Following the benchmark setup of Farenga et al., 800 training simulations on a $\lambda$3 FEM grid are used, with test rollouts extending from $\lambda$4 to $\lambda$5. AE-ViT achieves a mean relative rollout error of 0.0029 on $\lambda$6 and 0.0059 on $\lambda$7, versus 0.0123/0.6473 for DL-ROM and 0.0117/0.0229 for AE + 1D transformer, with only ~6M parameters compared to ~52M (DL-ROM) and ~38M (latent transformer). Plain ViT performs poorly (0.0997/0.2366). DL-ROM cannot extrapolate past its training window since it predicts entire trajectories conditioned on time. Rollout error grows approximately linearly, with dips consistent with the periodic post-transient solution.

An ablation isolating single enhancements against an unconditioned baseline (rollout error 0.4274 at $\lambda$8) shows that FiLM in transformer attention alone yields 0.005188, FiLM in encoder/decoder alone 0.005207, coordinate encoding alone 0.0575, and a parameter token alone 0.008613. The near-identical performance of feature-only and attention-only conditioning supports the paper's three-regime taxonomy and motivates the fully conditioned configuration actually deployed.

Navier-Stokes cylinder wake. The harder benchmark solves 2D incompressible flow past a circular obstacle in a pipe, with parameters comprising inflow perturbation amplitude $\lambda$9, obstacle position $16 \times 16 \times 256$0, and radius $16 \times 16 \times 256$1 (hence Reynolds number $16 \times 16 \times 256$2). Velocity components and pressure are learned jointly on a masked $16 \times 16 \times 256$3 grid, trained on two inlet periods (450 snapshots per simulation, 800 simulations) and rolled out over five periods. AE-ViT attains the best error on every component: e.g., $16 \times 16 \times 256$4 rollout error of 0.01725 at $16 \times 16 \times 256$5 and 0.0276 at $16 \times 16 \times 256$6, versus 0.0839/0.1391 for the best baseline (AE + 1D transformer); pressure errors reach 0.0999/0.1861 versus 0.9068/1.3438. This is the source of the headline claim of roughly fivefold error reduction. Pressure exhibits periodic error spikes attributed to phase sensitivity: since pressure in incompressible flow is determined globally through a Poisson constraint, small phase drifts in predicted vortex positions produce large pointwise pressure errors even when the flow structure is well captured.

Limitations and open questions

The authors are explicit about several constraints. The method requires solutions interpolated onto rectangular grids, making it unsuitable for domains that do not naturally admit such discretization—although geometric parameters (obstacle position and radius) are handled here via masking, and fluid-structure interaction with genuinely varying domains remains future motivation rather than demonstrated capability. Self-attention imposes quadratic cost in token count, bounding practical resolution; sparse or linear-complexity attention variants are acknowledged but not explored. Evaluation is restricted to in-distribution parameter rollouts; extrapolation beyond the calibrated parameter range is not addressed. The stability observation—that error accumulates approximately linearly over horizons 250 times longer than the scheduled sampling window—is empirical, and the paper leaves open the development of rigorous error and complexity bounds via neural network approximation theory. The kernel-regression interpretation of parametric attention likewise remains heuristic, since standard softmax attention is not covered by existing theoretical analyses of Fourier- or Galerkin-type attention.

Conclusion

AE-ViT demonstrates that jointly trained convolutional autoencoding and vision-transformer latent evolution, equipped with multi-stage parameter injection and coordinate channel encoding, can deliver latent-space efficiency with full-field fidelity on multi-component, scale-imbalanced parametric PDEs. The empirical evidence—fivefold rollout-error reduction on Navier-Stokes wake prediction, linear error accumulation far beyond the training window, and ablations confirming complementary contributions of feature-level and attention-level conditioning—substantiates the design choices within the stated in-distribution regime. The main unresolved issues are domain generality, attention scalability, out-of-distribution parameter generalization, and theoretical grounding of the observed long-horizon stability.