Explainable AI for Jet Tagging: A Comparative Study of GNNExplainer, GNNShap, and GradCAM for Jet Tagging in the Lund Jet Plane

Este trabajo evalúa y compara métodos de explicabilidad basados en perturbaciones, en valores de Shapley y en gradientes, adaptados a la representación del Plano de Jet de Lund, demostrando que estas técnicas correlacionan con éxito las predicciones de las redes neuronales con observables clásicos de la QCD y revelan cambios distintos en el enfoque entre los regímenes perturbativo y no perturbativo en la identificación de jets.

Lo esencial

La Gran Imagen: El Problema de la "Caja Negra"

Imagina que eres un detective tratando de resolver un crimen. Tienes un asistente de IA súper inteligente que puede observar una escena del crimen desordenada y decirte, con un 99% de precisión, exactamente quién es el criminal. Sin embargo, cuando le preguntas a la IA por qué piensa eso, simplemente dice: "Estoy seguro". No te muestra sus notas ni explica su razonamiento.

En el mundo de la física de partículas (específicamente en el Gran Colisionador de Hadrones), los científicos utilizan modelos de IA similares de "caja negra" para identificar chorros (jets). Los chorros son ráfagas de partículas diminutas creadas cuando los protones chocan entre sí. A veces estos chorros provienen de partículas comunes (como quarks o gluones), y a veces provienen de partículas raras y pesadas (como el bosón de Higgs o un quark top).

La IA es excelente para detectar la diferencia, pero los físicos están preocupados: ¿Está la IA realmente aprendiendo las leyes de la física, o simplemente está memorizando peculiaridades de la simulación informática utilizada para entrenarla? Si solo está memorizando la simulación, podría fallar al observar datos reales.

La Solución: El Mapa del "Plano Lund del Chorro"

Para solucionar esto, los investigadores decidieron dejar de mirar las partículas como un montón desordenado y empezar a verlas como un mapa.

Utilizaron algo llamado el Plano Lund del Chorro. Piensa en esto como un mapa topográfico de una cordillera montañosa.

• El eje X representa lo ancho que es el chorro de partículas.
• El eje Y representa cuánta energía tienen las partículas.

Cada "división" individual donde una partícula se rompe en dos partículas más pequeñas se plotea como un punto en este mapa. Debido a que este mapa se basa en las leyes reales de la física (Cromodinámica Cuántica), cada punto tiene un significado específico y conocido. Es como tener un mapa donde cada colina y valle corresponde a un evento geológico específico.

El Experimento: Tres Diferentes "Linternas"

Los investigadores tomaron tres tipos diferentes de modelos de IA (los "detectives") y proyectaron tres tipos diferentes de "linternas" (herramientas de IA Explicable) sobre ellos para ver qué partes del mapa estaban observando.

1. GNNExplainer (La linterna "¿Qué pasaría si?"): Esta herramienta pregunta: "Si oculto esta parte del mapa, ¿la IA sigue obteniendo la respuesta correcta?". Destaca las áreas más críticas al ver qué sucede cuando se eliminan.
2. GNNShap (La linterna "Parte Justa"): Esta herramienta utiliza un concepto matemático de la teoría de juegos para calcular exactamente cuánto "crédito" merece cada punto en el mapa por la decisión final. Es como dividir una cuenta de manera justa basándose en quién comió qué.
3. GradCAM (La linterna "Mapa de Calor"): Esta herramienta observa los "neuronas" internas de la IA disparándose y dibuja un mapa de calor que muestra qué áreas estaban más activas cuando la IA tomó una decisión.

Los Hallazgos: ¿Qué vio realmente la IA?

Los investigadores compararon las vistas de la "linterna" de la IA contra la física conocida del mapa. Esto es lo que encontraron:

1. La IA Aprendió la Física Real
Para las partículas pesadas (como el quark top o el bosón de Higgs), las linternas de la IA iluminaron correctamente las divisiones "duras" específicas en el mapa donde la partícula pesada se desintegró.

