Closure-Validated Circuit Discovery in Attention Heads: Co-activation Proposes, Ablation Disposes

Resumen Técnico: Descubrimiento de Circuitos Validados por Cierre en Cabezales de Atención

Declaración del Problema
La interpretabilidad mecánica se ha desplazado hacia la visión de grupos de componentes (subespacios, variedades, conjuntos de cabezales) en lugar de unidades individuales como el objeto básico de estudio. Un enfoque metodológico prevalente propone encontrar estos grupos mediante el agrupamiento (clustering) de estadísticas de coactivación (por ejemplo, utilizando modelos de Ising binarizados sobre características de autoencoders dispersos o cabezales de atención). Sin embargo, persiste una brecha crítica: ¿constituye un grupo de coactivación estadísticamente coherente un "circuito" en el sentido de carga funcional? Un circuito verdadero es un grupo de componentes cuya eliminación degrada causalmente el rendimiento del modelo en entradas donde se aplica la función del grupo.

El artículo investiga si las señales económicas —específicamente la pertenencia a clústeres de coactivación, la selectividad de atención al objetivo y la razón de participación— pueden servir como evidencia suficiente para el descubrimiento de circuitos, o si simplemente proponen candidatos que requieren validación causal. Los autores se preguntan si el agrupamiento por sí solo puede distinguir entre:

1. Circuitos Verdaderos: Grupos que implementan conjuntamente una función y son de carga funcional (load-bearing).
2. Comunidades Coactivas: Grupos que se activan juntos debido a impulsores ascendentes compartidos, pero contribuyen de manera redundante o marginal.
3. Comunidades Cancelables: Grupos que añaden señales de sesgo que la red compensa, de modo que su eliminación podría incluso mejorar el rendimiento.

Metodología
El estudio emplea un marco de "propuesta y validación" a través de tres modelos de lenguaje de escala 1B (Pythia 1B, OLMo 1B y el Mixture-of-Experts OLMoE-1B-7B) y dos distribuciones de entrada (inducción sintética y texto natural).

1. Propuesta (Agrupamiento):

   • Extracción de Señal: Para cada cabezal de atención, se extrae una señal sin plantillas (template-free): el peso máximo de atención asignado a cualquier posición de clave en una entrada.
   • Binarización: Las señales se binarizan en la mediana por cabezal, asegurando que cada cabezal se active en el 50% de los ejemplos por construcción, eliminando sesgos de tasa de activación.
   • Ajuste de Ising: Se ajusta un modelo de Ising de pares a las señales binarizadas utilizando verosimilitud pseudo para estimar las fuerzas de acoplamiento entre cabezales.
   • Recuperación de Comunidades: Se aplica agrupamiento espectral a la matriz de acoplamiento absoluto. El número de clústeres ($k$) se selecciona para maximizar el Índice de Rand Ajustado (ARI) contra una clasificación supervisada de las capacidades de los cabezales (por ejemplo, token anterior, inducción).
   • Selección de Candidatos: Se selecciona un sub-clúster específico basado en pureza, aislamiento (ratio de acoplamiento interno frente a externo) y restricciones de tamaño.

2. Validación (Prueba de Cierre):

   • Ablación: La comunidad seleccionada se ablaciona poniendo a cero las porciones de salida de los cabezales identificados antes de que entren en la corriente residual.
   • Comparación: El daño por ejemplo (cambio en la pérdida de entropía cruzada, precisión top-1 y logit del token objetivo) se compara contra cinco ablaciones de control aleatorias emparejadas del mismo tamaño.
   • Criterio de Veredicto: Un candidato pasa el cierre si el daño está en la dirección esperada (la pérdida aumenta, la precisión/logit disminuye) y es estadísticamente significativo (z-score > 1.8) respecto a la distribución de control aleatoria.

3. Proxies Extendidos:

   • El estudio también rastrea la selectividad de atención al objetivo y la razón de participación (PR) a través de los puntos de control de entrenamiento para ver si estas señales "económicas" correlacionan con la función de carga a lo largo del tiempo.

Resultados Clave

• Modelos Densos (Pythia 1B, OLMo 1B):

  • Las comunidades de coactivación identificadas tanto en distribuciones de texto sintético como natural pasan el cierre.
  • La ablación de estas comunidades aumenta consistentemente la pérdida y disminuye la precisión/logit del objetivo en comparación con los controles aleatorios.
  • OLMo 1B (Texto Natural): Produjo la señal más fuerte, con un z-score de +36.4$\sigma$ en pérdida y -50.9$\sigma$ en el logit del objetivo.
  • Pythia 1B (Texto Natural): Mostró una "firma de redundancia descendente". Aunque el logit del objetivo cayó significativamente (-52.5$\sigma$), el cambio agregado en la pérdida fue pequeño (+1.95$\sigma$). Esto indica que el modelo posee vías descendentes redundantes que reconstruyen la salida ablada, aislando la pérdida de entropía cruzada mientras la confianza de la predicción específica colapsa.

