Resumen Técnico: La estructura de la tarea revierte el codificado de estado por capas en modelos de secuencias

Planteamiento del problema

La interpretabilidad mecánica de los modelos de secuencias suele tratar los codificados de estado por capas como rasgos arquitectónicos fijos. El supuesto prevaleciente es que los modelos recurrentes (p. ej., LSTMs, GRUs) y los modelos de espacio de estados (SSMs como Mamba) concentran el estado legible en las capas tardías, mientras que los modelos basados en atención (Transformers) distribuyen la acumulación de estado a través de la profundidad. Este artículo desafía la visión de que estas estrategias de codificación de estado son invariantes arquitectónicas estables. Postula que la misma arquitectura revierte su perfil de codificación de estado por capas cuando la estructura computacional de la tarea cambia. Además, la literatura suele confundir dos posibles impulsores de estos patrones: la estructura algebraica (específicamente la conmutatividad) frente a la estructura computacional (actualizaciones de prefijo frente al seguimiento de una pila o stack). Los autores pretenden desentrazar estos factores y determinar si el codificado de estado es una propiedad de la arquitectura por sí sola o una interacción entre la arquitectura y la tarea.

Metodología

El estudio emplea una suite experimental controlada a través de cinco familias de arquitecturas: Transformers, Mamba, Mamba-2, LSTMs y GRUs.

Tareas:
Los autores utilizan tres tareas formales de seguimiento de estado y un benchmark semi-real:

1. Paridad: Una tarea de actualización de prefijo de un solo bit y conmutativa.
2. Composición de permutación $S_3$: Una tarea de actualización de prefijo no conmutativa de 6 clases. Esta sirve como discriminante crítico: si la conmutatividad algebraica impulsa el perfil de codificación, $S_3$ debería comportarse como Dyck; si la estructura computacional lo impulsa, $S_3$ debería comportarse como Paridad.
3. Dyck-$k$: Una tarea de seguimiento de pila de profundidad acotada (emparejamiento de paréntesis) con una profundidad máxima de 10.
4. Profundidad de código (Code-Depth): Una tarea de profundidad de anidamiento a nivel de token derivada de código Python en CodeSearchNet, que sirve como validación semi-real.

Configuración Experimental:

• Tareas Formales: Los experimentos utilizan modelos de 4 capas con $d_{model}=128$. Las arquitecturas están emparejadas por parámetros (aprox. 830K parámetros) siempre que es posible. El entrenamiento ocurre en una longitud de secuencia $L=40$, con evaluación en $L \in \{40, 60, 80, 100, 120\}$.
• Escala Preentrenada: El ajuste fino (fine-tuning) se realiza en Mamba-130M (24 capas) y Pythia-160M/410M (12/24 capas).
• Mediciones:
  • Sondeo Lineal (Linear Probing): Entrenamiento de clasificadores de regresión logística sobre estados ocultos congelados para medir la legibilidad del estado en cada capa.
  • Intervenciones Causales: Ablación de cero (proyectar fuera las direcciones de la sonda) y ablación de subespacio (eliminar las principales direcciones PCA o singulares) para probar la necesidad funcional.
  • Generalización OOD: Aplicación de ablaciones entrenadas en $L=40$ a longitudes no vistas ($L=80, 100, 120$).
  • Atribución de Componentes: Ablación del componente Conv1D de Mamba bajo un backend de PyTorch puro emparejado para aislar su contribución a través de las tareas.
  • Parcheo de Activación (Activation Patching): En Mamba preentrenado, el parcheo de activaciones de posición media de entradas limpias en entradas corruptas para medir la recuperación del logit.

Contribuciones Clave

1. Reversión Dependiente de la Tarea: El artículo demuestra que el patrón de "concentración tardía" frente a "construcción distribuida" no es un invariante arquitectónico. Se invierte dependiendo de la tarea:
   • Tareas de Actualización de Prefijo (Paridad, $S_3$): Los modelos recurrentes y los SSMs concentran el estado legible en la capa final, mientras que los Transformers lo construyen gradualmente a través de las capas.
   • Seguimiento de Pila Acotada (Dyck): El patrón se revierte. Los Transformers logran una legibilidad casi perfecta de inmediato (capas tempranas), mientras que los modelos recurrentes y los SSMs construyen la representación progresivamente.
