A Method to Simultaneously Facilitate All Jet Physics Tasks

The paper introduces "OmniLearn," a jet-physics foundation model that leverages a representation learned from a specific multiclass generation and classification task to simultaneously enhance the accuracy, precision, and speed of diverse downstream jet physics tasks across different datasets, collision systems, and analysis goals.

The Gist

Imagine you are trying to learn how to drive.

In the past, if you wanted to learn how to drive a sports car, you would take a specific driving school course just for sports cars. If you later wanted to learn how to drive a truck, you'd have to go back to school and take a completely new, separate course. If you wanted to learn how to drive in the rain or on ice, you'd need yet another specialized class.

This is how particle physics has worked for a long time. Physicists study "jets" (which are sprays of tiny particles created when protons smash together at high speeds, like in the Large Hadron Collider). Because these jets are incredibly complex—made of hundreds of particles moving in chaotic patterns—physicists have built separate, specialized computer programs (models) for every single job:

• One program to tell if a jet came from a top quark.
• Another to tell if it came from a gluon.
• Another to simulate what a jet should look like.
• Another to find "weird" jets that might be new physics.

Each program is trained from scratch, like a student starting with a blank notebook every time. This is slow, expensive, and inefficient.

The Big Idea: The "Universal Driving School"

The authors of this paper, Vinicius Mikuni and Benjamin Nachman, asked a simple question: "Can we build one 'Super-Student' that learns to drive everything at once, so it can help us with any driving task later?"

They created a system called OmniLearn. Think of OmniLearn not as a single driver, but as a universal foundation model.

Here is how it works, using simple analogies:

1. The "Swiss Army Knife" Brain

Instead of training a model just to classify jets (like a sports car driver), they trained OmniLearn to do two things simultaneously:

1. Classify: Look at a jet and guess what kind of particle made it (Top quark? Gluon?).
2. Generate: Look at a description of a jet and try to draw (simulate) a new one from scratch.

By forcing the computer to learn how to create jets and identify them at the same time, the model's "brain" learns the deep, fundamental rules of how jets work. It's like a student who learns the physics of engines and the rules of the road at the same time. They understand the essence of driving, not just the specific steps for one car.

2. The "Pre-Trained" Advantage

Once OmniLearn is trained on a massive dataset (100 million jets!), it has learned a "general representation" of jets. It knows what a jet "feels" like.

Now, imagine you want to solve a new problem:

• Task A: You need to analyze data from a different type of particle collider (like electron-proton collisions instead of proton-proton).
• Task B: You need to find a rare, weird jet that looks like a new particle (Anomaly Detection).
• Task C: You need to correct a simulation to match real-world data.

In the old way, you would start from zero. With OmniLearn, you just take the "pre-trained brain," give it a tiny bit of new information (fine-tuning), and it instantly becomes an expert at the new task. It's like taking a master chef who knows how to cook French cuisine and asking them to cook Italian. They don't need to learn what a stove is or how to chop onions; they just need to learn the specific Italian recipes.

What Did They Prove?

The paper shows that OmniLearn is a "Foundation Model" for jet physics. They tested it in many different scenarios, and it consistently won:

• Speed: It learned new tasks 3 to 3.5 times faster than starting from scratch. It's like the difference between a student who has to learn the alphabet before reading a book versus one who already knows how to read.
• Accuracy: It was often more accurate than the best specialized models, even when the new data looked very different from the training data (e.g., different detectors or collision types).
• Versatility: It worked for:
  • Tagging: Identifying what a jet is.
  • Generation: Creating fake jets for simulations.
  • Reweighting: Fixing simulations to match reality.
  • Anomaly Detection: Finding the "needle in the haystack" (new physics).

The "OmniLearn" Metaphor

Think of OmniLearn as a universal translator for the language of the universe.

• Before, if you wanted to translate a book from French to German, you hired a French-German translator. If you wanted to translate from French to Japanese, you hired a different person.
• OmniLearn is like a person who has mastered the structure of all languages. They can translate French to anything instantly because they understand the underlying grammar of communication, not just the specific words.

