Pre-Training for Simulation-Based Science: A Study on Jet Foundation Model Training Objectives

2026-06-12

AI Summary

This study investigates the design of pretraining objectives for foundation models in simulation-based scientific domains such as high-energy physics, aiming to effectively support both classification and generation downstream tasks. Leveraging the JetClass dataset, the authors systematically evaluate three pretraining strategies—supervised classification, flow matching for generation, and self-supervised masked particle modeling—and fine-tune the resulting models on top-quark jet classification and JetNet conditional generation tasks. The findings reveal that classification and generation objectives are largely orthogonal, making it difficult for a single pretraining objective to excel at both. Supervised classification achieves optimal performance when labels are abundant, whereas masked particle modeling substantially improves results in low-label regimes. Flow matching benefits generation tasks only when included during pretraining. Built upon the OmniLearned framework, this work demonstrates that jointly optimizing multiple pretraining objectives is key to enabling effective cross-task transfer.

Abstract

Foundation models (FMs) trained on large datasets and fine-tuned on downstream tasks have emerged as a powerful paradigm in AI for science. Industrial FMs are typically trained using self-supervision with masking due to the lack of labels. In many scientific domains, accurate simulations are plentiful and facilitate large, labeled datasets. This opens up new possibilities for pre-training. We present a systematic comparison of pre-training methods using the OmniLearned High Energy Physics FM framework. We test supervised classification, flow-matching generation, and self-supervised masked particle modeling. All models are pre-trained on the JetClass dataset and fine-tuned on two representative downstream tasks, top jet classification and JetNet conditional generation. Among other observations, for classification tasks, we find that pure classifier pre-training is optimal when downstream labels and model capacity are plentiful, but combining it with self-supervised masked particle modeling (MPM) is uniquely powerful in the low-finetuning label regime. Flow matching-based generative pre-training seems to provide little benefit for downstream classification, and interestingly, for downstream generation, we find that flow matching must be in the pre-training objective to see a significant finetuning advantage, hinting at the orthogonality of classification and generation tasks. That is, for a model to transfer to both generative and classification downstream tasks, it must be pre-trained on both. This study provides a template for controlled scaling analysis of pre-training objectives for foundation models in simulation-based sciences.

Problem

• foundation models
• pre-training objectives
• simulation-based science
• jet physics
• transfer learning

Innovation

• foundation models
• pre-training objectives
• masked particle modeling
• flow matching
• simulation-based science