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

El artículo presenta OmniLearned, una versión mejorada del modelo fundacional para física de jets que, gracias a una nueva arquitectura, un entrenamiento con más de mil millones de jets y un software accesible, establece el estado del arte en tareas de etiquetado y detección de anomalías, ampliando así el potencial de descubrimiento en experimentos de colisionadores.

Lo esencial

Imagina que el universo es una inmensa cocina donde las partículas subatómicas son los ingredientes. Cuando chocan dos protones a velocidades increíbles (como en el Gran Colisionador de Hadrones, o LHC), se produce una explosión de energía que se convierte en una "lluvia" de partículas. A esta lluvia la llamamos chorro (o jet en inglés).

El problema es que estos chorros son como platos complicados: a veces es un "pastel" hecho de un quark top, a veces es una "sopa" de gluones, y a veces es algo completamente nuevo que nunca hemos visto. Los físicos necesitan identificar qué es cada plato para descubrir nuevas leyes de la naturaleza, pero hay tantos platos y son tan complejos que es imposible aprender a cocinar (o clasificar) cada uno por separado.

Aquí es donde entra OmniLearned, el protagonista de este artículo.

¿Qué es OmniLearned? (El "Chef Maestro")

Piensa en OmniLearned no como un simple cocinero, sino como un Chef Maestro que ha pasado años comiendo y analizando mil millones de platos diferentes.

1. El entrenamiento masivo: Antes, los modelos de inteligencia artificial tenían que aprender a reconocer un "pastel de quark top" con solo unos pocos ejemplos. OmniLearned, en cambio, ha "comido" más de mil millones de chorros de todas las formas posibles. Ha visto la cocina en su totalidad.
2. La base de conocimiento: Gracias a esta cantidad masiva de datos, el Chef Maestro ha desarrollado un "gusto" y una intuición profunda sobre cómo se comportan las partículas. Ya no necesita memorizar recetas específicas; entiende la esencia de los ingredientes.
3. La adaptación rápida: Lo más increíble es que, una vez que este Chef Maestro ha aprendido la base, puedes pedirle que se especialice en algo muy específico (como detectar un tipo de quark raro o buscar un ingrediente secreto) y lo hace en un instante, con muy pocos ejemplos adicionales. Es como si un chef que conoce la cocina francesa pudiera aprender a hacer sushi en una tarde porque ya entiende la lógica de los ingredientes.

Las tres mejoras clave (La "Receta Actualizada")

Los autores han actualizado la "receta" de este Chef Maestro con tres trucos nuevos:

• Una visión más detallada: Antes, el Chef miraba solo la forma básica de los ingredientes. Ahora, OmniLearned también mira detalles finos, como la "identidad" de cada partícula (si es un electrón, un fotón, etc.) y de dónde vino. Es como si antes solo miraras el color de la fruta, y ahora pudieras olerla y sentir su textura.
• Atención a los vecinos: El modelo ha aprendido a mirar cómo interactúan las partículas entre sí. No solo ve una partícula aislada, sino que entiende cómo se relaciona con sus "vecinos" inmediatos en el chorro. Es como entender que en una fiesta, no solo importa quién está en la sala, sino con quién está hablando cada persona.
• Generación de realidad: Además de clasificar, OmniLearned puede "imaginar" o generar nuevos chorros. Si le dices "hazme un chorro que se vea como un quark top", el modelo puede crear uno virtual. Esto es útil para simular lo que debería pasar si no hubiera nada nuevo, y así comparar con la realidad para encontrar anomalías.

¿Qué logró este Chef? (Los resultados)

El papel muestra que OmniLearned es el mejor en tres pruebas importantes:

1. Reconocer al "Rey" (Quark Top): En una prueba estándar de la comunidad, OmniLearned identificó a los chorros de quarks top con una precisión que supera a todos los métodos anteriores. Es como si pudiera distinguir un pastel de chocolate de uno de vainilla incluso si están cubiertos de azúcar.
2. Cazar el "Invisible" (Etiquetado de sabor): En el experimento ATLAS, el modelo ayudó a distinguir chorros que vienen de quarks "b" (bottom) de otros tipos. Esto es crucial porque los quarks "b" son a menudo la firma de procesos físicos muy interesantes. OmniLearned lo hizo mejor que los expertos actuales, ahorrando tiempo y recursos.
3. Detectar lo "Raro" (Detección de anomalías): Esta es la parte más emocionante. Usando datos reales del experimento CMS, el modelo buscó algo que no estaba en su entrenamiento: la reaparición de un quark top en un canal donde no se esperaba.
   • La analogía: Imagina que tienes un álbum de fotos de paisajes normales. De repente, aparece una foto con un unicornio. Un sistema normal podría decir "no sé qué es esto". OmniLearned, al haber visto millones de paisajes, nota inmediatamente: "¡Eso no encaja! ¡Es un unicornio!".
   • El modelo logró encontrar la señal del quark top en los datos reales del CMS, demostrando que puede encontrar "agujas en el pajar" sin necesidad de saber exactamente qué agujas buscar de antemano.

