Resumen Técnico: Asignar y Sumar: Un Estudio Mecanístico de la Aritmética Composicional

Declaración del Problema

Los modelos de lenguaje de gran tamaño (LLM) demuestran una capacidad para componer habilidades para realizar tareas complejas no vistas durante el entrenamiento, sin embargo, los mecanismos específicos que permiten esta "generalización composicional" siguen siendo elusivos. Aunque la generalización se atribuye a menudo a la composición de mecanismos atómicos, la interacción entre distintos mecanismos —específicamente la vinculación de variables (recuerdo asociativo mediante cabezales de inducción) y la aritmética modular (representaciones estructuradas mediante características de Fourier)— no ha sido estudiada sistemáticamente en un entorno controlado.

Este artículo aborda la pregunta: ¿Cómo aprenden los transformadores a componer la vinculación de variables y los mecanismos de suma modular para generalizar más allá de la distribución de entrenamiento, y cómo emergen estas capacidades durante el entrenamiento?

Metodología

Los autores emplean una configuración experimental controlada que involucra un transformador pequeño y mecánicamente interpretable entrenado en una tarea sintética que combina la asignación de variables y la suma modular.

• Arquitectura del Modelo: Un transformador de 2 capas y un solo cabezal con $d_{model} = d_{head} = 128$ y $d_{mlp} = 512$. La arquitectura omite deliberadamente sesgos, capas de normalización y el MLP de la Capa-1 para simplificar el descubrimiento de circuitos. Utiliza incrustaciones de posición absoluta de estilo GPT-2 y enmascaramiento causal.
• Definición de la Tarea: El modelo es entrenado para realizar suma modular ($N=59$) donde los operandos pueden ser constantes o variables asignadas dinámicamente en el contexto (porاً ejemplo, c=1, b=17, a=42, b+19=?).
• Distribución de Datos: Las secuencias se tokenizan con relleno aleatorio para evitar la dependencia de posiciones absolutas. El conjunto de entrenamiento se particiona para evaluar la generalización:
  • 0-variables (0-var): Suma de constantes pura.
  • 1-variable (1-var) y 2-variables (2-var): Suma que involucra variables asignadas.
  • Conjuntos Excluidos: Aproximadamente el 30% de los pares de suma se eliminan del entrenamiento (por ejemplo, si 1+3 no es visto, se excluyen todas las secuencias que lo contienen). Además, ciertas variables se restringen de ciertas posiciones en el conjunto de entrenamiento para probar la generalización posicional.
• Técnicas de Análisis: El estudio utiliza herramientas de interpretabilidad mecánica (a través de la biblioteca TransformerLens) para analizar patrones de atención, vectores del flujo residual y matrices de pesos. También emplea sondas lineales y análisis de gradiente teórico para rastrear la dinámica del entrenamiento.

Contribuciones Clave

1. Modelo Tratable para la Aritmética Composicional: Los autores presentan un modelo simplificado que permite la interpretación mecánica de cómo la asignación de variables se compone con la aritmética modular.
2. Dinámica de Entrenamiento Desacoplada: El análisis revela que la aritmética modular se aprende antes y de forma independiente a la asignación de variables. El modelo logra una precisión casi perfecta en tareas de 0-var mientras aún está aprendiendo la vinculación de variables.
3. Explicación Mecanística de la Generalización: El artículo identifica que la generalización surge de la composición de dos circuitos específicos:
   • Un cabezal de inducción modificado en la Capa 1 que recupera las variables asignadas y redirige sus valores.
   • Un circuito de suma modular en la Capa 2 (específicamente la matriz de Valor de Salida y el MLP de la Capa 2) que realiza la aritmética sobre los valores recuperados.
4. Marco Teórico: Los autores proporcionan un análisis teórico que demuestra cómo la dinámica del gradiente facilita la composición de los circuitos de vinculación de variables y de suma modular, mostrando que el gradiente negativo en la matriz Query-Key de la Capa 2 conduce naturalmente a la formación de un bloque de memoria asociativa.

Resultados

Dinámicas de Entrenamiento y Fases

El proceso de entrenamiento exhibe tres fases distintas:

1. Aprendizaje de la Suma Modular: El modelo primero aprende a realizar la suma en secuencias sin variables (0-var), logrando una alta precisión antes de que la vinculación de variables se establezca por completo.
2. Aprendizaje de la Asignación de Variables: Posteriormente, el modelo aprende la estructura requerida para la asignación de variables. Un agudo "pico" en la precisión (alrededor del paso 5,300–5,400) coincide con la emergencia del circuito de enrutamiento de asignación de variables, donde la atención de la Capa 1 enruta correctamente las constantes hacia sus variables asignadas.
3. Fase de Refinamiento: Una segunda fase (pasos 14,500–16,000) implica el refinamiento del enrutamiento para secuencias "difíciles" (por ejemplo, variables específicas en posiciones excluidas), corrigiendo fallos iniciales al atender a ciertas posiciones de variables.

Capacidades de Generalización

• Generalización Posicional: El modelo generaliza a nuevos ordenamientos de variables (por ejemplo, predecir b+14 después de haber visto 14+b).
• Generalización Fuera de Distribación (OOD): El modelo predice con éxito pares de suma no vistos (por ejemplo, 1+3) que nunca fueron vistos en ninguna forma durante el entrenamiento, logrando una precisión del 99.4% en conjuntos de prueba curados.
• Requerimientos de Datos: La generalización a pares de constantes no vistos requiere entrenar en al menos ~25% de la distribución total ($f \approx 0.25$). La generalización a secuencias de 2-variables requiere una frecuencia relativa de ~20% ($r \approx 0.2$) de secuencias de 2-variables en los datos de entrenamiento.

Hallazgos Mecanísticos

• Representación Compartida: Los vectores del flujo residual para la suma con variables asignadas y la suma de constantes pura son casi idénticos (similitud de coseno $\approx 1$) antes del MLP de la Capa 2, lo que indica un mecanismo computacional compartido.
• Características de Fourier: Las neuronas del MLP de la Capa 2 exhiben preactivaciones que coinciden con la forma teórica de las características de Fourier ($\cos(\frac{2\pi k}{N}n + \alpha_k) + \cos(\frac{2\pi k}{N}m + \beta_k)$), confirmando el uso de mecanismos periódicos establecidos para la aritmética modular.
• Modificación del Cabezal de Inducción: La atención de la Capa 1 actúa como un "cabezal de token previo" para las constantes, atendiendo a la variable inmediatamente precedente. La matriz Query-Key de la Capa 2 forma entonces un fuerte patrón de identidad en las incrustaciones de las variables, permitiendo que el token = atienda directamente a las constantes correctas asociadas con las variables.

Significancia

El artículo sostiene que la generalización composicional en los transformadores puede ser una consecuencia natural de la composicionalidad de los mecanismos internos. Al demostrar que distintos mecanismos (cabezales de inducción y aritmética basada en Fourier) pueden aprenderse independientemente y luego componerse mediante el enrutamiento de la atención, este trabajo proporciona una base mecánica para comprender cómo los transformadores generalizan a tareas de razonamiento complejas.

Los autores mantienen la modestia respecto a su alcance, señalando que el estudio está limitado a un entorno pequeño y controlado para facilitar el análisis. No pretenden que estos hallazgos expliquen directamente todos los comportos de los LLM, sino que sugieren que ofrecen un marco para entender cómo surgen las capacidades de generalización en los transformadores, lo que podría motivar investigaciones futuras sobre tareas más complejas y leyes de escala.