|
|
--- |
|
|
tags: |
|
|
- walrus |
|
|
- foundation-model |
|
|
- physics |
|
|
- continuum-dynamics |
|
|
- transformer |
|
|
- PDE |
|
|
datasets: |
|
|
- polymathic-ai/shear_flow |
|
|
- polymathic-ai/gray_scott_reaction_diffusion |
|
|
- polymathic-ai/active_matter |
|
|
- polymathic-ai/turbulent_radiative_layer_2D |
|
|
- polymathic-ai/supernova_explosion_64 |
|
|
- polymathic-ai/turbulence_gravity_cooling |
|
|
- polymathic-ai/rayleigh_benard |
|
|
- polymathic-ai/planetswe |
|
|
- polymathic-ai/acoustic_scattering_inclusions |
|
|
- polymathic-ai/MHD_64 |
|
|
- polymathic-ai/rayleigh_taylor_instability |
|
|
- polymathic-ai/acoustic_scattering_discontinuous |
|
|
- polymathic-ai/acoustic_scattering_maze |
|
|
- polymathic-ai/helmholtz_staircase |
|
|
- polymathic-ai/viscoelastic_instability |
|
|
- BGLab/FlowBench |
|
|
license: mit |
|
|
--- |
|
|
|
|
|
# Walrus: A Cross-Domain Foundation Model for Continuum Dynamics |
|
|
|
|
|
[](https://opensource.org/licenses/MIT) |
|
|
[](https://github.com/PolymathicAI/walrus) |
|
|
[](https://arxiv.org/abs/2511.15684) |
|
|
|
|
|
Walrus is a large-scale **physics foundation model** capable of modeling a broad range of continuum dynamical systems. |
|
|
|
|
|
Walrus is trained jointly across **19 diverse physical domains** spanning: |
|
|
- astrophysics |
|
|
- geoscience |
|
|
- rheology |
|
|
- plasma physics |
|
|
- acoustics |
|
|
- classical fluids |
|
|
|
|
|
These systems have diverse boundary conditions and physical parameterizations. The model is optimized to serve as a **general-purpose surrogate** for physical simulation and a **strong initialization** for downstream fine-tuning on new PDE systems. |
|
|
|
|
|
--- |
|
|
|
|
|
# Model Description |
|
|
|
|
|
Walrus is a **1.3B-parameter space–time Transformer** trained autoregressively to predict the temporal evolution of physical fields. Walrus is trained to model the evolution of physical systems in space and time. A simulation snapshot at time t is written as u(t). |
|
|
|
|
|
We define the difference between two consecutive snapshots as: |
|
|
Δu(t+1) = u(t+1) − u(t) |
|
|
|
|
|
Given a short history of snapshots: |
|
|
U(t) = [u(t − τ + 1), ..., u(t)] |
|
|
|
|
|
The model predicts the next state using: |
|
|
u(t+1) ≈ u(t) + M(U(t)) |
|
|
|
|
|
### Key architectural components |
|
|
|
|
|
- **Adaptive-compute patch embedding** |
|
|
- Token count automatically balanced across resolutions |
|
|
- Enables mixing 2D and 3D datasets efficiently |
|
|
|
|
|
- **Patch Jittering** |
|
|
- A harmonic-analysis–motivated augmentation technique |
|
|
- Reduces aliasing and spectral artifacts |
|
|
- Improves long-horizon stability across 17/19 pretraining datasets |
|
|
|
|
|
- **Tensor-law–aware data augmentation** |
|
|
- 2D data embedded into 3D through plane rotations |
|
|
- Vector/tensor fields rotated with correct physical transformations |
|
|
|
|
|
- **Asymmetric normalization** |
|
|
- **Asymmetric normalization:** Walrus normalizes inputs by RMS over space-time and de-normalizes the predicted Δu using the RMS of Δ. |
|
|
|
|
|
--- |
|
|
|
|
|
# Pretraining Details |
|
|
|
|
|
Walrus is pretrained 19 physical datasets with: |
|
|
|
|
|
- **Loss**: Per-field normalized L1 loss |
|
|
- **Optimizer**: AdamW |
|
|
- **Batching**: System-uniform hierarchical sampling |
|
|
- **Time-striding**: Random stride (1–5) per training example |
|
|
- **Patch jitter range**: Uniform per-axis random offset |
|
|
- **Dimensional unification**: 2D fields embedded as thin 3D volumes |
|
|
|
|
|
The model was pretrained on 96 **NVIDIA H100 GPUs** using distributed HSDP (4 GPU per shard group) with sampling matching distribution structure for minimal deadweight loss. |
|
|
|
|
|
--- |
|
|
|
|
|
# Intended Use |
|
|
|
|
|
This pretrained checkpoint is suitable for: |
|
|
|
|
|
### ✔ Next-step prediction |
|
|
### ✔ Fast surrogate simulation |
|
|
### ✔ Autoregressive rollout of physical systems |
|
|
### ✔ Transfer learning to new physical settings |
|
|
|
|
|
# Resources |
|
|
|
|
|
Paper: https://arxiv.org/pdf/2511.15684 |
|
|
Github: https://github.com/PolymathicAI/walrus |
|
|
Tutorial: https://github.com/PolymathicAI/walrus/demo_notebooks |
|
|
|
|
|
Note, the training code in the repository is closely coupled with tools from [the Well](https://github.com/PolymathicAI/the_well), so |
|
|
it can be beneficial to format data to match that schema. If that's not possible, the tutorial does show how one would use the model |
|
|
without Well-formatted data. |
|
|
|
|
|
|
|
|
# Demonstrated downstream tasks |
|
|
|
|
|
We show the strong performance of Walrus by finetuning on a range of challenging downstream tasks as shown in the paper. |
|
|
Paths to access the finetuned walrus checkpoints for various downstream tasks is as follows: |
|
|
|
|
|
### PDEGym CE-RM: https://huggingface.co/polymathic-ai/walrus_ft_CE-RM/tree/main |
|
|
### PDEBench CNS Turbulent: https://huggingface.co/polymathic-ai/walrus_ft_CNS3D_64_Turb/tree/main |
|
|
### PDEBench CNS Random: https://huggingface.co/polymathic-ai/walrus_ft_CNS3D_128_Rand/tree/main |
|
|
### Flowbench FPOSkelenton: https://huggingface.co/polymathic-ai/walrus_ft_flowbench_skelenton/tree/main |
|
|
### The Well Postmerger Neutron Star: https://huggingface.co/polymathic-ai/walrus_ft_post_neutron_star_merger/tree/main |
|
|
### The Well Convective envelope RSG: https://huggingface.co/polymathic-ai/walrus_ft_convective_envelope_rsg/tree/main |
|
|
### PDEArena Conditioned Incompressible NS: https://huggingface.co/polymathic-ai/walrus_ft_pdearena_ins/tree/main |
|
|
### BubbleML 2.0 PoolBoil Subcooled: https://huggingface.co/polymathic-ai/walrus_ft_bubbleML_poolboil/tree/main |
|
|
|
|
|
|
|
|
Additional checkpoints not included in the Walrus collection on HF can be found [here](https://users.flatironinstitute.org/~polymathic/data/walrus_project_checkpoints/) though the endpoint is a bit finicky. |
|
|
|
|
|
More finetuning checkpoints will continue to be added to HF over time. |
