|
| 1 | +<!--  --> |
| 2 | +<!-- <img src="assets/walrus_logo.png" width="360"> --> |
| 3 | + |
| 4 | +# Walrus: A Cross-domain Foundation Model for Continuum Dynamics |
| 5 | +<div align="center"> |
| 6 | + |
| 7 | +[](https://opensource.org/licenses/MIT) |
| 8 | +[](https://pytorch.org/) |
| 9 | +[](https://arxiv.org/abs/2511.15684) |
| 10 | +[](https://huggingface.co/polymathic-ai/walrus) |
| 11 | + |
| 12 | +[Getting Started](#getting-started) • [Tutorials](#tutorials) • [Model Overview](#model-overview) |
| 13 | + |
| 14 | +</div> |
| 15 | + |
| 16 | +--- |
| 17 | + |
| 18 | +## Overview |
| 19 | +<div align="center"> |
| 20 | + <img src="assets/ArchitectureWIP.png" alt="Walrus schematic" width="600"> |
| 21 | +</div> |
| 22 | + |
| 23 | +This repo is built for training and evaluating Walrus, a multi-domain foundation model for continuum dynamics trained primarily on fluid-like behaviors. |
| 24 | +Walrus was trained on 19 different physical scenarios spanning 63 physical variables in both 2 and 3D. Walrus utilizes new tools for adaptive computation and improved stability |
| 25 | +in order to achieve accurate long-term rollouts while co-adapting sampling and distribution to improve training throughput despite handling varying dimensions, resolutions, and |
| 26 | +aspect ratios. |
| 27 | + |
| 28 | +## Getting Started |
| 29 | + |
| 30 | +### Installation |
| 31 | + |
| 32 | +Clone the repository and install locally. Requirements are documented in the pyproject.toml file. |
| 33 | + |
| 34 | +Most of the data used in experiments is from [the Well](https://github.com/PolymathicAI/the_well), |
| 35 | +so it may be easier to get started if you have the Well at least partially available. |
| 36 | + |
| 37 | +```bash |
| 38 | +git clone git@github.com:PolymathicAI/walrus.git |
| 39 | +cd walrus |
| 40 | +pip install . |
| 41 | +``` |
| 42 | +By default, this repository does not include all dependencies of non-Walrus models. To install optional dependencies, |
| 43 | +instead run: |
| 44 | + |
| 45 | +```bash |
| 46 | +pip install .[external_modes] |
| 47 | +``` |
| 48 | +## Running Experiments |
| 49 | + |
| 50 | +This project is orchestrated using [Hydra](https://github.com/facebookresearch/hydra) to give users modular access over |
| 51 | +various model components, datasets, and runtime options. All training was done in slurm environments. Example invocations |
| 52 | +for training, validating, and finetuning Walrus models can be found in [walrus/run_scripts](walrus/run_scripts). |
| 53 | + |
| 54 | +Most of these use relative paths and slurm. For example, one would launch the training script with: |
| 55 | +```bash |
| 56 | +cd /path/to/thisfolder/walrus/walrus |
| 57 | +sbatch run_scripts/pretrain_example_distributed_walrus.sh |
| 58 | +``` |
| 59 | +Local invocations also available for smaller training and validation runs. |
| 60 | + |
| 61 | +## Tutorials |
| 62 | + |
| 63 | +To help get you started, we include a set of demo notebooks. These are: |
| 64 | + |
| 65 | +- [Transforming data into Well Format](demo_notebooks/walrus_example_0_ConvertingDataIntoWellFormat.ipynb) - Since our repository largely uses structures from the Well, we include a guide |
| 66 | +to transforming data so it can be used easily within this repository. |
| 67 | +- [Running Walrus](demo_notebooks/walrus_example_1_RunningWalrus.ipynb) - This notebook walks through the basics of how to use Walrus both with Well data and with external data. |
| 68 | + |
| 69 | +If you already have a local copy of the Well, you can skip straight to example 1. However, if you do not, it's recommended to follow example 0 to get some data to use. We also include several example run scripts for training and validation, both using slurm and torchrun for distribution and just running single-GPU code. |
| 70 | + |
| 71 | +## Model Overview |
| 72 | + |
| 73 | +### Architecture |
| 74 | + |
| 75 | +Walrus uses and encoder-processor-decoder structure. Encoder/decoder are hMLPs/transposed hMLPs using stride modulation to dynamically adjust the |
| 76 | +internal resolution. The processor consists of blocks containing factorized space and time attention. |
| 77 | + |
| 78 | +### Patch Jittering |
| 79 | + |
| 80 | +Walrus suppressed the growth of long-run instabilities through the use of *patch jittering*. Patch jittering involves randomly translating the reference frame (with padding for boundaries) |
| 81 | +before each step. While the paper goes into more theoretical detail on why this works, the core idea is that the specific downsampling pattern leads to predictable accumulation |
| 82 | +of error and that randomizing this process can help alleviate this pathology. |
| 83 | + |
| 84 | +### Adaptive Compute |
| 85 | + |
| 86 | +To handle varying compute budgets and problem complexities, we also employ [stride modulation](https://arxiv.org/pdf/2507.09264) to allow users to adjust |
| 87 | +their downstream resolution. During pretraining, this was used to keep internal resolution fairly consistent (32/33 per dim in 2D, 16/17 in 3D). In this approach, the downsampling |
| 88 | +layers of the encoder/decoder will dynamically adjust their stride based on a target internal resolution. |
| 89 | + |
| 90 | +### Efficient Training |
| 91 | + |
| 92 | +Heterogeneous data can easily lead to training bottlenecks as many distribution primitives require regular syncing. Our approach |
| 93 | +adjusts downsampling to ensure consistent token counts, but also adjust context length and batch size to minimize gaps. This isn't |
| 94 | +enough to completely eliminate discrepencies, so we must also adjust our sampling strategy to be aware of where these dead cycles can emerge. |
| 95 | +We implement sampling such that when training with HSDP, all nodes within a sharding group are forced to sample the same data source. This |
| 96 | +balances batch diversity and efficiency. |
| 97 | + |
| 98 | +## License |
| 99 | +This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. |
| 100 | + |
| 101 | +## Contact |
| 102 | + |
| 103 | +- **Issues**: [GitHub Issues](https://github.com/PolymathicAI/AION/issues) |
| 104 | + |
| 105 | +For other queries, please reach out to the corresponding author: mmccabe@flatironinstitute.org. |
| 106 | + |
| 107 | +## Acknowledgements |
| 108 | + |
| 109 | +Walrus is built by [Polymathic AI](https://polymathic-ai.org/) as part of our mission of advancing the frontier of AI for scientific application. Polymathic AI gratefully acknowledges funding from the Simons Foundation and Schmidt Sciences, LLC. This work was performed with compute from the Scientific Computing Core, a |
| 110 | +division of the Flatiron Institute, a division of the Simons Foundation and from the National AI Research Resource Pilot, including support from NVIDIA |
| 111 | +and NVIDIA’s DGX Cloud product which includes the NVIDIA AI Enterprise Software Platform. |
| 112 | + |
| 113 | +## Citing Walrus |
| 114 | + |
| 115 | +Please use the following citation for Walrus: |
| 116 | + |
| 117 | +``` |
| 118 | +@misc{mccabe2025walruscrossdomainfoundationmodel, |
| 119 | + title={Walrus: A Cross-Domain Foundation Model for Continuum Dynamics}, |
| 120 | + author={Michael McCabe and Payel Mukhopadhyay and Tanya Marwah and Bruno Regaldo-Saint Blancard and Francois Rozet and Cristiana Diaconu and Lucas Meyer and Kaze W. K. Wong and Hadi Sotoudeh and Alberto Bietti and Irina Espejo and Rio Fear and Siavash Golkar and Tom Hehir and Keiya Hirashima and Geraud Krawezik and Francois Lanusse and Rudy Morel and Ruben Ohana and Liam Parker and Mariel Pettee and Jeff Shen and Kyunghyun Cho and Miles Cranmer and Shirley Ho}, |
| 121 | + year={2025}, |
| 122 | + eprint={2511.15684}, |
| 123 | + archivePrefix={arXiv}, |
| 124 | + primaryClass={cs.LG}, |
| 125 | + url={https://arxiv.org/abs/2511.15684}, |
| 126 | +} |
| 127 | +``` |
| 128 | + |
0 commit comments