Multi-Modal Foundation · polymathic-ai
WALRUS
A 1.3-billion-parameter space–time Transformer surrogate for multi-domain continuum dynamics simulations.
What WALRUS is
WALRUS is a cross-domain physics foundation model developed by Polymathic AI. Designed to model the temporal evolution of continuum dynamical systems governed by partial differential equations (PDEs), it serves as a fast general-purpose surrogate for physical simulation and as a base initialization for fine-tuning on new physical systems.
Traditional numerical solvers require extensive computational resources and are typically tailored to specific domains, geometries, or boundary conditions. WALRUS addresses this by pre-training across 19 diverse physical datasets spanning astrophysics, geoscience, rheology, plasma physics, acoustics, and classical fluids, allowing it to generalise across distinct boundary configurations and physical parameterisations.
The model operates autoregressively, predicting state differences over time from a brief history of simulation snapshots. By standardising diverse physical problems into a single space–time architecture, it supports zero-shot surrogate rollouts and efficient transfer learning on downstream mechanics and fluid tasks.
How it works
WALRUS is built as a 1.3-billion-parameter space–time Transformer that predicts the next snapshot state by approximating the temporal increment between consecutive simulation frames. Its input consists of a historical window of snapshots normalized using root-mean-square (RMS) metrics across space and time, whilst predicted updates are de-normalized using the RMS of the change increments. To ingest varying dimensions and resolutions into unified representations, it uses adaptive-compute patch embeddings and embeds two-dimensional physical fields as thin three-dimensional volumes.
Training uses an AdamW optimizer targeting a per-field normalized L1 loss. Key training procedures include patch jittering—a harmonic-analysis–motivated augmentation introducing random offsets to reduce spectral aliasing—tensor-law–aware data augmentation that applies proper physical rotations to vector and tensor fields, system-uniform hierarchical batch sampling, and random time-striding of one to five steps. Pre-training was executed across 96 NVIDIA H100 GPUs using distributed Hybrid Sharded Data Parallelism (HSDP).
Key facts
- Model architecture
- Space–time Transformer
- Parameter count
- 1.3 billion
- Hardware used for training
- 96 NVIDIA H100 GPUs
- Pre-training domains
- 19 physical domains (including astrophysics, geoscience, rheology, plasma physics, acoustics, classical fluids)
- Training objective
- Per-field normalized L1 loss
What you can use it for
- Fast Surrogate Simulation
- Accelerate time-stepping simulations by replacing standard numerical differential equation integrators with autoregressive neural surrogate rollouts.
- Downstream Fine-Tuning for Specialized PDEs
- Use pre-trained weights to adapt to specific scenarios such as compressible Navier–Stokes turbulence, neutron star post-mergers, or bubble boiling simulations.
- Cross-Dimensional Physical Field Forecasting
- Forecast multi-field state trajectories across both 2D and 3D continuum dynamics using adaptive patch embeddings.
Run it here
We list this model but cannot run it yet: no pipeline_tag — needs a human.
Everything on this page still applies. Running it needs a loader for this repository, not a bigger machine.
Strengths and limits
Good at
- Pre-trained across 19 distinct continuum dynamics datasets with varied boundary conditions and parameters.
- Employs patch jittering to suppress spectral aliasing and increase long-horizon rollout stability across most training domains.
- Integrates tensor-law–aware data augmentations that enforce physically consistent vector and tensor field transformations.
- Unifies two-dimensional and three-dimensional simulation data by embedding 2D fields into thin 3D volumes.
Watch out for
- Native training infrastructure and code are tightly coupled with the schema from the Well, requiring schema matching or workarounds.
- Rollout accuracy on custom physical configurations outside its pre-training regimes may require fine-tuning rather than zero-shot inference.
- External non-Hugging-Face checkpoint storage endpoints for additional fine-tuned models can be unstable.
Inputs and outputs
You provide
Sequence or text
A protein or nucleotide sequence, or free text — whichever this model was trained on.
You get back
Mean-pooled representation, one vector per input.
Licence
Check the licence before commercial use · mit
The mit licence carries terms worth reading before you depend on it.
Questions
Can I use WALRUS commercially?
Check the licence before commercial use. The mit licence carries terms worth reading before you depend on it.
How does WALRUS unify 2D and 3D datasets?
The model uses an adaptive-compute patch embedding that adjusts token counts across resolutions, and it embeds two-dimensional field data as thin three-dimensional slices rotated in 3D space.
What is the formulation for predicting the next time step?
Given a history of recent snapshots, WALRUS predicts the state update increment, computing the next state by adding this predicted difference to the current state.
What downstream benchmarks has the model been evaluated on?
It has been fine-tuned on benchmarks including PDEBench (turbulent and random compressible Navier–Stokes), PDEGym, FlowBench, PDEArena (incompressible Navier–Stokes), BubbleML 2.0 pool boiling, and astrophysical simulations from the Well.
How to cite it
arXiv:2511.15684
Written from WALRUS's own documentation (model card). Statements citing a figure that does not appear there are removed before publication. Licence and running details come from our registry.