arXiv preprint ·

One-Way Coupling via Patchwise Adaptive Normalization in Transformer Surrogates

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

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Machine Learning

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.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Machine Learning

Get the few papers worth your time in Machine Learning every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics