Resumen Técnico: Emergencia de Circuitos Mínimos para la Identificación del Objeto Indirecto en Transformers de Solo Atención

Declaración del Problema

La interpretabilidad mecánica busca la ingeniería inversa de los Grandes Modelos de Lenguaje (LLMs) para convertirlos en circuitos computacionales comprensibles para el ser humano. Sin embargo, la escala, las conexiones residuales y las no linealidades de los LLMs modernos a menudo oscurecen los mecanismos mínimos requeridos para tareas de razonamiento específicas. Mientras que trabajos previos (e.g., Wang et al., 2023) identificaron circuitos de atención complejos y de múltiples saltos para la Identificación del Objeto Indirecto (IOI) en modelos preentrenados como GPT-2 small, estos mecanismos surgen dentro de un entorno de optimización complejo que involucra el lenguaje natural y presiones multitarea. El desafío es determinar si tales capacidades de razonamiento pueden emerger de arquitecturas más simples y restringidas, y aislar los motivos computacionales fundamentales involucrados.

Metodología

Para abordar esto, los autores entrenaron modelos transformer de solo atención mínimos desde cero en una versión puramente simbólica de la tarea IOI. Este enfoque elimina los factores de confusión lingüística y aisla el núcleo del desafío de razonamiento relacional.

• Formulación de la Tarea: La tarea IOI consiste en predecir el objeto indirecto (IO) en una oración donde se introducen un sujeto (S) y un IO, y el sujeto se repite. El modelo debe predecir el nombre no repetido. Los autores abstrajeron esto en una secuencia simbólica de 6 tokens: <BOS> IO S1 S2 <MID> ?. El modelo predice el token en la posición <MID>, que debe ser el nombre que no se repitió antes del token <MID> (por ejemplo, distinguiendo entre plantillas "BABA" y "BAAB").
• Arquitectura del Modelo: El estudio utilizó transformers de solo atención con embeddings posicionales absolutos, omitiendo las MLPs y la normalización de capa para aislar los mecanismos de atención.
  • Configuración 1: Un modelo de una sola capa con dos cabezales de atención.
  • Configuración 2: Un modelo de dos capas con un solo cabezal de atención por capa.
  • Vocabulario: 8 tokens en total (6 nombres + <BOS> + <MID>).
  • Entrenamiento: Los modelos fueron entrenados en las 60 secuencias únicas posibles utilizando pérdida de entropía cruzada, optimizador AdamW y un programador de tasa de aprendizaje OneCycle.
• Técnicas de Análisis: Los autores emplearon descomposición del flujo residual, análisis espectral de las matrices de transición Query-Key (QK) y Output-Value (OV), intervenciones en embeddings (LogitLens) y ablaciones dirigidas (eliminación de embeddings posicionales o vías de composición específicas).

Resultados Clave

1. Fracaso de los Modelos de un Solo Cabezal

Un modelo de cero capas (bigrama) y un modelo de una sola capa y un solo cabezal fallaron en resolver la tarea. El modelo de un solo cabezal asignó una probabilidad aproximadamente igual a ambos nombres en la cláusula dependiente. El análisis reveló que un solo cabezal no puede realizar simultáneamente la "detección de referencia" (identificar el token correcto) y el "copiado" (propagar esa información a la posición de predicción); estos roles parecían funcionalmente incompatibles dentro de un solo cabezal de atención.

2. Éxito del Modelo de Dos Cabezales y Una Capa

Un modelo de una sola capa con dos cabezales de atención logró una precisión del 100%. El análisis mecánico reveló una división del trabajo altamente interpretable:

• Cabezal 1 (Aditivo): Atiende equitativamente a ambos nombres en la cláusula dependiente inicial. Su salida se alinea con la dirección de la suma de los embeddings de los nombres (Correcto + Incorrecto). Actúa como un agregador de señales referenciales.
• Cabezal 2 (Contrastivo): Atiende al sujeto de la cláusula principal y al nombre no repetido en la cláusula dependiente. Su salida se alinea con la dirección de la diferencia (Correcto – Incorrecto). Realiza la supresión contrastiva del alternativo incorrecto.
• Mecanismo: Los logits finales se computan mediante la combinación aditiva de estos dos cabezales. La señal aditiva del Cabezal 1 y la señal contrastiva del Cabezal 2 resultan en la cancelación del logit del token incorrecto y la amplificación del logit del token correcto.

