Spotlight

Walrus: A Cross-domain Foundation Model for Continuum Dynamics

Michael McCabe ⋅ Payel Mukhopadhyay ⋅ Tanya Marwah ⋅ Bruno Régaldo-Saint Blancard ⋅ François Rozet ⋅ Cristiana Diaconu ⋅ Lucas Meyer ⋅ Kaze Wong ⋅ Hadi Sotoudeh ⋅ Alberto Bietti ⋅ Irina Espejo ⋅ Rio Fear ⋅ Siavash Golkar ⋅ Tom Hehir ⋅ Keiya Hirashima ⋅ Geraud Krawezik ⋅ Francois Lanusse ⋅ Rudy Morel ⋅ Ruben Ohana ⋅ Liam Parker ⋅ Mariel Pettee ⋅ Jeff Shen ⋅ Kyunghyun Cho ⋅ Miles Cranmer ⋅ Shirley Ho

Abstract

Foundation models have transformed machine learning for language and vision, but achieving comparable impact in physical simulation remains a challenge. Data heterogeneity and unstable long-term dynamics inhibit learning from sufficiently diverse dynamics, while varying resolutions and dimensionalities challenge efficient training on modern hardware. Through empirical and theoretical analysis, we incorporate new approaches to mitigate these obstacles, including a harmonic-analysis–based stabilization method, load-balanced distributed 2D-3D training strategies, and compute-adaptive tokenization. Using these tools, we develop Walrus, a transformer-based foundation model developed primarily for fluid-like continuum dynamics. Walrus is pretrained on nineteen diverse scenarios spanning astrophysics, geoscience, rheology, plasma physics, acoustics, and classical fluids. Experiments show that Walrus outperforms prior foundation models on both short- and long-term prediction horizons on downstream tasks and across the breadth of pretraining data, while ablation studies confirm the value of our contributions to forecast stability, training throughput, and transfer performance over conventional approaches.

Lay Summary

Foundation models have been slow to have the same impact on physical simulation in large part because the massive size and diversity of high quality data in this space make it difficult to train on sufficiently diverse datasets. A key problem is that any solution to modeling problems in 3D/4D space needs to scale computationally in ways modeling lower dimensional problems doesn't. In our paper, we tackle some of these problems with solutions derived both from theory and experiments. Walrus is fundamentally built to be adaptable. The internal resolution is controllable, it's trained on a much more diverse data corpus than prior foundation models, and new tricks introduced here improve the stability of forecasts. The end result is a foundation model that our experiments show outperforms previous models across multiple time horizons and on a wide variety of problems.