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

TL;DR

The State-Exchange Attention (SEA) module is introduced, a novel transformer-based module enabling information exchange between encoded fields through multi-head cross-attention, and it is demonstrated that the SEA module alone can reduce errors by 97% for state variables that are highly dependent on other states of the system.

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

Figures (16)

• Figure 1: The ViT-based mesh autoencoder divides the domain into patches and pads them to ensure equal sizes. An MLP is then applied to reduce the dimensionality, and the MHSA mechanism provides global awareness to create spatially coherent reconstructions with minimal patch artifacts. These patches are subsequently flattened and adapted as tokens for the temporal model.
• Figure 2: (a) State-Exchange Attention (SEA) integrated Transformer model architecture, incorporating the additional modules of SEA and Time Invariant Parameter Injection (TIPI). Dashed lines represent the inclusion of additional fields. (b) Representation of the TIPI, designed to incorporate time-invariant information after the SEA module. (c) Schematic of the SEA module, illustrating how fields communicate through this module.
• Figure 3: Comparison of the rollout error for the cylinder flow dataset.
• Figure 4: Contour maps of the generated fields at Re=400, and time step of 250 where Von Karman vortex street is formed.
• Figure 6: Comparison of the contour maps of the volume fraction ($\alpha$), between the predictions and ground truth in the case of $\rho = 850$