Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Walrus: A Cross-Domain Foundation Model for Continuum Dynamics

Michael McCabe, Payel Mukhopadhyay, Tanya Marwah, Bruno Regaldo-Saint Blancard, Francois Rozet, Cristiana Diaconu, Lucas Meyer, Kaze W. K. 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

TL;DR

Walrus introduces a cross-domain foundation model for continuum dynamics by extending transformer architectures with stability and efficiency mechanisms tailored to heterogeneous 2D/3D physics data. Key contributions include patch jittering, 2D–3D data augmentation, adaptive-compute tokenization, and topology-aware sampling, enabling stable long-horizon forecasting across 19 physical scenarios. Empirical results show Walrus achieving substantial improvements over prior foundation models on both short- and long-horizon tasks and across diverse domains, with 3D performance highlighting practical viability for real-world simulations. The work emphasizes the importance of representational diversity in pretraining and proposes scalable training strategies to handle multi-resolution, multi-physics data, marking a substantive step toward robust, cross-domain physics emulation.

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 and 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. Code and weights are released for community use.

Walrus: A Cross-Domain Foundation Model for Continuum Dynamics

TL;DR

Walrus introduces a cross-domain foundation model for continuum dynamics by extending transformer architectures with stability and efficiency mechanisms tailored to heterogeneous 2D/3D physics data. Key contributions include patch jittering, 2D–3D data augmentation, adaptive-compute tokenization, and topology-aware sampling, enabling stable long-horizon forecasting across 19 physical scenarios. Empirical results show Walrus achieving substantial improvements over prior foundation models on both short- and long-horizon tasks and across diverse domains, with 3D performance highlighting practical viability for real-world simulations. The work emphasizes the importance of representational diversity in pretraining and proposes scalable training strategies to handle multi-resolution, multi-physics data, marking a substantive step toward robust, cross-domain physics emulation.

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 and 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. Code and weights are released for community use.

Paper Structure

This paper contains 42 sections, 13 equations, 43 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Related Work
  4. Walrus Model
  5. Architecture
  6. Patch Jittering
  7. Walrus Training
  8. Augmenting into a Combined 2D-3D Space
  9. Efficient Multi-Task Training
  10. Experiments
  11. Downstream Performance
  12. Cross-domain Analysis
  13. Impact of Pretraining Strategy
  14. Conclusion
  15. Model Details
  16. ...and 27 more sections

Figures (43)

  • Figure 1: Walrus is a modern transformer incorporating novel stabilizing techniques and recent adaptive-compute methods to learn from a highly diverse set of physical dynamics. Walrus takes as input a short sequence of snapshots and predicts the next step in the sequence.
  • Figure 2: Patch jittering (middle) reduces the accumulation of high frequency artifacts (top) allowing for more stable long-term forecasts.
  • Figure 3: All raw data is projected into 3D by appending singleton dimensions and zero-padding tensor-valued fields prior to applying symmetry-preserving augmentations resulting in the originally final input data being embedded in arbitrary axis-aligned directions in 3D
  • Figure 4: Tokenization and distribution strategies are carefully balanced to ensure each copy of Walrus receives similarly sized buckets of tokens within each synchronous block, minimizing deadweight and maximizing throughput.
  • Figure 5: Visualizing the prediction of a finetuned Walrus for 3D research-frontier level simulations.
  • ...and 38 more figures