OmniLearned Framework: Foundation Model in Jet Physics

• OmniLearned is a foundation-model framework defined for jet physics, combining a PET v2 backbone and multi-dataset pre-training to deliver state-of-the-art results.
• The framework employs advanced pre-training with joint classification and generative objectives using over one billion jet samples and a robust software ecosystem.
• OmniLearned demonstrates versatility by achieving high performance in top tagging, anomaly detection, and cross-domain transfers to neutrino, cosmology, and molecular dynamics.

OmniLearned is a foundation-model framework for jet physics in which a single, very large pretrained neural network, together with a software and data ecosystem, learns a reusable representation of jets and is then adapted to downstream tasks such as tagging, generative modeling, and anomaly detection. Introduced as a major upgrade to OmniLearn, the framework combines a Point Edge Transformer v2 (PET v2), joint classification and flow-matching objectives, training on over one billion jets, and well-documented software for accessing datasets and models; later studies reuse the same framework to compare pre-training objectives systematically and to transfer learned representations from collider jets to neutrino interactions, cosmology, and molecular dynamics (Bhimji et al., 28 Oct 2025, Elsharkawy et al., 12 Jun 2026).

1. Scope, lineage, and conceptual position

OmniLearned is presented as a foundation-model framework for all tasks involving jet physics. In that formulation, the central object is not a single benchmark model but a reusable infrastructure: a pretrained backbone, a family of task heads, a multi-dataset training corpus, and software for downstream adaptation (Bhimji et al., 28 Oct 2025). It directly extends OmniLearn, preserving the earlier emphasis on physically motivated, permutation-invariant architectures operating on jet constituents, while upgrading the model architecture, the training corpus, and the accompanying software stack.

The framework is explicitly organized around three new elements relative to OmniLearn: updates to the model architecture and training, using over one billion jets used for training, and providing well-documented software for accessing all datasets and models (Bhimji et al., 28 Oct 2025). Its reference downstream demonstrations are top-quark jet tagging with a community Delphes-based benchmark dataset, b-tagging with ATLAS full simulation, and anomaly detection with CMS experimental data, with the paper reporting state-of-the-art performance in each case (Bhimji et al., 28 Oct 2025).

A later study recasts the same line of work as the OmniLearned High Energy Physics FM framework, emphasizing a shared point-edge-transformer body, modular task heads, and a standardized fine-tuning protocol for downstream jet classification and generation (Elsharkawy et al., 12 Jun 2026). In that formulation, OmniLearned becomes both an operational model family and an experimental platform for controlled studies of pre-training objectives, scaling, and transfer in simulation-based science.

2. PET v2 architecture and jet representation

OmniLearned represents jets as point clouds of constituents. In the 2025 collider framework, up to 150 particles are retained per jet, each described by the four kinematic features $\Delta\eta$, $\Delta\phi$, $\log p_T$, and $\log E$, with optional auxiliary per-particle information such as PID and vertex information when available (Bhimji et al., 28 Oct 2025). The architecture is designed to respect the unordered-set structure of constituents and to capture both local and global particle-particle correlations.

The backbone is PET v2, a Point Edge Transformer that combines local attention layers and global transformer layers. Its local message construction is explicitly physics-aware. For a particle $x_i$ and its neighbor $x_j$, OmniLearned uses

$$f(x_i, x_j) = \big[ x_i - x_j,\, \log m(x_i,x_j),\, \log \Delta R(x_i,x_j),\, \log(\min\{p_{T i},p_{T j}\}\,\Delta R(x_i,x_j)) \big],$$

where $m(x_i,x_j)$ is the invariant mass of the summed four-vectors and $\Delta R(x_i,x_j)$ is their angular separation (Bhimji et al., 28 Oct 2025). Local aggregation is performed with a local transformer block over the $k=10$ nearest neighbors rather than a simple average, while global multi-head attention is augmented with a bias derived from the same pairwise physics-inspired features (Bhimji et al., 28 Oct 2025).

Input encoding is modular. PID is handled by a lookup-table embedding, vertex information by a small MLP, and absent auxiliary channels are set to zero so that the same pretrained weights can be used across heterogeneous datasets (Bhimji et al., 28 Oct 2025). For generative tasks, the time variable is encoded with Fourier features followed by an MLP and appended as an additional point in the point cloud, rather than being added uniformly to every constituent embedding (Bhimji et al., 28 Oct 2025).

