Dissecting Jet-Tagger Through Mechanistic Interpretability

Este trabajo aplica técnicas de interpretabilidad mecanicista a un clasificador de jets Particle Transformer, revelando que un circuito disperso de seis cabezas que se basa en representaciones derivadas de correladores de energía y en una arquitectura específica de fuente-relé-lectura puede recuperar el rendimiento completo de clasificación del modelo, al tiempo que demuestra que el descenso de gradiente descubre naturalmente características de subestructura de jets con significado físico.

Lo esencial

La Gran Imagen: Abriendo la Caja Negra

Imagina a un detective altamente cualificado (un programa informático llamado Particle Transformer) entrenado para examinar una escena del crimen caótica (un "chorro" de partículas creado en un colisionador de partículas) y decidir: "¿Fue esto un Quark Top (el sospechoso) o simplemente ruido de fondo aleatorio (QCD)?".

Durante mucho tiempo, supimos que el detective era increíblemente bueno resolviendo el caso, pero no sabíamos cómo funcionaba. Era una "caja negra". Este artículo es como contratar a un equipo forense para abrir el cerebro del detective, mapear exactamente qué neuronas se están activando y explicar la lógica paso a paso que utilizan para llegar a un veredicto.

El Cerebro del Detective: Un Equipo de Especialistas

Los investigadores descubrieron que el detective no utiliza todo su cerebro para resolver el caso. En su lugar, depende de un equipo diminuto y eficiente de solo seis especialistas (de 16 disponibles) para realizar el 97% del trabajo pesado. Llaman a este equipo el "Circuito".

Así es como funciona este equipo de seis personas, usando una analogía de una carrera de relevos:

1. El Explorador (Fuente Primaria): Un especialista en la primera capa del cerebro actúa como explorador. Esta persona no busca directamente a los "malos". En su lugar, escanea la multitud buscando el "ruido de fondo" (partículas suaves que colisionan). Al entender el ruido, preparan el escenario para todos los demás. Son la persona más importante; si se les elimina, el equipo pierde casi toda su capacidad para resolver el caso.
2. El Segundo Explorador (Fuente Secundaria): Otro especialista en la primera capa ayuda al Explorador. Son muy similares al Explorador pero se centran en detalles ligeramente diferentes.
3. Los Corredores de Relevos (Capa Media): Tres especialistas en las capas medias actúan como corredores. Toman la información de los Exploradores y buscan algo específico: pares de partículas pesadas y energéticas. En el mundo de la física de partículas, un Quark Top decae en un "bosón W", que luego se divide en dos partículas pesadas. Estos corredores son expertos en detectar estos pares pesados.
   • Descubrimiento Crucial: Aunque se supone que el detective debe encontrar un "Quark Top" (que es una estructura de 3 partes), estos corredores en realidad solo buscan el "bosón W" (una estructura de 2 partes). El artículo sugiere que el detective encontró un atajo: "Si puedo encontrar el pesado bosón W de 2 partes, puedo estar bastante seguro de que es un Quark Top". Es como un detective resolviendo un asesinato encontrando el arma del crimen, en lugar de intentar reconstruir toda la escena del crimen.
4. El Juez (Lectura): Un especialista en la capa final actúa como juez. No mira las partículas directamente. En su lugar, toma los informes de los Corredores de Relevos, los resume y toma la decisión final: "Culpable" (Quark Top) o "No Culpable" (Fondo).

El Momento "¡Ajá!": No es una Nueva Idea, Solo un Nuevo Lenguaje

Uno de los hallazgos más sorprendentes del artículo se refiere a cuándo toma el detective su decisión.

Normalmente, pensamos que el detective recopila pistas capa por capa y de repente grita: "¡Sé quién lo hizo!" al final. Sin embargo, los investigadores descubrieron que el detective en realidad conoce la respuesta casi inmediatamente (después de la primera capa de escaneo).

Entonces, ¿por qué el paso final parece tan dramático?

• La Analogía: Imagina que el detective tiene la respuesta escrita en un código secreto (un idioma diferente) en su primera capa. El paso final no es "pensar" ni "encontrar nuevas pistas"; es simplemente traducir ese código secreto al inglés claro para que el juez final pueda leerlo.
• El artículo llama a esto una "Rotación de Base". La información ya estaba allí; solo necesitaba ser rotada a la orientación correcta para ser entendida por la salida final.

¿Qué Aprendió Realmente el Detective?

