[Paper Review] Learning the language of QCD jets with transformers

Thorben Finke, M. Krämer+2 more|arXiv (Cornell University)|Mar 13, 2023

Computational Physics and Python Applications4 citations

TL;DR

This paper proposes a transformer-based auto-regressive density estimation framework for high-energy jet physics by modeling jet constituents as discrete 'words' in a language-like structure. By training on QCD and top jets, the model learns complex, high-dimensional probability distributions, enabling high-fidelity data generation that closely matches real jet observables and is indistinguishable from real data even by advanced classifiers.

ABSTRACT

Transformers have become the primary architecture for natural language processing. In this study, we explore their use for auto-regressive density estimation in high-energy jet physics, which involves working with a high-dimensional space. We draw an analogy between sentences and words in natural language and jets and their constituents in high-energy physics. Specifically, we investigate density estimation for light QCD jets and hadronically decaying boosted top jets. Since transformers allow easy sampling from learned densities, we exploit their generative capability to assess the quality of the density estimate. Our results indicate that the generated data samples closely resemble the original data, as evidenced by the excellent agreement of distributions such as particle multiplicity or jet mass. Furthermore, the generated samples are difficult to distinguish from the original data, even by a powerful supervised classifier. Given their exceptional data processing capabilities, transformers could potentially be trained directly on the massive LHC data sets to learn the probability densities in high-energy jet physics.

Motivation & Objective

• To address the challenge of high-dimensional density estimation in jet physics, where traditional methods struggle with the curse of dimensionality and loss of correlations.
• To explore whether transformers—successful in NLP—can model jet constituents as discrete tokens to learn complex, auto-regressive probability distributions.
• To enable both likelihood evaluation and high-quality generative sampling of jets with variable multiplicity, avoiding the limitations of fixed-size models like normalizing flows.
• To assess the quality of learned densities through generative sampling and classifier-based evaluation, ensuring realism and statistical fidelity.
• To demonstrate that the model captures key physical features of QCD and top jets, including jet mass, multiplicity, and angular distributions.

Proposed method

• Discretize jet constituent properties—transverse momentum $p_T$, $\Delta\eta$, and $\Delta\phi$—into discrete bins to form a 'vocabulary' of particle states.
• Order jet constituents by $p_T$ to introduce a basic grammar, enabling auto-regressive modeling of the joint density as a product of conditional probabilities.
• Train a standard transformer encoder to predict the probability of each constituent given all prior ones in the sequence, using maximum likelihood estimation.
• Use top-k and full-probability sampling strategies to generate synthetic jet samples from the learned density.
• Evaluate generative quality using observables like multiplicity, $p_T$, $\Delta\eta$, and jet mass, comparing distributions between real and generated data.
• Train a ParticleNet classifier to distinguish real from generated jets, using the AUC score as a proxy for distributional fidelity.

Experimental results

Research questions

• RQ1 Can transformers effectively learn the high-dimensional probability density of QCD and top jets by treating jet constituents as discrete 'words' in a sequence?
• RQ2 How well do the generated jet samples reproduce key physical observables such as multiplicity, $p_T$, $\Delta\eta$, and jet mass compared to real data?
• RQ3 To what extent can a supervised classifier distinguish between real and generated jets, indicating the realism of the learned density?
• RQ4 Does the model generalize well to top jets, which have more complex internal structures than QCD jets?
• RQ5 How do different sampling strategies (top-k vs. full distribution) affect the realism and physical consistency of generated samples?

Key findings

• The generated QCD jet samples show excellent agreement with real data in distributions of particle multiplicity, transverse momentum, and angular variables, with minimal statistical differences.
• The ParticleNet classifier achieves an AUC of approximately 0.95 when distinguishing real from top-k sampled QCD jets, indicating that only about 1% of generated samples are easily separable from real data.
• Sampling from the full probability distribution leads to a slight oversampling of high-multiplicity and high-mass jets, introducing unphysical tails and increasing the fraction of distinguishable samples to ~5%.
• For top jets, the model reproduces one-dimensional distributions ($\log(p_T)$, $\Delta\eta$, $m^{\text{jet}}$) with fidelity comparable to QCD jets, though multiplicity is less accurately captured due to limited training on high-multiplicity events.
• The AUC for classifying top jet samples from real data is ~0.7, indicating slightly better separation than for QCD jets, likely due to the more complex structure of top jets and higher variance in sampling.
• The results suggest that more extensive training on larger datasets with higher constituent counts is necessary to further improve the modeling of top jets.