jBOT: Semantic Jet Representation Clustering Emerges from Self-Distillation

AI summary

Overview

Research area: Machine learning for high-energy physics (self-supervised representation learning applied to CERN LHC jet data), with methodological roots in computer-vision self-distillation.

Technical level: Intermediate to Advanced. The paper assumes familiarity with transformer architectures, self-supervised learning objectives, and basic collider physics concepts such as jets, transverse momentum, and jet tagging.

Scope: The paper introduces jBOT, a self-distillation pre-training method adapted from the iBOT computer-vision framework to tokenized jet constituents, and evaluates whether the resulting frozen embeddings support classification and anomaly detection on the JetNet dataset.

What This Paper Is About

Supervised jet taggers at the LHC are trained from labeled simulations, but labels are expensive and generic unlabeled collision data is abundant. The authors adapt iBOT — a self-distillation method originally built for images — so that a teacher-student transformer pair can be pre-trained on unlabeled jets and learn useful structure without any class labels. The goal is to test whether such pre-training produces semantically clustered jet representations that can then be probed or fine-tuned for classification and anomaly detection.

Key Contributions

1. Introduces jBOT, a self-distillation pre-training framework for jets built on iBOT, combining a particle-level (same-view) distillation loss on masked particle tokens with a jet-level (cross-view) distillation loss on the [CLS] token, plus a KoLeo regularizer for feature-space diversification.

2. Adapts masking to physics: instead of masking a fixed number of particles, jBOT uses a momentum-aware scheme that masks particles until the cumulative transverse momentum of masked particles reaches a target ratio (uniformly sampled from 0 to 50% per view).

3. Demonstrates emergent semantic clustering: t-SNE projections of the [CLS] embedding during pre-training show compact clusters separating different jet classes form progressively from random initialization, even though no labels are used.

4. Shows downstream utility in two regimes: frozen-embedding probes and fine-tuned classifiers for five-class classification and top tagging, and distance-based anomaly detection using a model pre-trained on QCD jets (q, g) only.

Main Findings

• Emergent clustering without labels: The [CLS] embedding, visualized with 2D t-SNE, forms progressively more compact class-separating clusters during pre-training, starting from random initialization. Attention visualizations from the [CLS] query in the last transformer block show different heads attending to different particle groups within a jet.

• No representation collapse: Learning curves on the five-class data show steadily decreasing losses that plateau. The teacher entropy stabilizes slightly below the student entropy and both remain well below the bound log(d_proj) = log(32) ≈ 3.47, and above zero.

• Frozen five-class accuracy: k-NN (k = 30) on frozen features reaches 71.02 ± 0.58% (jBOT-S) and 70.53 ± 0.57% (jBOT-B); linear probes reach 67.43 ± 0.56% (jBOT-S) and 69.42 ± 0.46% (jBOT-B). The authors characterize this as around 70%, compared with 76% for a supervised model.

• Fine-tuning beats training from scratch, most at low label counts: Five-class accuracy: jBOT-S (10%) 74.25 ± 0.56% vs Sup.-S (10%) 72.51 ± 0.52%; jBOT-S (100%) 75.26 ± 0.64% vs Sup.-S (100%) 74.31 ± 0.56%; jBOT-B (10%) 74.99 ± 0.53% vs Sup.-B (10%) 73.79 ± 0.56%; jBOT-B (100%) 76.43 ± 0.61% vs Sup.-B (100%) 76.04 ± 0.57%.

• Top tagging gains: Frozen k-NN reaches 87.77 ± 0.37% (jBOT-S) and 87.93 ± 0.35% (jBOT-B). Fine-tuned: jBOT-S (10%) 88.07 ± 0.47% vs Sup.-S (10%) 87.23 ± 0.40%; jBOT-S (100%) 88.75 ± 0.40% vs Sup.-S (100%) 88.52 ± 0.47%; jBOT-B (10%) 88.62 ± 0.38% vs Sup.-B (10%) 87.84 ± 0.51%; jBOT-B (100%) 89.11 ± 0.29% vs Sup.-B (100%) 88.99 ± 0.35%.