• Analogía: Si estás buscando un tipo específico de árbol en un bosque, la IA no solo adivinó; señaló correctamente la forma única de las hojas y la corteza. El estudio demuestra que la IA no está solo adivinando; ha aprendido la estructura real de cómo se desintegran estas partículas pesadas.

2. La "Anomalía de la QCD" (El Bosque Brumoso)
Para las partículas comunes (chorros de QCD), la linterna de la IA no se centró en un punto específico. En su lugar, iluminó todo el mapa, especialmente las áreas "suaves" y "anchas".

• Analogía: Imagina tratar de identificar un pino genérico. No hay una sola rama única que lo defina; es la forma general y la forma en que se dispersan las agujas. La IA se dio cuenta correctamente que, para estos chorros comunes, la respuesta yace en el patrón completo, no solo en un lugar especial. Los investigadores llaman a esto una "anomalía de fidelidad", pero en realidad es una señal de que la IA entiende la física perfectamente.

3. Herramientas Diferentes para Trabajos Diferentes
El estudio encontró que ninguna "linterna" única funciona mejor para todos los modelos de IA.

• Para algunos modelos de IA, la herramienta de "Parte Justa" (GNNShap) fue la mejor para encontrar las divisiones duras.
• Para otros, la herramienta de "Mapa de Calor" (GradCAM) fue mejor para ver los patrones más amplios.
• Conclusión: No puedes simplemente elegir una herramienta de explicación y usarla para siempre. Necesitas emparejar la herramienta con el modelo de IA específico que estás utilizando.

4. El Efecto "Impulso"
Los investigadores observaron chorros moviéndose a diferentes velocidades (baja energía vs. alta energía). Descubrieron que a medida que los chorros se movían más rápido, el enfoque de la IA se volvía más nítido y concentrado en las divisiones duras específicas, tal como predice la física que debería suceder.

La Conclusión

El artículo concluye que los clasificadores modernos de chorros de IA no son solo cajas negras mágicas. Cuando proyectas la luz correcta sobre ellos, puedes ver que han aprendido genuinamente las leyes de la física. Saben dónde se desintegran las partículas pesadas y entienden la diferencia entre un evento pesado específico y un chorro genérico de partículas.

Esto es algo importante porque significa que los científicos pueden confiar más en estas herramientas de IA cuando se utilizan para buscar nuevas partículas desconocidas en experimentos futuros. La IA no está solo memorizando un libro de texto; está haciendo física.

Resumen técnico

A continuación se presenta un resumen técnico detallado del artículo "IA Explicable para la Etiquetación de Jets: Un Estudio Comparativo de GNNExplainer, GNNShap y GradCAM para la Etiquetación de Jets en el Plano de Jet de Lund".

1. Planteamiento del Problema

A pesar del rendimiento de vanguardia de las Redes Neuronales de Grafos (GNN) como ParticleNet y Particle Transformer en la identificación de jets (chorros colimados de partículas) en el Gran Colisionador de Hadrones (LHC), sus procesos de toma de decisiones siguen siendo "cajas negras".

• La Brecha: Los observables clásicos de la subestructura de jets (por ejemplo, la subjetividad $N$-subjetosa $\tau_{21}$, las funciones de correlación de energía) se derivan de los principios fundamentales de la Cromodinámica Cuántica (QCD) y son calculables analíticamente. En contraste, las redes neuronales aprenden representaciones de características complejas y opacas.
• El Riesgo: Si un modelo aprende artefactos de las simulaciones de Monte Carlo (MC) (por ejemplo, modelos específicos de hadronización o respuestas del detector) en lugar de física fundamental, su rendimiento se degradará en datos de colisiones reales.
• El Desafío: Los métodos existentes de IA Explicable (XAI) no se han aplicado sistemáticamente a las GNN basadas en el Plano de Jet de Lund (LJP). El LJP ofrece una oportunidad única porque cada nodo en su representación gráfica corresponde a una división de partones físicamente significativa, permitiendo una comparación directa entre la atención del modelo y la teoría de la QCD.

