Patchwise adaptive normalization improves transformer PDE surrogates efficiency
One-Way Coupling via Patchwise Adaptive Normalization in Transformer Surrogates
Machine Learning
Summary
Physical systems often have factors that affect their behavior but aren’t changed by it; these are called one-way couplings. The authors present a method called Patchwise Adaptive Normalization (PAdaNorm) that helps transformer models better understand these factors by dividing them into patches. This technique improves prediction accuracy and keeps important physical rules intact over time. It also makes predictions faster when the conditions don’t change. Tests across several physical systems show PAdaNorm performs as well or better than existing methods.
What this means in practice
- •For engineering simulation teams: Improve accuracy and speed of physics-based simulations that use transformer models with static environmental factors.
- •For fluid dynamics modelers: Maintain physical constraints more reliably in long-term flow simulations around fixed obstacles using transformer surrogates.
Authors
Emmanuel Gottlob, Hugo Thimonier, Ludovico Buizza, August Martens, Alexandre Gavaudan, Thiago Guzella
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.