2. Desentrañar la Álgebra y la Computación: Al introducir la tarea no conmutativa $S_3$, los autores muestran que el perfil de codificación sigue la estructura computacional (actualización de prefijo vs. pila) en lugar de la estructura algebraica (conmutatividad). $S_3$ se agrupa con Paridad en todas las arquitecturas, refutando la hipótesis de que la conmutatividad es el principal motor del perfil de codificación.
3. Desacoplamiento de Legibilidad–Causalidad: El estudio revela que la capa donde el estado es más linealmente legible (mejor capa de la sonda) no siempre coincide con la capa que es causalmente necesaria para la tarea. Este desacoplamiento es particularmente pronunciado en modelos preentrenados, donde las representaciones son más redundantes y distribuidas.
4. Dependencia de la Tarea a Nivel de Componente: Incluso dentro de una sola arquitectura (Mamba), la importancia de componentes específicos varía según la tarea. Por ejemplo, el componente Conv1D es crítico para $S_3$ y Paridad, pero no tiene un efecto significativo en el rendimiento de Dyck, lo que desafía la noción de un "mecanismo" único para una arquitectura.

Resultados Clave

• Nivel de Comportamiento: Las líneas base recurrentes generalizan perfectamente en Paridad y $S_3$ y superan a los SSMs y Transformers en Dyck. El orden de rendimiento de $S_3$ (Recurrente > SSM > Transformer) refleja el de Paridad, no el de Dyck.
• Perfiles de Sondeo (Probing):
  • Paridad/$S_3$: Mamba y LSTMs muestran un salto brusco en la precisión de la sonda en la capa final ($L_3$ en modelos de 4 capas). Los Transformers muestran un aumento gradual desde $L_1$ hasta $L_3$.
  • Dyck: Los Transformers alcanzan una alta precisión en $L_0$ o $L_1$. Mamba y LSTMs muestran un aumento gradual, alcanzando la saturación solo en capas posteriores.
  • Confirmación de $S_3$: $S_3$ sigue el perfil de Paridad (concentración tardía para RNNs/SSMs, gradual para Transformers), confirmando la hipótesis de la estructura computacional.
• Intervenciones Causales:
  • En modelos formales de 4 capas, la ablación de la mejor dirección de la sonda causa caídas significativas en la precisión, y estas direcciones siguen siendo necesarias en longitudes de secuencia más largas (OOD).
  • En Modelos Preentrenados, ocurre una división:
    • Pythia (Dyck): Existe un fuerte cuello de botella en las capas medias (la ablación en $L_6$–$L_7$ reduce la precisión en un ~81% en 160M). La mejor capa de la sonda ($L_{11}$) no es la más críticamente causal.
    • Mamba (Paridad/Dyck/$S_3$): La capa final es altamente legible, pero la ablación de una sola dirección causa caídas insignificantes (<0.2%). Sin embargo, el parcheo de activación de posición media en la capa final recupera ~97–98% de la brecha del logit, lo que indica que el estado es suficiente pero está distribuido en muchas direcciones, no capturado por una única sonda lineal.
• Atribución de Conv1D: Eliminar Conv1D de Mamba causa una gran caída en $S_3$ (-12.64 pp) y Paridad (-6.83 pp), pero no una caída significativa en Dyck (+4.04 pp, no significativa). Esto indica que la importancia del componente depende de la tarea.

Significado y Reivindicaciones

El artículo argumenta que preguntar "¿cómo codifica la arquitectura X el estado?" es incompleto sin especificar la estructura computacional de la tarea. Los autores afirman que:

1. Las firmas mecánicas son propiedades conjuntas: Las estrategias de codificación de estado no son rasgos arquitectónicos fijos, sino que emergen de la interacción entre los sesgos inductivos del modelo y los requisitos computacionales de la tarea (actualización de prefijo vs. seguimiento de pila).
2. Cautela metodológica: Extrapolar hallazgos mecánicos de un régimen de tareas (p. ej., actualizaciones de prefijo) a otro (p. ej., seguimiento de pila acotada) conlleva el riesgo de confundir firmas específicas de la tarea con invariantes arquitectónicos.
3. Legibilidad vs. Causalidad: El sondeo identifica dónde la información está linealmente disponible, pero esto no siempre mapea con dónde se produce el cuello de botella computacional, especialmente en modelos preentrenados de gran escala donde la redundancia es alta.
4. Álgebra vs. Computación: En el contexto de estas tareas formales, la estructura computacional (actualización de prefijo secuencial vs. pila jerárquica) es un mejor predictor de los perfiles de codificación por capas que las propiedades algebraicas como la conmutatividad.

Los autores concluyen que sus hallazgos se mantienen a través de tareas formales, escalas preentrenadas (hasta 410M de parámetros) y benchmarks de código semi-reales, lo que sugiere una interacción robusta entre la estructura de la arquitectura y la de la tarea en el modelado de secuencias.