HEP-JEPA: A foundation model for collider physics

Jai Bardhan, Radhikesh Agrawal, Abhiram Tilak, Cyrin Neeraj, Subhadip Mitra

Keywords: self supervised learning, particle physics, high energy physics, jepa, foundational model

Abstract: 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. Moreover, this work lays the foundation for future endeavours in world models for HEP, where unified representations of heterogeneous collider data could revolutionize our approach to uncovering new physics.

Submission Number: 56