Resumen Técnico: Los Transformers Representan Linealmente Modelos del Mundo Altamente Estructurados

Enunciado del Problema

El artículo investiga si los transformers entrenados en trazas de razonamiento secuencial desarrollan "modelos del mundo" internos que reflejen la estructura algebraica subyacente del dominio de la tarea, en lugar de meramente su presentación superficial. Mientras que trabajos anteriores (por ejemplo, Li et al. [2023], Nanda et al. [2023]) demostraron que los modelos entrenados en Othello aprenden espontáneamente representaciones lineales de la ocupación del tablero (una característica por celda), sigue sin estar claro si esto se mantiene para tareas con diferentes restricciones estructurales. Específicamente, los autores preguntan: en una tarea combinatoria de satisfacción de restricciones como el Sudoku, donde la validez está determinada por subestructuras superpuestas (filas, columnas y cajas de 3×3) en lugar de celdas aisladas, ¿representa el modelo el estado celda por celda o subestructura por subestructura?

Metodología

Los autores entrenaron un transformer causal de solo decodificador de 8 capas (8 cabezas de atención, $d_{modelo}=576$) en 2.7 millones de trazas de resolución de Sudoku generadas por un solver estilo Norvig con retroceso. El modelo alcanzó una precisión del 98.4% por celda y del 97.5% por cuadrícula en un conjunto de prueba retenido.

Para analizar el cálculo interno, los autores emplearon una metodología de "sondeo y parche":

1. Sondeo Lineal: Se entrenaron clasificadores de regresión logística sobre las activaciones del flujo residual en el token [clues_end] (que marca el final de las pistas iniciales) para decodificar conceptos específicos. Se probaron tres familias de sondeos:
   • Estado de celda: Predecir el dígito en una celda específica (81 sondeos).
   • Candidatos de celda: Predecir dígitos candidatos válidos para una celda específica (729 sondeos).
   • Estado de subestructura: Predecir si un dígito está presente en una subestructura específica (fila, columna o caja) (243 sondeos).
2. Intervención Causal (Parcheo de Activación): Para verificar si las representaciones decodificadas se utilizan causalmente, los autores restaron direcciones de sondeo específicas del flujo residual y midieron la caída resultante en los logits del token objetivo.
3. Análisis de Atención: Los autores descompusieron las cabezas de atención en circuitos de Consulta-Clave (QK) y Salida-Valor (OV) para determinar qué información atienden las cabezas y cómo modifican el flujo residual.
4. Lente de Logits y Atribución Directa de Logits (DLA): Estas herramientas se utilizaron para rastrear cómo evolucionan las predicciones a través de las capas y para atribuir cambios en los logits a componentes específicos (cabezas o neuronas).
5. Ablación de Neuronas: Se identificaron y ablaron neuronas específicas en la capa final MLP para probar su papel causal en la toma de decisiones.

Contribuciones y Resultados Clave

1. Emergencia de un Modelo del Mundo de Subestructuras

Contrario a la representación "celda por celda" encontrada en los modelos de Othello, el transformer de Sudoku construye un modelo del mundo organizado en torno al álgebra de restricciones del rompecabezas.

• Separabilidad Lineal: Los sondeos lineales para "¿está presente el dígito $d$ en la subestructura $S$?" alcanzan una precisión de coincidencia exacta perfecta (1.0) en capas medias (Capas 4–7) en todos los 243 pares (subestructura, dígito).
• Limitaciones a Nivel de Celda: En contraste, los sondeos que predicen el estado de celdas individuales se estabilizan en una precisión de ~80%. Los sondeos para candidatos de celda alcanzan una precisión del 100%, pero exhiben alineación estructural: los vectores de candidatos para celdas que comparten una subestructura están altamente correlacionados, lo que sugiere que son combinaciones lineales de características a nivel de subestructura.
• Estabilidad: Esta representación de subestructura emerge en la Capa 4 y permanece geométricamente estable a lo largo de la traza de resolución, según lo confirmado por experimentos de transferencia entre capas y entre posiciones.
• Uso Causal: El parcheo de activación confirma que estas direcciones son causalmente necesarias; sobrescribirlas para indicar que un dígito está presente en una subestructura hace que la predicción del modelo para esa colocación colapse.