The model scales from Small to Large. The released configurations are Small with 8 blocks, 8 heads, latent dimension 128, and about 3M parameters; Medium with 12 blocks, 16 heads, latent dimension 512, and about 58M parameters; and Large with 28 blocks, 32 heads, latent dimension 1024, and about 423M parameters (Bhimji et al., 28 Oct 2025). The later controlled-study variant formalizes the same family as a shared PET body with pluggable classifier, flow-matching generator, and masked particle modeling heads, making the architecture explicitly multi-objective (Elsharkawy et al., 12 Jun 2026).

3. Pre-training objectives, data ecosystem, and software

OmniLearned pre-training is deliberately multi-task. In the 2025 framework, the shared backbone is trained jointly for jet flavor classification, sample or domain classification, and flow-matching generative modeling, so that the representation must simultaneously encode discriminative jet structure and support particle-level generation (Bhimji et al., 28 Oct 2025). The sample-classification branch is particularly important because it allows the framework to incorporate jets without truth flavor labels, including real CMS data, while still forcing the latent space to model experiment- and sample-level structure (Bhimji et al., 28 Oct 2025).

The pre-training corpus comprises 1,057.7 million training jets, 101.8 million validation jets, and 67.6 million test jets drawn from JetClass, JetClass2, Aspen Open Jets, ATLAS Top Tagging, H1 DIS, CMS QCD, and CMS BSM samples (Bhimji et al., 28 Oct 2025). The label space includes 200 physics classes and 10 sample/domain classes (Bhimji et al., 28 Oct 2025). Training uses the Perlmutter supercomputer, 32–512 GPUs, a global batch size of 4096, 3 full passes over the pre-training set, a cosine learning-rate schedule from $\Delta\phi$0 to $\Delta\phi$1, and the Lion optimizer with $\Delta\phi$2 and $\Delta\phi$3 (Bhimji et al., 28 Oct 2025).

A later study extends this setup into a controlled comparison of three pre-training heads: a classifier head, a flow-matching generator head, and a masked particle modeling (MPM) head (Elsharkawy et al., 12 Jun 2026). When all three are active, the total pre-training loss is

$\Delta\phi$4

with $\Delta\phi$5 and $\Delta\phi$6 (Elsharkawy et al., 12 Jun 2026). That study reports a sharp task dependence of objective choice: pure classifier pre-training is optimal when downstream labels and model capacity are plentiful, Classifier+MPM is uniquely powerful in the low-finetuning label regime, and flow matching must be in the pre-training objective to see a significant finetuning advantage on downstream generation; for a model to transfer well to both classification and generation, it must be pre-trained on both (Elsharkawy et al., 12 Jun 2026).

The software stack is part of the framework rather than an afterthought. OmniLearned provides a unified PyTorch codebase, automatic download and conversion scripts for the major datasets, pretrained checkpoints for multiple model sizes, and task-specific scripts for classification, generation, and anomaly detection (Bhimji et al., 28 Oct 2025). In that sense, “framework” refers as much to reproducible access and cross-dataset interoperability as to model architecture.

4. Reference collider applications

The framework’s canonical demonstrations are top tagging, flavor tagging, and anomaly detection. On the community Delphes-based top-tagging benchmark, OmniLearned-l achieves Acc 0.944, AUC 0.9880, background rejection $\Delta\phi$7 at $\Delta\phi$8, and $\Delta\phi$9 at $\log p_T$0; the corresponding OmniLearned-m value at $\log p_T$1 is $\log p_T$2, and OmniLearned-s reaches $\log p_T$3 (Bhimji et al., 28 Oct 2025). These numbers exceed the earlier OmniLearn and the strongest non-foundation competitors reported in the same benchmark.

