One-Way Coupling via Patchwise Adaptive Normalization in Transformer Surrogates

Emmanuel Gottlob; Hugo Thimonier; Ludovico Buizza; August Martens; Alexandre Gavaudan; Thiago Guzella

cs.LGOct 10, 2026

One-Way Coupling via Patchwise Adaptive Normalization in Transformer Surrogates

Authors: Emmanuel Gottlob, Hugo Thimonier, Ludovico Buizza, August Martens, Alexandre Gavaudan, Thiago Guzella

Organizations: Mistral AI · ASML The Netherlands B.V.

Abstract

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

Panel of Figure 1 — Figure 1: Patchwise Adaptive Normalization . The spatially varying conditioning fields $\mathbf{C}$ are partitioned into patches aligned with the input tokenization. Each patch is independently mapped through a shared MLP $f_{\theta}$ to produce a local conditioning vector $\mathbf{c}_{s}$ . Each token $\mathbf{h}_{s}$ , in every frame of the history window, is then modulated by the scale, shift, and gate derived from $\mathbf{c}_{s}$ (conditioning vector of the corresponding patch).

Panel of Table 1 — Table 1: Long-horizon validation vRMSE ( $\downarrow$ , $T{\in}[21{:}60]$ ) for baselines and PAdaNorm spatial-only, reported as mean $\pm$ 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 (

\downarrow

T{\in}[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
Patch strategy	#Params	$T{\in}[21{:}60]$	Overall	$T{\in}[21{:}60]$	Overall	$T{\in}[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{\times}$	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 (

\downarrow

) 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{\times}

. All metrics are reported at the checkpoint with the best

T{\in}[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 (

\downarrow

) 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

\pm

sample standard deviation over three seeds (1234, 21, 67), each interval read at that seed’s best

T{\in}[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 (

{\sim}

23 M). ‡ Parameter-matched In-context controls targeting the per-block PAdaNorm parameter count (

{\sim}

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.

Panel of Table 5 — Table 5: Best validation vRMSE ( $\downarrow$ ) for baselines and PAdaNorm spatial-only on FlowBench FPO ( Tali et al., 2025 ) (2D Navier–Stokes, $1024{\times}256$ , 60-step autoregressive rollout). ‡ Parameter-matched In-context controls targeting the PAdaNorm parameter count ( ${\sim}38.6$ M). Each entry is the mean $\pm$ sample standard deviation over three seeds (1234, 21, 67). Per seed the three intervals are read at that seed’s best $T{\in}[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.

Panel of Table 6 — Table 6: Best validation vRMSE ( $\downarrow$ ) for baselines and PAdaNorm spatial-only on CO 2 –water flow in porous media ( Abdellatif et al., 2026 ) ( $512{\times}512$ , $T_{\mathrm{in}}{=}4$ , 60-step autoregressive rollout). ‡ Parameter-matched In-context controls targeting the PAdaNorm parameter count ( ${\sim}37.7$ M). Each entry is the mean $\pm$ sample standard deviation over three seeds (1234, 21, 67). Per seed the three intervals are read at that seed’s best $T{\in}[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.

Panel of Table 7 — Table 7: Mean absolute error on rock pixels ( $\downarrow$ , z-score normalized units) on CO 2 –water flow in porous media, averaged over $T{\in}[21{:}60]$ . All fields must be zero within the rock pixels. Each entry is the mean $\pm$ sample standard deviation over three seeds (1234, 21, 67), each read at its best $T{\in}[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 $\mathrm{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 $\lVert\mathbf{u}\rVert$ . 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 ).

Conditioning	Spatial modulation	Temporal modulation	Axis-specific	Block-specific
Spatial-only	✓	✗	–	✓
Temporal-only	✗	✓	–	✓
Temporal (shared proj.)	✓	✓	✗	✓
Temporal + Spatial	✓	✓	✓	✓
Fully shared	✓	✓	✗	✗
Fully shared (spatial)	✓	✗	–	✗

Appendix

Table 8: PAdaNorm conditioning variants.

Data
Datasets	acoustic scattering: maze, inclusions, discontinuous ( $256{\times}256$ )
Split	official The Well train / valid splits
Joint training	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{\times}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	$\log_{10}$ 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{\times}512$
Split (simulations)	500 train / 124 valid, split by base geometry
Split procedure	100 base geometries permuted with seed 42; 80 train / 20 valid
Trajectory length	100 steps; 96 sliding windows per trajectory
State fields	$\alpha_{\mathrm{water}},p,u_{x},u_{y}$ (4 channels), z-score normalized
Conditioning channels	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 .

Panel of Table 12 — 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 $\Phi_{\mathrm{step}}=\Phi_{\mathrm{fwd}}-\Phi_{\mathrm{static}}$ and $60\,\Phi_{\mathrm{step}}+\Phi_{\mathrm{static}}$ . App. B.5 gives the counting convention and the derivations.

Panel of Figure 5 — 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\times P$ state patches. The width of the orange region is $k$ for the sliding window, $P+L_{\mathrm{conv}}(k-1)$ for the convolutional encoder, and $mP$ for coarse conditioning.

PDE Surrogate Modeling

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

See all 331 papers

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs
Mar 3, 2026Pengyu Lai, Yixiao Chen, Dewu Yang +3PDE Surrogate Modeling Neural Operators
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
Physics Transformer: Tailoring Transformer for General PDE Prediction
Jul 27, 2026Guoze Sun, Rui Zhang, Jiankai Tang +4Transformer PDE Surrogate Modeling
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China · School of Mechanics and Engineering Science, Peking University, Beijing, China
When Attention Beats Fourier: Multi-Scale Transformers for PDE Solving on Irregular Domains
May 8, 2026Brandon Yee, Pairie Koh, Jack Rodriguez +1Neural PDE Solvers PDE Solving
Physics Lab, Yee Collins Research Group

Mar 3, 2026J/K move · Enter open · S save

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs

Pengyu Lai, Yixiao Chen, Dewu Yang, Rui Wang+2

School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China

PDE Surrogate Modeling Neural Operators Physics-Informed Neural Operators Transformer Attention Attention Mechanisms

Open paper AskPDF arXiv

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.

Abstract

PDE Surrogate Modeling

Explore similar work

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs

Physics Transformer: Tailoring Transformer for General PDE Prediction

When Attention Beats Fourier: Multi-Scale Transformers for PDE Solving on Irregular Domains