Why Does This Matter?

In particle physics, data is messy, simulations are slow, and new discoveries are rare.

• Old Way: "Let's spend 6 months training a new AI to find this specific new particle."
• OmniLearn Way: "Let's take our universal AI, give it a quick 2-day refresher course, and it's ready to go."

This approach saves massive amounts of computing power and time. It allows physicists to ask more questions and get answers faster. The authors have made their code public, meaning any physicist can now use this "universal jet brain" to accelerate their own research.

In short: They built a "super-model" that learned the general rules of particle jets so well that it can now help solve almost any problem in jet physics, faster and better than ever before.

Technical Summary

1. Problem Statement

Jet physics in high-energy collider experiments involves analyzing complex, high-dimensional data structures (jets composed of many particles). While deep learning has revolutionized this field, current innovations typically proceed in parallel, with separate models trained for specific tasks (e.g., top-tagging, quark/gluon discrimination, anomaly detection, or generative modeling).

• The Challenge: There is no unified approach that leverages a single learned representation to improve all downstream jet physics tasks simultaneously.
• The Gap: Existing "foundation models" in particle physics often rely on self-supervised learning (e.g., masking data) which may not align with actual analysis goals. Furthermore, transfer learning has been applied to single downstream tasks but not broadly across the diverse landscape of jet physics (different collision systems, detectors, and task types).

2. Methodology: OmniLearn

The authors propose OmniLearn, a "jet-physics foundation model" designed to learn a general, transferable representation of jets that can accelerate and improve accuracy across classification, generation, and other physics tasks.

Core Architecture: Point-Edge Transformer (PET)

The backbone of OmniLearn is a Point-Edge Transformer (PET), which treats jets as unordered point clouds of particles.

• Input Representation: Jets are represented as sets of particles with kinematic features (momentum, energy) and particle identification (PID).
• Feature Handling: To handle datasets with varying feature availability (e.g., some lack PID or vertex info), the model employs Feature Drop (randomly zeroing out PID features with $p=0.2$) during training, forcing the network to learn robust representations regardless of missing inputs.
• Local Geometry: The model uses Dynamic Graph Convolution (DGCNN) to create local edge features based on $k$-nearest neighbors in $(\eta, \phi)$ space, capturing local geometric correlations before passing data to the transformer.
• Transformer Blocks: The core uses multi-head attention with LayerScale to improve training stability and convergence.
• Time Conditioning: Since the model uses a diffusion process, it is conditioned on a time parameter $t$. The time embedding is designed to be zero at $t=0$ (clean data) to ensure the classifier head works effectively on unperturbed data.

Dual-Head Strategy

OmniLearn utilizes a shared PET body with two distinct heads:

1. Classifier Head:
   • Takes particle embeddings and global jet kinematics (mass, $p_T$, $\eta$, multiplicity).
   • Uses a trainable Class Token (similar to CLS tokens in NLP) to summarize particle information for classification.
   • Outputs class probabilities.
2. Generator Head:
   • Based on Diffusion Models. It learns to reverse a perturbation process to generate particle distributions.
   • Uses a Diffusion Token to condition the generation on jet type (class labels) and time.
   • Employs Layer Drop (randomly zeroing class labels with $p=0.1$) to encourage the model to learn both general and specialized representations, mimicking classifier-free guidance.

Training Strategy

• Dataset: Trained on the JetClass dataset (100 million jets) containing 10 different jet classes (e.g., top, W, Z, gluon, quark) simulated with MADGRAPH5, Pythia8, and Delphes.
• Loss Function: A composite loss combining:
  • Cross-Entropy ($L_{class}$): For classification accuracy.
  • Velocity Prediction ($L_{gen}$): For the diffusion generation task (predicting the score function).
  • Smear Loss ($L_{class\_smear}$): Cross-entropy calculated on perturbed (noisy) inputs to improve robustness.