3. Análisis Espectral

El análisis espectral de las matrices QK y OV confirmó la especialización funcional:

• Cabezal 1: Exhibe una dinámica de atención relativamente neutral con una mezcla de autovalores amplificadores y supresores, actuando como un mecanismo de "paso" o copiado.
• Cabezal 2: Muestra una fuerte dinámica supresiva (autovalores negativos dominantes en QK) y una mezcla de contribuciones aditivas y rotacionales (contrastivas) en el circuito OV, restando efectivamente la contribución del token incorrecto.

4. Rol de los Embeddings Posicionales

Cuando se eliminaron los embeddings posicionales, el modelo aún logró un ~70% de precisión, apoyándose en relaciones contextuales semánticas. Sin embargo, con los embeddings posicionales, el modelo alcanzó el 100% de precisión y se desacopló perfectamente en los roles aditivo-contrastivos. Las pistas posicionales simplificaron el paisaje de optimización, permitiendo que los cabezales se especializaran de forma más limpia.

5. Modelo de Dos Capas y Un Solo Cabezal

Un modelo de dos capas con un solo cabezal por capa también resolvió la tarea. A diferencia de la especialización paralela del modelo de dos cabezales, esta arquitectura dependió de la composición a través de las capas. Los estudios de ablación sobre las composiciones de Q, K y V mostraron que el modelo depende fuertemente de la composición Q (caída de precisión de aprox. 100% cuando se ablaciona) y la composición V (caída de aprox. 93%), lo que indica que la segunda capa lee y transforma la información agregada de la primera capa para resolver la tarea.

Contribuciones

El artículo realiza tres contribuciones principales:

1. Suficiencia Mínima: Demuestra que un modelo de solo atención de una capa y dos cabezales es suficiente para resolver perfectamente la tarea IOI en una plantilla fija, una tarea que anteriormente se pensaba requería circuitos complejos de múltiples saltos en modelos más grandes.
2. Descubrimiento de Circuitos Mecanísticos: Revela un circuito mínimo donde la solución se computa mediante una combinación aditiva directa de salidas de cabezales de atención especializados (uno aditivo, uno contrastivo), contrastando con los complejos pipelines de múltiples saltos encontrados en GPT-2 small.
3. Perspectiva de la Dinámica de Entrenamiento: Argumenta que la complejidad de los circuitos en modelos grandes y ampliamente preentrenados puede derivar de presiones multitarea, mientras que el entrenamiento restringido a una tarea puede revelar mecanismos más parsimoniosos e interpretables.

Significancia y Reivindicaciones

Los autores afirman que su trabajo proporciona un banco de pruebas controlado para sondear las bases computacionales del razonamiento en transformers. Al eliminar la complejidad del lenguaje natural y los objetivos de preentrenamiento, demuestran que:

• El entrenamiento específico de la tarea induce circuitos altamente interpretables y mínimos.
• La complejidad de "caja negra" de los LLMs puede no ser inherente a la tarea de razonamiento en sí, sino un subproducto del paisaje de optimización en configuraciones de gran escala y multitarea.
• Los motivos computacionales específicos identificados (circuitos aditivo-contrastivos y composición por capas) ofrecen una base para entender cómo modelos más grandes podrían implementar razonamientos similares, sugiriendo que la lógica fundamental de la resolución de IO es más simple de lo observado previamente en modelos preentrenados.

El artículo mantiene la modestia respecto a su alcance, reconociendo que estos hallazgos están limitados por estructuras de secuencias rígidas y aún no exploran cómo estos circuitos mínimos escalan a variaciones en la longitud de la secuencia, sintaxis dinámica o la dinámica del desarrollo durante el entrenamiento del circuito.