2. Mecanismos de Atención Alineados con Subestructuras

Las cabezas de atención de capas medias (Capas 4–6) están especializadas para operar en subestructuras específicas.

• Enrutamiento: Las cabezas de atención concentran su enfoque en filas, columnas, cajas, bandas o pilas específicas.
• Aplicación de Restricciones: Los circuitos OV de estas cabezas implementan un mecanismo simétrico de verificación de restricciones. Si un dígito está presente en la subestructura atendida, la cabeza suprime los logits para ese dígito en las celdas relevantes; si está ausente, promueve el dígito.
• Localización: La ablación causal de cabezas específicas (por ejemplo, una cabeza que atiende a las columnas 4–6) aumenta los logits de colocaciones ilegales solo dentro de esa subestructura específica, dejando otras subestructuras sin afectar.

3. El Circuito del "Único Desnudo"

Los autores identificaron un circuito disperso y monosémico en la capa final MLP responsable de comprometerse con colocaciones de "único desnudo" (celdas con exactamente un candidato válido).

• Estado Pre-MLP: Antes del MLP final, el flujo residual ya clasifica el dígito correcto en primer lugar en el 98.77% de los casos únicos de único desnudo.
• Neuronas Dedicadas: El MLP final contiene 91 neuronas dedicadas que se activan exclusivamente cuando su celda asociada tiene exactamente un candidato. Estas neuronas son monosémicas y no se activan para otras cantidades de candidatos.
• Mecanismo de Amplificación: Cuando una neurona de único desnudo se activa, aplica un impulso uniforme y masivo (+3.00 logits) a todos los tokens de colocación para esa celda específica. Dado que el dígito correcto ya estaba clasificado primero por las capas medias, esta amplificación uniforme preserva el liderazgo y concentra la probabilidad softmax en el dígito correcto.
• Confirmación Causal: Ablar estas neuronas reduce la probabilidad del único desnudo objetivo en ~0.59 sin afectar otros únicos desnudos simultáneos en el mismo tablero, confirmando su papel como la puerta final de compromiso.

Significado y Afirmaciones

El artículo afirma demostrar que la geometría de los modelos del mundo emergentes está moldeada por el álgebra de restricciones del dominio en lugar de su descomposición superficial. En el Sudoku, donde la validez se define por la intersección de filas, columnas y cajas, el modelo aprende a representar el estado en términos de estas subestructuras.

Además, el estudio proporciona una cuenta algorítmica completa, mecanicista y de extremo a extremo de cómo un transformer resuelve una tarea de razonamiento combinatorio:

1. Representación: Las capas medias construyen un modelo del mundo lineal y alineado con subestructuras.
2. Propagación de Restricciones: Cabezas de atención especializadas aplican restricciones suprimiendo candidatos inválidos dentro de esas subestructuras.
3. Decisión: Un conjunto disperso de neuronas monosémicas en la capa final detecta estados lógicos específicos (únicos desnudos) y amplifica la respuesta correcta precalculada.

Los autores concluyen que las herramientas de interpretabilidad mecanicista pueden revertir exitosamente estos complejos procesos de razonamiento, revelando que los circuitos de decisión resultantes son dispersos, monosémicos y completamente interpretables.

Limitaciones

Los autores señalan que estos hallazgos son específicos de una única arquitectura de modelo (8 capas, 8 cabezas) entrenada en un único tipo de rompecabezas (Sudoku). No afirman que los modelos del mundo de subestructuras o los circuitos de único desnudo aparecerán necesariamente en modelos de diferentes tamaños o en aquellos entrenados en trazas sin retroceso. Además, el análisis de cómo se seleccionan los patrones de atención (circuitos QK) no está detallado completamente, y la evidencia causal depende de un protocolo de parcheo específico.