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

Benjamin Nachman, Wahid Bhimji, Ibrahim Elsharkawy, Vinicius Mikuni, Joschka Birk, Gregor Kasieczka

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 12 functions out of this paper's own repositories and ran 7 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

Repositories linked to this paper

• ibrahimEls/PretrainingForScience

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

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 selfsupervision 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 selfsupervised masked particle modeling. All models are pre-trained on the JetClass dataset and finetuned 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 selfsupervised 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.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2606.14870")
get_code_for_paper("2606.14870")
have("2606.14870")

Connect an agent — have() is free.