2. Metodología

Los autores desarrollaron un marco integral para evaluar y comparar tres paradigmas distintos de XAI adaptados a las GNN basadas en LJP.

A. Representación de Datos: El Plano de Jet de Lund (LJP)

En lugar de tratar los jets como imágenes o nubes de puntos genéricas, el estudio utiliza el Plano de Jet de Lund, un espacio de fase 2D definido por:

• $\ln(k_T)$: Logaritmo del momento transversal de la división.
• $\ln(1/\Delta R)$: Logaritmo del inverso del ángulo de división.
• Construcción del Grafico: Cada nodo representa una emisión específica de partones en la historia de la radiación del jet. Esto permite que los métodos de explicabilidad asignen puntuaciones de importancia a eventos de división física específicos.

B. Modelos y Métodos de Explicabilidad

El estudio evalúa tres arquitecturas:

1. LundNet: Una GNN que utiliza EdgeConv en un grafo $k$-NN estático construido a partir del LJP.
2. ParticleNet: Una GNN que utiliza EdgeConv Dinámico (recalculando los vecinos en cada capa).
3. Particle Transformer: Una arquitectura basada en Transformer que utiliza auto-atención sobre los constituyentes del LJP.

Estas se emparejan con tres métodos de XAI:

1. GNNExplainer (basado en perturbación): Optimiza una máscara para maximizar la información mutua entre el subgrafo y la predicción.
2. GNNShap (basado en valores de Shapley): Utiliza la teoría de juegos cooperativos para asignar valores de importancia justos a las aristas basándose en su contribución marginal.
3. GradCAM (basado en gradientes): Calcula mapas de activación de clase ponderados por gradientes para identificar nodos/características importantes.

C. Marco de Evaluación

Los autores introdujeron una evaluación informada por la física más allá de las métricas estándar de fidelidad:

1. Máscaras de Verdad de MC: Se construyeron máscaras de explicación de verdad fundamental trazando los partones de dispersión dura a través del proceso de radiación (por ejemplo, identificando la estructura de 3 puntas para quarks Top o la estructura de 2 puntas para bosones de Higgs).
2. Métricas de Fidelidad ($Fid^+, Fid^-$): Medición de cuánto cae la predicción del modelo cuando se eliminan o retienen aristas importantes.
3. Correlación de Subestructura: Cálculo de correlaciones de rango de Spearman entre la importancia de nodos asignada por XAI y observables clásicos ($\tau_{21}, \tau_{32}, C_2, C_3$).
4. Estratificación por $p_T$: El análisis se realizó en tres intervalos de momento transversal ($500\text{--}700$ GeV, $800\text{--}1000$ GeV e inclusivo) para estudiar la transición entre regímenes no perturbativos y perturbativos.

3. Contribuciones Clave

1. Primera Comparación Sistemática Multi-Método: El artículo proporciona la primera comparación rigurosa de métodos de XAI basados en perturbación, teoría de juegos y gradientes específicamente para GNN que operan sobre la representación físicamente fundamentada del LJP.
2. Verdad Fundamental Informada por la Física: La introducción de máscaras de explicación de verdad de MC permite una verificación cuantitativa de si el modelo se centra en las topologías de desintegración físicas correctas (por ejemplo, divisiones duras para la señal frente a radiación suave para el fondo).
3. Descubrimiento de la "Anomalía de Fidelidad": Los autores identificaron un comportamiento contra intuitivo pero físicamente significativo donde eliminar aristas "importantes" de los jets de QCD no degrada la confianza del modelo, a diferencia de los jets de señal.
4. Emparejamiento Arquitectura-Método: El estudio demuestra que ningún método único de XAI es universalmente mejor; la elección óptima depende de la arquitectura neuronal subyacente.

4. Resultados Clave

A. Aprendizaje del Modelo y Alineación con la Física

