Resumen Técnico: Corrección de la Localización de Circuitos Basada en Gradientes mediante la Retropropagación Consciente de las Interacciones

Planteamiento del Problema

La localización de circuitos tiene como objetivo identificar el subconjunto específico de componentes del modelo (por ejemplo, cabezales de atención, neuronas, conexiones) responsables de comportamientos específicos en los Grandes Modelos de Lenguaje (LLMs) para permitir la interpretabilidad mecánica. La mayoría de los métodos existentes operan bajo el supuesto de que los componentes actúan de forma independiente, estimando la importancia perturbando cada componente de forma aislada.

Sin embargo, los componentes de las redes neuronales interactúan, e ignorar estas interacciones conduce a una estimación errónea sistemática de la importancia de los componentes. Los autores identifican una interacción específica y problemática denominada autoreparación de la atención (attention self-repair). En este fenómeno, cuando múltiples posiciones poseen pesos de atención elevados y contienen información similar (vectores de valor), la ablación de un único puntaje de atención tiene un efecto mínimo en la salida. Esto ocurre porque la función de activación softmax redistribuye el peso hacia las posiciones redundantes restantes para compensar. En consecuencia, los métodos de explicación estándar, tanto basados en perturbaciones como en gradientes, subestiman la importancia de los componentes que contribuyen a estos puntajes de atención, lo que a menudo resulta en gradientes evanescentes para puntajes influyentes.

Metodología: Modificaciones de Interacción de Gradientes (GIM)

El artículo introduce las Modificaciones de Interacción de Gradientes (GIM), un método de atribución basado en gradientes diseñado para estimar el impacto de perturbar múltiples componentes interactuantes simultáneamente. GIM modifica la retropropagación estándar a través de tres técnicas complementarias:

1. Gradientes de Softmax con Ajuste de Temperatura (TSG):

   • Motivación: Abordar la autoreparación de la atención, donde los gradientes estándar se desvanecen porque la ganancia al aumentar un peso de softmax se cancela con la pérdida al suprimir otros cuando los valores son similares.
   • Mecanismo: Durante la retropropagación, TSG vuelve a calcular el softmax con una temperatura elevada ($T > 1$) antes de calcular los gradientes. Esto aplana la distribución de probabilidad, reduciendo los pesos grandes y aumentando los pequeños. Al ampliar el conjunto de posiciones "significativas", se relaja el supuesto de que las contribuciones son idénticas en estas posiciones. Esto asegura que las diferencias de gradiente por pares ($g_j - g_k$) no se desvanezcan, permitiendo que la señal del gradiente sobreviva y aproxime el efecto de la perturbación simultánea.

2. Congelación de la Capa de Normalización (Layernorm Freeze):

   • Motivación: La normalización de capa introduce una dinámica de suma cero donde la ablación de un componente reduce la magnitud de la activación, causando que el factor de normalización amplifique los componentes restantes (una forma de autoreparación).
   • Mecanismo: El factor de normalización ($\sigma$) se trata como una constante durante la retropropagación. Esto mide la contribución directa de cada componente antes de que ocurra el reescalado de la normalización, evitando la subestimación de la importancia causada por la amplificación compensatoria.

3. Norma de Gradiente:

   • Motivación: Las operaciones multiplicativas (por ejemplo, $A = QK^T$ donde tanto $Q$ como $K$ dependen de la entrada $X$) conducen a una doble cuenta en los gradientes estándar de primer orden. Perturbar $X$ afecta a $Q$ y $K$ simultáneamente, pero las reglas de la cadena estándar suman sus efectos individuales, sobreestimando el cambio.
   • Mecanismo: Los gradientes de las interacciones multiplicativas se dividen por el número de entradas multiplicativas (por ejemplo, dividir por 2 para $Q \times K$). Esto se alinea con la "regla uniforme" y corrige la sobrecuenta de las contribuciones en la multiplicación de atención-valor, la multiplicación de consulta-clave y la multiplicación de proyección de puerta del MLP.

Contribuciones Clave

1. Marco GIM: La introducción de un método de atribución basado en gradientes que explica explícitamente las interacciones de características, logrando un rendimiento de vanguardia en la pista de localización de circuitos del Benchmark de Interpretabilidad Mecánica (MIB).
2. Identificación de la Autoreparación de la Atención: La identificación formal y la demostración de que la autoreparación de la atención (una interacción tipo compuerta OR) causa que los gradientes se desvanezcan para los puntajes de atención influyentes, provocando una subestimación sistemática de la importancia de los componentes.
3. Superioridad en la Atribución de Características: Demostración de que GIM supera a los métodos basados en gradientes existentes en la atribución de características a través de diversos modelos y tareas.

Resultados

• Localización de Circuitos: En el MIB, GIM supera significativamente a los métodos existentes (incluyendo EAP, EAP-IG, Hybrid-ens e IFR) en múltiples modelos (GPT-2, Qwen-2.5, Gemma-2, Llama-3.1) y tareas (Identificación de Objeto Indirecto, Aritmética, MCQA, ARC). Logra el mayor Ratio de Rendimiento de Circuito (CPR) promedio, particularmente en tareas que requieren razonamiento de múltiples pasos.
• Atribución de Características: A través de siete LLMs ajustados por instrucciones y seis conjuntos de datos (que cubren QA, verificación de hechos, sentimiento y detección de odio), GIM generalmente logra puntuaciones de comprehensividad y suficiencia más altas en comparación con GradientXInput, Integrated Gradients, DeepLIFT y TransformerLRP.
• Validación de la Autoreparación: El análisis empírico confirma que la autoreparación de la atención es prevalente en los LLMs reales, ocurriendo en promedio entre 65 y 1200 veces por entrada. La ablación conjunta de dos grandes puntajes de atención produce un efecto significativamente mayor que la suma de sus ablaciones individuales, validando el análisis teórico.
• Estudio de Ablación: Eliminar cualquier modificación única (TSG, Layernorm Freeze o Gradient Norm) de GIM degrada el rendimiento en la mayoría de los pares modelo-conjunto de datos, confirmando que cada componente contribuye significativamente a la fidelidad del método.

Significado

El artículo afirma que GIM permite un análisis mecánico más fiel de los LLMs al corregir la subestimación sistemática de la importancia de los componentes causada por las interacciones de características. Al modelar explícitamente cómo interactúan los componentes —específicamente a través de la redistribución de softmax y el reescalado de normalización—, GIM proporciona una imagen más precisa de los mecanismos internos que impulsan el comportamiento del modelo. Los autores posicionan a GIM no solo como una mejora incremental, sino como una corrección necesaria al "supuesto de independencia" prevalente en las herramientas de interpretabilidad actuales, permitiendo a los investigadores identificar circuitos que anteriormente estaban ocultos por los gradientes evanescentes. El método se publica como un paquete de Python de código abierto para facilitar una adopción más amplia.