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