Full event interpretation with machine-learning-based particle-flow reconstruction in the CMS detector

Este artículo presenta un algoritmo de flujo de partículas basado en aprendizaje automático (MLPF) de vanguardia implementado en GPUs dentro del marco de trabajo de CMS, el cual unifica la reconstrucción de eventos en un único modelo mientras logra un rendimiento físico comparable a los métodos estándar, una resolución de energía de chorro significativamente mejorada y una reducción de cinco veces en el tiempo de inferencia.

Lo esencial

Imagina el Gran Colisionador de Hadrones (LHC) en el CERN como un enorme y veloz choque de partículas. Cuando dos protones colisionan, no solo se rompen; explotan en una lluvia caótica de miles de diminutos e invisibles fragmentos. El detector CMS es una cámara gigante y de alta tecnología que intenta tomar una foto de esta explosión. Su trabajo es averiguar exactamente qué es cada uno de los fragmentos (¿es un fotón? ¿un electrón? ¿un trozo de un protón?) y con qué velocidad se mueve.

Durante años, el CMS ha utilizado un "libro de recetas" llamado el algoritmo de Flujo de Partículas (PF por sus siglas en inglés). Piensa en el antiguo algoritmo PF como un equipo de detectives humanos intentando resolver un rompecabezas. Observan pistas de diferentes partes de la cámara (el rastreador, los calorímetros) y utilizan una larga lista de reglas estrictas y escritas a mano para conectar los puntos. "Si una trayectoria se ve así y una masa de energía se ve de esa otra forma, deben ser la misma partícula". Funciona bien, pero es lento, rígido y requiere mucha calibración manual.

Este artículo presenta a un detective nuevo y más inteligente: MLPF (Flujo de Partículas Aprendido por Aprendizaje Automático).

El Nuevo Detective: Una Red Neuronal

En lugar de seguir un libro de reglas rígido, el MLPF es como un estudiante que ha leído millones de libros de texto de física y ha visto millones de explosiones simuladas. Utiliza un tipo de inteligencia artificial llamado Transformer (la misma tecnología detrás de los modelos de lenguaje avanzados).

• Cómo aprende: El equipo alimentó a esta IA con millones de colisiones "simuladas". Le mostraron los datos brutos (las trayectorias y las masas de energía) y le dijeron: "Esto es lo que realmente sucedió en la simulación". La IA aprendió a reconocer patrones y correlaciones que las reglas humanas podrían pasar por alto.
• Cómo piensa: En lugar de revisar las pistas una por una, la IA observa la explosión completa de una sola vez. Entiende cómo cada pieza del rompecabezas se relaciona con cada una de las demás piezas simultáneamente.

Las Grandes Victorias

1. Es Mucho Más Rápido (El Velocista)
El viejo detective (el PF estándar) se ejecuta en procesadores de computadora estándar (CPUs) y tarda unos 110 milisegundos en analizar una colisión. Eso es como tardar mucho tiempo en ordenar una baraja de cartas.
El nuevo detective de IA (MLPF) se ejecuta en una tarjeta gráfica especializada (GPU), que está diseñada para este tipo de tareas pesadas. Termina el mismo trabajo en solo 20 milisegundos. Eso es una aceleración de 5 veces. Es como cambiar el ordenar cartas a mano por el uso de una máquina de alta velocidad. Esta velocidad es crucial porque el LHC se está volviendo más concurrido y necesitan procesar más colisiones en menos tiempo.

2. Es Más Preciso (El Tirador de Precisión)
Debido a que la IA aprendió de tantos ejemplos, obtiene los detalles mejor que el viejo libro de reglas.

• Resolución de Energía de Jets: En física, los "jets" son chorros de partículas que actan como un paquete único. El artículo encontró que, para jets de tamaño medio, la nueva IA mide su energía con un 10–20% más de precisión que el método antiguo. Imagina intentar pesar una bolsa de manzanas; el método antiguo podría fallar por unos pocos onzas, mientras que el nuevo método es preciso hasta el gramo.
• Partículas Neutras: Es particularmente buena detectando "hadrones neutros" (partículas que no tienen carga eléctrica y son difíciles de rastrear), encontrando más de ellas sin cometer más errores.

3. Es Flexible (El Camaleón)
Las reglas antiguas fueron construidas para condiciones específicas del detector. Si el detector cambia o la energía de la colisión cambia, las reglas a menudo deben ser reescritas. La IA, sin embargo, aprendió los principios de la física de partículas. El artículo muestra que, incluso cuando la probaron con datos de un año o nivel de energía ligeramente diferente (que no había visto durante el entrenamiento), funcionó bien. Generaliza, lo que significa que puede adaptarse a nuevas situaciones sin necesidad de una revisión completa.

