Mixture-of-Experts Graph Transformers for Interpretable Particle Collision Detection

Resumen Técnico: Transformadores de Grafos con Mezcla de Expertos para la Detección Interpretable de Colisiones de Partículas

Planteamiento del Problema
El Gran Colisionador de Hadrones (LHC) en el CERN genera volúmenes vastos de datos complejos de colisiones de alta energía, lo que requiere técnicas analíticas sofisticadas para la clasificación de eventos e identificación de objetos. Aunque las Redes Neuronales de Grafos (GNN) han demostrado un rendimiento de vanguardia al representar las colisiones como grafos, su naturaleza de "caja negra" limita su explicabilidad. En la física de altas energías (HEP), donde establecer la confianza en descubrimientos novedosos es primordial, la incapacidad de interpretar los procesos de toma de decisiones es una limitación crítica. Los métodos de explicabilidad post-hoc existentes (por ejemplo, PGExplainer, GNNExplainer) a menudo carecen de robustez y no logran proporcionar una comprensión integral del funcionamiento interno de los modelos. Este trabajo aborda el desafío de desarrollar modelos de aprendizaje automático que alcancen una alta precisión predictiva mientras integran explicabilidad intrínseca directamente en la arquitectura.

Metodología
Los autores proponen una arquitectura novedosa, el Transformador de Grafos con Mezcla de Expertos (MGT), diseñada para clasificar eventos de colisión simulados del experimento ATLAS. La tarea consiste en distinguir eventos de señal de Supersimetría (SUSY) poco comunes de los fondos dominantes del Modelo Estándar (SM) (específicamente la producción de $t\bar{t}$ y $Wt$).

1. Representación de Datos: Cada evento de colisión se modela como un grafo totalmente conectado $G$ con 6 o 7 nodos. Los nodos representan partículas reconstruidas (jets, jets con etiquetado b, leptones y energía transversal faltante) y contienen hasta 6 características cinemáticas (por ejemplo, momento transversal $p_T$, pseudorapidez $\eta$, ángulo azimutal $\phi$, cuantil de etiquetado b y masa).
2. Arquitectura: El modelo se basa en el marco de trabajo del Transformador de Grafos (GT).
   • Transformador de Grafos: Utiliza mecanismos de autoatención de múltiples cabezales para actualizar las representaciones de los nodos, lo que permite al modelo visualizar mapas de atención y cuantificar las relaciones entre partículas.
   • Mezcla de Expertos (MoE): La red de alimentación hacia adelante (FFN) estándar dentro de las capas del Transformador es reemplazada por un bloque MoE. Este bloque consiste en múltiples subredes especializadas ("expertos") y un mecanismo de compuerta (gating). El enrutador asigna dinámicamente cada nodo a un subconjunto de expertos basándose en patrones aprendidos.
   • Entrenamiento: El modelo emplea un mecanismo de compuerta top-K con ruido y una pérdida de equilibrio de carga para asegurar la utilización uniforme de los expertos y evitar el entrenamiento de solo un pequeño subconjunto de ellos.
3. Mecanismo de Explicabilidad: La arquitectura proporciona explicabilidad intrínseca a través de dos canales principales:
   • Mapas de Atención: La visualización de los pesos de atención revela en qué nodos (partículas) se enfoca el modelo al realizar las predicciones.
   • Especialización de Expertos: El análisis de la red de compuerta identifica qué expertos se activan para tipos de nodos específicos, revelando cómo el modelo descompone el problema en tareas especializadas.

Contribuciones Clave

• Explicabilidad Intrínseca: A diferencia de los métodos post-hoc, este enfoque integra la explicabilidad en el núcleo de la arquitectura del modelo, permitiendo rastrear directamente las predicciones hacia mecanismos de atención y activaciones de expertos específicos.
• Interpretación Alineada con la Física: El estudio demuestra que los mecanismos internos del modelo se alinean con principios físicos conocidos. Por ejemplo, los mapas de atención muestran un alto enfoque en los b-jets ($b_1, b_2$) y la energía faltante para los eventos de señal, lo cual es consistente con la desintegración de los bosones de Higgs y los candidatos a materia oscura.
• Patrones de Especialización de Expertos: El análisis revela que los expertos de MoE se especializan naturalmente en distintos tipos de partículas (por ejemplo, expertos específicos que se enfocan en pares de leptones/energía o en b-jets), ofreciendo una visión granular del proceso de toma de decisiones del modelo.

Resultados
El modelo fue evaluado en un conjunto de datos de 800,000 eventos simulados por Monte Carlo (400,000 de señal, 400,000 de fondo) utilizando diez ejecuciones con diferentes semillas aleatorias.

• Rendimiento: El modelo MGT logró un rendimiento superior comparado con las arquitecturas de referencia (MLP, GCN y el GT estándar). Alcanzó una precisión de 0.852 ± 0.0005 y un AUC de 0.929 ± 0.0039, superando al GT estándar (0.849 de precisión, 0.928 AUC) y excediendo significativamente al GCN (0.750 de precisión, 0.832 AUC).
• Hallazgos de Interpretabilidad:
  • Análisis de Atención: En los eventos de señal clasificados correctamente, la atención se concentró fuertemente en los nodos de $b_1$, $b_2$ y energía, reflejando la física de la desintegración de Higgs y la energía faltante de la materia oscura. En los eventos de fondo, la atención estuvo más distribuida.
  • Análisis de Expertos: Los expertos en la primera capa mostraron una especialización distinta (por ejemplo, los Expertos 2 y 4 se enfocaron en leptones y energía, mientras que los Expertos 1 y 5 se enfocaron en b-jets). Las capas más profundas mostraron mayor solapamiento, sugiriendo un nivel de abstracción más alto en la combinación de características.
  • Estudios de Ablación: Los autores encontraron que la pérdida de equilibrio de carga es crítica para la especialización de los expertos; sin ella, los expertos no logran especializarse. La configuración óptima involucró 6 expertos con enrutamiento top-2.

Significancia y Reivindicaciones
El artículo afirma que el modelo MGT equilibra con éxito un alto rendimiento predictivo con una explicabilidad robusta, ofreciendo una herramienta transparente para el análisis de datos de HEP. Al vincular las predicciones con características informadas por la física a través de mapas de atención y especialización de expertos, el modelo fomenta la confianza en los descubrimientos impulsados por IA. Los autores enfatizan que la explicabilidad es tan crucial como la precisión cuando se explora la posible nueva física. Concluyen que, si bien el estudio actual se centra en una búsqueda específica de SUSY, el marco demuestra la viabilidad de las GNN intrínsecamente explicables para datos científicos complejos, allanando el camino para metodologías de IA más fiables. Los autores mantienen la modestia respecto a la escalabilidad, señalando que la generalización a conjuntos de datos más grandes y tareas de física más amplias requiere de una validación adicional.