Walrus: A Cross-Domain Foundation Model for Continuum Dynamics

AI summary

Overview

• Research area: Machine learning for physical simulation; specifically, foundation models for continuum dynamics (fluids, plasmas, and related fields governed by partial differential equations).
• Technical level: Advanced. The paper assumes familiarity with transformer architectures, tokenization, distributed training, and numerical simulation of physical systems.
• Scope: The authors build and evaluate Walrus, a transformer-based foundation model trained jointly on nineteen heterogeneous physical simulation datasets, demonstrating improved short- and long-horizon forecasting over prior foundation models.

What This Paper Is About

Foundation models have reshaped language and vision, but the same recipe has not translated cleanly to physical simulation. The core obstacles are that simulation data is wildly heterogeneous (different physics, units, resolutions, and dimensionalities), and that autoregressive rollouts of learned models tend to drift or blow up over long horizons. This paper tackles both problems, aiming to produce a single pretrained model that can be fine-tuned or adapted across many fluid-like continuum domains rather than training separate models per problem.

Key Contributions

1. A harmonic-analysis-based stabilization method that improves the numerical stability of long-term autoregressive forecasts, addressing the drift and divergence that plague learned simulators.
2. Load-balanced distributed 2D and 3D training strategies that let the model train efficiently across mixed-dimensional and mixed-resolution data on modern hardware without stalling on poorly balanced work.
3. Compute-adaptive tokenization, a scheme that adjusts how physical fields are tokenized based on available compute, enabling efficient training across varying resolutions and dimensionalities.
4. Walrus itself, a transformer foundation model pretrained on nineteen diverse scenarios spanning astrophysics, geoscience, rheology, plasma physics, acoustics, and classical fluids, with code and weights released publicly.

Main Findings

• Superior forecasting across horizons: Walrus outperforms prior foundation models on both short-term and long-term prediction tasks on downstream problems, not just on the domains it saw most often.
• Breadth of transfer: Strong performance holds across the full span of pretraining data, indicating the model generalizes across physically distinct regimes rather than specializing in one.
• Stability gains confirmed: Ablations show the harmonic-analysis stabilization materially improves forecast stability compared with conventional approaches.
• Throughput gains confirmed: The load-balanced distributed training strategies are shown to improve training efficiency, supporting the scale of the pretraining run.
• Transfer gains confirmed: Ablation studies attribute measurable improvements in downstream transfer performance to the paper's proposed techniques rather than to scale alone.

Methodology in Plain English

The authors first diagnose why foundation models struggle in physics: data from different simulation codes disagree on units, grids, and dimensionality, and small per-step errors compound into catastrophic divergence during long rollouts. Rather than accept these as fixed constraints, they design targeted fixes. They use ideas from harmonic analysis, essentially decomposing fields into frequency components and controlling how the model handles them, to keep long rollouts from drifting. To make training practical, they distribute work across many GPUs in a way that balances 2D and 3D data so no device sits idle, and they let the tokenization scheme adapt to the compute budget, so higher-resolution data does not force proportionally more expensive training. With these pieces in place, they assemble a large corpus of nineteen simulation scenarios covering very different physics, train a single transformer on all of it, and then test how well it predicts held-out and downstream dynamics. Ablation experiments isolate how much each design choice contributes.

Why This Matters

This work is an argument that physical simulation can have its own foundation model moment, and it supplies concrete engineering to get there. It shifts the field from bespoke, single-domain neural surrogates toward a shared pretrained backbone that practitioners adapt to their own problem, which could substantially lower the cost of building fast simulators.

Real-world applications include:

• Climate and weather modeling: Fast surrogate models for atmospheric and oceanic dynamics, where long-horizon stability is the binding constraint.
• Fusion energy: Predicting plasma behavior in tokamaks, where simulations are expensive and instabilities are the central concern.
• Astrophysics: Modeling turbulence, accretion, and cosmological fluid structure at scales where traditional simulation is prohibitively costly.
• Engineering and manufacturing: Aerodynamic design, acoustics, and rheology for materials processing, where rapid iteration on fluid behavior drives product decisions.

Industry relevance spans any sector that relies on computational fluid dynamics or plasma simulation: aerospace, energy, automotive, materials, and semiconductor manufacturing. A single adaptable model that reduces reliance on domain-specific retraining has direct cost and time-to-answer implications for these groups.

Future Directions

• Beyond fluid-like dynamics: Extending the approach to solid mechanics, elastodynamics, and other continuum regimes not covered by the nineteen pretraining scenarios.
• Scaling laws for physics: Determining whether the language-model-style relationship between compute, data, and performance holds for physical simulation, and where it breaks.
• Uncertainty and conservation guarantees: Quantifying predictive uncertainty and enforcing physical invariants such as conservation laws inside the model rather than only in the tokenization.
• Fusion with classical solvers: Hybrid schemes where the learned model accelerates or corrects a numerical solver, rather than replacing it wholesale.

Target Audience

Machine learning researchers working on scientific computing and neural surrogates, computational physicists and engineers who need fast simulation tools, and practitioners in climate, fusion, astrophysics, or fluid engineering evaluating whether foundation models are mature enough for their domain. Readers without a background in transformers or PDE-constrained simulation will find the paper demanding, but the conceptual framing of data heterogeneity and rollout stability is accessible to a broader technical audience.

Authors’ 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.