Vision Transformers and Graph Neural Networks for Charged Particle Tracking in the ATLAS Muon Spectrometer

Este artículo presenta dos enfoques basados en aprendizaje automático para mejorar el seguimiento de muones en el Espectrómetro de Muones de ATLAS: el uso de Redes Neuronales de Grafos para la rechazo de fondos que acelera la reconstrucción en un 15 %, y una prueba de concepto con Vision Transformers que logra una reconstrucción aproximada ultra rápida en 2,3 ms con una eficiencia del 98 %.

Lo esencial

¡Claro que sí! Imagina que el experimento ATLAS en el CERN es como un estadio de fútbol gigante donde ocurren millones de partidos de "fútbol de partículas" cada segundo. El objetivo es encontrar a los jugadores especiales (los muones) que son clave para descubrir nuevos secretos del universo.

El problema es que, en el futuro, el estadio se llenará tanto de gente (colisiones) que será imposible distinguir a los jugadores reales de los espectadores que gritan, los globos que rebotan o el ruido de las cámaras. A esto lo llamamos "ruido de fondo".

Este documento presenta dos nuevas ideas, como si fueran dos tipos de detectives inteligentes (basados en Inteligencia Artificial), para ayudar a ATLAS a encontrar a los jugadores reales más rápido y mejor.

Aquí tienes la explicación sencilla:

1. El Contexto: El Estadio se Llena Demasiado

Actualmente, en cada "golpe" de las partículas, hay unas 60 colisiones simultáneas. Pero en el futuro (después de 2030), habrá hasta 200 colisiones a la vez.

• El problema: El detector de muones (el "campo de juego" exterior) se llena de tanta basura (ruido) que los ordenadores tradicionales tardan mucho en limpiar el campo y encontrar a los jugadores. Es como intentar encontrar una aguja en un pajar, pero el pajar está creciendo a velocidad luz.

2. La Primera Solución: El "Filtro de Red Neuronal" (GNN)

La analogía: Imagina que tienes una caja llena de miles de piezas de LEGO mezcladas. Algunas forman un castillo (el jugador real) y otras son solo basura suelta.

• Cómo funciona el método antiguo: El ordenador revisa pieza por pieza, una por una, para ver si encajan. Es lento.
• La nueva idea (GNN): En lugar de mirar pieza por pieza, el detector agrupa las piezas en "paquetes" (llamados "Cubos de Muones"). La Inteligencia Artificial actúa como un jefe de obra experto que mira el paquete entero y dice: "¡Ese paquete es basura, tíralo!" antes de intentar armar el castillo.
• El resultado: Al tirar la basura primero, el ordenador tiene que trabajar con mucho menos material.
  • Ahorro de tiempo: El proceso se hizo un 15% más rápido (de 255 milisegundos a 217).
  • Ventaja: No pierde ningún jugador real, solo tira la basura. Es como tener un tamiz muy fino que deja pasar solo lo importante.

3. La Segunda Solución: El "Ojo de Águila" (Vision Transformers)

La analogía: Imagina que en lugar de armar el castillo pieza por pieza, tienes una foto aérea del estadio tomada por un dron súper inteligente.

• Cómo funciona: En lugar de buscar conexiones locales, este sistema (llamado Vision Transformer o ViT) mira toda la foto de golpe, como un humano que ve un patrón de inmediato.
  • Paso 1 (El Filtro): El dron identifica rápidamente qué puntos de la foto son "ruido" y los borra. Reduce el estadio de 6,900 puntos a solo 55 puntos relevantes. ¡Es como si el dron hiciera desaparecer a 99% de la multitud!
  • Paso 2 (El Reconocimiento): Con solo esos 55 puntos limpios, el dron dibuja las líneas que conectan a los jugadores.
• La magia: Este sistema es tan rápido que puede hacer todo el trabajo en 2.3 milisegundos usando una tarjeta gráfica de ordenador normal (como las que usan los gamers), en lugar de necesitar superordenadores gigantes.
• El resultado: Encuentra al 98% de los jugadores reales y los conecta correctamente, aunque a veces le cuesta un poco calcular la velocidad exacta del jugador (eso es lo que siguen mejorando).

¿Por qué es importante esto?

Imagina que el sistema de seguridad del estadio tiene que decidir en una fracción de segundo qué videos guardar para la historia.