• Redescubrimiento de la QCD: Las tres arquitecturas aprendieron con éxito características canónicas de subestructura de QCD.
  • Jets Top ($t \to bW \to bq\bar{q}'$): Los modelos se centraron en la estructura de división dura de 3 puntas.
  • Jets Higgs ($H \to c\bar{c}$): Los modelos se centraron en la región de caída de masa de 2 puntas.
  • Jets QCD: Los modelos identificaron correctamente la ausencia de escalas duras, confiando en la radiación suave integrada de gran ángulo (cola de Sudakov).
• Fuerza de Correlación: El Particle Transformer mostró la correlación más fuerte ($|\rho| \approx 0.5\text{--}0.6$) entre sus pesos de atención internos y los observables clásicos (por ejemplo, $\tau_{21}$), lo que sugiere que los mecanismos de atención codifican características cinemáticas de pares de manera más directa que las GNN de paso de mensajes locales.

B. Comportamientos Específicos del Método

• GNNShap: Más adecuado para Particle Transformer y LundNet. Actúa como un "filtro" para ramificaciones duras y colineales, asignando un peso alto a las aristas específicas responsables de la masa de la señal.
• GradCAM: Más adecuado para ParticleNet. Tiende a ponderar la actividad de fragmentación de bajo $k_T$ (radiación suave), capturando la ruta completa de paso de mensajes en lugar de solo las aristas causales.
• GNNExplainer: Altamente efectivo para revelar el papel del intercambio de momento ($\ln z$) y la extensión angular, a menudo destacando el patrón completo de radiación.

C. La Anomalía de Fidelidad (Jets QCD)

• Para jets de señal (2 puntas/3 puntas), eliminar las $k$ aristas más importantes destruye la predicción ($Fid^+ \to 1$), como se esperaba.
• Para jets QCD, eliminar las $k$ aristas más importantes resulta en ninguna caída significativa en la predicción ($Fid^+ \to 0$).
• Interpretación: Esto no es un fallo del explicador, sino una validación de la factorización de la QCD. Los jets de QCD carecen de una escala dura discriminante única; su clasificación depende del patrón colectivo de emisiones suaves. Por lo tanto, ningún subgrafo pequeño es suficiente para definir la clase.

D. Dependencia de $p_T$

• A medida que aumenta el momento transversal ($p_T$), las divisiones duras se vuelven más colimadas.
• En regímenes de alto $p_T$, el acuerdo entre diferentes métodos de XAI generalmente mejora para los jets de señal, ya que la señal dura se vuelve más distinta del fondo suave. Sin embargo, para jets de 2 puntas altamente impulsados, el límite entre la señal y el fondo suave se vuelve ambiguo, causando cierta divergencia en las interpretaciones de los métodos.

5. Significado y Perspectivas

• Confianza en el ML: El estudio proporciona evidencia cuantitativa de que los etiquetadores modernos de aprendizaje profundo no solo están explotando artefactos de simulación, sino que están aprendiendo la física fundamental de la subestructura predicha por la QCD.
• Cambio Metodológico: Establece que la evaluación multi-método es esencial. Confiar en una sola técnica de XAI puede llevar a interpretaciones incompletas o sesgadas de lo que ha aprendido un modelo.
• Aplicaciones Futuras:
  • Detección de Anomalías: Utilizar XAI para localizar poblaciones inesperadas en el plano de Lund para búsquedas más allá del Modelo Estándar (BSM).
  • Diseño de Arquitecturas: Utilizar mapas de explicación para podar entradas o diseñar etiquetadores más eficientes (por ejemplo, restringiendo las entradas a regiones de Lund de alta importancia).
  • Validación Basada en Datos: El marco está listo para aplicarse a datos reales del LHC (ATLAS/CMS) para validar los etiquetadores de ML directamente contra mediciones experimentales en lugar de simulaciones.

En conclusión, este artículo cierra la brecha entre el aprendizaje automático de alto rendimiento y la física de partículas teórica, demostrando que con la representación adecuada (Plano de Jet de Lund) y un XAI riguroso, las redes neuronales pueden interpretarse como aprendiendo el lenguaje de la QCD.