Detection vs. Execution: Single-Bucket Probes Miss Half the Mamba-2 State Sink

Resumen Técnico: Detección vs. Ejecución en los Sumideros de Estado de Mamba-2

Planteamiento del Problema
La interpretabilidad mecánica opera frecuentemente bajo la suposición de que las sondas que identifican una firma representacional específica (por ejemplo, alta activación en tokens de frontera) también identifican el circuito de computación específico que la ejecuta. Este artículo desafía esa suposición en el contexto de los Modelos de Espacio de Estados (SSM) selectivos, específicamente Mamba-2. Los autores investigan el "sumidero de estado" (state sink), análogo al "sumidero de atención" (attention sink) en los Transformers, donde ocurre una activación desproporcionada de la compuerta condicionada al token ($\Delta$) en los tokens de frontera (BOS/Inicio de Secuencia y saltos de línea). El problema central es si las sondas de un solo cubo (single-bucket probes) —que clasifican unidades basándose en la activación de una sola clase— localizan con precisión el circuito causal responsable de la gestión del estado, o si pasan por alto un conjunto de unidades mucho mayor y funcionalmente distinto que comparten la misma firma representacional.

Metodología
El estudio emplea una rigurosa cuadrícula de intervención que involucra 12,096 celdas (6,060 evaluadas) a través de cuatro puntos de control de modelos: Mamba-1 (130M, 2.8B) y Mamba-2 (1.3B, 2.7B), junto con Pythia (160M–2.8B) como línea base de control arquitectónico.

• Clasificación de Especialistas: Las unidades (canales en Mamba-1, cabezales en Mamba-2) se clasifican como "especialistas" si su media de activación $\Delta$ pre-softplus en una clase de token específica (BOS, salto de línea, etc.) supera un umbral ($\tau=0.5$) relativo a un conjunto de referencia. Esto genera especialistas de "un solo cubo" (una clase) y "duales" (dos clases).
• Intervenciones Causales: Los autores realizan estudios de ablación utilizando seis tipos de intervención (ej. gate_zero, delta_zero, gate_one) sobre conjuntos específicos de unidades. Miden el diferencial en la Negativa Log-Verosimilitud (NLL) en los primeros ocho tokens después de un token de frontera, comparando el conjunto de especialistas contra complementos aleatorios de igual tamaño.
• Comparación de Granularidad: El estudio contrasta Mamba-1 (granularidad de canal, proyecciones $\Delta$ independientes) con Mamba-2 (granularidad de cabezal, proyecciones $\Delta$ compartidas por cabezal) para aislar efectos arquitectónicos.
• Validación: Los resultados se validan en dos corpora (wikitext-2, Pile-10k), múltiples escalas y una tarea de recuperación de contexto largo (RULER Needle-in-a-Haystack).

Contribuciones Clave y Resultados

1. Causalidad Condicionada a la Granularidad:

   • En Mamba-1 (granularidad de canal), los especialistas de BOS de un solo cubo portan la masa causal del sumidero de estado. Su ablación provoca una degradación significativa del rendimiento, y la sonda rastrea con precisión el circuito de ejecución.
   • En Mamba-2 (granularidad de cabezal), las sondas de un solo cubo recuperan solo una pequeña "capa de ejecución" (~5% de los cabezales). Sin embargo, un conjunto mucho mayor de "cabezales duales" (27–35% de los cabezales), que comparten la misma firma representacional (alta activación tanto en BOS como en saltos de línea), es omitido por la clasificación de un solo cubo.

2. Dicotomía Detección vs. Ejecución:

   • Capa de Ejecución: Los especialistas de BOS de un solo cubo son la capa de ejecución causal. Su ablación colapsa la precisión de recuperación de 1.00 a 0.00 en las tareas RULER NIAH y aumenta significamente la NLL en los objetivos de contexto BOS y saltos de línea.
   • Capa de Detección: Los cabezales duales exhiben una alta similitud representacional con los especialistas (similitud de coseno $\approx 0.89$ en BOS vs. salto de línea) pero muestran efectos causales sustancialmente más débiles bajo ablación. De hecho, la ablación de los cabezales duales a menudo tiene efectos insignificantes o incluso negativos en el rendimiento en comparación con la línea base, lo que indica que son una "capa de detección representacional" en lugar de la capa de ejecución funcional.
   • Sondas Engañosas: El artículo demuestra que la similitud representacional no implica equivalencia funcional. Las sondas de un solo cubo subestiman sistemáticamente el sustrato causal en un factor de 3 a 7 veces en tamaño de conjunto en la granularidad de cabezal de Mamba-2.

3. Impulsor Arquitectónico:

   • La asimetría se atribuye a la arquitectura de Mamba-2, donde la proyección $\Delta$ se comparte entre canales dentro de un cabezal. Esto obliga a los cabezales a multiplexar entre clases de frontera, creando la gran población de cabezales duales.
   • Un experimento de control que consistió en el agrupamiento aleatorio de canales en Mamba-1 para igualar la relación cabezal-clase de Mamba-2 no logró reproducir el dominio de los cabezales duales, lo que implica que la proyección $\Delta$ compartida por cabezal es el impulsor, y no la granularidad del sustrato per se.

4. Línea Base de Arquitectura Cruzada:

   • Los cabezales de atención de Pythia (Transformer) muestran selectividad de un solo cubo pero carecen del dominante conjunto de cabezales duales encontrado en Mamba-2. Esto confirma que la división detección/ejecución es específica de la mecánica de espacio de estados de la arquitectura Mamba-2, no una propiedad universal de los modelos con granularidad de cabezal.

Significado y Reivindicaciones
El artículo sostiene que el movimiento estándar de "etiqueta-a-locus" en la interpretabilidad mecánica —asumir que las unidades de mayor activación identificadas por una sonda son el circuito causal— es sistemáticamente poco fiable en la granularidad de cabezal de M la Mamba-2.

• Necesidad de Re-auditoría: Los hallazgos de interpretabilidad localizados en los cabezales de Mamba-2 mediante sondas de un solo cubo requieren una re-auditoría bajo reglas de agregación de múltiples clases.
• Descomposición Funcional: El sumidero de estado en Mamba-2 se descompone en dos capas funcionales distintas: una pequeña capa de ejecución (especialistas de un solo cubo) y una gran capa de detección (cabezales duales). Separarlas requiere ablación condicional a la clase en lugar de confiar en la similitud de coseno condicional a la clase.
• Ancla Activa: A diferencia de los sumideros de atención en los Transformers, que pueden ser pasivos o suprimibles, el sumidero de estado de Mamba-2 es un ancla activa donde cabezales específicos son necesarios para la recuperación de contexto largo y la gestión de fronteras de estado.

Los autores concluyen que, si bien las sondas son efectivas para el análisis representacional, pueden fallar al identificar el verdadero circuito causal en arquitecturas con proyecciones compartidas y perfiles de activación multiplexados, lo que hace necesaria la verificación basada en intervenciones para distinguir entre detección y ejecución.