Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models (2401.13537v3)

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.

• The paper introduces a novel self-supervised learning framework (MPM) that predicts masked particle features in collider physics data.
• The methodology employs VQ-VAE tokenization and a permutation invariant design to pre-train foundation models for diverse HEP tasks.
• The results demonstrate improved jet classification and robust generalization across unseen classes and domains, cutting reliance on extensive labeled datasets.

Overview of Masked Particle Modeling on Sets

The paper "Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models" introduces a novel self-supervised learning (SSL) approach specifically designed for high energy physics (HEP) data. The authors propose a strategy called Masked Particle Modeling (MPM), which is aimed at learning generic, transferable, and reusable representations from unordered sets of particles in collider physics experiments. This strategy draws inspiration from masked modeling techniques successfully applied in other domains, such as NLP and Computer Vision (CV).

Goals and Methodology

The primary objective of MPM is to construct large foundation models for HEP that can be pre-trained in a self-supervised manner and subsequently fine-tuned for various downstream tasks such as jet classification. Contrary to traditional supervised learning approaches, which depend heavily on labeled data and are susceptible to domain overfitting, MPM seeks to leverage unlabeled data for robust and domain-generalizable feature extraction.

In the MPM framework, particles within a jet are masked, and the training model aims to predict the properties of the masked particles based on the information from the unmasked particles. This is analogous to the masked LLMing used in models like BERT but adapted to unordered, continuous data typical in HEP. Various strategies for masking, ordering, and predicting particle features are investigated within this method.

Tokenization and Permutation Invariance

A salient challenge tackled by the authors is adapting the discretization of continuous features and ensuring permutation invariance for unordered sets of particles. To create discrete tokens from continuous particle features, the authors utilize a Vector Quantized Variational Autoencoder (VQ-VAE). The paper investigates different tokenization techniques, including direct binning of features and using k-means clustering, to evaluate their effectiveness in pre-training.

Moreover, the model also addresses the permutation invariance of particle sets by experimenting with ordering strategies. The paper finds that ordering particles exclusively in the prediction head significantly improves performance while maintaining the essential permutation invariance in the backbone model. This subtle yet critical adjustment ensures that the model can adapt to various downstream tasks without an implicit sequence bias.

Fine-Tuning and Evaluation

The pre-trained models are evaluated on several downstream tasks to measure their effectiveness:

• In-context Classification: The model is fine-tuned and tested on the same JetClass dataset used for pre-training. The results demonstrate substantial performance improvements, particularly with small labeled datasets, underscoring the effectiveness of SSL pre-training.
• Out-of-context Classification: To evaluate the generalizability of the learned representations, the model is pre-trained on a subset of classes and fine-tuned on new, unseen classes. The model continues to exhibit strong performance, indicating that the backbone has learned generalizable features useful across different particle jet categories.
• Out-of-domain Classification: The model's adaptability to different datasets (RODEM) further validates its generalizability. Even with domain shifts, the pre-trained models show superior performance compared to models trained from scratch, highlighting their potential to mitigate domain-related discrepancies in HEP data analytics.

Additionally, the paper explores the utility of weakly supervised learning by using "noisy" labels in fine-tuning. The pre-trained models achieve significant improvements in performance compared to fully supervised models trained from scratch, suggesting their practical applicability in real-world scenarios where clean labels might not always be available.

Implications and Future Directions

The implications of adopting MPM in HEP are multifaceted. This approach can significantly reduce reliance on large labeled datasets, which are often expensive and time-consuming to generate. Moreover, the potential to mitigate domain shifts by pre-training on real experimental data while fine-tuning on simulated data opens new avenues for robust model deployment in high energy physics.

Future work could involve scaling the MPM framework with larger models and data sets, thus enhancing the representation learning capability further. Additionally, exploring other SSL techniques and their synergistic incorporation with MPM could yield more refined and powerful foundation models for HEP.

Furthermore, this paper sets a precedent for cross-domain application of SSL techniques, suggesting that methodologies successful in NLP and CV can be adapted for scientific data with appropriate modifications. This could spur innovation not just within HEP but across various scientific disciplines that deal with large, complex, and unlabeled datasets.

Conclusion

This paper presents a comprehensive approach to self-supervised learning in high energy physics, demonstrating the feasibility and advantages of masked particle modeling. The results underscore the potential of SSL to revolutionize data analysis in HEP, providing a pathway for more efficient and generalizable models capable of addressing the complexities inherent to high energy particle data.