Resumen Técnico: Poda de Nodos Multigranular para el Descubrimiento de Circuitos Causales

1. Planteamiento del Problema

El descubrimiento de circuitos causales busca identificar subredes mínimas dentro de los Modelos de Lenguaje Extensos (LLMs) que impulsan causalmente comportamientos específicos. Aunque los métodos existentes han avanzado en la interpretabilidad mecánica, enfrentan dos limitaciones primarias:

1. Escalabilidad Computacional: Los enfoques actuales dependen predominantemente de la poda de aristas (eliminación de conexiones entre componentes). Dado que el número de aristas en un transformador crece cuadráticamente con el número de nodos ($O(|N|^2)$), la poda a nivel de arista se vuelve computacionalmente intratable a medida que aumenta el tamaño del modelo. Además, estos métodos suelen requerir el almacenamiento de representaciones del flujo residual y activaciones intermedias desentrelazadas, lo que conduce a un uso de memoria prohibitivo.
2. Limitaciones de Granularidad: Los métodos existentes suelen operar a una granularidad gruesa, tratando cabezales de atención o bloques MLP enteros como unidades atómicas. Esto pasa por alto la posibilidad de que neuronas (nodos) específicas dentro de estos componentes sean las responsables de los comportos objetivo, lo que resulta en circuitos recuperados que omiten estructuras finas y retienen componentes innecesarios.

2. Metodología

Los autores proponen un marco de poda de nodos multigranular que descubre circuitos mínimos a través de múltiples niveles de granularidad dentro de un único paso de optimización unificado.

2.1 Diseño de Granularidad Jerárquica

En lugar de podar aristas entre componentes gruesos, el marco aprende máscaras sobre nodos en cuatro niveles distintos:

• Nivel de Bloque: Bloques de Atención completos y Bloques MLP.
• Nivel de Cabezal: Cabezales de Atención individuales.
• Nivel de Neurona: Neuronas individuales dentro de las capas de Atención y las dos capas lineales de los bloques MLP (MLP1/Keys y MLP2/Values).

2.2 Configuración de Poda Causal de Dos Flujos

Para identificar la importancia causal sin introducir cambios en la distribución (un riesgo con la ablación de cero simple), el método emplea una pasada hacia adelante de dos flujos:

• Flujo Limpio (Clean Stream): El modelo procesa una entrada diseñada para eludir un comportamiento específico (por ejemplo, predecir un año > 50).
• Flujo Corrupto (Corrupted Stream): El modelo procesa una versión mínimamente perturbada de la entrada (por ejemplo, invertir pronombres de género o cambiar el año de inicio) para eludir una predicción diferente.
• Interpolación: Para cada nodo $n_i$, una máscara aprendible $m_i \in [0, 1]$ interpola entre la activación limpia ($h^{clean}$) y la activación corrupta ($h^{corrupted}$):
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  La optimización busca encontrar máscaras que permitan al modelo producir la salida correcta (coincidiendo con el comportamiento limpio) mientras se maximiza la dispersión (sparsity).

2.3 Optimización y Parametrización de Máscaras

• Distribución Hard Concrete: Las máscaras se parametrizan utilizando la distribución Hard Concrete para permitir la optimización basada en gradientes mientras se induce un comportamiento aproximadamente binario (0 o 1).
• Objetivo Unificado: El objetivo de entrenamiento combina tres componentes:
  1. Pérdida de Fidelidad ($L_{faith}$): Una divergencia KL que asegura que la distribución de salida del circuito podado coincida con la del modelo base.
  2. Pérdida de Tarea ($L_{task}$): Una pérdida de ranking basada en margen que asegura que el circuito distinga correctamente entre las opciones de verdad fundamental y las de distracción.
  3. Regularizador de Dispersión ($L_0$): Fomenta que las máscaras sean cero, minimizando el tamaño del circuito.
• Consistencia Jerárquica: Tras el entrenamiento, las máscaras se binarizan con una restricción que asegura que si un nodo padre (por ejemplo, un bloque MLP) es podado, todos sus nodos hijos (por ejemplo, neuronas dentro de ese bloque) también sean desactivados.

3. Contribuciones Clave

