Skip to content

pedrocurvo/mspt

Repository files navigation

MSPT

MSPT: Efficient Large-Scale Physical Modeling via Parallelized Multi-Scale Attention

MSPT Parallelized Multi-Scale Attention overview

arXiv AlphaXiv HF license submodule-enabled

MSPT is a scalable neural PDE solver based on Parallelized Multi-Scale Attention (PMSA). It combines local patch attention with global supernode communication, enabling strong accuracy and efficient scaling on PDE and CFD benchmarks.

Abstract

A key scalability challenge in neural solvers for industrial-scale physics simulations is efficiently capturing both fine-grained local interactions and long-range global dependencies across millions of spatial elements. We introduce the Multi-Scale Patch Transformer (MSPT), an architecture that combines local point attention within patches with global attention to coarse patch-level representations. To partition the input domain into spatially coherent patches, we employ ball trees, which handle irregular geometries efficiently. This dual-scale design enables MSPT to scale to millions of points on a single GPU. We validate MSPT on standard PDE benchmarks (elasticity, plasticity, fluid dynamics, porous flow) and large-scale aerodynamic datasets (ShapeNet-Car, Ahmed-ML), achieving state-of-the-art accuracy with substantially lower memory footprint and computational cost.

Links

  • Paper: https://arxiv.org/abs/2512.01738
  • AlphaXiv: https://www.alphaxiv.org/abs/2512.01738
  • Project website: ./docs (GitHub Pages content)
  • Checkpoints collection: https://huggingface.co/collections/pedrocurvo/mspt

Repository Layout

  • Neural-Solver-Library-MSPT/: codebase for training/evaluation (models, scripts, configs)
  • docs/: project website source (GitHub Pages)

Models

  • MSPT.py: combined implementation containing both structured and unstructured MSPT variants.
  • MSPTStr.py: structured MSPT variant (separate file).
  • MSPTUn.py: unstructured MSPT variant (separate file).

Warning

MSPT.py is the reference file used in Neural-Solver-Library-MSPT, and it contains both structured and unstructured models in one place.
MSPTStr.py and MSPTUn.py are the same model family split into separate files, with minor implementation variations.

Setup

This repository uses a git submodule.

git clone --recurse-submodules https://github.com/pedrocurvo/mspt.git
cd mspt

If already cloned without submodules:

git submodule update --init --recursive

Install dependencies:

cd Neural-Solver-Library-MSPT
pip install -r requirements.txt

Reproducing Experiments

Run benchmark scripts from Neural-Solver-Library-MSPT.

Example (MSPT on Plasticity):

bash ./scripts/StandardBench/plasticity/MSPT.sh

Additional examples:

bash ./scripts/StandardBench/plasticity/Transolver.sh
bash ./scripts/StandardBench/plasticity/Galerkin_Transformer.sh
bash ./scripts/StandardBench/pipe/U_NO.sh

Available StandardBench datasets:

  • airfoil
  • darcy
  • elasticity
  • ns
  • pipe
  • plasticity

Notes

  • Main training and evaluation pipeline lives in Neural-Solver-Library-MSPT/.
  • The website in docs/ is aligned with project figures and benchmark results.

Citation

@misc{curvo2025msptefficientlargescalephysical,
      title={MSPT: Efficient Large-Scale Physical Modeling via Parallelized Multi-Scale Attention},
      author={Pedro M. P. Curvo and Jan-Willem van de Meent and Maksim Zhdanov},
      year={2025},
      eprint={2512.01738},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2512.01738},
}