OmniLearned: A Foundation Model Framework for All Tasks Involving Jet Physics

This paper introduces OmniLearned, a major upgrade to the OmniLearn foundation model framework for jet physics that leverages a billion-jet training dataset and improved architecture to achieve state-of-the-art performance in top-quark tagging, b-tagging, and anomaly detection, thereby significantly enhancing the discovery potential of collider experiments.

The Gist

Imagine you are trying to teach a computer to recognize different types of clouds in the sky. You could show it a picture of a cumulus cloud and say, "This is a cumulus," then show it a stratus and say, "This is a stratus." But what if you wanted the computer to recognize every kind of cloud, predict how they move, and spot a strange, never-before-seen cloud formation that might signal a tornado?

To do that, you wouldn't just show it a few pictures. You'd need to show it billions of cloud photos, let it study the physics of how water vapor behaves, and let it learn the "essence" of what a cloud is before you ever ask it to identify a specific type.

This is exactly what the paper "Foundation Model Framework for All Tasks Involving Jet Physics" is about, but instead of clouds, it's about particle physics jets.

What is a "Jet"?

In particle colliders (like the Large Hadron Collider), scientists smash particles together at nearly the speed of light. When a heavy particle (like a top quark) breaks apart, it doesn't just vanish; it sprays out a shower of smaller particles. To our detectors, this spray looks like a single, messy blob of energy. Physicists call this a "jet."

For decades, scientists have tried to build computer programs to look at these messy blobs and say, "Ah, that's a top quark!" or "That's just a boring background particle." The problem is, there are so many different jobs to do with jets, and training a new computer program for every single job takes forever and requires massive amounts of data.

The Old Way vs. The New Way

The Old Way (The Specialist):
Imagine hiring a different chef for every meal. You hire a sushi chef for dinner, a pizza chef for lunch, and a baker for breakfast. Each chef is great at their specific task, but they have to start from scratch every time. If you want a new dish, you have to hire a new chef and train them from zero. This is how particle physics used to work: a new AI model for every new experiment.

The New Way (The Foundation Model):
This paper introduces OmniLearned, which is like hiring a Master Chef who has tasted every dish in the world.

1. The Training: They fed this Master Chef a massive library of over 1 billion different "jet" recipes (simulated data). The chef didn't just memorize the recipes; they learned the fundamental rules of cooking (the physics of how particles interact).
2. The Result: Now, this chef is a "Foundation Model." They have a deep, universal understanding of jets.
3. The Magic: When you need a specific dish (like identifying a top quark), you don't hire a new chef. You just give the Master Chef a quick, 5-minute briefing on what you want. Because they already know the basics, they can adapt instantly and perfectly.

What Makes OmniLearned Special?

The authors upgraded their previous model (called OmniLearn) with three major improvements:

1. A Bigger Brain (Architecture): They tweaked the computer's "brain" (the neural network) to pay better attention to the tiny details inside the jet, like how particles are arranged relative to each other. It's like upgrading from a magnifying glass to a high-powered microscope.
2. More Data (1 Billion Jets): The previous model was trained on 100 million examples. The new one saw 10 times more. In the world of AI, more data usually means a smarter, more reliable model.
3. Open Source (The Cookbook): They didn't just keep the secret recipe. They released the software, the data, and the instructions to the whole world. Anyone can now download this "Master Chef" and use it for their own experiments.

What Did They Prove?

To show off their new Master Chef, they tested it on three very different "kitchen tasks":

• Task 1: The Top Quark Hunt (The "Top Chef" Challenge)
  They asked the model to find top quarks in a standard test dataset. The OmniLearned model didn't just pass; it crushed the competition, beating every other existing AI model by a wide margin. It found the signal with much less "noise" (background junk).

• Task 2: The Flavor Tagging (The "Ingredient" Challenge)
  In real experiments (using data from the ATLAS detector), they asked the model to distinguish between jets made of "bottom" quarks, "charm" quarks, and regular light quarks. This is crucial for finding new physics. OmniLearned was significantly better at this than the current best methods used by the ATLAS collaboration, even though it had to learn a new "language" (the specific detector data) very quickly.

