README.md

---
license: mit
tags:
- particle-physics
- lhc
- atlas-open-data
- flow-matching
- riemannian-flow-matching
- event-generation
- arxiv:2607.16144
---

# ShellFlow — pretrained models

Pretrained checkpoints for the paper **[Learning Standard Model structure from LHC data with Riemannian flow matching](https://arxiv.org/abs/2607.16144)** (arXiv:2607.16144).

ShellFlow is a transformer generative model of LHC collision events, trained with Riemannian flow matching on a product manifold (R × S², plus a log-mass coordinate for massive objects) using ATLAS Open Data. Code: https://github.com/pomidori/ShellFlow

## Models

All models share the same backbone (d_model 128, 6 layers, 4 heads, ~6.0M parameters) and were trained on the union of the ATLAS Open Data `1LMET30` and `2to4lep` skims (events with ≤ 8 reconstructed particles, de-duplicated by run/event number).

| File | Model | Chart | Aux head | Training |
|---|---|---|---|---|
| `union_1LMET30_2to4lep/model.ckpt` | Main paper model | on-shell (manifold) | yes | 30 epochs, 717k steps |
| `e1_v0_manifold_aux/model.ckpt` | Ablation E1-V0 | on-shell (manifold) | yes | 2 epochs, 48k steps |
| `e1_v1_euclidean_aux/model.ckpt` | Ablation E1-V1 | free (Euclidean) | yes | 2 epochs, 48k steps |
| `e1_v2_manifold_noaux/model.ckpt` | Ablation E1-V2 | on-shell (manifold) | no | 2 epochs, 48k steps |
| `e1_v3_euclidean_noaux/model.ckpt` | Ablation E1-V3 | free (Euclidean) | no | 2 epochs, 48k steps |

Checkpoint format: PyTorch Lightning (weights + EMA weights + hyperparameters; optimizer state stripped). The hyperparameters carry `architecture: legacy_type_film`, which selects the paper architecture in the ShellFlow code. Distribution matching (a training-time reweighting) is disabled in the stored hyperparameters so the checkpoints load without any dataset present; it has no effect on generation.

## Usage

Install the [ShellFlow code](https://github.com/pomidori/ShellFlow), then:

```python
from huggingface_hub import hf_hub_download
from shellflow.training.riemannian_DiTs_lightning import RiemannianDiTsModule

path = hf_hub_download("pomidori73/ShellFlow", "union_1LMET30_2to4lep/model.ckpt")
module = RiemannianDiTsModule.load_from_checkpoint(path, map_location="cpu").eval()
model = module.ema_model  # paper results use the EMA weights
```

Generation is conditioned on the event composition (particle types, charges, detector extras, MET). The simplest generation harness mirrors the training-time monitoring callback:

```python
import torch
from shellflow.training.riemannian_DiTs_callback import RiemannianDiTsSamplePlotCallback

cb = RiemannianDiTsSamplePlotCallback(val_dataset=None, collator=None, n_gen_steps=200)
with torch.no_grad():
    kinematics = cb._generate_batch_rfm(model, batch, module)  # (B, N, 4) = (px, py, pz, E) in GeV
```

where `batch` is a dict with `particle_type` (B, N) int64 codes {0=pad, 1=jet, 10=largeRjet, 11=e, 13=mu, 15=tau, 22=photon}, `particle_charge` (B, N) int64 in {0, 1, 2}, `particle_extra` (B, N, 13) float, `particle_mask` (B, N) 0/1, and optionally `met_vector` (B, 2). For physical samples, draw the conditioning from real events (see the dataset pipeline in the code repository).

## Citation

```bibtex
@article{shellflow2026,
  title   = {Learning Standard Model structure from LHC data with Riemannian flow matching},
  author  = {Kato, Midori and Urqu{\'i}a-Calder{\'o}n, Kevin A. and Timiryasov, Inar and Ruchayskiy, Oleg},
  journal = {arXiv preprint arXiv:2607.16144},
  year    = {2026},
  url     = {https://arxiv.org/abs/2607.16144}
}
```