MATEY: multiscale adaptive foundation models for spatiotemporal physical systems

TL;DR

MATEY introduces adaptive tokenization and decoupled spatiotemporal attention schemes in vision transformers to efficiently model multiscale features in physical systems, improving accuracy and efficiency especially with limited data.

Contribution

The paper proposes novel adaptive tokenization schemes and spatiotemporal attention methods for vision transformers applied to physical systems modeling.

Findings

Adaptive tokenization improves accuracy without increasing sequence length.

Fully decoupled axial attention is less efficient than coupled schemes.

Pretrained models outperform from-scratch training, especially with limited data.

Abstract

Accurate representation of the multiscale features in spatiotemporal physical systems using vision transformer (ViT) architectures requires extremely long, computationally prohibitive token sequences. To address this issue, we propose two adaptive tokenization schemes that dynamically adjust patch sizes based on local features: one ensures convergent behavior to uniform patch refinement, while the other offers better computational efficiency. Moreover, we present a set of spatiotemporal attention schemes, where the temporal or axial spatial dimensions are decoupled, and evaluate their computational and data efficiencies. We assess the performance of the proposed multiscale adaptive model, MATEY, in a sequence of experiments. The results show that adaptive tokenization schemes achieve improved accuracy without significantly increasing the length of the token sequence. Compared to a full spatiotemporal attention scheme or a scheme that decouples only the temporal dimension, we find that fully decoupled axial attention is less efficient and expressive, requiring more training time and model weights to achieve the same accuracy. Finally, we demonstrate in two fine-tuning tasks featuring different physics that models pretrained on PDEBench data outperform the ones trained from scratch, especially in the low data regime with frozen attention.