---
title: 'OmniLearned Framework: Foundation Model in Jet Physics'
url: https://www.emergentmind.com/topics/omnilearned-framework
type: topic
---

# OmniLearned Framework: Foundation Model in Jet Physics

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 [2510.24066] [2606.14870].

## 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 [2510.24066]. 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** [2510.24066]. 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 [2510.24066].

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 [2606.14870]. 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 [2510.24066]. 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 [2510.24066]. 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 [2510.24066].

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 [2510.24066]. 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 [2510.24066].

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 [2510.24066]. 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 [2606.14870].

## 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 [2510.24066]. 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 [2510.24066].

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 [2510.24066]. The label space includes **200 physics classes** and **10 sample/domain classes** [2510.24066]. 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 **$1\times10^{-5}$** to **$10^{-6}$**, and the **Lion** optimizer with $\beta_1=0.95$ and $\beta_2=0.98$ [2510.24066].

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 [2606.14870]. When all three are active, the total pre-training loss is
$$
\begin{aligned}
\mathcal{L}
&=
\mathcal{L}_{\mathrm{class}}
+
\mathcal{L}_{\mathrm{gen}}
+
\mathcal{L}_{\mathrm{MPM}} \\
&\quad+
\alpha(t)^2\,\mathcal{L}_{\mathrm{class}^{\mathrm{gen,\,perturb}}}
+
w_{\rm MPM}\,\mathcal{L}_{\mathrm{class}^{\mathrm{MPM,\,perturb}}},
\end{aligned}
$$
with $\alpha(t)\equiv 1-t$ and $w_{\rm MPM}=1-f_m=0.6$ [2606.14870]. 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 [2606.14870].

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 [2510.24066]. 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 **$688\pm 9$** at $\epsilon_S=0.5$, and **$3486\pm157$** at $\epsilon_S=0.3$; the corresponding **OmniLearned-m** value at $\epsilon_S=0.3$ is **$3208\pm176$**, and **OmniLearned-s** reaches **$2637\pm128$** [2510.24066]. 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** [2510.24066]. 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 [2510.24066]. 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 [2510.24066].

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 [2510.24066]. 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 [2510.24066]. 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 [2510.24066].

## 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-$Q^2$ simulated and real $pp$ and $ep$ 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** [2604.12364].

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 [2512.24422]. On CAMELS-SAM, OmniCosmos reports test-set $R^2$ values of **$0.87\pm0.01$** for $\Omega_m$ and **$0.92\pm0.01$** for $\sigma_8$, compared with **$0.83\pm0.02$** and **$0.89\pm0.01$** for the same architecture trained from scratch [2512.24422]. On Quijote, the corresponding values are **$0.849\pm0.003$** and **$0.871\pm0.003$**, 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 [2512.24422]. The paper describes this as the first demonstration that a collider-physics model generalizes across scientific fields [2512.24422].

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 [2601.10791]. 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%** [2601.10791]. 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 [2601.10791].

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 [2604.12364] [2512.24422].

## 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 [2412.06693] | “OmniLearned evaluation core” | Static Builder + Dynamic Data Flow |
| Capybara-OMNI [2504.12315] | General omni-modal training paradigm | Unified LLM backbone + plugin encoders + staged alignment |
| LessonL [2505.23946] | “OmniLearned”-style multi-agent system | Lesson solicitation, banking, selection, adjustment |
| OmniLV [2504.04903] | Omni-learned low-level vision system | Separate text/visual encoders + shallow feature control |
| RoboOmni [2510.23763] | End-to-end proactive robot framework | Perceiver–Thinker–Talker–Executor |
| Omni-AutoThink [2512.03783] | 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 [2412.06693]. **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” [2504.12315]. **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 [2505.23946].

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 [2504.04903]. **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 [2510.23763]. **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 [2512.03783].

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 [2606.14870]. 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 [2512.24422] [2601.10791]. In neutrino transfer, all models retain some regression bias at high $q_3$, and the transfer paper explicitly notes this as a remaining methodological issue [2604.12364]. 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.

Source: https://www.emergentmind.com/topics/omnilearned-framework