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

Dieses Paper stellt OmniLearned vor, ein verbessertes Framework für ein Fundamentmodell in der Jet-Physik, das durch eine aktualisierte Architektur, das Training mit über einer Milliarde Jets und offene Software den State-of-the-Art bei Aufgaben wie Top-Quark-Tagging, b-Tagging und Anomalieerkennung erreicht und damit das Entdeckungspotenzial von Teilchenbeschleuniger-Experimenten erweitert.

Das Wesentliche

Titel: Der „OmniLearned"-Superheld für die Teilchenphysik

Stellen Sie sich vor, Sie sind ein Detektiv in einer riesigen, chaotischen Bibliothek. Diese Bibliothek ist ein Teilchenbeschleuniger (wie der LHC am CERN). In dieser Bibliothek werden Billionen von Büchern (Teilchenkollisionen) pro Sekunde produziert. Die meisten dieser Bücher sind langweilige Standardgeschichten (das sogenannte „Hintergrundrauschen" der Natur). Aber manchmal, ganz selten, taucht ein Buch auf, das eine völlig neue, aufregende Geschichte erzählt – vielleicht eine Entdeckung einer neuen Physik.

Das Problem für die Detektive (die Physiker) ist: Um diese seltenen neuen Geschichten zu finden, müssen sie erst lernen, wie die alten, bekannten Geschichten aussehen. Dafür brauchen sie extrem viel Übungsmaterial. Früher mussten sie für jede einzelne Art von „neuer Geschichte" (z. B. ein Top-Quark oder ein Higgs-Boson) einen neuen, spezialisierten Detektiv von Grund auf ausbilden. Das war wie ein Schüler, der für jeden neuen Fall ein ganz neues Lehrbuch schreiben musste.

Die Lösung: Ein universelles Genie

In diesem Papier stellen die Autoren OmniLearned vor. Man kann sich das wie einen universellen Super-Detektiv vorstellen, der nicht für einen einzigen Fall trainiert wurde, sondern für alle Fälle gleichzeitig.

Hier ist die einfache Erklärung, wie das funktioniert, mit ein paar kreativen Vergleichen:

1. Der riesige Trainingscampus (1 Milliarde Jets)

Früher trainierten diese KI-Modelle mit etwa 100 Millionen Beispielen. Das ist wie ein Schüler, der 100 Bücher liest.
OmniLearned wurde jedoch mit über einer Milliarde Beispielen trainiert. Das ist, als würde dieser Schüler nicht nur 100 Bücher lesen, sondern die gesamte Bibliothek durchwühlen. Er hat Millionen von verschiedenen „Geschichten" (Teilchenjets) gesehen: von einfachen Kollisionen bis hin zu komplexen Zerfällen. Durch diese massive Menge an Wissen hat er ein tiefes, intuitives Verständnis dafür entwickelt, wie die Welt der Teilchen funktioniert.

2. Das Gehirn des Detektivs (Die Architektur)

Das Herzstück von OmniLearned ist eine spezielle KI-Architektur, die sie PET v2 nennen.

• Der alte Ansatz: Früher schauten die Detektiven nur auf die groben Umrisse eines Teilchens.
• Der neue Ansatz: OmniLearned schaut sich jedes Teilchen wie einen einzelnen Buchstaben in einem Wort an. Es versteht nicht nur, welche Buchstaben da sind, sondern auch, wie sie zueinander stehen (Abstand, Energie, Masse). Es ist wie ein Meisterleser, der sofort erkennt, ob ein Satz grammatikalisch korrekt ist oder ob ein Wort „falsch" geschrieben wurde (was auf eine neue Physik hindeuten könnte).

3. Die drei großen Aufgaben (Was kann der Super-Detektiv?)

Das Papier zeigt, dass dieser eine Super-Detektiv drei verschiedene Jobs gleichzeitig erledigen kann, ohne dass man ihn jedes Mal neu erfinden muss:

• Job A: Der Klassifizierer (Top-Quark-Jagd)
  Stellen Sie sich vor, Sie müssen in einem Haufen von roten und blauen Murmeln die seltenen goldenen Murmeln finden. OmniLearned wurde so trainiert, dass er sofort erkennt: „Aha, das hier ist ein Top-Quark-Jet!" Er ist in dieser Aufgabe besser als alle bisherigen Spezialisten, weil er sein riesiges Vorwissen nutzt, um winzige Details zu erkennen, die andere übersehen.

• Job B: Der Fein-Tuner (b-Tagging)
  Manchmal muss man zwischen sehr ähnlichen Dingen unterscheiden, wie zwischen einem echten Diamanten und einem sehr guten Glasimitat (z. B. b-Quarks vs. andere Teilchen). Hier zeigt OmniLearned, dass er nicht nur „grob" unterscheiden kann, sondern extrem präzise ist. Er kann sogar Teile seines Gehirns umfunktionieren, um nicht nur das ganze Teilchen, sondern sogar die einzelnen Spuren (wie Fußabdrücke) zu analysieren. Das spart enorme Rechenzeit.

• Job C: Der Anomalie-Detektor (Die Suche nach dem Unbekannten)
  Das ist das Coolste: Was passiert, wenn wir nicht wissen, wonach wir suchen? Was, wenn es eine völlig neue Art von Teilchen gibt, die in keinem Lehrbuch steht?
  Hier nutzt OmniLearned seine Fähigkeit, zu generieren. Er hat so viele normale Teilchen gesehen, dass er weiß, wie ein „normales" Teilchen aussehen müsste. Wenn er nun Daten sieht, die von diesem Idealbild abweichen, sagt er: „Moment mal! Das hier passt nicht in mein Muster!"
  In einem Test mit echten Daten vom CMS-Experiment hat OmniLearned erfolgreich die Top-Quarks wiederentdeckt, obwohl er nicht explizit dafür trainiert wurde, sie zu finden. Er hat sie einfach als „etwas Besonderes" erkannt, das aus dem Raster fiel.

Warum ist das so wichtig?

Stellen Sie sich vor, Sie müssten für jede neue Krankheit einen neuen Arzt ausbilden, der Jahre studiert. Das ist teuer und langsam.
OmniLearned ist wie ein Allgemeinmediziner mit einem Genie-Verstand, der sofort in jede Spezialklinik gehen und dort arbeiten kann, weil er die Grundlagen der menschlichen Anatomie (die Physik der Teilchen) so perfekt beherrscht.

• Effizienz: Man muss nicht für jede neue Frage ein neues Modell von Grund auf trainieren.
• Entdeckungspotenzial: Da das Modell so gut versteht, was „normal" ist, findet es leichter das „Unnormale" – und genau dort liegen die neuen physikalischen Entdeckungen.
• Offenheit: Die Autoren haben den Code und die Daten (die „Bibliothek") für alle zugänglich gemacht, damit jeder diesen Super-Detektiv nutzen kann.

Fazit:
Dieses Papier ist ein großer Schritt in Richtung einer neuen Ära der Teilchenphysik. Anstatt viele kleine, spezialisierte Werkzeuge zu haben, haben wir jetzt einen einzigen, mächtigen, flexiblen und extrem gut ausgebildeten „Allrounder", der uns hilft, die Geheimnisse des Universums schneller und genauer zu entschlüsseln.

Technische Zusammenfassung

Titel: Foundation Model Framework für alle Aufgaben der Jet-Physik

Autoren: Wahid Bhimji, Chris Harris, Vinicius Mikuni, Benjamin Nachman

---

1. Problemstellung

In der Hochenergiephysik sind hadronische Jets (Kaskaden von Teilchen aus Quarks und Gluonen) allgegenwärtig. Die Analyse ihrer komplexen Struktur ist entscheidend für die Entdeckung neuer Physik.

• Herausforderung: Derzeitige Deep-Learning-Ansätze (z. B. Jet-Tagger) erfordern für jede spezifische Aufgabe (z. B. Top-Quark-Identifikation, b-Tagging, Anomalieerkennung) große, maßgeschneiderte Trainingsdatensätze. Das Sammeln und Labeln ausreichender Daten für jede einzelne Aufgabe ist ressourcenintensiv und limitiert die Entdeckungspotenziale.
• Bisherige Ansätze: Viele Forschungsarbeiten zu „Foundation Models" in der Teilchenphysik nutzen tokenisierte Modelle mit selbstüberwachtem Lernen (ähnlich wie LLMs). Diese nutzen jedoch oft nicht die Kontinuität der Daten oder die in Simulationen verfügbaren Labelinformationen voll aus.

2. Methodik: OmniLearned Framework

Das Paper stellt OmniLearned vor, ein signifikantes Upgrade des vorherigen OmniLearn-Modells. Es handelt sich um ein generisches Transformer-basiertes Framework, das als Fundament für diverse Downstream-Aufgaben dient.

A. Architektur-Updates (PET v2)

Das Rückgrat des Modells ist eine Weiterentwicklung des Point Edge Transformer (PET), genannt PET v2:

• Input-Features: Reduktion auf ein minimales, aber effektives Set von Eingaben: $(\Delta\eta, \Delta\phi, \log p_T, \log E)$. Vorherige Standardisierungsschritte wurden entfernt, da die logarithmische Darstellung bereits zentriert ist.
• Erweiterte Information: Das Modell kann nun zusätzliche Informationen wie Teilchenidentifikation (PID) und Vertex-Informationen integrieren. Diese werden separat eingebettet (PID via Lookup-Tabelle, Vertex via MLP) und dynamisch hinzugefügt oder auf Null gesetzt (ähnlich Dropout), um Robustheit gegenüber fehlenden Daten zu gewährleisten.
• Lokale Aufmerksamkeit: Statt einfacher Mittelwertbildung über Nachbarn werden Attention-Mechanismen verwendet. Die Interaktionsfunktionen zwischen Teilchen und ihren $k$-nächsten Nachbarn enthalten physikalisch inspirierte Terme (invariante Masse, $\Delta R$, $p_T$-Produkte).
• Globale Aufmerksamkeit: Die Attention-Matrix wird durch physikalische Bias-Terme angereichert. Anstelle von Standard-Layer-Normalisierung wird eine lernbare hyperbolische Tangens-Funktion verwendet, um die Stabilität bei tiefen Netzwerken zu erhöhen.
• Aufgaben-spezifische Köpfe (Heads):
  • Klassifikation: Statt eines einzigen „Class-Token" werden fünf lernbare Tokens verwendet, die verschiedene Aspekte der Jet-Information zusammenfassen (Multi-Head-Ansatz).
  • Generierung: Ein Diffusions-Head (bzw. Flow-Matching) generiert Jet-Partikel.
• Verlustfunktion: Eine Kombination aus Klassifikationsverlust (Cross-Entropy) und Generierungsverlust. Neu ist die Einführung eines separaten Klassifikationsziels für „Sample Classification", um auch Datensätze ohne spezifische physikalische Labels (z. B. experimentelle Daten) im Pre-Training zu nutzen.
• Flow Matching: Der Generierungsprozess wurde von reinen Diffusionsmodellen auf Flow Matching umgestellt, was zu einer besseren Generierungsqualität und schnelleren Konvergenz führt.

B. Der „1-Billion-Datensatz"

Ein zentraler Baustein ist die Zusammenführung und Vereinheitlichung von über 1 Milliarde Jets aus verschiedenen Quellen in einem einzigen Format:

• Quellen: JetClass, JetClass2, Aspen Open Jets, ATLAS Top Tagging, H1 DIS (Deep Inelastic Scattering) und CMS QCD/BSM (Beyond Standard Model) Daten.
• Vielfalt: Der Datensatz umfasst 210 Klassen, darunter verschiedene Jet-Flavours, experimentelle Daten (CMS, ATLAS, H1) und neue Physik-Szenarien (z. B. Vektorartige Quarks, SUSY).
• Training: Das Modell wurde auf dem Perlmutter-Supercomputer trainiert (bis zu 512 GPUs) mit einem globalen Batch-Size von 4096 über 3 Epochen.

3. Wichtige Beiträge

1. OmniLearned Framework: Ein skalierbares, einheitliches Framework für Jet-Physik, das Pre-Training auf riesigen, heterogenen Datensätzen mit Feinabstimmung (Fine-Tuning) für spezifische Aufgaben kombiniert.
2. Skalierungsgesetze: Demonstration, dass größere Modelle (bis zu 460M trainierbare Parameter) durch das Pre-Training auf 1 Milliarde Jets signifikant bessere Ergebnisse erzielen als kleinere Modelle oder Modelle, die von Grund auf neu trainiert werden.
3. Software-Ökosystem: Bereitstellung einer umfassenden Software-Package, die den Zugriff auf alle Trainingsdaten, Testdaten und vortrainierten Modelle (Small, Medium, Large) vereinfacht.
4. Flexibilität: Das Modell kann Aufgaben übernehmen, für die es nicht ursprünglich gedacht war (z. B. Nutzung des Generierungs-Heads für die Klassifikation von Track-Ursprüngen im ATLAS Flavor-Tagging).

4. Ergebnisse

Das Framework wurde an drei repräsentativen Aufgaben getestet und erreichte in allen Fällen den State-of-the-Art (SOTA):

• Jet-Klassifikation (Top-Quark & Quark/Gluon):

  • Auf dem Community-Benchmark-Datensatz für Top-Quark-Tagging übertrifft OmniLearned (insbesondere die großen Modelle) alle bestehenden Methoden (z. B. ParticleNet, ParT, L-GATr).
  • Bei 30% Signal-Effizienz erreicht das große Modell eine Untergrund-Ablehnung (Background Rejection) von 3486, verglichen mit ~2894 für den vorherigen SOTA (L-GATr-f.t.).
  • Ähnliche Verbesserungen wurden für die Quark/Gluon-Trennung erzielt.

• Flavor-Tagging (b- und c-Tagging mit ATLAS-Daten):

  • OmniLearned wurde auf ATLAS-Public-Simulationsdaten für Flavor-Tagging feinabgestimmt.
  • Das Modell übertrifft das aktuelle ATLAS-Standardmodell (GN2) deutlich.
  • b-Tagging: Die Ablehnung von Light-Jets verbessert sich um mehr als 50% bei 70% b-Effizienz.
  • c-Tagging: Die Ablehnung von Light-Jets verbessert sich um 30–50%.
  • Ein wichtiger Befund: Die Feinabstimmung ist extrem effizient; das Medium-Modell erreicht Konvergenz mit weniger als der Hälfte der Updates im Vergleich zum Training von Grund auf.

• Anomalieerkennung (CMS Open Data):

  • Anwendung auf reale CMS-Daten (Aspen Open Jets), um Top-Quarks (ca. 0,1% der Jets) in einem reinen QCD-Hintergrund wiederzufinden.
  • Strategie 1 (Cathode-ähnlich): Nutzung des generativen Heads zur Hintergrundmodellierung und des Klassifikators zur Anomalie-Score-Bestimmung. OmniLearned konnte das Signal mit einer Signifikanz von über 15$\sigma$ entdecken, während Modelle, die von Grund auf neu trainiert wurden, das Signal nicht fanden.
  • Strategie 2 (Zero-Shot): Direkte Nutzung der während des Pre-Trainings gelernten Klassen (z. B. 3-Prong-Zerfälle) als Anomalie-Score ohne Feinabstimmung. Auch dies führte zur erfolgreichen Entdeckung des Signals.

5. Bedeutung und Ausblick

• Paradigmenwechsel: OmniLearned etabliert Foundation Models als neuen Standard in der Jet-Physik. Es zeigt, dass ein einziges, großskaliertes Pre-Training die Notwendigkeit für viele spezialisierte Modelle reduziert und gleichzeitig die Leistungsfähigkeit steigert.
• Entdeckungspotenzial: Durch die Fähigkeit, auch mit ungelabelten experimentellen Daten zu lernen und Anomalien ohne spezifische Signal-Annahmen zu finden, erweitert das Framework das Potenzial für Entdeckungen in aktuellen und zukünftigen Kollaborationen (LHC, zukünftige Collider).
• Allgemeingültigkeit: Die Autoren betonen, dass die Methodik und die vortrainierten Gewichte potenziell über die Jet-Physik hinaus anwendbar sind, z. B. für die Analyse ganzer Ereignstopologien oder anderer Kollisionssysteme.

Zusammenfassend stellt OmniLearned einen Meilenstein dar, der die Skalierbarkeit von Transformer-Architekturen mit physikalisch fundierten Datenrepräsentationen verbindet, um die Grenzen der jetzigen Analysemethoden in der Teilchenphysik zu überwinden.