On the ATLAS public flavor-tagging dataset, OmniLearned is adapted to track-level inputs and reuses the generative head as a per-track auxiliary classifier. For b-tagging at 70% b-jet efficiency, OmniLearned-m reports rejection 66.5 against c-jets and 1853 against light jets, compared with 45.5 and 1097 for GN2 (Bhimji et al., 28 Oct 2025). For c-tagging at 30% c-jet efficiency, OmniLearned-m reports rejection 24.8 against b-jets and 235 against light jets, compared with 21.1 and 166 for GN2 (Bhimji et al., 28 Oct 2025). The same study also notes that larger OmniLearned models converge in less than half as many updates as training PET v2 from scratch on the same task (Bhimji et al., 28 Oct 2025).

The anomaly-detection application uses CMS Open Data in a CATHODE-style setup. A sideband-trained mass model and a jet-level conditioning model provide the conditioning variables, and OmniLearned then generates particle-level QCD background jets in the signal region; a classifier trained to distinguish real signal-region data from generated background defines the anomaly score (Bhimji et al., 28 Oct 2025). In this setting, pretrained OmniLearned models show clear excesses, with significances above discovery threshold for appropriate score thresholds, whereas corresponding models trained from scratch show only marginal sensitivity (Bhimji et al., 28 Oct 2025). The paper also introduces a more direct strategy: using ratios of pretrained class scores, such as a generic 3-prong versus QCD score, as a “zero-shot” anomaly observable. For the boosted-top signal considered there, this direct use of pretrained classes yields a higher significance than the CATHODE-style method with fine-tuning (Bhimji et al., 28 Oct 2025).

5. Cross-domain transfer and scientific generalization

A distinctive feature of OmniLearned is that it has been used as a vehicle for transfer across scientific domains. In neutrino interactions, a PET2-based OmniLearned model pretrained on diverse high-$\log p_T$4 simulated and real $\log p_T$5 and $\log p_T$6 collisions is transferred to MINERvA neutrino–nucleus scattering for available-energy regression and binary classification of charged-current pion final states. The transfer paper reports that pretrained OmniLearned models consistently outperform similarly sized models trained from scratch, achieving better overall performance at the same compute budget and at the same number of training steps, and interprets this as evidence for detector-agnostic inference (Krzmanc et al., 14 Apr 2026).

In cosmology, OmniCosmos adapts the small OmniLearned PET-v2 model to halo point clouds from CAMELS-SAM and Quijote by replacing collider-specific local pairwise features with geometric features in 3D position space and reusing the backbone and pooling heads (Mikuni et al., 30 Dec 2025). On CAMELS-SAM, OmniCosmos reports test-set $\log p_T$7 values of $\log p_T$8 for $\log p_T$9 and $\log E$0 for $\log E$1, compared with $\log E$2 and $\log E$3 for the same architecture trained from scratch (Mikuni et al., 30 Dec 2025). On Quijote, the corresponding values are $\log E$4 and $\log E$5, again above the scratch baseline, and the paper states that benchmark-level performance can be reached with less than 10% of the available simulations for some tasks (Mikuni et al., 30 Dec 2025). The paper describes this as the first demonstration that a collider-physics model generalizes across scientific fields (Mikuni et al., 30 Dec 2025).

In molecular dynamics, OmniMol adapts OmniLearned to small-molecule machine-learned interatomic potentials by preserving the PET backbone and the interaction-matrix attention bias, while changing the encoders, pairwise interaction features, and output heads to molecular forms (Elsharkawy et al., 15 Jan 2026). The paper states that OmniLearned was pretrained on a diverse set of one billion particle jets, and reports strong low-data benefits from this initialization. On 100k oMol molecules, OmniMol-s-d-pt improves over OmniMol-s-d by +12.3% for energy and +19.5% for forces, while OmniMol-m-d-pt improves over OmniMol-m-d by +29.4% and +26.9% (Elsharkawy et al., 15 Jan 2026). After only two passes over the 4M-molecule dataset, the medium model shows a +54.6% and +56.9% advantage in energy and force MAE relative to its non-pretrained counterpart (Elsharkawy et al., 15 Jan 2026).

Taken together, these transfers support a specific interpretation of the framework: OmniLearned appears to learn point-cloud inductive biases that remain useful across large differences in energy scale, detector technology, underlying physics processes, and even scientific discipline. That interpretation is explicit in the neutrino and cosmology transfer papers, and a plausible implication is that the framework’s most portable components are its set-based representation, local-plus-global interaction structure, and pairwise physics-aware attention biases (Krzmanc et al., 14 Apr 2026, Mikuni et al., 30 Dec 2025).