En resumen

Este artículo presenta OmniLearned, una inteligencia artificial que ha aprendido de mil millones de ejemplos para entender la física de partículas como nunca antes.

Es como tener un super-sensor que, tras ver todo el universo de colisiones, puede:

• Clasificar lo que ve mejor que nadie.
• Aprender nuevas tareas rápidamente.
• Detectar cosas extrañas o nuevas que los humanos podrían pasar por alto.

Esto no solo mejora lo que sabemos hoy, sino que abre la puerta a descubrir nueva física en los datos del futuro, haciendo que los experimentos del LHC sean mucho más potentes y eficientes. Es un paso gigante hacia una comprensión más profunda de cómo está construido nuestro universo.

Resumen técnico

Resumen Técnico: OmniLearned

1. Planteamiento del Problema

En la física de colisionadores, los jets hadrónicos (resultados de quarks y gluones de alta energía) son ubicuos y contienen información crucial sobre el origen de las partículas y las propiedades de la fuerza fuerte. Tradicionalmente, los investigadores han diseñado observables físicos específicos para caracterizar esta estructura compleja. Sin embargo, el enfoque moderno basado en aprendizaje profundo (deep learning) enfrenta un desafío crítico: la necesidad de grandes volúmenes de datos etiquetados para entrenar modelos de última generación para cada tarea específica (como la identificación de quarks top, b-tagging o detección de anomalías).

Obtener suficientes ejemplos de jets simulados de alta fidelidad (simulaciones completas tipo Geant) para cada tarea es computacionalmente costoso y limitante. Los modelos fundamentales (foundation models) previos en física de partículas se han centrado en modelos tokenizados entrenados mediante aprendizaje auto-supervisado (similares a los LLMs), pero a menudo no aprovechan completamente la continuidad de los datos ni la información de etiquetas disponible en las simulaciones.

2. Metodología: OmniLearned

El artículo presenta OmniLearned, una actualización mayor del modelo anterior OmniLearn. Este nuevo marco se basa en tres pilares fundamentales:

A. Arquitectura del Modelo (PET v2)
El núcleo del modelo es una mejora de la arquitectura Point Edge Transformer (PET), ahora versión 2 (PET v2), diseñada para manejar jets como nubes de puntos con un número variable de partículas. Las mejoras clave incluyen:

• Entradas Optimizadas: Se reduce el conjunto de características de entrada a un mínimo esencial pero informativo: $(\Delta\eta, \Delta\phi, \log p_T, \log E)$, eliminando la necesidad de normalización previa.
• Codificación de Información Adicional: Se incorporan bloques de incrustación (embedding) separados para la identificación de partículas (PID) y la información de vértices, permitiendo que el modelo funcione tanto con datos etiquetados como sin ellos.
• Capas de Atención Local y Global:
  • Se introducen términos de interacción inspirados en la física entre partículas y sus $k$-vecinos más cercanos (incluyendo masa invariante y distancia $\Delta R$).
  • Se reemplaza el promedio simple por una atención ponderada aprendible (usando bloques de transformador), lo que permite al modelo aprender la importancia de cada vecino y ser robusto ante variaciones en la multiplicidad de partículas.
  • En la atención global, se añaden sesgos inspirados en la física a la matriz de atención y se utiliza una normalización de capa basada en tangente hiperbólica aprendible para mayor estabilidad.
• Función de Pérdida Híbrida: El modelo se entrena simultáneamente para clasificación y generación.
  • Para la generación, se adopta un enfoque de Flow Matching (en lugar de solo difusión), prediciendo el campo de velocidad $u = \epsilon - x$, lo que mejora la calidad de generación y la convergencia.
  • Se introduce una tarea de clasificación de muestras para datos sin etiquetas de verdad (como datos experimentales reales), clasificando el origen del dataset.

B. Conjunto de Datos (1 Mil Millones de Jets)
OmniLearned se entrena con un conjunto de datos masivo de más de 1.000 millones de jets (10 veces más que el modelo anterior). Este conjunto unificado incluye:

• Datos de simulación rápida (Delphes) y simulación completa (Geant4/Geant3).
• Múltiples experimentos: ATLAS, CMS, H1.
• Diversidad de procesos: Jets de QCD, quarks top, bosones W/Z, Higgs, y nuevos procesos de física más allá del modelo estándar (BSM).
• Se crearon 210 clases en total (200 sabores de jets específicos y 10 etiquetas de dataset).

C. Software y Accesibilidad
Se ha desarrollado un paquete de software unificado que permite el acceso, descarga y formato estandarizado de todos los conjuntos de datos y modelos pre-entrenados, facilitando la reutilización por la comunidad.

3. Contribuciones Clave

1. Escalabilidad: Demostración exitosa de la escalabilidad de los transformadores en física de jets, entrenando con 10 veces más datos que la versión anterior.
2. Flexibilidad de Arquitectura: Capacidad de manejar datos con y sin etiquetas de verdad, y de integrar información adicional (PID, vértices) de forma dinámica.
3. Reutilización de Cabezas (Heads): Demostración de que las cabezas de generación pueden reutilizarse para tareas de clasificación (ej. predicción del origen de trazas en b-tagging), y viceversa.
4. Aplicación a Datos Reales: Primera aplicación de un modelo fundamental a datos experimentales reales del CMS para la detección de anomalías.

4. Resultados

El modelo se evaluó en tres tareas representativas, superando el estado del arte (SOTA) en todos los casos:

• Clasificación de Jets (Top y Quark vs. Gluón):

  • En el conjunto de datos de referencia de la comunidad para top tagging, OmniLearned (especialmente las versiones medium y large) superó significativamente a modelos anteriores como ResNeXt, ParticleNet y ParT.
  • Logró una mayor eficiencia de rechazo de fondo para una eficiencia de señal fija del 30% (ej. rechazo de fondo de ~3486 para la versión large frente a ~2789 del mejor modelo anterior).
  • Se observó una fuerte correlación entre el tamaño del modelo y el rendimiento tras el fine-tuning.

• Flavor Tagging (b-tagging y c-tagging) con ATLAS:

  • Utilizando simulaciones completas de ATLAS, OmniLearned superó al modelo actual de referencia de ATLAS (GN2).
  • Mejoras notables en el rechazo de fondos: >50% de mejora en el rechazo de jets ligeros para b-tagging y mejoras de hasta un 50% en c-tagging.
  • El fine-tuning de OmniLearned fue un 50% más rápido que entrenar un modelo desde cero, requiriendo menos de la mitad de las actualizaciones de red.

• Detección de Anomalías con Datos del CMS:

  • Se aplicó a datos abiertos del CMS (2016) para redescubrir el quark top (que constituye ~0.1% de los jets) en el canal de jets individuales de alto $p_T$.
  • Estrategia 1 (Supervisada): Se utilizó la capacidad generativa para modelar el fondo en las regiones laterales (sidebands) y un clasificador para identificar anomalías en la región de señal. OmniLearned logró una significancia estadística de descubrimiento (>5$\sigma$) donde los modelos entrenados desde cero fallaron.
  • Estrategia 2 (No supervisada/Zero-shot): Se utilizó directamente la salida de las clases pre-entrenadas (relación entre desintegraciones de 3 puntas y QCD) sin fine-tuning adicional, logrando también identificar la señal, demostrando la capacidad de generalización del modelo.

5. Significado e Impacto

El trabajo de OmniLearned representa un cambio de paradigma en el análisis de datos de colisionadores:

• Eficiencia Computacional: Reduce drásticamente la necesidad de generar simulaciones completas costosas para cada nueva tarea, permitiendo el fine-tuning de un modelo fundamental pre-entrenado.
• Descubrimiento de Nueva Física: Al mejorar la capacidad de detección de anomalías en todo el espacio de fases, aumenta el potencial de descubrimiento en experimentos pasados, presentes y futuros (como el HL-LHC).
• Reproducibilidad y Comunidad: La liberación de un marco de software unificado y modelos pre-entrenados estandariza el acceso a la tecnología de IA en física de altas energías, fomentando la colaboración y la comparación justa entre algoritmos.
• Generalización: La metodología sugiere que estos modelos fundamentales podrían extenderse más allá de los jets, abarcando topologías completas de eventos y diferentes sistemas de colisión.

En conclusión, OmniLearned establece un nuevo estándar de rendimiento en la física de jets, combinando arquitecturas avanzadas de transformadores, entrenamiento a gran escala y una filosofía de modelo fundamental que maximiza la utilidad de los datos de simulación y experimentales.