---
title: 'JetClass Dataset: HEP ML Benchmark'
url: https://www.emergentmind.com/topics/jetclass-dataset
type: topic
---

# JetClass Dataset: HEP ML Benchmark

JetClass is a large-scale, publicly available multi-class jet tagging dataset specifically designed to advance machine learning (ML) techniques in high-energy physics (HEP). Comprising 100 million simulated jets equally distributed over ten physically motivated categories, JetClass presents a foundational benchmark for transformer-based and set-based approaches to jet classification, anomaly detection, and generative modeling. The dataset’s construction, physics content, high-level feature engineering, and role in model pre-training have set new standards for empirical rigor and statistical power in the HEP ML domain [2202.03772].

## 1. Dataset Construction and Physics Processes

JetClass was generated by simulating proton-proton collisions, primarily using MadGraph5_aMC@NLO for hard-process generation, Pythia 8 for parton showering and hadronization, and Delphes for fast detector simulation with CMS-like parameters [2202.03772, 2406.06638]. Jets are clustered from stable final-state particles using the anti-$k_T$ algorithm with radius parameter $R=0.8$. Typical jet selection requires $500~\mathrm{GeV} < p_T^\mathrm{jet} < 1000~\mathrm{GeV}$ and $|\eta^{\mathrm{jet}}|<2$, ensuring high-purity, high-$p_T$ jets [2202.03772, 2312.00123].

The dataset’s ten classes correspond to key Standard Model processes observed at the LHC:

- **Background (QCD):** light-quark- and gluon-initiated jets (q/g)
- **Higgs boson decays:** $H \rightarrow b \bar{b}$, $H \rightarrow c \bar{c}$, $H \rightarrow g g$, $H \rightarrow 4q$, $H \rightarrow \ell \nu qq'$
- **Top quark decays:** $t \rightarrow bqq'$, $t \rightarrow b \ell \nu$
- **W and Z boson decays:** $W \rightarrow qq'$, $Z \rightarrow qq$

Each class contributes exactly 10 million jets, ensuring strict balance across splits for unbiased training and benchmarking [2406.06638]. Jets are labeled at generator level by matching final-state decay products within $\Delta R < 0.8$ of the jet axis, except for the q/g background [2202.03772].

## 2. Input Features, Representation, and Preprocessing

A defining aspect of JetClass is its granular, constituent-level representation. Each jet is encoded as an unordered set (or "cloud") of up to $N$ constituents, where $N$ typically ranges from 10–100, depending on event structure. For each constituent, an array of real-valued and discrete features captures spatial, kinematic, and identification information [2202.03772, 2312.00123].

**Per-particle features:**  
- Kinematic: $\Delta\eta = \eta_\text{const} - \eta_\text{jet}$, $\Delta\phi = \phi_\text{const} - \phi_\text{jet}$, $\log p_T$, $\log E$, $\log(p_T/p_{T,\text{jet}})$, $\log(E/E_\text{jet})$, $\Delta R = \sqrt{\Delta\eta^2 + \Delta\phi^2}$
- Charge and PID: $\mathrm{charge}\in\{-1, 0, +1\}$, one-hot flags for $\mathrm{isChargedHadron}$, $\mathrm{isNeutralHadron}$, $\mathrm{isPhoton}$, $\mathrm{isElectron}$, $\mathrm{isMuon}$
- Impact parameter: $d_0$, $d_z$, $\sigma_{d_0}$, $\sigma_{d_z}$ (displacement and uncertainties, set to zero for neutrals)

**Per-pair features (Lund plane parametrization):**  
For each pair of constituents, features include $\ln\Delta R$, $\ln k_T$ (with $k_T = \min(p_{T,i}, p_{T,j}) \Delta R$), $\ln z$ (with $z = \min(p_{T,i}, p_{T,j})/(p_{T,i}+p_{T,j})$), and $\ln m^2$ (with $m^2 = (E_i+E_j)^2 - |\vec{p}_i+\vec{p}_j|^2$). These features are crucial for transformer-based architectures incorporating pairwise attention [2202.03772, 2406.06638].

**Preprocessing:**  
Most features are used in raw or log-transformed form, with no dataset-level normalization; activations are normalized within network architectures via LayerNorm or BatchNorm [2406.06638, 2202.03772]. For generative modeling, all continuous features can be standardized to zero mean and unit variance, and constituents are zero-padded or masked up to $N=128$ [2312.00123].

## 3. Dataset Splits, Storage, and Access

JetClass is distributed as ROOT TTree objects with 41 branches, each encoding per-jet, per-particle, or per-pair features [2406.06638]. The standardized division is:

| Split         | Jets per class    | Total jets   | Proportion         |
|---------------|-------------------|--------------|--------------------|
| Training      | 10 million        | 100 million  | 83.3%              |
| Validation    | 0.5 million       | 5 million    | 4.2%               |
| Test          | 2 million         | 20 million   | 12.5%              |

