Induction Heads Interpolate N-Grams

Resumen Técnico: Las cabezas de inducción interpolan N-gramas

Planteamiento del problema
Las cabezas de inducción, identificadas como el circuito principal para el aprendizaje en contexto (ICL) en los transformadores, se caracterizan tradicionalmente en el límite de "atención dura". En este régimen, el circuito realiza una coincidencia de patrones exacta y un conteo de máxima verosimilitud (ML). Aunque es mecánicamente atractivo, esta visión es estadísticamente incompleta para dependencias de orden superior: las coincidencias exactas de $k$-gramas suelen ser demasiado raras en muestras finitas como para soportar una predicción fiable, y el estimador de ML asigna una probabilidad de cero a las continuaciones no vistas. Los modelos de lenguaje clásicos abordan esto mediante mecanismos de suavizado y retroceso (por ejemplo, Jelinek–Mercer, Katz back-off), que redistribuyen la masa de probabilidad hacia eventos no vistos e interpolan a través de longitudes de contexto. La cuestión central abordada por este trabajo es si las cabezas de inducción implementan una clase de estimadores más rica que el simple conteo estricto, específicamente capaz de realizar dicho suavizado.

Metodología
Los autores analizan transformadores entrenados en cadenas de Markov de orden-$k$, un entorno que equilibra la trazabilidad analítica con los compromisos de muestra finita inherentes a la predicción basada en contexto.

1. Marco Teórico: El estudio utiliza una arquitectura de transformador "desentrelazada" (eliminando los MLP y utilizando la concatenación para los flujos residuales) para facilitar un análisis exacto. Los autores construyen un transformador específico de dos capas con Codificación Posicional Relativa (RPE) y demuestran que implementa un estimador de coincidencia de contexto suave (soft context-matching estimator).
2. Construcción del Mecanismo:
   • Capa 1 (Cabezas de Copia): Utiliza $k$ cabezas de atención para copiar tokens desde desfases relativos específicos ($t-1, \dots, t-k$), construyendo una representación del contexto actual.
   • Capa 2 (Coincidencia de Contexto): Utiliza una única cabeza de atención para comparar el contexto de la consulta (query) contra todos los contextos candidatos en la historia. Las puntuaciones de atención están diseñadas para computar un producto interno alineado por desplazamiento, ponderando las coincidencias según el número de tokens superpuestos (solapamiento de Hamming) en lugar de requerir coincidencias exactas.
   • Token BOS: Los autores incorporan un token de Inicio de Secuencia (BOS) como un objetivo de atención independiente de la secuencia para introducir pseudo-conteos aditivos.
3. Análisis Estadístico: Los autores derivan el estimador de forma cerrada implementado por este circuito y lo comparan con los predictores de Bayes-óptimos bajo dos prioris de tarea:
   • Prior de Dirichlet Independiente: Donde los contextos no están correlacionados.
   • Prior de Dirichlet Jerárquico: Donde las transiciones de orden superior están centradas en sus sufijos de orden inferior, creando correlaciones entre contextos que comparten un sufijo.

Contribuciones Clave
El artículo establece que las cabezas de inducción implementan dos mecanismos de suavizado complementarios cuando operan en escalas de peso de atención finitas:

1. Coincidencia de Contexto Suave (Interpolación):

   • El circuito agrega contribuciones tanto de coincidencias de contexto exactas como parciales.
   • Las coincidencias se ponderan exponencialmente por su tamaño de solapamiento ($|M|$).
   • Esto induce una interpolación dependiente de los datos a través de órdenes de contexto, análoga al suavizado de Jelinek–Mercer. A diferencia del suavizado JM clásico, que utiliza pesos globales fijos, los pesos del transformador dependen de los conteos realizados en la secuencia específica.
   • El estimador unifica modelos de mezcla (como JM) y estrategias de retroceso (como Katz): cuando las coincidencias exactas son raras, el modelo realiza automáticamente un "retroceso" (back-off) hacia coincidencias parciales de orden inferior sin reglas explícitas.

2. BOS como Pseudo-conteos Aditivos:

   • El token BOS actúa como un sumidero fijo e independiente de la secuencia para la atención.
   • Este mecanismo permite al modelo añadir pseudo-conteos constantes a los conteos de transición específicos del contexto.
   • Teóricamente, esto recupera el suavizado add-$\alpha$ (por ejemplo, el suavizado de Laplace) en el límite de grandes pesos de atención, proporcionando una explicación a nivel de circuito de cómo las convenciones arquitectónicas implementan la regularización tipo prior.

3. Validación Empírica:

   • Los autores entrenan tanto transformadores desentrelazados como estándar (totalmente parametrizados) en tareas de cadenas de Markov.
   • Los resultados muestran que los modelos entrenados recuperan los patrones de atención predichos:
     • Sin BOS, dependen de la coincidencia de contexto suave para interpolar entre órdenes.
     • Con BOS, utilizan pseudo-conteos para alcanzar un rendimiento cercano al de Bayes-óptimo bajo prioris independientes.
     • Bajo prioris jerárquicos, donde existen evidencias estructuradas, los modelos superan a las líneas base de suavizado fijo mediante el aprovechamiento del mecanismo de interpolación.
   • El análisis mecánico confirma que los transformadores estándar entrenados aprenden las estructuras de peso específicas (por ejemplo, matrices de desplazamiento por bloques) predichas por la teoría.

Resultados

• Divergencia KL: En entornos donde el suavizado add-$\alpha$ es Bayes-óptimo (Dirichlet Independiente), los transformadores entrenados con BOS igualan al predictor Bayes-óptimo. En entornos jerárquicos, donde existe evidencia estructurada, el mecanismo de interpolación permite a los transformadores superar a las líneas base de suavizado fijo.
• Patrones de Atención: Las visualizaciones de los modelos entrenados confirman el circuito de dos etapas: la Capa 1 realiza la copia de desfase específico y la Capa 2 realiza la coincidencia suave, distribuyendo la masa de atención entre coincidencias completas y parciales según el solapamiento.
• Pesos de Desfase Asimétricos: En tareas jerárquicas, los modelos entrenados aprenden a asignar mayores pesos a las coincidencias en el token más reciente (desfase 1) en comparación con los tokens más antiguos (desfase 2), lo cual es consistente con la estructura del prior donde los tokens recientes son mejores predictores del siguiente estado.

Significado
El artículo cierra la brecha entre la interpretabilidad mecánica de las cabezas de inducción y el suavizado estadístico clásico. Demuestra que los transformadores no solo realizan un conteo no regularizado; más bien, aprenden a regularizar la estimación en contexto. Al interpolar dinámicamente a través de órdenes de contexto y utilizar tokens arquitectónicos (BOS) para el suavizado aditivo, las cabezas de inducción implementan procedimientos estadísticos sofisticados que se adaptan a la distribución de los datos, explicando su robustez en las tareas de aprendizaje en contexto.