HEP-JEPA: A foundation model for collider physics using joint embedding predictive architecture
This work addresses classification tasks in high-energy physics, such as jet analysis at colliders, but is incremental as it applies an existing self-supervised method to a new domain.
The authors tackled the problem of classifying jets in high-energy particle colliders by developing a transformer-based foundation model using a self-supervised Joint Embedding Predictive Architecture, pre-trained on 100M jets and showing competitive performance on tasks like top tagging and quark-gluon differentiation compared to state-of-the-art models.
We present a transformer architecture-based foundation model for tasks at high-energy particle colliders such as the Large Hadron Collider. We train the model to classify jets using a self-supervised strategy inspired by the Joint Embedding Predictive Architecture. We use the JetClass dataset containing 100M jets of various known particles to pre-train the model with a data-centric approach -- the model uses a fraction of the jet constituents as the context to predict the embeddings of the unseen target constituents. Our pre-trained model fares well with other datasets for standard classification benchmark tasks. We test our model on two additional downstream tasks: top tagging and differentiating light-quark jets from gluon jets. We also evaluate our model with task-specific metrics and baselines and compare it with state-of-the-art models in high-energy physics. Project site: https://hep-jepa.github.io/