Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Resumen Técnico: Evidencia Causal de Representaciones de Pila en el Modelado de Lenguajes de Contador mediante Transformers

Planteamiento del Problema
Si bien la interpretabilidad mecánica ha identificado con éxito características latentes en modelos de lenguaje (LM) entrenados en lenguajes formales, persiste una brecha crítica en establecer el papel causal de estas características. Trabajos previos, como los de Tiwari et al. (2025), demostraron que los transformers entrenados en lenguajes de contador (p. ej., Dyck-1, Shuffle-k) aprenden representaciones consistentes con una estructura de pila subyacente. Sin embargo, estos estudios no determinaron si estas representaciones son meramente correlacionadas con la tarea o si son causalmente necesarias para la computación del modelo. Este artículo aborda la pregunta: ¿Están las representaciones de pila aprendidas involucradas causalmente en la capacidad del transformer para predecir el siguiente token?

Metodología
El estudio emplea un flujo de trabajo experimental de tres etapas que involucra entrenamiento, sondeo e intervención causal:

1. Entrenamiento del Modelo: Se entrena un transformer de solo codificador para la predicción del siguiente token en lenguajes Shuffle-k (donde $k \in \{2, 4, 6, 8\}$). El Shuffle-k se define como el entrelazado de $k$ lenguajes Dyck-1 independientes, lo que requiere que el modelo rastree la profundidad de $k$ pilas distintas simultáneamente. La arquitectura del modelo consiste en una sola capa con 4 cabezales de atención, una dimensión de embedding de 32 y una dimensión oculta de 64. El modelo logra una precisión de validación perfecta para la época 25, asegurando que todos los análisis posteriores se realicen sobre un modelo totalmente convergido.
2. Sondeo de la Profundidad de la Pila: Se entrenan sondas lineales para predecir la profundidad de pilas específicas ($c_i(t)$) a partir de los estados ocultos del modelo ($h^{(l)}_t$) en la capa $l$. Para asegurar que las sondas capturen información estructural genuina en lugar de correlaciones espurias, se entrena una sonda de control sobre etiquetas perturbadas. La métrica Selectividad se define como la diferencia de precisión entre la sonda experimental y la sonda de control. Una alta selectividad indica que la profundidad de la pila está genuinamente codificada en el espacio oculto.
3. Intervención Causal (Ablación):
   • Extracción de Dirección: La matriz de pesos $W$ de la sonda se descompone mediante la Descomposición en Valores Singulares (SVD). El vector singular derecho superior, $w$, se identifica como la dirección principal correspondiente a la representación de la profundidad de la pila.
   • Ablación: Durante la inferencia, el componente del estado oculto alineado con $w$ se elimina utilizando la operación $h' = h - \alpha(w \cdot h)w$, donde $\alpha \in [0, 1]$ representa la fuerza de la ablación.
   • Control: Direcciones aleatorias son abladas de manera similar para servir como línea base.
   • Evaluación: Los estados ablados se pasan directamente al cabezal del decodificador (sin volver a ejecutar las capas internas del transformer) para medir el impacto en la Precisión Posicional ($A_{pos}$) y la Precisión Secuencial ($A_{seq}$).

Resultos Clave

• Existencia de la Representación: Las sondas lineales extraen con éxito la información de la profundidad de la pila con alta selectividad (p. ej., 85.8% para $k=6$, Pila 2), confirmando que el modelo aprende representaciones consistentes con una estructura de pila.
• Necesidad Causal: La ablación dirigida de la dirección principal de la pila ($w$) causa un colapso dramático en el rendimiento del modelo. A medida que la fuerza de la ablación $\alpha$ aumenta de 0 a 1, la Precisión Secuencial cae a casi 0%.
• Especificidad: La ablación de direcciones aleatorias no tiene un efecto significativo en la precisión, confirmando que el colapso del rendimiento es específico a la eliminación de la representación de la pila y no un artefacto general de la manipulación de vectores.
• Sensibilidad de la Métrica: El estudio observa que, si bien la precisión posicional disminuye, la precisión secuencial colapsa de forma más aguda. Esto se atribuye a la naturaleza compuesta de los errores en la generación de secuencias, donde un solo token incorrecto invalida toda la cadena.

Significancia y Reivindicaciones
El artículo afirma proporcionar evidencia empírica sólida de que las representaciones de pila en los transformers entrenados en lenguajes de contador no son solo artefactos emergentes, sino que son causalmente necesarias para la computación interna del modelo y la predicción del siguiente token.

Al demostrar que la eliminación de una dirección de representación específica y aprendida conduce a un fallo casi total, los autores argumentan que el transformer utiliza explícitamente estas estructuras tipo pila para resolver las tareas de razonamiento lógico requeridas por los lenguajes formales. Esto mueve el campo más allá del análisis correlacional (sondeo) hacia una comprensión mecánica causal, validando que el comportamiento de "caja negra" del modelo depende de primitivas computacionales específicas e identificables.

Direcciones Futuras
Los autores sugieren que el trabajo futuro debería explorar:

• Intervenir a mitad de la red pasando estados ablados o editados a través de la red de alimentación hacia adelante (feed-forward) para observar cambios en los resultados posteriores.
• Investigar el papel de diferentes arquitecturas en la emergencia, redundancia y causalidad de estas representaciones.
• Comprender cómo los transformers manejan la manipulación de otras primitivas computacionales emergentes como los contadores.