Los investigadores también verificaron qué tipo de "física" aprendió el detective. Compararon las notas internas del detective con las fórmulas físicas estándar utilizadas por expertos humanos.

• El Resultado: El detective ignoró las fórmulas complejas de 3 partes que los humanos suelen usar. En su lugar, descubrió y prefirió naturalmente fórmulas más simples de 2 partes (llamadas Correladores de Energía).
• La Conclusión: La computadora no necesitó que un humano le dijera: "¡Busca el bosón W!". Descubrió por sí misma que encontrar el decaimiento pesado de 2 partes era la forma más fácil y confiable de resolver el acertijo. Redescubrió una verdad física significativa simplemente intentando ganar el juego.

Resumen

Este artículo demuestra que podemos tomar una IA compleja y moderna utilizada en física de altas energías y realizarle una ingeniería inversa para encontrar un circuito simple y lógico en su interior.

1. Es eficiente: Un pequeño equipo de 6 "neuronas" hace casi todo el trabajo.
2. Es lógico: El equipo sigue un camino claro: Explorar el ruido $\rightarrow$ Relevar los pares pesados $\rightarrow$ Juzgar el resultado.
3. Es inteligente: La IA descubrió que resolver un subproblema más simple (encontrar el bosón W de 2 partes) es la mejor manera de resolver el gran problema (encontrar el Quark Top).
4. Es una traducción: El paso final de la IA es simplemente traducir su conocimiento temprano y secreto en una respuesta final, no un nuevo descubrimiento.

Los autores concluyen que las herramientas que utilizamos para entender los modelos de lenguaje de IA (como los Chatbots) funcionan perfectamente bien para entender la IA en física de partículas, revelando que estas máquinas pueden aprender verdades físicas profundas por sí mismas.

Resumen técnico

Resumen Técnico: Diseccionando Jet-Tagger a través de la Interpretabilidad Mecanística

Enunciado del Problema
Las arquitecturas de aprendizaje profundo, particularmente el Transformador de Partículas (ParT), han logrado un rendimiento de vanguardia en tareas de etiquetado de jets, como distinguir las desintegraciones hadrónicas del quark top de los jets de fondo de QCD. Sin embargo, los mecanismos computacionales internos mediante los cuales estos modelos toman decisiones de clasificación permanecen opacos. Aunque trabajos anteriores han utilizado métodos de atribución post-hoc (por ejemplo, valores de Shapley, mapas de saliencia) o visualización de atención, estos enfoques identifican qué entradas importan, pero no explican cómo la red las combina ni aíslan el sub-red causal mínima (circuito) responsable del comportamiento. Este artículo pretende cerrar esta brecha aplicando el conjunto completo de herramientas de interpretabilidad mecanística —originalmente desarrollado para modelos de lenguaje natural— a un clasificador de física de jets.

Metodología
Los autores entrenaron un Transformador de Partículas pequeño (4 capas de atención de partículas, 4 cabezas por capa, ~1.3M de parámetros) en un subconjunto del conjunto de datos de referencia para el etiquetado de quarks top (señal: $t \to Wb \to q\bar{q}b$; fondo: quarks ligeros/gluones). El análisis emplea una serie de técnicas de intervención y sondeo:

1. Ablación Cero: Establecer sistemáticamente la salida de cabezas de atención individuales en cero para medir su importancia estructural mediante la caída en la diferencia media de logits.
2. Parcheo de Trayectorias (Path Patching): Una intervención causal donde la salida de una cabeza específica en una entrada "limpia" se sustituye en una entrada "corrupta" (utilizando reemplazo de partículas dentro del lote o permutación de todo el jet). Esto aísla los efectos directos y los efectos de trayectoria (flujo de información) entre cabezas.
3. Lente de Logits y Sondas por Capa: La lente de logits estándar proyecta las representaciones intermedias a través de la última cabeza de clasificación entrenada. Para resolver la desalineación de bases, los autores también entrenaron sondas de regresión logística por capa en la representación de cada capa para determinar la accesibilidad lineal real de la información de clase.
4. Sondeo Lineal: Entrenar modelos de regresión Ridge para predecir observables clásicos de subestructura de jets (por ejemplo, $N$-subjettiness, Funciones de Correlación de Energía) a partir del flujo residual en varias profundidades para caracterizar el contenido físico de las representaciones internas.

Contribuciones y Resultados Clave

