Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models

TL;DR

Masked Particle Modeling (MPM) introduces a self-supervised, permutation-invariant pre-training scheme for unordered jet particle sets by masking a subset of particles and predicting their discrete token identities from a VQ-VAE codebook, using a transformer encoder with eight blocks and a hidden size of $1024$. Jets are treated as $X=\{x_i ight\}_{i=1}^N$ and token targets come from a 512-entry codebook, enabling context-aware predictions that capture jet-level structure. The approach is evaluated on the JetClass dataset (100M jets across 10 classes) with domain-shift tests on RODEM and includes analyses of ordering, discretization, and three fine-tuning strategies (fixed backbone, fine-tuned, from scratch). Results show data-efficient transfer to unseen classes and cross-domain data, and demonstrate potential benefits of pre-training on real data to mitigate simulation-to-data discrepancies, suggesting a scalable path toward HEP foundation models.

Abstract

We propose masked particle modeling (MPM) as a self-supervised method for learning generic, transferable, and reusable representations on unordered sets of inputs for use in high energy physics (HEP) scientific data. This work provides a novel scheme to perform masked modeling based pre-training to learn permutation invariant functions on sets. More generally, this work provides a step towards building large foundation models for HEP that can be generically pre-trained with self-supervised learning and later fine-tuned for a variety of down-stream tasks. In MPM, particles in a set are masked and the training objective is to recover their identity, as defined by a discretized token representation of a pre-trained vector quantized variational autoencoder. We study the efficacy of the method in samples of high energy jets at collider physics experiments, including studies on the impact of discretization, permutation invariance, and ordering. We also study the fine-tuning capability of the model, showing that it can be adapted to tasks such as supervised and weakly supervised jet classification, and that the model can transfer efficiently with small fine-tuning data sets to new classes and new data domains.

Figures (8)

• Figure 1: The proposed model and training scheme for a FM for jets. A jet is represented as a set of particles, each a list of features, and some particles are replaced by a learnable vector and passed through a transformer encoder. Training aims to predict the discrete token identity, defined by the encoder of a pre-trained VQ-VAE, of the masked particles. The MPM Backbone in is a transformer encoder in this work. The positional embedding shown in box A is used only in the prediction head to preserve the permutation invariance in the backbone; it is required in the model to break the degeneracy from when $m_1=m_2$.
• Figure 2: The reconstruction of the jet mass and transverse momentum using the decoded output of the VQ-VAE trained to quantize the particles in a jet.
• Figure 3: Accuracy of different training strategies as a function of the number of labelled training samples. Accuracy is calculated on the ten classes in the JetClass dataset. The average and standard deviation over 5 trainings is shown in solid lines and uncertainty bands, respectively. Models with frozen pretrained backbone weights during fine-tuning are "Fixed", and those with updated weights are "Fine-tuned".
• Figure 4: Accuracy of different training strategies as a function of labelled training samples. Accuracy is calculated on four classes, held out from all pre-trained models, in the JetClass dataset. The average and standard deviation over 5 trainings is shown in solid lines and uncertainty bands, respectively. Models with frozen pretrained backbone weights during fine-tuning are "Fixed", and those with updated weights are "Fine-tuned".
• Figure 5: The QCD rejection at $50\%$ top-jet efficiency evaluated on (left) the RODEM test set and (right) the JetClass test set, as a function of the size of the RODEM data set used for fine-tuning. All models are pre-trained on JetClass. The average and standard deviation of rejection over 5 trainings is shown in solid lines and uncertainty bands, respectively. Models with frozen pretrained backbone weights during fine-tuning are "Fixed", and those with updated weights are "Fine-tuned".