ShellFlow — pretrained models

Pretrained checkpoints for the paper Learning Standard Model structure from LHC data with Riemannian flow matching (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, then:

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:

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

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