• Antes: Tardaba mucho en limpiar la imagen, por lo que perdía momentos importantes o necesitaba ordenadores carísimos.
• Ahora: Con estas nuevas herramientas de IA, el sistema puede:
  1. Limpiar la basura casi al instante.
  2. Encontrar a los jugadores en tiempo real, incluso cuando el estadio está lleno al máximo.
  3. Hacerlo con hardware más barato (tarjetas gráficas de consumo).

En resumen

Los científicos de ATLAS están enseñando a sus ordenadores a pensar de dos formas nuevas:

1. Como un experto en agrupar: Para descartar la basura rápidamente y trabajar menos.
2. Como un artista que ve el cuadro completo: Para encontrar patrones complejos en milisegundos, sin perderse en los detalles pequeños.

Esto asegura que, cuando llegue la era del "Gran Colisionador de Alta Luminosidad" (cuando todo esté mucho más lleno y ruidoso), ATLAS seguirá siendo capaz de encontrar los secretos más raros del universo sin quedarse atascado en el tráfico de datos. ¡Es como pasar de buscar una aguja en un pajar a tener un imán que atrae solo las agujas!

Resumen técnico

Resumen Técnico: Vision Transformers y Redes Neuronales de Grafos para el Rastreo de Partículas Cargadas en el Espectrómetro de Muones de ATLAS

1. El Problema

La reconstrucción de partículas cargadas (específicamente muones) en el experimento ATLAS del Gran Colisionador de Hadrones (LHC) enfrenta desafíos críticos con la llegada de la era de Alta Luminosidad (HL-LHC) post-2030.

• Aumento de la Pile-up: La densidad de colisiones por cruce de paquetes aumentará de $\langle\mu\rangle \approx 60$ (Run 3) hasta $\langle\mu\rangle \approx 200$.
• Alta Ocupación y Ruido: Esto incrementa drásticamente la ocupación dentro del Espectrómetro de Muones, generando una relación señal-ruido extremadamente baja (hasta 0.6% en algunas regiones) debido a partículas secundarias, ruido electrónico y radiación cósmica.
• Restricciones de Latencia: El sistema de disparo (Trigger), específicamente el Event Filter (EF), requiere algoritmos robustos y rápidos para seleccionar eventos físicos interesantes a partir de una tasa de entrada de 1 MHz.
• Limitaciones del Algoritmo Base: El algoritmo de reconstrucción actual (basado en transformada de Hough y ajuste $\chi^2$) tarda aproximadamente 255 ms por evento en un solo hilo de CPU, lo cual es insuficiente para las condiciones futuras sin optimizaciones. Además, el espectrómetro integra cinco tecnologías de subdetectores heterogéneas en un campo magnético inhomogéneo, complicando la asignación de "hits" a trayectorias.

2. Metodología

El trabajo propone dos enfoques basados en aprendizaje automático (ML) para abordar estos desafíos, ambos centrados en la eliminación de "hits" de fondo antes de la reconstrucción principal.

A. Enfoque 1: Redes Neuronales de Grafos (GNN) para Rechazo de Fondo

• Concepto: Integrar GNNs en la cadena de reconstrucción existente (no-ML) para filtrar "hits" de fondo.
• Estructura de Datos: En lugar de usar hits individuales (computacionalmente costoso), se construyen grafos dinámicos a partir de "Cubos de Muones" (Muon Buckets), que son agrupaciones de hits espaciales.
• Arquitectura: Se utiliza una arquitectura basada en EdgeConv. Los nodos son los cubos válidos y las conexiones (bordes) se definen por proximidad espacial entre sectores adyacentes o idénticos.
• Implementación: El modelo se integra en el framework Athena de ATLAS. Un desafío técnico fue la implementación de operadores EdgeConv personalizados para soportar el formato ONNX dentro del entorno C++ de ATLAS.
• Objetivo: Clasificar los cubos como señal, múltiples segmentos o ruido para reducir la entrada de datos antes de la búsqueda de patrones.

B. Enfoque 2: Vision Transformers (ViT) para Reconstrucción de Extremo a Extremo

