MATEY: multiscale adaptive foundation models for spatiotemporal physical systems

ICLR 2025 Conference Submission12946 Authors

Keywords: Foundation model; vision transformer, physical systems; adaptive tokenization; decoupled spatiotemporal attentions; computational fluid dynamics

TL;DR: We present MATEY, a multiscale adaptive foundation model for spatiotemporal physical systems, featuring adaptive tokenization and decoupled spatiotemporal attention schemes, and demonstrate the effectiveness of pretraining in two fine-tuning cases.

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 an adaptive tokenization scheme which dynamically adjusts the token sizes based on local features.
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 achieves improved accuracy without significantly increasing token sequence length, but the improvement deteriorates in more complex data configurations.
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.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 12946