From particle clouds to tokens: building foundation models for particle physics

This work presents OmniJet-$\alpha$, the first multi-task foundation model for particle physics in the context of the Large Hadron Collider (LHC) at CERN. In contrast to natural language, particle jet data is represented by point-cloud-like objects, requiring a different type of encoding strategy to make it suitable for auto-regressive generation. We introduce a comprehensive set of evaluation methods to investigate the encoding of particles into a discrete set of tokens. These methods guide us to adopt a more precise tokenization method compared to previous strategies, and we provide insights into how a rather small set of 8192 tokens can accurately represent a complex data space spanned by three continuous physical features (the momenta of the particles). Moreover, we showcase the efficacy of transfer learning between an unsupervised task (jet generation) and a common supervised task (jet tagging). This integration of disparate tasks and the successful transfer learning between them marks a significant advancement in the development of foundation models for particle physics. The code and the model checkpoint are publicly available on GitHub, with the link provided following the review process.