Class frequencies are identical in all splits, obviating the need for further weighting or oversampling [2406.06638, 2202.03772]. Data ingestion for ML frameworks is enabled by direct loading into (x, U) tensor pairs (per-particle, per-pair), or as fixed-size tensors plus masks for set-based equivariant models [2407.08682, 2312.00123].

## 4. Applications: Discriminative and Generative Modeling

JetClass underpins multiple research directions in jet tagging:

- **Supervised classification:** The original Particle Transformer (ParT) achieved accuracy 0.861 and AUC 0.9877 on the JetClass test set, outperforming ParticleNet and plain transformer baselines. Performance scales with training size, with notable gains when moving from 2M to 100M jets [2202.03772].
- **Transfer learning:** Pre-training on JetClass followed by downstream fine-tuning (e.g., on top-tagging or quark–gluon discrimination) delivers superior accuracy, rejection power, and data efficiency compared to training from scratch [2202.03772, 2407.08682].
- **Self-supervised and contrastive learning:** JetClass can serve as an unlabeled corpus for training encoders via contrastive losses (e.g., JetCLR) with physics-motivated augmentations (random translation in $\eta$–$\phi$, particle dropout, momentum jitter, global rotation) [2408.09343].
- **Generative modeling:** Permutation-equivariant continuous normalizing flows (CNFs), trained with flow-matching loss and conditioned on jet type and axis kinematics, produce high-fidelity synthetic jets capturing both kinematics and discrete PID/displacement information [2312.00123].

## 5. Rationale, Selection, and Robustness

The design of JetClass addresses several key goals for large-scale HEP ML benchmarks:

- **Inclusive physics coverage:** Ten classes sample the main LHC resonance processes, allowing studies of both signal vs background and subtype discrimination.
- **Realistic detector modeling:** The use of full detector simulation (in baseline versions) approximates tracker resolution, calorimeter smearing, and efficiency effects encountered in CMS-like environments [2202.03772, 2406.06638].
- **Feature completeness:** Particle flow–like feature vectors include all information needed for state-of-the-art ML models, from kinematics to particle type and trajectory.
- **Balanced splits and strict selection:** Jets are not only class-balanced, but also pass standard $p_T$ and $|\eta|$ cuts and containment requirements to ensure statistical control and reproducibility.
- **Augmentation for self-supervised robustness:** Physics-inspired perturbations are critical for learning representations invariant to pileup, detector resolution, and calibration uncertainties, as demonstrated in contrastive SSL setups [2408.09343].
- **No hand-tuned normalization:** Relying instead on architecture-level normalization (LayerNorm, BatchNorm) lets models adapt optimally to the varied dynamical range of input features [2406.06638].

## 6. Benchmarks and Empirical Performance

Quantitative performance of supervised architectures trained on JetClass demonstrates the utility of both the dataset’s scale and its physics-driven structure:

| Model           | Accuracy | AUC    | Notable rejection (example)          |
|-----------------|----------|--------|--------------------------------------|
| PFN (DeepSets)  | 0.772    | 0.9714 | –                                    |
| P-CNN (DeepAK8) | 0.809    | 0.9789 | –                                    |
| ParticleNet     | 0.844    | 0.9849 | $t_\mathrm{had}@99.5\%$: $1.12\times10^4$ |
| ParT            | 0.861    | 0.9877 | $t_\mathrm{had}@99.5\%$: $1.59\times10^4$ |

ParT improves background rejection by up to $\times3$ for $t$-jets, $\times2$ for $H \rightarrow bb$, etc. Increasing training set size consistently boosts both accuracy and AUC across all benchmark architectures [2202.03772].

For downstream transfer, pre-trained ParT models attain higher performance in top, Higgs, and quark/gluon discrimination compared to state-of-the-art alternatives, including LorentzNet, particularly when leveraging all PID inputs [2202.03772, 2407.08682].

A plausible implication is that JetClass’s scale and constituent-level detail render it suitable not only for jet tagging, but also for advanced representation learning and anomaly detection in collider physics [2202.03772, 2312.00123, 2408.09343].

## 7. Position Relative to Other Jet Datasets

JetClass occupies a foundational position in the modern ML-for-HEP landscape:

- It is approximately two orders of magnitude larger than earlier public datasets (e.g., JetNet).
- Its constituent-level information (including impact parameters and PID) surpasses the limited feature sets of prior datasets.
- JetClass-II [2405.12972], developed subsequently, expands drastically to 188 classes and includes more exotic signatures, refined reweighting protocols, and is designed for foundational multi-class pre-training (e.g., in the Sophon architecture). JetClass, however, remains the standard for large-scale, balanced, 10-class classification and transfer learning tasks in collider jet physics.

JetClass has thus become the reference standard for both benchmarking and developing new ML models (transformers, equivariant flows, contrastive SSL) that target realistic, high-granularity collider environments. Since its release, JetClass has catalyzed progress in model design, empirical jet physics, and scientific discovery pipelines [2202.03772, 2408.09343, 2406.06638].

Source: https://www.emergentmind.com/topics/jetclass-dataset