• Identificación de un Circuito Escaso de Seis Cabezas: Mediante ablación cero y parcheo de trayectorias, los autores identificaron un circuito mínimo de seis cabezas de atención que recupera el 97.3% del AUC del modelo completo. Este circuito es significativamente más performante que subconjuntos de seis cabezas muestreados aleatoriamente (situándose en el percentil 96 de una distribución de referencia aleatoria).
• Estructura Causal (Fuente-Relé-Lectura): El circuito exhibe una jerarquía causal clara:
  • Fuente Primaria ($L0H1$): Una sola cabeza en la primera capa de atención de partículas que actúa como la fuente causal principal. Por sí sola recupera el 88.6% del AUC del modelo completo y exhibe una "super-recuperación" en el parcheo de trayectorias. Atiende preferentemente a pares de partículas suaves y colineales, proporcionando una normalización contextual.
  • Fuente Secundaria ($L0H2$): Una cabeza en la misma capa con un espacio representacional similar al de $L0H1$ pero un rol causal distinto, contribuyendo una señal complementaria.
  • Cabezas Relé ($L1H0, L1H1, L1H3$): Un grupo de cabezas en la segunda capa que atienden selectivamente a la subestructura de pares duros (alta masa invariante, alto $k_T$). Su función es condicional a la señal aguas arriba de las cabezas fuente.
  • Cabeza de Lectura ($L3H3$): Una sola cabeza en la cuarta capa que agrega las señales retransmitidas.
• Rotación de Bases vs. Ganancia de Información: El análisis estándar de lente de logits sugirió que la información de clase solo emerge en el primer bloque de atención de clase ($Cls0$), mostrando un salto de un AUC de 0.111 a 0.973. Sin embargo, las sondas entrenadas por capa revelaron que la señal discriminadora de clases ya es linealmente accesible con un AUC $\approx$ 0.97 en la primera capa de atención de partículas ($L1$). El salto dramático en $Cls0$ se interpreta, por tanto, no como una generación de nueva información, sino como una rotación de bases que alinea la señal latente con la base de la cabeza de clasificación final.
• Contenido Físico: Codificación de 2 puntas sobre 3 puntas: El sondeo lineal del flujo residual contra observables clásicos reveló que el modelo codifica preferentemente observables de correlador de energía de 2 puntas (por ejemplo, $D^{(\beta=1)}_2$) sobre observables de 3 puntas (por ejemplo, $C^{(\beta=1)}_3, N^{(\beta=1)}_3$), a pesar de que la tarea es el etiquetado de top de 3 puntas. El modelo factoriza implícitamente el problema, centrándose en identificar la desintegración hadrónica del bosón $W$ (una subestructura de 2 puntas) en lugar de la topología completa de 3 cuerpos. Esta preferencia persiste incluso después de residualizar contra la masa del jet.
• Base de Correlador de Energía: El flujo residual está alineado preferentemente con la base del Correlador de Energía sobre la base de $N$-subjettiness, lo que sugiere que el modelo ha redescubierto estructuras físicamente significativas a través del descenso de gradiente sin supervisión explícita.

Significado y Afirmaciones
El artículo afirma demostrar que los métodos de interpretabilidad mecanística desarrollados para modelos de lenguaje natural pueden transferirse con éxito a clasificadores de física de jets. Los hallazgos indican que:

1. El descenso de gradiente puede redescubrir aspectos físicamente significativos del etiquetado de jets (específicamente la factorización de la desintegración del top en un subproblema de bosón $W$) sin restricciones físicas explícitas en la arquitectura.
2. La estructura de circuito fuente-relé-lectura identificada aquí puede ser un patrón característico para Transformadores en dominios de física, distinto de la estructura específica de los modelos de lenguaje.
3. El estudio destaca la necesidad de estrategias de corrupción en la variedad (on-manifold) para el parcheo de trayectorias en conjuntos de datos de física cinemáticamente estrechos, ya que la corrupción fuera de la variedad (Gaussiana) puede conducir a incompatibilidades estructurales con las formulaciones estándar de puntuación de recuperación.

Los autores permanecen modestos respecto a la universalidad de sus hallazgos, señalando que el circuito específico de seis cabezas se identificó en un modelo pequeño y que modelos más grandes podrían poseer estructuras de circuito más ricas. También reconocen que las sondas lineales proporcionan una cota inferior del contenido de información, ya que las codificaciones no lineales no se detectan.