1. Marco Multigranular: La introducción de un marco de optimización unificado que aprende máscaras simultáneamente a través de bloques de transformadores, cabezales de atención, bloques MLP y neuronas individuales.
2. Escalabilidad y Eficiencia: El método evita la escala cuadrática de la poda de aristas y la necesidad de almacenar activaciones intermedias para todas las aristas. Reduce el uso de memoria GPU en un factor de 3× a 11× en comparación con métodos previos, permitiendo el descubrimiento de circuitos en Llama 3.1-8B usando una sola GPU de 40GB.
3. Descubrimiento de Grano Fino: El enfoque identifica circuitos significativamente más compactos que las líneas base, logrando específicamente una alta dispersión en los niveles de bloque MLP y de neurona, algo que los métodos anteriores no lograron alcanzar.

4. Resultados Experimentales

El método fue evaluado en tres tareas estándar de descubrimiento de circuitos (Identificación de Objeto Indirecto, Pronombres de Género y Mayor Que) utilizando modelos GPT-2, GPT-2-XL y Llama.

4.1 Dispersión y Compactación

• Poda de MLP y Neurona: El método propuesto supera significativamente a las líneas base de poda de aristas (EAP, EP) y a la poda de pesos no estructurada (DiscoGP). En el escenario menos favorable, podó un 33.34% más de bloques MLP y un 59.8% más de neuronas que los métodos anteriores.
• Estructuras Dependientes de la Tarea: Los circuitos descubiertos exhiben topologías distintas según la tarea:
  • IOI: Dominado por la computación distribuida de MLP; la atención es mayormente podada, excepto en las capas finales.
  • Pronombres de Género (GP): Altamente localizado y discontinuo; la computación de MLP se restringe a unas pocas capas, mientras que la atención es dispersa.
  • Mayor Que (GT): Exhibe una dispersión extrema con bloques contiguos de capas inactivas y patrones de ejecución tipo "salto" (skip-like).

4.2 Fidelidad y Consistencia

• Desempeño: Los circuitos podados mantienen una alta precisión de la tarea (a menudo >95%) y una baja divergencia KL con respecto al modelo completo.
• Consistencia: Experimentos a través de cinco semillas aleatorias muestran que los circuitos descubiertos (métricas de fidelidad y perfiles de dispersión) son altamente estables.
• Análisis Cualitativo: En Llama 3.2-1B, se encontró que los cabezales de atención supervivientes en la tarea IOI desempeñan roles funcionales específicos (IO-Movers, S-Inhibitors), reflejando los circuitos descubiertos manualmente en GPT-2, lo que valida la capacidad del método para recuperar conocimientos mecánicos.

4.3 Eficiencia de Cómputo

• Memoria: El método requiere al menos 3× menos memoria que las líneas base de poda de aristas (por ejemplo, 6.27 GB frente a 72.79 GB para EAP en GPT-2).
• Tiempo: El entrenamiento se completa en menos de 6 minutos, significativamente más rápido que los métodos de poda de aristas (aprox. 50 minutos), aunque ligeramente más lento que el método EAP de una sola pasada (que sacrifica tiempo por un uso masivo de memoria).

5. Significado y Reivindicaciones

El artículo afirma que su principal importancia radica en superar los cuellos de botella de escalabilidad y granularidad de la actualidad. Al cambiar de la poda de aristas a la poda de nodos multigranular, los autores demuestran que:

• Muchas neuronas retenidas por los métodos de granularidad gruesa son innecesarias y pueden eliminarse con un impacto insignificante en el desempeño.
• Los modelos más grandes (por ejemplo, Llama 8B) soportan una poda igual de fuerte o más fuerte que los modelos más pequeños, lo que sugiere que el aumento de escala introduce redundancia en lugar de requerir circuitos específicos de la tarea más densos.
• El método proporciona un camino computacionalmente factible para descubrir circuitos causales de grano fino en modelos a gran escala, revelando que diferentes comportamientos son codificados por patrones distintos de actividad de atención, MLP y neuronas.

Los autores reconocen una limitación: el marco identifica qué nodos son necesarios, pero no recupera explícitamente la estructura de interacción (aristas) entre los nodos supervivientes. Sugieren que una dirección futura es combinar la selección de nodos con la recuperación de aristas para inferir estructuras de interacción.