SEA: State-Exchange Attention for High-Fidelity Physics Based Transformers

TL;DR

This paper introduces the State-Exchange Attention (SEA) module, a transformer-based approach that improves the accuracy of physics-based predictions in dynamical systems by enabling multidirectional information exchange between fields, significantly reducing rollout errors.

Contribution

The paper presents the novel SEA module with cross-field attention and an efficient mesh autoencoder, achieving state-of-the-art error reduction in physics-based transformers for dynamical systems.

Findings

Outperforms existing models with 88-91% error reduction

SEA module reduces errors by 97% in highly dependent state variables

Demonstrates superior rollout accuracy over baseline models

Abstract

Current approaches using sequential networks have shown promise in estimating field variables for dynamical systems, but they are often limited by high rollout errors. The unresolved issue of rollout error accumulation results in unreliable estimations as the network predicts further into the future, with each step's error compounding and leading to an increase in inaccuracy. Here, we introduce the State-Exchange Attention (SEA) module, a novel transformer-based module enabling information exchange between encoded fields through multi-head cross-attention. The cross-field multidirectional information exchange design enables all state variables in the system to exchange information with one another, capturing physical relationships and symmetries between fields. Additionally, we introduce an efficient ViT-like mesh autoencoder to generate spatially coherent mesh embeddings for a large number of meshing cells. The SEA integrated transformer demonstrates the state-of-the-art rollout error compared to other competitive baselines. Specifically, we outperform PbGMR-GMUS Transformer-RealNVP and GMR-GMUS Transformer, with a reduction in error of 88% and 91%, respectively. Furthermore, we demonstrate that the SEA module alone can reduce errors by 97% for state variables that are highly dependent on other states of the system. The repository for this work is available at: https://github.com/ParsaEsmati/SEA