• Fine-tuning uses layer-wise learning rate decay (LLRD) with decay factors scanned over {0.6, 0.65, 0.7, 0.75, 0.8} and base learning rates over {4×10⁻⁵, 8×10⁻⁵, 2×10⁻⁴, 4×10⁻⁴, 8×10⁻⁴, 2×10⁻³}, batch size 1024, 100 epochs.

• Anomaly detection is competitive with reconstruction baselines: Using a model pre-trained on q and g jets only and tested on W, Z, and t as anomalies (k = 30, GMM K = 4, cosine temperature τ = 0.05):

  • jBOT-S (k-NN): W 0.8072 ± 0.0028, Z 0.8355 ± 0.0025, t 0.8356 ± 0.0029, combined 0.8261 ± 0.0028.
  • jBOT-S (Cosine): W 0.8064 ± 0.0028, Z 0.8355 ± 0.0027, t 0.8388 ± 0.0028, combined 0.8269 ± 0.0027.
  • jBOT-S (Mahalanobis): W 0.7431 ± 0.0030, Z 0.7821 ± 0.0029, t 0.8620 ± 0.0026, combined 0.7957 ± 0.0037.
  • jBOT-S (GMM): W 0.8062 ± 0.0040, Z 0.8204 ± 0.0040, t 0.8197 ± 0.0049, combined 0.8155 ± 0.0038.
  • Baselines from Ref. [50]: CNNAE (0.6886 W, 0.7247 Z, 0.8962 t, 0.7700 combined), GNNAE (0.7558, 0.7805, 0.8917, 0.8195), LGAE (0.7489, 0.7909, 0.8669, 0.8313).

• Bimodal QCD anomaly score distribution: The QCD score distribution is bimodal because of the separation between q and g clusters that emerges from pre-training, which is resolved by metrics sensitive to local structure.

• Stated caveat on the label-efficiency claim: Models fine-tuned on 10% labeled data were already exposed to the full training set during unlabeled pre-training, so the authors frame this as a probe of label efficiency rather than a controlled disjoint-data comparison; that comparison is left to future work.

• Stated caveat on the dataset and hyperparameters: The JetNet dataset (880k simulated jets) is described as modest in the context of foundation model training, and pre-training hyperparameters were "not rigorously optimized but are reasonably chosen."

Methodology in Plain English

The setup. A jet is represented as up to 30 constituent particles, each with four features: relative pseudorapidity, relative azimuthal angle, relative transverse momentum fraction, and a boolean "valid" flag. Jet-level features such as jet pT, η, and φ are deliberately excluded because their distributions differ between classes and could bias classification away from genuine substructure information.

Two views per jet. Each jet is augmented twice to form a positive pair, using three physics-motivated augmentations taken from JetCLR-style practice: uniform rotation of particles about the jet axis (angle sampled uniformly from −π to π), Gaussian smearing of each particle's η and φ with variance Λ_QCD/pT where Λ_QCD = 100 MeV, and collinear splitting of particles with conserved transverse momentum, applied to jets with fewer than 30 valid particles.

Tokenization. Each particle's feature vector is linearly embedded into a d_model-dimensional space, and a learnable [CLS] token is prepended, giving N_p + 1 tokens processed by a ViT-style transformer encoder. A projection head (an MLP) maps encoded tokens into a d_proj-dimensional space where distillation losses are computed.

Teacher-student distillation. The teacher and student share architecture and initialization and train jointly; only the student receives gradients. The teacher's weights are an exponential moving average of the student's (τ_EMA cosine-scheduled from 0.996 to 1), and teacher outputs are re-centered to their batch mean with centering momentum τ_c = 0.9. Temperature-scaled softmax with 0.04 (teacher) and 0.1 (student) sharpens the targets.

Two losses plus a regularizer. The particle-level loss replaces masked particle tokens with a learnable [MASK] token in the student's views and has the student predict the teacher's outputs for those masked positions within the same view. The jet-level loss matches the teacher [CLS] from one view to the student [CLS] from the other view, in both directions. A KoLeo regularizer (scale λ = 0.01) applied to the L2-normalized student [CLS] from one view encourages a diverse spread of examples in the embedding. The total loss is the sum of the two distillation losses plus λ times the KoLeo loss.