• Task 3: The Anomaly Detector (The "Alien Cloud" Challenge)
  This is the coolest part. They asked the model to find something it had never seen before. They used real data from the CMS experiment and asked the AI to find the "needle in the haystack."

  • The Trick: They didn't tell the AI what to look for. They just asked it to find anything that looked weird compared to the "normal" background.
  • The Result: The AI successfully found the top quarks hidden in the data, even though it wasn't explicitly told to look for them. It essentially said, "Hey, these blobs look different from the rest; let's investigate." This proves the model can be used for discovery, not just classification.

Why Does This Matter?

Think of the universe as a giant, complex puzzle. For years, physicists have been trying to solve it with small, specialized tools. OmniLearned is like handing them a universal key.

• Speed: It takes less time and computing power to solve new problems.
• Discovery: It can spot strange new particles that we didn't even know to look for.
• Collaboration: By making the code and data public, they are letting scientists all over the world use this powerful tool to accelerate their own discoveries.

In short, this paper says: "We built a super-smart AI that has studied the entire history of particle collisions. Now, instead of training a new AI for every new question, we can just ask this one master AI, and it will help us find the answers faster and better than ever before."

Technical Summary

1. Problem Statement

High-energy physics (HEP) experiments, particularly at colliders like the LHC, generate vast amounts of data involving hadronic jets (collimated sprays of particles resulting from quarks and gluons). Characterizing these jets is crucial for tasks such as identifying specific particles (e.g., top quarks, $b$-quarks) or detecting anomalies (new physics).

The primary challenges addressed in this paper are:

• Data Scarcity for Specific Tasks: Training state-of-the-art machine learning (ML) models for every specific jet physics task requires massive amounts of labeled data, which is often computationally expensive to generate via full simulations (e.g., Geant4).
• Fragmented Approaches: Previous foundation models in HEP relied heavily on tokenized, self-supervised learning (similar to Large Language Models), which often fails to fully utilize the continuous nature of jet data or available simulation labels.
• Generalization: Existing models often struggle to generalize across different detectors, collision systems, or tasks without extensive retraining from scratch.

2. Methodology: The OmniLearned Framework

The authors introduce OmniLearned, a major upgrade to their previous "OmniLearn" model. It is a foundation model designed to learn a universal representation of jets that can be fine-tuned for diverse downstream tasks.

A. Architecture Improvements (PET v2)

The backbone of the model is an upgraded Point Edge Transformer (PET v2), which treats jets as point clouds with variable particle counts. Key architectural modifications include:

• Input Features: Reduced to a minimal, physically meaningful set: $(\Delta\eta, \Delta\phi, \log p_T, \log E)$. Pre-processing standardization was removed as the log-scale inputs naturally center around zero.
• Handling Auxiliary Information: The model now natively supports Particle Identification (PID) and vertex information.
  • PID is embedded via a lookup table (similar to token embeddings in LLMs).
  • Vertex information is encoded via MLPs.
  • These are added to kinematic features; if absent, the corresponding encoding is zeroed out, allowing the model to handle datasets with or without these features seamlessly.
• Local Attention: Replaced simple averaging with learnable weighted attention over $k$-nearest neighbors ($k=10$). The interaction terms include physics-inspired features like invariant mass and $\Delta R$ between particles.
• Global Attention: Modified the attention matrix to include physics-inspired bias terms (similar to local interactions) and replaced standard layer normalization with a learnable hyperbolic tangent operation for improved stability in deep networks.
• Task Heads: Instead of a single classification token, the model uses five learnable tokens added to the particle set. These tokens summarize different aspects of the jet, acting similarly to multi-head attention mechanisms.

B. Training Strategy & Loss Function

• Dataset Scale: The model is pre-trained on over 1 billion jets (10x the size of the previous JetClass dataset), combining data from JetClass, JetClass2, Aspen Open Jets, ATLAS Top Tagging, H1 DIS, and CMS QCD/BSM simulations.
• Multi-Task Learning: The model is trained jointly on:
  1. Jet Classification: Distinguishing between 210 classes (200 jet flavors + 10 dataset identifiers).
  2. Jet Generation: Using a Flow Matching objective (an alternative to Diffusion models) to generate jet constituents. The loss minimizes the error in predicting the velocity field $u = \epsilon - x$.
  3. Sample Classification: A specific task to classify which dataset a jet comes from, helping the model learn dataset-invariant features.