• Mixture-of-Experts (OLMoE-1B-7B):

  • Fallo Marginal: En texto natural, el modelo de Ising marginal colapsó (ARI $\approx$ 0.006), fallando en recuperar cualquier estructura.
  • Recuperación Condicional a la Ruta: Estratificar los ejemplos por patrones de enrutamiento (usando k-means en los pesos de enrutamiento) y ajustar Isings separados permitió recuperar una señal estadísticamente significativa (ARI 0.191, +3.05$\sigma$ por encima de un nulo de partición aleatoria).
  • Fallo de Cierre: A pesar de la recuperación estadística, la comunidad descubierta falló el cierre en la dirección incorrecta. Ablar la comunidad mejoró la pérdida (z = -5.64$\sigma$ en el estrato de descubrimiento, -12.94$\sigma$ fuera de él). Los cabezales actuaron como ruido modulado por la ruta; eliminarlos ayudó al modelo, particularmente en entradas fuera de ese estrato de enrutamiento específico.

• Eje de Entrenamiento Proxies:

  • Función sin Forma: En OLMo 1B, los cabezales BOS (inicio de secuencia) eran de carga funcional (pasaron el cierre) a los 2B tokens antes de que su patrón de selectividad de atención-a-BOS se formara (la selectividad era cercana a cero).
  • Forma sin Función: En Pythia 1B, los cabezales de token anterior mostraron alta selectividad y alta razón de participación al 1B de tokens, pero fallaron el cierre (la ablación no dañó el rendimiento, y en algunos casos ayudó).
  • Este desacople bidireccional demuestra que los patrones de atención y la función son constructos distintos, no meramente diferentes sensibilidades de un mismo fenómeno.

Reporte Multimétrica
El artículo argumenta que la pérdida de entropía cruzada por sí sola es una métrica poco fiable para el cierre. Puede verse inflada por el colapso de la varianza de control (como se vio en OLMo texto natural) o enmascarada por la redundancia descendente (como en Pythia texto natural). Los autores recomiendan un estándar multimétrico:

1. Dirección: Todas las métricas deben moverse en la dirección del daño.
2. Logit del Objetivo: El indicador primario de impacto causal.
3. Precisión: Evidencia corroborativa.
4. Pérdida: Un piso conservador, sujeto a la holgura de agregación.

Afirmaciones y Significancia
El artículo concluye con cuatro afirmaciones específicas:

1. Propuesta vs. Descubrimiento: El agrupamiento de coactivación no supervisado es un método de propuesta de circuitos, no un método de descubrimiento. Identifica comunidades candidatas, pero solo la ablación causal (cierre) las valida como de carga funcional.
2. Éxito en Densos: En modelos densos de escala 1B, las comunidades de coactivación en textos sintéticos y naturales pueden proponer exitosamente circuitos de carga funcional que sobreviven al cierre.
3. Fallo en MoE: En arquitecturas MoE, incluso la recuperación estadística condicional a la ruta (que pasa un test nulo cuidadoso) es insuficiente como evidencia de un circuito. Las comunidades recuperadas pueden fallar el cierre en la dirección opuesta (mejorando el rendimiento tras su eliminación).
4. Divergencia de Métricas: La brecha entre los patrones de atención (selectividad/PR) y la función es bidireccional; un cabezal puede tener un patrón agudo sin función, y función sin un patrón.

Alcance y Limitaciones
Los autores declaran explícitamente que no pretenden que estos resultados se generalicen a las variedades de características de los SAE (que utilizan criterios de validación distintos), modelos más grandes, otras arquitecturas (como LLaMA, Mistral) u otras distribuciones de entrada (como código). También señalan que el fallo de "dirección incorrecta" en MoE es específico de la arquitectura y el protocolo de enrutamiento probados, y que la afinidad de Ising utilizada es una elección para la comparabilidad, no necesariamente la óptima para circuitos discretos. El estudio enfatiza que el resultado negativo en MoE es esencial, ya que demuestra que la coherencia estadística no garantiza la carga funcional causal, especialmente en arquitecturas complejas.