FM4NPP/PP_collision

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 (accepted to ICLR 2026). Developed at Brookhaven National Laboratory.

• 📄 Paper: https://arxiv.org/abs/2508.14087 · OpenReview
• 💻 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.

# 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

@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}
}