jBOT: Semantic Jet Representation Clustering Emerges from Self-Distillation

TL;DR

This work addresses jet representation learning for LHC data by adapting the iBOT self-distillation framework to tokenized jets, enabling SSL pre-training on unlabeled jets. It introduces jBOT, which learns from two augmented views using per-particle masking, a CLS-based global objective, and a KoLeo regularizer within a teacher-student EMA setup. The learned embeddings exhibit emergent semantic clustering without labels and support effective anomaly detection via simple distance-based metrics, while fine-tuning yields improvements over scratch-based supervised training, especially with limited labeled data. The approach demonstrates practical benefits for jet tagging and anomaly detection and points toward scalable SSL applications in high-energy physics.

Abstract

Self-supervised learning is a powerful pre-training method for learning feature representations without labels, which often capture generic underlying semantics from the data and can later be fine-tuned for downstream tasks. In this work, we introduce jBOT, a pre-training method based on self-distillation for jet data from the CERN Large Hadron Collider, which combines local particle-level distillation with global jet-level distillation to learn jet representations that support downstream tasks such as anomaly detection and classification. We observe that pre-training on unlabeled jets leads to emergent semantic class clustering in the representation space. The clustering in the frozen embedding, when pre-trained on background jets only, enables anomaly detection via simple distance-based metrics, and the learned embedding can be fine-tuned for classification with improved performance compared to supervised models trained from scratch.

Figures (9)

• Figure 1: Schematic diagram of the jBOT pre-training method. A teacher-student architecture is used with a backbone encoder and a projection head, where stop-gradient is applied to the teacher network, whose weights are an EMA of the student network weights. Starting from an input jet, two augmented views are generated; in each view, each particle is embedded into a token space, and a [CLS] token is prepend. Both views are passed to the student and teacher networks. The student network processes distorted views where some of the particle tokens are masked and replaced by a learnable [MASK] token, while the teacher network processes the full views. Same-view and cross-view distillation losses are computed in the projection space, and the KoLeo loss is computed on the student [CLS] embedding from only one of the two views.
• Figure 2: Downstream tasks are performed using the [CLS] embedding from the pre-trained student encoder.
• Figure 3: One example jet per class, where each circle represents a particle: the circle center is at the particle location and the radius is proportional to its $p_{\text{T}}$. Upper row: input jets. Lower row: augmented jets with masking shown in red (e.g., $\sim$30% of the jet $p_{\text{T}}$).
• Figure 4: Learning curves for jBOT-B when pre-trained on five-class data. Upper left to right: total loss, $\mathcal{L}_{\texttt{[CLS]}}$, and $\mathcal{L}_{\text{Part}}$. Lower left to right: $\mathcal{L}_{\text{KoLeo}}$, entropy, and centering norm.
• Figure 5: Attention weights from the last transformer block in jBOT-S pre-trained on all five classes, obtained using the [CLS] token as the query attending to the particle tokens. Top row: one example input jet per class (each circle represents a particle: the circle center is at the particle location, the radius is proportional to its $p_{\text{T}}$, and the edge alpha is uniform across all particles here). Other rows: attention weights per head for the same input jets, shown with the same drawing style as the input jets, but with the attention weight represented by the edge alpha (higher edge alpha indicates larger attention weight).