• Concepto: Un prototipo de reconstrucción completa (detección, asignación y estimación de parámetros) utilizando arquitecturas de visión por computadora.
• Arquitectura: Adaptación de Mask2Former, un modelo de segmentación panóptica de vanguardia.
• Innovaciones Clave:
  • Tokens de Hits: Cada hit individual del detector se trata como un "token" en lugar de parches de imagen.
  • Atención Flash Ventaneada (Windowed Flash Attention): Para manejar la alta dimensionalidad y la naturaleza dispersa de los datos, se utiliza una atención limitada a una ventana de tamaño $W$ (512 hits), escalando la complejidad a $O(W \times N)$ en lugar de $O(N^2)$.
  • Prioridad Física: Los hits se ordenan por ángulo azimutal ($\phi$) y se asume que solo los hits cercanos en $\phi$ pertenecen a la misma trayectoria.
  • Pipeline:
    1. Fase de Filtrado: Clasificación binaria (señal/ruido) usando atención global ventaneada.
    2. Fase de Rastreo: Uso de consultas aprendidas (learnable queries) y atención cruzada para asignar hits a trayectorias y regesar parámetros ($p_T, \eta, \phi, Q$).

3. Contribuciones Clave

• Integración Híbrida: Demostración exitosa de integrar GNNs en una cadena de reconstrucción física tradicional (Athena) sin introducir sesgos negativos.
• Prototipo End-to-End: Primera aplicación de arquitecturas Transformer modernas (Mask2Former) para el rastreo de muones en el Espectrómetro de ATLAS, tratando el problema como una tarea de segmentación de imágenes.
• Optimización de Hardware: Validación de que estas arquitecturas pueden ejecutarse en hardware de consumo (GPUs de gama media) con latencias extremadamente bajas, abriendo la puerta a implementaciones en tiempo real más económicas.
• Manejo de Heterogeneidad: El enfoque ViT demuestra capacidad para procesar datos de múltiples tecnologías de detectores y geometrías complejas sin necesidad de una ingeniería de características manual exhaustiva.

4. Resultados

A. Resultados del GNN (Filtro de Cubos):

• Velocidad: Reducción del tiempo de reconstrucción por evento en un 15% (de 255 ms a 217 ms) en un entorno con $\langle\mu\rangle=200$.
• Eficiencia: Rechazo de aproximadamente el 97% de los cubos de fondo a $\langle\mu\rangle=60$.
• Calidad: Se preservó la eficiencia y precisión de la reconstrucción de señales en todo el rango cinemático, sin sesgo negativo.

B. Resultados del Vision Transformer (ViT):

• Fase de Filtrado:
  • AUC (Área bajo la curva ROC) de 0.9997.
  • A un punto de trabajo de 99% de eficiencia de señal, la pureza de los hits aumenta de 0.6% a 66.5%, rechazando el 99.7% del fondo.
  • Reducción de la ocupación de ~6,900 hits a ~55 hits por evento.
• Fase de Rastreo:
  • Eficiencia de detección de señal: 98.0% (para $p_T \ge 5$ GeV).
  • Eficiencia de doble coincidencia (Double Matching): 94.59% (eficiencia y pureza de asignación > 50%).
  • Precisión de carga: 96.35% de exactitud en la clasificación del signo de la carga.
  • Velocidad de Inferencia:
    • En GPU de consumo (RTX 3090): 2.3 ms por evento (lote de 200).
    • En GPU de alto rendimiento (H100): 0.9 ms por evento.
• Limitaciones: La regresión de parámetros precisos de la trayectoria y la clasificación de carga aún no son competitivas con el algoritmo base tradicional, y la latencia individual (lote 1) está limitada por la sobrecarga de lanzamiento de kernels de GPU.

5. Significado e Impacto

Este trabajo es fundamental para el futuro del experimento ATLAS en la era HL-LHC:

1. Viabilidad del Trigger: Demuestra que las técnicas de ML avanzadas pueden cumplir con las estrictas restricciones de latencia del Event Filter, reduciendo la carga computacional y permitiendo procesar eventos de alta ocupación.
2. Paradigma de Diseño: Cambia el enfoque de algoritmos heurísticos secuenciales a arquitecturas de atención global que pueden resolver problemas combinatorios complejos de manera más eficiente.
3. Sostenibilidad y Costo: La capacidad de ejecutar estos modelos en GPUs de consumo (costo ~1.5k CHF) sugiere que la infraestructura de disparo podría volverse más escalable y económica, reduciendo la dependencia de clusters de CPU masivos.
4. Futuro de la Investigación: Establece una base para integrar la asignación de hits y la estimación de parámetros directamente en el sistema de disparo, y abre la puerta a la detección de partículas de vida larga (LLP) y a la optimización mediante compilación de modelos (TorchScript, etc.) para reducir aún más la latencia.

En resumen, el estudio valida que los Transformers y las GNNs no solo son teóricamente aplicables a la física de altas energías, sino que ofrecen mejoras tangibles en velocidad y eficiencia para la reconstrucción de muones en condiciones extremas de ruido y ocupación.