Training scale. Two model sizes: jBOT-S (d_model = 32, 2 transformer blocks, 4 attention heads) and jBOT-B (d_model = 64, 4 blocks, 6 heads); projection dimension is d_model/2. Dropout rate 20%, GELU activations, AdamW, learning rate 5×10⁻⁴ × (batch size/256) with linear warm-up for 10 epochs, batch size 1024, weight decay 10⁻⁴. Implemented in TensorFlow and Keras.

Three pre-training datasets, three downstream tasks. Pre-training uses (1) all five classes, (2) q, g, and t, or (3) q and g, feeding five-class classification, top tagging, and anomaly detection respectively. Splits are roughly 80/10/10% with balanced classes. Pre-training runs for 100 epochs on a single Nvidia A100 GPU — for the base model on the five-class set with 700k examples, around 11 minutes per epoch and 18 hours total.

Downstream evaluation. The projection head is removed, and unaugmented, unmasked test jets pass through the student encoder. Frozen features are probed with a k-NN classifier (k = 30) and a linear classifier. Fine-tuning attaches an MLP classifier head to the [CLS] token and trains the encoder with LLRD. For anomaly detection, the encoder is frozen after pre-training on q and g only, the training set is mapped to [CLS] vectors as an in-distribution reference set, and four scores (k-NN distance, cosine similarity, Mahalanobis distance, GMM) measure how far test examples sit from that reference bulk. Ten jBOT-S models are pre-trained with identical configuration and the best-performing one is reported.

Why This Matters

Research impact. The paper adds a self-distillation option to a small but growing set of self-supervised methods for HEP (contrastive approaches such as JetCLR and DarkCLR, masked particle modeling, J-JEPA, MACK, RINO), and it demonstrates that class structure can emerge in jet embeddings without any labels. It also connects the HEP SSL literature to the broader foundation-model effort in the field, and shows that simple distance metrics on a frozen self-supervised embedding are a viable alternative to reconstruction-based autoencoders for anomaly detection.

Real-world applications:

• Jet tagging for LHC analyses: improved identification of boosted W, Z, and top jets supports precision Standard Model measurements and searches for new physics.
• Anomaly detection for new-physics searches: a model pre-trained only on QCD jets provides a signal-agnostic way to flag unusual jets without committing to a specific signal hypothesis.
• Label-efficient workflows: because the largest gains appear when labeled data is limited (10% fine-tuning sets), the approach suits scenarios where simulation or calibration labels are scarce while unlabeled collision data is plentiful.
• Cross-domain robustness: the related-work discussion notes SSL methods such as MACK and RINO aimed at reducing the simulation-to-real performance gap caused by mismodeling, a persistent operational concern for deployed taggers.

Industry relevance. The methodological backbone — masked-token distillation with a teacher-student EMA, momentum-aware masking, and KoLeo regularization — is generic self-supervised representation learning. Practitioners outside physics working on sparse, variable-length, unordered point sets (particle clouds, point clouds, event streams) may find the momentum-weighted masking scheme and the dual local/global distillation objective directly transferable. The open-source release (https://github.com/hftsoi/jbot) in TensorFlow/Keras makes the method easy to adapt.

Future Directions

• Scaling up pre-training on much larger unlabeled datasets, since the authors explicitly describe JetNet as modest in size for foundation-model-scale training.
• More robust anomaly detection strategies, including dedicated augmentations and combination of the learned embedding with other scoring approaches; the authors note the large freedom in metric and hyperparameter choice and the need for optimization scans.
• A controlled disjoint-data experiment, separating pre-training and fine-tuning datasets, to move beyond the current comparison in which 10%-fine-tuned models were pre-trained on the full training set.
• Fine-tuning on more diverse labeled data, and extension of the framework to similar physics objects at the LHC, as part of ongoing foundation-model efforts in the field.

Target Audience

Machine learning researchers working on self-supervised learning for scientific data; high-energy physicists developing jet taggers and anomaly detection methods for LHC analyses; and practitioners interested in adapting computer-vision self-distillation frameworks (iBOT, DINO, DINOv2) to sparse, variable-length, unordered point sets. Readers without a physics background can follow the machine learning content, but the dataset description assumes familiarity with jet physics terminology.

Authors’ abstract

Self-supervised learning, in the context of foundation model training, 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.