README.md

---
license: lgpl
library_name: pytorch
tags:
- physics
- particle-physics
- nuclear-physics
- foundation-model
- mamba
- state-space-model
- track-reconstruction
- arxiv:2508.14087
---

# FM4NPP: Foundation Models for Nuclear and Particle Physics

Pretrained state space model (Mamba / Mamba2) checkpoints for particle physics,
from the paper **[Foundation Models for Particle Physics](https://arxiv.org/abs/2508.14087)**
(accepted to ICLR 2026). Developed at Brookhaven National Laboratory.

- 📄 Paper: https://arxiv.org/abs/2508.14087 · [OpenReview](https://openreview.net/forum?id=qaI3cLFsiX)
- 💻 Code: https://github.com/FM4NPP/PP_collision
- 📊 Dataset (TPCpp-10M): https://doi.org/10.5281/zenodo.16970029

## Model description

FM4NPP pretrains state space models (Mamba, Mamba2) on simulated proton–proton
collision data from a Time Projection Chamber (TPC), then transfers the learned
representations to downstream tasks such as track reconstruction (track finding).

Two model sizes are provided:

**Mamba 5M**

| Parameter | Value |
|---|---|
| `embed_dim` | 256 |
| `num_layers` | 12 |
| `d_state` | 16 |
| `d_conv` | 4 |
| `expand` | 2 |

**Mamba2 5M**

| Parameter | Value |
|---|---|
| `embed_dim` | 256 |
| `num_layers` | 12 |
| `d_state` | 128 |
| `headdim` | 64 |
| `ngroups` | 1 |

## Intended use

These checkpoints are intended for research on foundation models in nuclear and
particle physics: pretraining on unlabeled TPC data and fine-tuning for downstream
tasks like track reconstruction, particle classification, and noise tagging.

## How to use

Full training and inference code is in the
[GitHub repository](https://github.com/FM4NPP/PP_collision).

```bash
# Environment
conda create -n fm4npp python=3.10
conda activate fm4npp
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install mamba-ssm causal-conv1d triton
pip install pyyaml numpy scipy tqdm mmap-ninja

# Download the checkpoint from this repo
huggingface-cli download <your-username>/<your-repo> --local-dir ./checkpoints

# Fine-tune for track reconstruction (point pretrained_ckpt at the downloaded file)
python train/downstream/track_finding_trainer.py \
    --yaml_config=scripts/configs/mamba_tracking.yaml \
    --config=mamba_5m_downstream \
    --run_num=run0
```

## Training data

Pretrained on **TPCpp-10M**: 10M simulated proton–proton collision events (100 files,
~118.5 GB compressed, NumPy `.npz`) for pretraining, plus 70k/13k/7k labeled
train/val/test events for downstream tasks. Each spacepoint carries 30D features
(position, momentum, energy, time, detector metadata).

Dataset: https://doi.org/10.5281/zenodo.16970029 ·
Dataset paper: https://www.sciencedirect.com/science/article/pii/S2352340925011060

## Citation

```bibtex
@article{park2025fm4npp,
  title={FM4NPP: A Scaling Foundation Model for Nuclear and Particle Physics},
  author={Park, David and Li, Shuhang and Huang, Yi and Luo, Xihaier and Yu, Haiwang and Go, Yeonju and Pinkenburg, Christopher and Lin, Yuewei and Yoo, Shinjae and Osborn, Joseph and others},
  journal={arXiv preprint arXiv:2508.14087},
  year={2025}
}

@article{tpcpp10m2025,
  title={TPCpp-10M: Simulated proton-proton collisions in a Time Projection Chamber for AI Foundation Models},
  author={Li, Shuhang and Huang, Yi and Park, David and Luo, Xihaier and Yu, Haiwang and Go, Yeonju and Pinkenburg, Christopher and Lin, Yuewei and Yoo, Shinjae and Osborn, Joseph and Roland, Christof and Huang, Jin and Ren, Yihui},
  journal={arXiv preprint arXiv:2509.05792},
  year={2025}
}
```