⚛️ phenomenology

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

This paper introduces Masked Particle Modeling (MPM), a self-supervised pre-training framework that learns permutation-invariant representations of unordered particle sets in high energy physics by reconstructing masked particles via vector quantization, demonstrating its effectiveness as a foundation model for diverse downstream tasks like jet classification and domain transfer.

Original authors: Tobias Golling, Lukas Heinrich, Michael Kagan, Samuel Klein, Matthew Leigh, Margarita Osadchy, John Andrew Raine

Published 2026-08-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tobias Golling, Lukas Heinrich, Michael Kagan, Samuel Klein, Matthew Leigh, Margarita Osadchy, John Andrew Raine

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

High energy physics is the study of the universe's most fundamental building blocks, conducted by smashing particles together at incredible speeds and watching what flies out. The detectors in these experiments do not see single, clean objects; instead, they record chaotic, unordered collections of thousands of tiny fragments, known as particles, that stream out from the collision. For decades, scientists have used artificial intelligence to make sense of this data, but these computer programs were usually trained like students memorizing flashcards: they were fed vast amounts of labeled examples and told exactly what to look for. This approach works well for specific tasks, but it is inefficient. It requires massive amounts of labeled data for every new question, and it struggles when the data changes slightly, such as when moving from computer simulations to real-world experimental results. The field is now looking toward a different kind of intelligence, one that learns by observing the world on its own, without needing a teacher to point out every detail. This is the realm of foundation models, systems that learn general patterns from huge, unlabeled datasets so they can be adapted to many different tasks later.

A team of researchers has taken a significant step toward building such a foundation model specifically for high energy physics. They developed a new method called masked particle modeling, which teaches a computer to understand the structure of a particle collision by playing a game of "fill in the blanks." Imagine a jet of particles as a scattered group of people in a room. In this method, the computer is shown the room with several people hidden behind a curtain. Its job is to guess who those hidden people are and what they are doing, based solely on the behavior and position of the people it can still see. To make this possible, the researchers had to solve two major problems. First, unlike words in a sentence which have a fixed order, particles in a jet have no specific sequence; they are a jumbled set. Second, the properties of these particles, like their speed and direction, are continuous numbers, not discrete words from a dictionary. The researchers created a system that converts these continuous numbers into a set of discrete "tokens," similar to how a digital image is broken down into pixels, allowing the computer to treat the particles as a vocabulary it can learn to predict.

The researchers tested this approach using a massive dataset containing one hundred million simulated particle jets, representing ten different types of particle collisions. They trained their model to hide a portion of the particles in each jet and then predict the identity of the missing ones. To do this, they used a specialized tool that translates the complex, continuous features of the particles into a simpler, discrete code. The model learned to understand the relationships between the visible particles to reconstruct the hidden ones. The results showed that this self-taught model developed a deep, general understanding of particle physics. When the researchers tested it on tasks it had never seen before, such as identifying specific types of jets or distinguishing between different particle sources, the model performed remarkably well. It was particularly effective when the amount of labeled data available for the new task was small. In these cases, the model that had been pre-trained on the massive, unlabeled dataset vastly outperformed models trained from scratch, proving that the initial "fill in the blanks" exercise had taught it useful, transferable skills.

The study also explored how well this knowledge could travel across different types of data. The researchers trained their model on one set of simulated data and then tried to apply it to a different dataset generated with different simulation settings, which mimics the challenge of moving from theory to real-world experiments. The pre-trained model adapted quickly to this new environment, maintaining high performance even with very little new labeled data. This suggests that the model learned fundamental physical patterns rather than just memorizing the quirks of a specific simulation. Furthermore, the researchers demonstrated that the model could be used in "weakly supervised" scenarios, where the data is messy and the labels are imperfect. By leveraging the knowledge gained from the initial pre-training, the model could still distinguish between signal and background noise with high accuracy, a crucial capability for real-world experiments where perfect data is rare.

The findings indicate that this approach offers a powerful new way to analyze high energy physics data. By learning from vast amounts of unlabeled data, these models can reduce the reliance on expensive, perfectly labeled simulations. The researchers found that the model's ability to generalize was strongest when they treated the particles as an unordered set and used a specific type of prediction loss that focused on classifying the identity of the missing particles rather than just guessing their exact numerical values. They also discovered that the order in which particles were presented to the model mattered only during the prediction phase, not during the core learning phase, preserving the natural symmetry of the data. While the work was conducted using simulated data, the authors suggest that this method could eventually be applied directly to real experimental data, allowing scientists to learn from the raw output of particle colliders without needing to first translate it into a simulation. This would represent a major shift in how physics data is processed, moving from task-specific training to a more robust, general understanding of the subatomic world.

Technical Summary: Masked Particle Modeling on Sets for High Energy Physics Foundation Models

Problem Statement

High Energy Physics (HEP) data analysis has traditionally relied on supervised learning models trained on specific tasks and limited labeled datasets. While Foundation Models (FMs) based on self-supervised learning (SSL) have revolutionized Natural Language Processing (NLP) and Computer Vision (CV), their application to HEP remains nascent. The HEP community faces unique challenges in adapting these paradigms:

Data Structure: HEP data, specifically particle jets, consists of unordered sets of particles, unlike the sequential nature of text or the grid structure of images.
Feature Continuity: Particle features (e.g., momentum, direction) are continuous, contrasting with the discrete vocabulary used in language models.
Label Scarcity & Domain Shift: Supervised learning often requires extensive simulated data, which may not perfectly match real experimental data, leading to domain shift uncertainties. There is a need for models that can learn generic representations from vast amounts of unlabeled data to improve downstream task performance with smaller labeled datasets.

Methodology: Masked Particle Modeling (MPM)

The authors propose Masked Particle Modeling (MPM), a self-supervised pre-training scheme adapted for unordered sets of continuous inputs. The methodology draws inspiration from Masked Language Modeling (MLM) in NLP (e.g., BERT) and masked image modeling (e.g., BEiT).

Core Framework

Input Representation: A jet is treated as a set $X = \{x_i\}$ of $N$ particles, where each $x_i$ is a vector of continuous features.
Masking Strategy: A subset of particles in each jet is masked and replaced with a learnable vector.
Tokenization (Discretization): To handle continuous features and enable classification-based loss functions, the authors employ a Vector Quantized Variational Autoencoder (VQ-VAE).
- A pre-trained VQ-VAE encoder maps continuous particle features to the nearest vector in a finite codebook.
- The index of the codebook vector serves as a discrete "token."
- This allows the model to predict the identity (token index) of masked particles rather than regressing continuous values directly.
- Alternative: K-means clustering is also explored as a simpler, context-independent tokenization method.
Model Architecture:
- Backbone: A Transformer encoder (based on Normformer) processes the set of particles.
- Permutation Invariance: Since jets are unordered sets, the backbone must be permutation equivariant. However, using identical mask vectors for all masked particles creates a redundancy where the model cannot distinguish between masked positions.
- Ordering Strategy: To break this degeneracy without violating permutation invariance for downstream tasks, the authors introduce ordering only at the input to the masked prediction head. The particles are ordered by transverse momentum ( $p_T$ ) and assigned learned positional embeddings only for the prediction step. The backbone itself remains permutation invariant.
Training Objective: The model minimizes a cross-entropy loss to predict the discrete token identity of the masked particles based on the unmasked context.

Fine-Tuning Strategies

The paper evaluates three fine-tuning approaches for downstream tasks (e.g., jet classification):

Fixed Backbone: The pre-trained backbone is frozen; only a linear classification head is trained.
Fine-Tuned: Both the backbone and the classification head are updated.
From Scratch: A model with the same architecture is trained from random initialization on the downstream task (supervised baseline).

Key Contributions and Experimental Results

1. Tokenization and Ordering Design

Experiments on the JetClass dataset (100M jets, 10 classes) revealed:

Ordering: Ordering particles only at the prediction head (not in the backbone) yielded the highest accuracy (56.8%), outperforming ordering the backbone inputs or no ordering. This breaks symmetry for precise per-particle prediction while preserving the backbone's permutation invariance.
Discretization vs. Regression: Using discrete tokens (via VQ-VAE) with a classification loss significantly outperformed direct regression of continuous features. Regression tends to learn the mean of the posterior, failing to capture multi-modal distributions, whereas classification learns the full posterior over tokens.
VQ-VAE vs. K-means: VQ-VAE tokens outperformed K-means tokens, likely due to the VQ-VAE's ability to capture context-dependent features. However, K-means provided reasonable performance with easier training.

2. Downstream Task Performance

In-Context Learning (Same Classes): On the JetClass test set, pre-trained models significantly outperformed "from scratch" models when labeled data was scarce (<10k jets). As labeled data increased, supervised models converged to the performance of fine-tuned pre-trained models.
Out-of-Context Learning (New Classes): The model was pre-trained on 6 classes (Higgs and QCD) and fine-tuned on 4 unseen classes (Top, W, Z). Pre-trained models again outperformed supervised baselines with limited data, demonstrating that the learned representations are generic and transferable to unseen particle types.
Out-of-Domain Learning (New Datasets): Pre-training on JetClass and fine-tuning on RODEM (a different dataset with different detector simulation and jet clustering parameters) showed strong data efficiency. Pre-trained models maintained high performance on the RODEM dataset with small labeled sets, suggesting robustness to domain shifts.
Weakly Supervised Classification: The model was tested on a task using "noisy labels" (mixing signal and background events). Pre-trained backbones significantly improved the significance of separating pure signal and background datasets compared to models trained from scratch, even when only the linear head was fine-tuned.

Significance and Claims

The paper claims that Masked Particle Modeling (MPM) is a viable and effective strategy for building Foundation Models in High Energy Physics. Its primary significance lies in:

Adaptation of SSL to Sets: It successfully generalizes masking strategies from NLP and CV to unordered sets of continuous data, addressing the specific structural challenges of HEP.
Data Efficiency: The pre-trained "backbone" provides a powerful initialization that allows models to achieve high performance on downstream tasks with significantly smaller labeled datasets compared to training from scratch.
Generalizability: The learned representations are not limited to the specific classes or datasets seen during pre-training. They transfer effectively to new particle classes and new data domains (simulations with different parameters).
Path to Real-Data Pre-training: The authors highlight the potential to pre-train directly on unlabeled real experimental data (which is abundant) and fine-tune on smaller simulated datasets. This approach could mitigate domain shift uncertainties arising from imperfect simulations, a critical issue in HEP.

The work positions MPM as a foundational step toward large-scale, self-supervised learning in HEP, suggesting that scaling up pre-training data and model sizes could further overcome domain adaptation challenges.