⚛️ high-energy experiments

SuperJEPA: Pretraining with a Supervision-Augmented Joint-Embedding Predictive Architecture for High Energy Physics

The paper introduces SuperJEPA, a joint-embedding predictive architecture augmented with supervised objectives, which demonstrates improved downstream jet-tagging performance across various model sizes and dataset scales, even in high energy physics contexts where large labeled simulation datasets are already available.

Original authors: Dylan Sheldon Rankin, Ho Fung Tsoi

Published 2026-10-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dylan Sheldon Rankin, Ho Fung Tsoi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the high-energy physics of particle colliders, scientists do not watch particles collide in real time; they watch the aftermath. When protons smash together at nearly the speed of light, they shatter into sprays of new particles called jets. These jets are the fingerprints of the original collision, carrying clues about the fundamental forces and particles that created them. To understand what happened in a collision, researchers must sort these jets into categories, a process known as jet tagging. For decades, the standard way to teach computers to do this sorting has been to show them millions of simulated examples where the answer is already known. Because physicists can generate these labeled simulations with perfect accuracy, they have traditionally believed that the most powerful way to train an artificial intelligence is simply to feed it as many of these labeled examples as possible, letting the computer learn the rules by rote.

However, a new approach challenges the idea that labeled data is the only path to intelligence. In a recent study, researchers at the University of Pennsylvania explored whether a different kind of learning—one where the computer tries to predict missing parts of a picture based on what it can see—could improve performance even when plenty of labeled data exists. They developed a method called SuperJEPA, which combines the traditional labeled training with a predictive task. The researchers found that this hybrid approach allowed their models to learn better representations of the particle jets, leading to more accurate sorting in various tests. The results suggest that even in a field where perfect labels are available, asking a computer to guess the missing pieces of a puzzle can help it understand the whole picture more deeply.

The researchers began by taking a massive dataset of simulated particle jets, containing millions of examples from ten different types of particle interactions. They built a computer model designed to look at a jet and identify what kind of particle created it. To test their new method, they split the training process into two parts. First, they used a technique where the model was shown a jet with a significant portion of its particles hidden away. The model had to look at the visible particles and predict the properties of the hidden ones. This forced the model to learn the underlying structure of the jet, understanding how the particles relate to one another, rather than just memorizing the final answer. Second, they combined this predictive task with the traditional method of showing the model the full jet and telling it the correct label.

They tested this combined approach, which they named SuperJEPA, against a standard model trained only on labeled data. They ran these tests across different sizes of computer models and different amounts of training data, ranging from a few hundred thousand jets to twenty million. In every scenario, the SuperJEPA model outperformed the standard model. The improvement was most noticeable when the model was asked to identify types of particles it had not seen during its initial training. For instance, when the model was trained on a specific set of four particle types and then asked to identify ten different types, the SuperJEPA model was significantly better at recognizing the new, unseen categories. This suggests that the predictive task helped the model learn a more flexible understanding of particle jets, one that could adapt to new situations more easily than simple memorization.

The study also looked at how the model performed when its internal settings were locked in place versus when they were allowed to change during the final testing phase. In both cases, the model trained with the predictive method performed better. The researchers observed that the model learned to organize the information about the particles in a way that made the different types of jets easier to distinguish. While the standard model tended to produce similar-looking internal representations for different jets, the SuperJEPA model created distinct patterns for each type, even when the jets looked very similar on the surface. This ability to find subtle differences proved crucial when the models were tested on datasets with different energy levels or particle counts than the ones they were trained on.

One of the most interesting findings was how the amount of labeled data influenced the best way to train the model. When the researchers tested the model on a dataset that was very different from the training data, they found that a balance between the predictive task and the labeled task worked best. However, when the training data and the test data were very similar, using a stronger emphasis on the labeled data yielded the best results. This indicates that the ideal training strategy depends on how much the model needs to generalize to new situations. The researchers also noted that their method worked well even with smaller amounts of training data, suggesting that the predictive task helps the model learn more efficiently from fewer examples.

The work does not claim to have solved all the problems in particle physics, nor does it suggest that labeled data is no longer useful. Instead, it offers a refined tool for a field that already has powerful methods. The study shows that adding a layer of predictive reasoning to the training process can squeeze out extra performance, making the models more robust and versatile. This is particularly important as experiments at particle colliders continue to produce more data than ever before. By improving how computers learn to sort these complex sprays of particles, scientists can hope to uncover rarer and more subtle signals of new physics hidden within the noise of high-energy collisions. The findings provide a clear path forward for developing the next generation of artificial intelligence tools in high-energy physics, proving that even in a world of perfect labels, the act of prediction remains a powerful teacher.

Technical Summary: SuperJEPA for High Energy Physics

Problem Statement
Foundation models in natural language and computer vision have demonstrated that self-supervised learning (SSL) objectives, such as next-token prediction, can yield powerful representations from vast unlabeled datasets. However, High Energy Physics (HEP), particularly collider physics, differs significantly: it possesses large quantities of high-fidelity, labeled simulation data. Consequently, the field has historically relied on supervised pretraining for tasks like jet tagging. Previous research in HEP has struggled to demonstrate significant performance gains from SSL when large, accurate labeled datasets are available. The central question addressed by this work is whether self-supervised objectives can improve upon supervised pretraining alone in an environment where labeled data is already abundant.