• Loss Function:
  $$L = CE(y, y_{pred}) + \|u - u_{pred}\|^2 + \alpha^2 CE(y, \hat{y}_{pred})$$
  Where the terms represent supervised classification, flow matching generation, and classification on perturbed (noisy) inputs, respectively.

C. Software & Reproducibility

The authors released a unified software package that automatically downloads, formats, and provides access to all training datasets and pre-trained checkpoints (Small, Medium, Large models).

3. Key Contributions

1. Scale: The first foundation model for jet physics trained on >1 billion jets, leveraging the scaling laws of transformers.
2. Architecture: Introduction of PET v2, which incorporates physics-informed interaction terms, handles missing auxiliary data (PID/vertex) gracefully, and utilizes flow matching for high-fidelity generation.
3. Unified Framework: A single pre-trained model capable of addressing classification, generation, and anomaly detection across different detectors (ATLAS, CMS, H1) and collision systems (pp, ep).
4. Open Science: Full release of code, data pipelines, and pre-trained weights (3M, 58M, and 423M parameter models).

4. Results

The paper evaluates OmniLearned on three representative tasks:

A. Jet Classification (Top Tagging & Quark/Gluon)

• Benchmark: Tested on community datasets (Delphes-based) for top quark tagging and quark/gluon separation.
• Performance: OmniLearned (especially the Medium and Large variants) surpassed all state-of-the-art models (including ParticleNet, ParT, and L-GATr) trained from scratch.
• Fine-tuning: Fine-tuning OmniLearned required significantly fewer epochs and computational resources than training from scratch while achieving higher accuracy (e.g., AUC improvements and better background rejection at fixed signal efficiency).

B. Flavor Tagging (ATLAS $b$-tagging)

• Task: Distinguishing $b$-jets, $c$-jets, light jets, and $\tau$-jets using full ATLAS simulation data.
• Innovation: The authors repurposed the generation head of the foundation model to perform track classification (an auxiliary task), demonstrating the flexibility of the architecture.
• Performance: OmniLearned outperformed the ATLAS Collaboration's current state-of-the-art GN2 architecture.
  • $b$-tagging: Improved light-jet rejection by >50% and $\tau$-jet rejection by a factor of 2 compared to GN2.
  • Efficiency: Fine-tuning was 50% faster than training a medium model from scratch.

C. Anomaly Detection (CMS Open Data)

• Task: Rediscovering the top quark (a known signal) in CMS Open Data as a proxy for discovering new physics. The signal constitutes only ~0.1% of the data.
• Strategy 1 (Generative): Used OmniLearned to model the Standard Model background in sideband regions and generate background samples in the signal region. A classifier then distinguished data from the generated background.
  • Result: Successfully identified the top quark signal with high significance (up to ~15$\sigma$), whereas models trained from scratch failed to detect the signal.
• Strategy 2 (Direct): Used the pre-trained class probabilities directly (without fine-tuning) by comparing the probability of "3-prong" decays (generic top-like topology) against QCD.
  • Result: Also successfully identified the anomaly, proving the model's ability to generalize without task-specific retraining.

5. Significance and Outlook

• Paradigm Shift: OmniLearned demonstrates that foundation models can replace the "train-from-scratch" paradigm in HEP, drastically reducing the computational cost of developing new analysis tools.
• Discovery Potential: By improving background rejection and anomaly detection sensitivity, the framework enhances the discovery potential of past, current, and future collider experiments.
• Generalizability: The success in applying a model trained on simulated data to real experimental data (CMS Open Data) and across different detectors (ATLAS, H1) suggests that the learned representations capture fundamental physical properties of jets rather than detector-specific artifacts.
• Future Work: The authors suggest the methodology could extend beyond jets to full event topologies and different collision systems, potentially revolutionizing how particle physics data is analyzed.

In summary, OmniLearned establishes a new standard for jet physics analysis by combining massive-scale pre-training, physics-informed neural architectures, and a versatile framework capable of handling classification, generation, and anomaly detection with superior performance and efficiency.