MATEY: multiscale adaptive foundation models for spatiotemporal physical systems

Accurate representation of the multiscale features in spatiotemporal systems using vision transformer (ViT) architectures requires extremely long, computationally prohibitive token sequences. To address this issue, we propose an adaptive tokenization scheme which dynamically adjusts the token sizes based on local features. Moreover, we introduce a set of spatiotemporal attention schemes, built on the axial attention approach in which a full attention mechanism is decoupled into attentions in the axial dimensions. We assess the performance of the proposed multiscale adaptive model, MATEY, in a sequence of experiments. The results show that adaptive tokenization is up to eight times more cost-efficient. Compared to a full spatiotemporal attention scheme, we find that decoupled attention requires more training time and tokens 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.