La Prueba del Mundo Real

El equipo no solo probó esto en simulaciones por computadora; de hecho, lo ejecutaron con datos reales recolectados por el detector CMS en 2024. Compararon la salida de la IA contra el método estándar en datos de colisiones reales. Los resultados fueron casi idénticos en términos de los resultados físicos, demostrando que la IA está lista para el mundo real.

Por Qué Esto Importa (Según el Artículo)

El artículo afirma que este es un gran paso adelante para el futuro del LHC. A medida que el colisionador se actualice para manejar colisiones aún más congestionadas (una fase llamada "LHC de Alta Luminosidad"), los métodos antiguos basados en reglas se volverán demasiado lentos y complejos de gestionar.

El algoritmo MLPF demuestra que podemos reemplazar las reglas de física complejas y diseñadas a mano con un modelo de IA único y unificado que es:

• Más rápido (ejecutándose eficientemente en GPUs modernas).
• Más inteligente (mejorando la precisión de las mediciones).
• Escalable (listo para las enormes cargas de datos del futuro).

En resumen, el experimento CMS está actualizando sus "ojos" de un par de detectives humanos siguiendo una lista de verificación a una IA superinteligente que ve el panorama completo instantáneamente, permitiendo a los físicos ver más profundamente en los secretos del universo.

Resumen técnico

Resumen Técnico: Interpretación de Eventos Completos con Reconstrucción de Flujo de Partículas Basada en Aprendizaje Automático en el Detector CMS

Planteamiento del Problema
El algoritmo de Flujo de Partículas (PF, por sus siglas en inglés) es la piedra angular de la reconstrucción de eventos en el experimento CMS en el LHC del CERN, con el objetivo de proporcionar una descripción global, a nivel de partícula, de las colisiones mediante la reconciliación de señales del rastreador, los calorímetros y los sistemas de muones. La implementación estándar de PF se basa en heurísticas motivadas por la física y métodos iterativos basados en reglas (por ejemplo, vinculación basada en la proximidad) para asociar los impactos del detector con las partículas. Si bien estos enfoques son efectivos, enfrentan desafíos en cuanto a escalabilidad y adaptabilidad, particularmente a medida que la complejidad del detector y el pileup (interacciones simultáneas de protones-protones) aumentan hacia la era del LHC de Alta Luminosidad (HL-LHC). Las tareas de agrupamiento tradicionales suelen exhibir una alta complejidad computacional ($O(N^3)$ o $O(N \log N)$) y requieren heurísticas diseñadas a mano que son difíciles de optimizar para el hardware moderno. Además, el marco actual de PF no es nativamente adecuado para una ejecución eficiente en Unidades de Procesamiento Gráfico (GPU), lo que limita el potencial de aceleración para la reconstrucción de eventos completos.

Metodología
El artículo presenta la implementación de un algoritmo de reconstrucción de Flujo de Partículas basado en Aprendizaje Automático (MLPF) dentro del marco de software de CMS. Este enfoque reemplaza múltiples pasos modulares de reconstrucción con un único modelo unificado entrenado directamente con datos simulados.

• Arquitectura: El algoritmo MLPF utiliza una arquitectura basada en transformers que emplea mecanismos de autoatención (self-attention). A diferencia de los enfoques previos de redes neuronales de grafos (GNN) probados en CMS, este modelo utiliza capas de transformer aceleradas por el núcleo FLASHATTENTION para mitigar la escala cuadrática de memoria y computación típicamente asociada con los mecanismos de atención. El modelo toma vectores de características de trayectorias (tracks) y cúmulos de calorímetro reconstruidos como entrada y predice todas las partículas del estado final simultáneamente en un solo paso de inferencia.
• Objetivo de Entrenamiento: Un componente crítico es la definición de un "objetivo de entrenamiento a nivel de partícula". Este objetivo se construye para incluir todas las partículas que dejan señales detectables en el detector (ya sea directamente o vía productos de decaimiento), distinguiéndolas de las partículas de "verdad" (truth) generadas por el generador de eventos (PYTHIA), que pueden quedar fuera de la aceptación. El objetivo incluye una etiqueta binaria para el origen del pileup, aunque la mitigación del pileup se realiza posteriormente mediante el algoritmo PUPPI.
• Función de Pérdida: El modelo se entrena de extremo a extremo utilizando una función de pérdida multitarea que consta de tres componentes:
  1. Clasificación Binaria ($L_{cls-binary}$): Determina si una trayectoria o un cúmulo de entrada es el elemento primario para una partícula objetivo.
  2. PID Multiclase ($L_{cls-PID}$): Clasifica el tipo de partícula (fotón, electrón, muón, hadrón cargado, hadrón neutro) utilizando pérdida focal (focal loss) para abordar el desequilibrio de clases.
  3. Regresión ($L_{reg}$): Predice los componentes del cuadrimomento ($p_T, \eta, \phi, E$). Los objetivos de regresión se transforman (razones logarítmicas respecto a la entrada primaria) para mejorar la estabilidad numérica, con un peso adicional dependiente de $p_T$ para enfatizar las partículas de alta energía.