• Scale: The model has ~1.3 million trainable parameters in the PET body, with additional parameters in the heads. It is trained on 100M jets using 128 GPUs.

3. Key Contributions

1. Unified Foundation Model: Demonstrates that a single model trained on a specific multi-class generation and classification task can serve as a foundation for all major jet physics tasks.
2. Supervised Pre-training: Unlike many foundation models that use self-supervision, OmniLearn leverages the unique advantage of particle physics: ab initio simulations to create large, labeled datasets for supervised pre-training.
3. Cross-Domain Generalization: Proves the model's ability to transfer knowledge across:
   • Jet Types: Top tagging, quark/gluon discrimination.
   • Detectors: From fast simulation (Delphes) to full detector simulation (CMS Open Data/Geant4).
   • Collision Systems: From LHC ($pp$) to HERA ($ep$) deep inelastic scattering.
   • Tasks: Classification, conditional generation, likelihood ratio estimation, and anomaly detection.

4. Results

The paper evaluates OmniLearn across nine different datasets and tasks, comparing it against models trained from scratch and state-of-the-art baselines (e.g., ParticleNet, ParT, FPCD).

• Classification Performance:

  • Top Tagging: OmniLearn achieves an AUC of 0.9872, surpassing the fine-tuned ParT model (0.9877 is slightly higher but OmniLearn matches it closely) and significantly outperforming models trained from scratch.
  • Quark/Gluon: Achieves an AUC of 0.9159, outperforming the scratch-trained PET classifier and matching fine-tuned baselines.
  • Training Efficiency: OmniLearn converges 3x faster than training from scratch, reaching optimal performance in just 3 epochs compared to 10+ for scratch models.

• Generalization Across Detectors (CMS Open Data):

  • When applied to real CMS Open Data (full Geant4 simulation), OmniLearn outperforms a scratch-trained classifier in all metrics (AUC, Accuracy, Background Rejection) and converges 2x faster.

• Generalization Across Collision Systems (HERA DIS):

  • Applied to electron-proton collision data (H1 experiment), a domain vastly different from the LHC training data.
  • While the task is difficult (low AUC ~0.57 due to subtle differences between simulations), OmniLearn converges 3.5x faster than training from scratch, highlighting its data efficiency.

• Generative Tasks:

  • On the JetNet dataset (30 and 150 particles), OmniLearn matches or improves upon the performance of specialized generative models (FPCD, EPiC-GAN) in metrics like Wasserstein distance and coverage.
  • It requires 20–30% fewer training epochs to converge.

• Likelihood Ratio Estimation & Anomaly Detection:

  • Reweighting: OmniLearn successfully reweights high-dimensional distributions (OmniFold task), showing better agreement with "data" (Herwig) than scratch models.
  • Anomaly Detection (CWoLa): In the LHC Olympics resonant anomaly detection challenge, OmniLearn detects signals with an initial significance of $S/\sqrt{B} \sim 2$ (600 injected events), whereas previous methods required $S/\sqrt{B} \sim 5$ (1500 events). This represents a significant sensitivity gain in low-data regimes.

5. Significance

• Paradigm Shift: OmniLearn challenges the notion that jet physics requires a "one-task-one-model" approach. It establishes that a single, well-trained foundation model can act as a universal feature extractor for the field.
• Data Efficiency: The model is particularly effective in low-data regimes (e.g., anomaly detection with limited signal events or new detector configurations), where it outperforms large models trained from scratch.
• Practical Utility: By reducing training time and improving accuracy across diverse tasks, OmniLearn lowers the barrier for high-precision analyses. It is particularly valuable for future facilities (like the EIC) where data may be scarce or simulation costs are high.
• Open Science: The authors have made the code and trained model publicly available, fostering reproducibility and further development in the particle physics community.

In conclusion, OmniLearn demonstrates that by leveraging large-scale supervised pre-training on simulated data, a single neural network can learn a robust, general representation of jet physics that accelerates and enhances virtually every downstream analysis task, from classification to anomaly detection.