Methodology: SuperJEPA
The authors introduce SuperJEPA (Supervision-Augmented Joint-Embedding Predictive Architecture), a method that combines a Joint-Embedding Predictive Architecture (JEPA) with a supervised objective. The architecture consists of four main components:

Encoder: A Transformer-based encoder processes jet constituents. Unlike previous HEP JEPA implementations that use an Exponential Moving Average (EMA) teacher network or stop-gradients, SuperJEPA employs a single, fully trainable encoder for both context and target views. To handle datasets with varying feature sets (e.g., JetClass vs. JetNet), the encoder uses a feature-group embedding scheme where specific feature groups can be masked during training to ensure robustness.
Preprocessing and Masking: Jet constituents are reclustering into subjets using the Cambridge/Aachen algorithm. For each jet, a "target" region (a subset of subjets) is masked, and the remaining "context" is used as input. The target fraction is chosen to contain a specific portion of the jet's transverse momentum.
Predictor: A Transformer decoder predicts the representations of the masked target constituents using only the context representations. Crucially, the predictor receives no information about the positions, features, or number of target constituents; it must infer the hidden structure solely from the context. The output is an unordered set of predicted vectors.
Loss Functions: The total pretraining loss ( $\mathcal{L}_{total}$ $L_{t o t a l}$ ) is a weighted sum of three components:
- Prediction Loss ( $\mathcal{L}_{pred}$ ): The predicted target representations are matched against the actual target representations (projected via a token projector) using optimal transport (Sinkhorn matching) to handle the unordered nature of the data.
- SIGReg Loss ( $\mathcal{L}_{SIGReg}$ ): A regularization term applied to the residuals of the projected tokens to prevent representation collapse, replacing the need for an EMA teacher.
- Classification Loss ( $\mathcal{L}_{CE}$ ): A standard cross-entropy loss applied to the [CLS] token to enforce supervised learning. The weight of this term is controlled by a hyperparameter $\alpha$ .

The model is pretrained on the JetClass dataset (125M jets total, with subsets used for training). After pretraining, the predictor, projector, and classifier heads are discarded. The encoder is then evaluated on downstream tasks (JetClass, JetNet, Top Tagging) either with frozen weights or via fine-tuning.

Key Contributions

Novel Architecture: This is the first application of the teacher-free LeJEPA formulation to HEP. It utilizes a single encoder and SIGReg for stability, avoiding the complexity of EMA teachers.
Supervision Augmentation: The paper proposes and validates a hybrid training objective that augments predictive self-supervision with supervised classification, rather than treating them as mutually exclusive.
Feature Robustness: The implementation includes specific mechanisms (feature-group dropout) to allow models trained on high-dimensional datasets (JetClass) to generalize to lower-dimensional datasets (JetNet, Top Tagging).

Results
The authors evaluated SuperJEPA across multiple model sizes (Small: 3 blocks, 128 dim; Medium: 4 blocks, 256 dim) and pretraining dataset sizes (400k, 4M, and 20M jets).

Performance Gains: SuperJEPA consistently outperformed equivalent supervised-only baselines across various jet-tagging tasks. The gains were observed for both frozen and fine-tuned encoders.
Hyperparameter Sensitivity: The weight of the supervised loss ( $\alpha$ $α$ ) significantly impacted performance.
- On the JetClass dataset (where pretraining and downstream tasks matched), the strongest supervision ( $\alpha=1.0$ ) yielded the best results.
- On the JetNet dataset (where the momentum scale and features differed significantly from JetClass), weaker supervision ( $\alpha=0.1$ ) generally performed better, suggesting the optimal balance depends on the domain shift between pretraining and downstream tasks.
Generalization: SuperJEPA models showed improved ability to separate classes not seen during pretraining (e.g., leptonic top quarks) compared to supervised models, particularly when using attention-pooling mechanisms on particle embeddings.
Representations: Analysis of particle embeddings ( $z_i$ ) revealed that SuperJEPA models retained more structured, class-dependent variation in particle tokens compared to supervised models, which tended to produce less structured embeddings. However, the global [CLS] embeddings showed similar class separation capabilities across both methods.

Significance and Claims
The paper claims that predictive self-supervision remains beneficial even when large labeled simulation samples are available, provided it is used to supplement rather than replace supervised objectives. The results suggest that the "next-token prediction" style objective helps the encoder learn the structural relationships within jets that pure classification might overlook.

The authors modestly note that while their results hold for datasets up to 20M jets, further study is required to determine if these conclusions scale to datasets in the hundreds of millions or billions. They also acknowledge that the optimal balance between self-supervision and supervision is dataset-dependent and that their current work focuses on jet tagging, leaving particle identification and event classification for future investigation. The work motivates the inclusion of predictive self-supervision in future large-scale particle-physics foundation models.

More like this