[Paper Review] Dataset for flavour tagging R&D

I. Ochoa, S. B. Klein+2 more|arXiv (Cornell University)|Aug 20, 2024

Electron and X-Ray Spectroscopy Techniques5 citations

TL;DR

The paper improves masked particle modeling (MPM) for jet physics by removing tokenization, upgrading the decoder, and exploring diverse reconstruction tasks and set-to-set generation to pre-train foundation-model–style backbones on jet data. It introduces MPMv2 and set-to-set flow-matching, demonstrating strong downstream performance including OOD tasks.

ABSTRACT

In this work, we significantly enhance masked particle modeling (MPM), a self-supervised learning scheme for constructing highly expressive representations of unordered sets relevant to developing foundation models for high-energy physics. In MPM, a model is trained to recover the missing elements of a set, a learning objective that requires no labels and can be applied directly to experimental data. We achieve significant performance improvements over previous work on MPM by addressing inefficiencies in the implementation and incorporating a more powerful decoder. We compare several pre-training tasks and introduce new reconstruction methods that utilize conditional generative models without data tokenization or discretization. We show that these new methods outperform the tokenized learning objective from the original MPM on a new test bed for foundation models for jets, which includes using a wide variety of downstream tasks relevant to jet physics, such as classification, secondary vertex finding, and track identification.

Motivation & Objective

• Motivate foundation-model style pre-training for high-energy physics using unlabeled jet data.
• Develop improved masked particle modeling (MPMv2) without VQVAE tokenization.
• Evaluate multiple reconstruction tasks including conditional generative approaches.
• Propose set-to-set flow-matching as a competitive pre-training paradigm for jets.

Proposed method

• Modify MPM to remove repeated masked tokens and use a full transformer decoder.
• Provide position-encoding only among masked elements to avoid trivializing the task.
• Investigate five continuous-feature reconstruction tasks and one categorical task (particle ID).
• Explore alternative targets: K-Means tokenization, CNF (conditional normalizing flow), flow-matching (CFM), and set-to-set flow-matching (SSFM).
• Use JetClass and BTag datasets simulated with Delphes to pre-train and evaluate backbone representations.
• Ablation study comparing decoder types, added features, and training settings (Table 1).

Figure 1 : A comparison of the original MPM encoder-decoder setup (left) and the new model configuration (right). The new model includes multiple reconstruction tasks, swaps the MLP decoder for a transformer, and only encodes the reduced set.

Experimental results

Research questions

• RQ1 Does removing tokenization and using a stronger decoder improve MPM performance over the original MPMv1?
• RQ2 Are alternative reconstruction targets (CNF, K-Means, flow-based methods) competitive with VQVAE tokenization for MPM pre-training?
• RQ3 How do the improved backbones perform on in-distribution, weakly supervised, and out-of-distribution jet tasks?
• RQ4 Can set-to-set flow-matching provide a viable pre-training paradigm for unordered jet constituents?
• RQ5 What is the impact of extended training, mask rate adjustments, and additional features on downstream tasks?

Key findings

• MPMv2 with a transformer decoder and reduced input cardinality substantially improves classification accuracy over MPMv1.
• Switching to MAE-style decoding and limiting positional encoding yields better performance and lower GPU memory usage.
• Adding impact-parameter features and particle IDs further boosts accuracy (examples: regression 62.2 to 80.4; k-means 70.2 to 83.0 after additions).
• Fully transformer-based decoder (MAE) achieves regression 79.2 and k-means 81.4 in ablations.
• Longer training, deeper decoder, and 40% mask rate yield best results: regression 83.3 and k-means 84.0.
• Pre-trained backbones outperform random initialization across in-distribution, weakly supervised, and out-of-distribution tasks, showing good generalization.

Figure 2 : A schematic overview of the SSFM model.