6. Broader usage as a general systems blueprint

Several adjacent papers use “OmniLearned” not as the name of the PET-based jet model itself, but as an explicit blueprint, archetype, or evaluation core for unified systems. This broader usage is interpretive rather than nominal, but it is recurrent across the cited literature.

Work	How the OmniLearned idea is used	Core mechanism
OmniEvalKit (Zhang et al., 2024)	“OmniLearned evaluation core”	Static Builder + Dynamic Data Flow
Capybara-OMNI (Ji et al., 10 Apr 2025)	General omni-modal training paradigm	Unified LLM backbone + plugin encoders + staged alignment
LessonL (Liu et al., 29 May 2025)	“OmniLearned”-style multi-agent system	Lesson solicitation, banking, selection, adjustment
OmniLV (Pu et al., 7 Apr 2025)	Omni-learned low-level vision system	Separate text/visual encoders + shallow feature control
RoboOmni (Wang et al., 27 Oct 2025)	End-to-end proactive robot framework	Perceiver–Thinker–Talker–Executor
Omni-AutoThink (Yang et al., 3 Dec 2025)	Adaptive reasoning framework for omni models	Adaptive SFT + Adaptive GRPO

Across these works, the common abstractions are unusually consistent. OmniEvalKit defines a unified, ultra-lightweight evaluation architecture organized as Static Builder + Dynamic Data Flow, explicitly presented as an “OmniLearned evaluation core” for multilingual, multidomain, and multimodal assessment (Zhang et al., 2024). Capybara-OMNI describes an omni-modal MLLM built from a unified language backbone, plugin encoders for image, video, and audio, and progressive staged training; the paper explicitly states that this methodology can serve as a general “OmniLearned Framework” (Ji et al., 10 Apr 2025). LessonL is described as an archetypal OmniLearned-style system in which multiple agents generate, store, select, and update compact “lessons” rather than weights, using a lesson solicitation–banking–selection mechanism (Liu et al., 29 May 2025).

The same pattern appears in other modalities. OmniLV presents a universal multimodal multi-task framework for low-level vision trained on about 40 million pairs covering over 100 sub-tasks, and its design principles—separate encoders, projection-addition for visual prompting, and shallow feature control—are explicitly summarized as guidance for an omni-learned low-level framework (Pu et al., 7 Apr 2025). RoboOmni treats vision, speech, environmental audio, language, and action as tokens in a joint sequence within a Perceiver–Thinker–Talker–Executor architecture, and is presented as a prototype of an OmniLearned framework for robotics (Wang et al., 27 Oct 2025). Omni-AutoThink likewise frames adaptive multimodal reasoning as a general OmniLearned training scheme based on Adaptive SFT and Adaptive GRPO, with a single unified policy over text-only, text-audio, text-visual, and text-audio-visual inputs (Yang et al., 3 Dec 2025).

This broader usage suggests an abstract definition that is wider than jet physics: an OmniLearned framework is a system with a shared representation core, task- or modality-specific adapters or heads, staged learning or adaptation, and a unified evaluation or control interface. The suggestion is consistent across the multimodal, robotic, evaluation, and multi-agent papers, although only the jet-physics lineage uses OmniLearned as the proper name of the model family.

Open questions remain. Within jet physics, later work shows that classification and generation are partially orthogonal objectives, so downstream coverage depends on explicitly including both during pre-training (Elsharkawy et al., 12 Jun 2026). In cross-domain transfer, the advantage over scratch can shrink as downstream data increase, as reported for Quijote in cosmology and for large-scale molecular training (Mikuni et al., 30 Dec 2025, Elsharkawy et al., 15 Jan 2026). In neutrino transfer, all models retain some regression bias at high $\log E$6, and the transfer paper explicitly notes this as a remaining methodological issue (Krzmanc et al., 14 Apr 2026). These results place OmniLearned in a characteristic foundation-model regime: broad reuse is real, but objective design, domain shift, and scaling behavior remain central research variables rather than solved engineering details.