• Integración: El modelo entrenado se exporta al formato Open Neural Network Exchange (ONNX) y se integra en el software offline de CMS a través de la interfaz ONNXRUNTIME. Esto permite que el modelo reemplace el módulo de reconstrucción PF estándar, dejando intactos los módulos de rastreo (tracking) ascendentes y los módulos de agrupamiento de chorros (jet clustering) descendentes.

Contribuciones Clave

1. Primera Reconstrucción de ML Validada con Datos: Este trabajo representa la primera aplicación de un flujo de reconstrucción de eventos completos basado en ML validado con datos recolectados por un experimento de colisionador de hadrones (CMS Run 3, 2023–2024).
2. Modelo de Transformer Unificado: La introducción de una arquitectura basada en transformers que realiza una reconstrucción de eventos completos aprendible, generalizando a través de las condiciones del detector y las energías de colisión sin depender de heurísticas iterativas.
3. Aceleración por GPU: El despliegue exitoso del algoritmo en GPUs, demostrando que la reconstrucción de eventos globales complejos puede transferirse de las CPU a los aceleradores, logrando aceleraciones significativas.
4. Validación del Rendimiento: Una evaluación exhaustiva que muestra que el enfoque de ML logra un rendimiento físico comparable a, y en algunas métricas superior al, algoritmo PF estándar.

Resultados
El rendimiento del algoritmo MLPF fue evaluado utilizando muestras simuladas (eventos de quark top-antiquark y QCD multijet) y un pequeño subconjunto de datos de Run 3.

• Rendimiento Físico:
  • Resolución de Energía de Chorro (Jet): En eventos simulados de quark top-antiquark bajo condiciones de Run 3, la resolución de energía de los chorros mejoró entre un 10–20% para chorros con momento transversal ($p_T$) entre 30–100 GeV en comparación con el PF estándar.
  • Identificación de Partículas: MLPF logró una eficiencia mejorada para hadrones neutros manteniendo la misma tasa de identificación errónea que el PF estándar. Para hadrones cargados y fotones, MLPF mostró una mayor eficiencia con una tasa de identificación errónea ligeramente superior, la cual puede ser ajustada.
  • Momento Transversal Faltante ($p_T^{miss}$): Las distribuciones de $p_T^{miss}$ reconstruidas fueron generalmente consistentes con el PF estándar, aunque MLPF reconstruyó un espectro algo más duro en las colas.
  • Validación de Datos: En los datos de Run 3, las distribuciones de asimetría de dijet y de $p_T$ del jet líder reconstrucidas por MLPF resultaron compatibles con las del PF estándar, confirmando la capacidad del modelo para producir salidas físicas realistas.
• Rendimiento Computacional:
  • Tiempo de Inferencia: En una GPU Nvidia L4, el algoritmo MLPF logró un tiempo de inferencia mediano de 20 ms por evento, comparado con aproximadamente 110 ms para la reconstrucción PF estándar de CMS ejecutándose en CPUs.
  • Escalabilidad: El uso de FLASHATTENTION permitió distribuciones de tiempo de ejecución estables y estrechas incluso para eventos complejos, con soporte para 64 flujos de inferencia paralelos en una sola GPU. El algoritmo demostró una mejora en la escalabilidad del tiempo de ejecución con el aumento de la complejidad del evento en comparación con el PF iterativo estándar.

Significancia
El artículo afirma que el algoritmo MLPF demuestra con éxito que el aprendizaje automático puede reemplazar la reconstrucción de flujo de partículas tradicional, basada en heurísticas, con un modelo unificado que ofrece un rendimiento físico comparable o mejor y una escalabilidad computacional superior. El trabajo establece las bases para la futura reconstrucción de eventos en el HL-LHC, donde el aumento del pileup y la complejidad del detector demandarán soluciones eficientes aceleradas por GPU. Al validar el enfoque con datos reales, los autores demuestran que la reconstrucción basada en ML no es meramente un ejercicio de simulación, sino un camino viable para el análisis de física a nivel de producción. El artículo enfatiza que, si bien los resultados actuales se limitan a las condiciones de Run 3, la metodología proporciona un flujo de trabajo escalable para la actualización de alta luminosidad y futuros colisionadores, permitiendo potencialmente la programación diferenciable para la optimización de detectores.