Many physical systems exhibit one-way coupling: exogenous fields such as medium density or obstacle geometry influence the dynamics without being affected by them. We investigate how explicitly representing these fields can improve the accuracy and computational efficiency of transformer-based PDE surrogates. Building on spatially adaptive normalization, we introduce Patchwise Adaptive Normalization (PAdaNorm), which partitions exogenous fields into patches aligned with the state tokenization. We evaluate PAdaNorm on datasets from three different domains, using a Walrus-inspired transformer architecture. Across the three PDE families, PAdaNorm is on par with or up to 16% better than parameter-matched baselines in mean long-horizon variance-scaled RMSE on 60-step autoregressive rollouts. We also find that PAdaNorm better maintains hard physical constraints over long rollouts, such as preventing flow into impenetrable solid regions in porous media. When the conditioning is static, PAdaNorm caches the conditioning embeddings once and reuses them across the entire rollout, reducing inference FLOPs per step by up to 36% compared with parameter-matched baselines. An ablation over eight conditioning variants shows that modulating a single attention axis (temporal or spatial) suffices and that a lightweight shared MLP across blocks remains competitive with per-block projections.
Figures & tables
Figure 1: Patchwise Adaptive Normalization . The spatially varying conditioning fields C are partitioned into patches aligned with the input tokenization. Each patch is independently mapped through a shared MLP fθ to produce a local conditioning vector cs . Each token hs , in every frame of the history window, is then modulated by the scale, shift, and gate derived from cs (conditioning vector of the corresponding patch).
Table 1: Long-horizon validation vRMSE ( ↓ , T∈[21:60] ) for baselines and PAdaNorm spatial-only, reported as mean ± standard deviation over three seeds. Best mean in bold and second best underlined within each dataset. ‡ Parameter-matched In-context controls. The Joint result reports vRMSE averaged over the three datasets, for a single model trained on all three simultaneously. GFLOPs denotes one forward pass at batch size one; cached rollout costs are reported in Table 12 .
Conditioning
#Params
Maze
Inclusions
Discontinuous
Joint
Spatial-only
19.1 M
0.3716
0.4094
0.2352
0.8015
Temporal-only
19.1 M
0.2996
0.7522
0.3864
0.9520
Temporal (shared proj.)
19.1 M
0.3352
0.9115
0.4749
1.0266
Temporal + Spatial
23.1 M
0.3203
0.7745
0.3732
0.9325
Fully shared
15.6 M
0.3447
0.4819
0.3553
0.7899
Fully shared (spatial)
15.6 M
0.3522
0.5263
0.2657
0.8628
Table 2: PAdaNorm conditioning mode : long-horizon validation vRMSE ( ↓ , T∈[21:60] ) on three acoustic scattering datasets from The Well ( Ohana et al., 2024 ) . We highlight the best metric in bold and underline the second best. The Joint column reports vRMSE averaged over the three datasets, each contributing 200 validation windows, for a single model trained on all three simultaneously. All entries are single runs (seed 1234); the three-seed results for spatial-only are in Tab. 1 .
Patch strategy
#Params
Maze
Inclusions
Discontinuous
T∈[21:60]
Overall
T∈[21:60]
Overall
T∈[21:60]
Overall
Base
19.1 M
0.3716
0.2938
0.4094
0.3054
0.2352
0.1732
Sliding-window overlap
19.1 M
0.3770
0.2982
0.3461
0.2629
0.2291
0.1688
Convolutional overlap ∗
20.8 M
0.2927
0.2236
0.4076
0.3037
0.2511
0.1852
Coarse 2×
19.3 M
0.4457
0.3374
0.5770
0.4311
0.3566
0.2604
Table 3: Effect of conditioning patch strategy for spatial-only PAdaNorm on validation vRMSE ( ↓ ) across three acoustic scattering datasets (60-step autoregressive rollout). We compare the standard (base) patching against three alternatives: sliding-window overlap, convolutional overlap, and coarse 2× . All metrics are reported at the checkpoint with the best T∈[21:60] score. We highlight the best metric in bold and underline the second best. All entries are single runs (seed 1234).
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Table 4: Best validation vRMSE ( ↓ ) across conditioning modes on three acoustic scattering datasets from The Well ( Ohana et al., 2024 ) (60-step autoregressive rollout). All models share similar hyperparameters where relevant. Rows that state a spread are the mean ± sample standard deviation over three seeds (1234, 21, 67), each interval read at that seed’s best T∈[21:60] checkpoint; rows without one are marked as single-seed and are not directly comparable to them. We highlight the best metric in bold and underline the second best. GFLOPs is the batch-1 forward pass, as in Tab. 12 . † Parameter-matched In-context controls use the same total parameter count as the PAdaNorm temporal+spatial variants ( ∼ 23 M). ‡ Parameter-matched In-context controls targeting the per-block PAdaNorm parameter count ( ∼ 19 M). The Joint column reports vRMSE averaged over the three datasets, each contributing 200 validation windows, for a single model trained on all three simultaneously.
Table 5: Best validation vRMSE ( ↓ ) for baselines and PAdaNorm spatial-only on FlowBench FPO ( Tali et al., 2025 ) (2D Navier–Stokes, 1024×256 , 60-step autoregressive rollout). ‡ Parameter-matched In-context controls targeting the PAdaNorm parameter count ( ∼38.6 M). Each entry is the mean ± sample standard deviation over three seeds (1234, 21, 67). Per seed the three intervals are read at that seed’s best T∈[21:60] checkpoint, so each row describes one model per seed. GFLOPs is the batch-1 forward pass, reproduced from Tab. 12 . We highlight the best mean in bold and underline the second best.
Table 6: Best validation vRMSE ( ↓ ) for baselines and PAdaNorm spatial-only on CO 2 –water flow in porous media ( Abdellatif et al., 2026 ) ( 512×512 , Tin=4 , 60-step autoregressive rollout). ‡ Parameter-matched In-context controls targeting the PAdaNorm parameter count ( ∼37.7 M). Each entry is the mean ± sample standard deviation over three seeds (1234, 21, 67). Per seed the three intervals are read at that seed’s best T∈[21:60] checkpoint, so each row describes one model per seed. GFLOPs is the batch-1 forward pass, reproduced from Tab. 12 . We highlight the best mean in bold and underline the second best.
Table 7: Mean absolute error on rock pixels ( ↓ , z-score normalized units) on CO 2 –water flow in porous media, averaged over T∈[21:60] . All fields must be zero within the rock pixels. Each entry is the mean ± sample standard deviation over three seeds (1234, 21, 67), each read at its best T∈[21:60] checkpoint, as in Tab. 6 . ‡ Parameter-matched In-context controls. We highlight the best mean in bold and underline the second best.
Figure 2: Autoregressive rollout on FlowBench FPO (pressure channel, Reynolds number Re=364 ). Rows: ground truth, the parameter-matched wide In-context baseline ( D=512 , 7 blocks), and PAdaNorm spatial-only ( D=384 , 10 blocks); each rollout is generated by the seed-1234 checkpoint from Tab. 1 . The obstacle is overlaid in gray.
Figure 3: Typical autoregressive rollout on acoustic scattering with inclusions (pressure). Rows: ground truth, the parameter-matched wide In-context baseline ( D=384 , 6 blocks), the strongest In-context baseline on this dataset in Tab. 1 , and PAdaNorm spatial-only ( D=256 , 10 blocks); each rollout is generated by a seed-21 checkpoint from the runs whose metrics enter that table. The inclusions are outlined in gray. The window is the one whose PAdaNorm -vs-baseline error gap is closest to the median over the candidates considered, so it is neither a best- nor a worst-case draw.
Figure 4: Typical autoregressive rollout on CO 2 –water flow in porous media ( Abdellatif et al., 2026 ) , showing the velocity magnitude ∥u∥ . The inset in the bottom-right corner of each panel magnifies the region outlined in white. The rows show the ground truth, In-context (deep) — the depth-matched parameter control — and PAdaNorm spatial-only. Throughout the rollout, PAdaNorm enforces the hard constraint of not having gas or water permeating the rock cells better than the parameter-matched In-context baseline, consistent with its lower error on rock pixels (Tab. 7 ).
the three datasets concatenated, sampled in proportion to size
State fields
pressure and velocity (3 channels), z-score normalized
Conditioning channels
density and speed of sound (2 channels)
Excluded trajectories
one trajectory (index 347) in inclusions, whose simulation diverges
Appendix
Table 9: Configuration for the acoustic scattering experiments from The Well ( Ohana et al., 2024 ) in Table 1 . All runs share the settings in the upper block and differ only in the lower one; head dimension is held at 128 throughout, so width and head count move together. Normalization uses The Well’s ZScoreNormalization with statistics computed on the training split. GNOT baselines use separate architectures and training settings, given in App. B.3.1 .
Data
Dataset
FlowBench FPO, 2D Navier–Stokes, 1024×256
Split (simulations)
881 train / 219 valid, stratified by geometry type
Trajectory length
242 steps; 236 sliding windows per trajectory
State fields
u,v,p (3 channels), z-score normalized
Conditioning channels
mask, SDF, Reynolds number (broadcast as a uniform 2D field)
Reynolds preprocessing
log10 then z-score normalized
Appendix
Table 10: Configuration for the FlowBench FPO ( Tali et al., 2025 ) experiments in Table 1 . All runs share the settings in the upper block and differ only in the lower one; conditioning inputs are normalized with statistics computed on the training split. GNOT baselines use separate architectures and training settings, given in App. B.3.1 .
Data
Dataset
CO 2 –water flow in porous media, 512×512
Split (simulations)
500 train / 124 valid, split by base geometry
Split procedure
100 base geometries permuted with seed 42; 80 train / 20 valid
binary pore mask (1 = pore, 0 = rock) and normalized x , y coordinates in [−1,1] (3 channels)
Appendix
Table 11: Configuration for the CO 2 –water flow in porous media ( Abdellatif et al., 2026 ) experiments in Table 1 . All runs share the settings in the upper block and differ only in the lower one; conditioning inputs are normalized with statistics computed on the training split. GNOT baselines use separate architectures and training settings, given in App. B.3.1 .
Table 12: Forward-pass and rollout cost of each conditioning strategy, for the models reported in Tab. 1 . Batch size 1, one multiply–add counted as two FLOPs. cond. (static) is the part of the conditioning cost that does not depend on the state and is computed once per rollout; cond. (state-dep.) is the part recomputed at every step, such as the cross-attention queries. The rollout columns are accordingly Φstep=Φfwd−Φstatic and 60Φstep+Φstatic . App. B.5 gives the counting convention and the derivations.
Figure 5: Conditioning patching strategies. Three approaches for conditioning patch design. Orange and blue regions show the input area contributing to two adjacent conditioning tokens, dashed squares mark the corresponding state patches, and the grid shows the P×P state patches. The width of the orange region is k for the sliding window, P+Lconv(k−1) for the convolutional encoder, and mP for coarse conditioning.
Partial differential equations (PDEs) are fundamental for modeling complex physical systems, yet classical numerical solvers face prohibitive computational costs in high-dimensional and multi-scale regimes. While Transformer-based neural operators have emerged as powerful data-driven alternatives, they conventionally treat all discretized spatial points as uniform, independent tokens. This monolithic approach ignores the intrinsic scale separation of physical fields, applying computationally prohibitive global attention that redundantly mixes smooth large-scale dynamics with high-frequency fluctuations. Rethinking Transformers through the lens of complex dynamics, we propose DynFormer, a novel dynamics-informed neural operator. Rather than applying a uniform attention mechanism across all scales, DynFormer explicitly assigns specialized network modules to distinct physical scales. It leverages a Spectral Embedding to isolate low-frequency modes, enabling a Kronecker-structured attention mechanism to efficiently capture large-scale global interactions with reduced complexity. Concurrently, we introduce a Local-Global-Mixing transformation. This module utilizes nonlinear multiplicative frequency mixing to implicitly reconstruct the small-scale, fast-varying turbulent cascades that are slaved to the macroscopic state, without incurring the cost of global attention. Integrating these modules into a hybrid evolutionary architecture ensures robust long-term temporal stability. Extensive memory-aligned evaluations across four PDE benchmarks demonstrate that DynFormer achieves up to a 95% reduction in relative error compared to state-of-the-art baselines, while significantly reducing GPU memory consumption. Our results establish that embedding first-principles physical dynamics into Transformer architectures yields a highly scalable, theoretically grounded blueprint for PDE surrogate modeling.
Pengyu Lai, Yixiao Chen, Dewu Yang +3
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
Transformer architectures have attracted increasing attention for solving partial differential equations (PDEs), owing to their flexibility in handling irregular discretizations and their ability to capture long-range physical dependencies. However, unlike discrete language tokens or fixed-resolution image patches, observed physical fields are finite samples of underlying infinite-dimensional functions. Consequently, effectively applying Transformers to PDEs requires a tokenizer that respects the functional nature of physical fields and constructs physically expressive tokens from arbitrary discretizations.To this end, we propose \methodname{Physics Transformer}, a function-projection-based Transformer architecture for physical field prediction. Physics Transformer treats a physical field as a continuous function and partitions its discretization into locality-preserving spatial patches. Within each patch, it dynamically learns a set of adaptive local basis functions and projects the sampled field onto these bases to obtain compact physics tokens. The resulting tokens capture diverse latent physical states while preserving fine-scale spatial structures, enabling efficient global interaction through factorized attention across space and physical states. The projected representation further supports efficient decoding at arbitrary query locations. Extensive experiments on diverse benchmarks, ranging from two-dimensional PDE dynamics to industrial-scale three-dimensional CFD simulations, demonstrate that Physics Transformer accurately captures fine-grained physical structures and achieves state-of-the-art predictive performance. These results establish function projection as a practical and effective foundation for designing Transformer architectures for PDE solving.
Guoze Sun, Rui Zhang, Jiankai Tang +4
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China · School of Mechanics and Engineering Science, Peking University, Beijing, China
We study the problem of architecture selection for deep learning models trained to solve partial differential equations (PDEs), asking when transformer-based architectures with learned attention outperform Fourier-domain neural operators. We introduce the Multi-Scale Attention Transformer (\msat{}), a deep learning architecture that encodes spatiotemporal solution histories as token sequences and trains end-to-end via a composite supervised objective with optional physics-informed regularization terms. We conduct a comprehensive empirical evaluation against nine baselines -- including physics-informed neural networks (PINNs), neural operators (FNO, DeepONet, GNOT), and state-space models (Mamba-NO) -- across five benchmark problems from the PINNacle suite, using identical train/test splits and reference data for all methods. \msat{} achieves state-of-the-art generalization on complex geometry problems (Lrel2=0.0101 on Heat2D-CG, a 3.7× improvement over FNO) at 34s total inference vs.\ 120,812s for Mamba-NO. Ablation studies over the physics regularization component reveal a precise inductive bias tradeoff: physics priors reduce test error on diffusion-dominated problems but degrade generalization on chaotic and recirculating-flow regimes, directly characterizing the prior misspecification boundary. Approximation error bounds as a function of domain boundary complexity κ provide a theoretical basis for these empirical findings and a principled rule for architecture selection.