Resumen Técnico: Inducción en Ambas Direcciones

Planteamiento del Problema

Mientras que los mecanismos internos de los transformadores autorregresivos (AR) han sido extensamente estudiados a través de la lente de la interpretabilidad mecánica, la mecánica operativa de los Modelos de Lenguaje de Difusión (DLM) permanece en gran medida inexplorada. Los DLM generan texto mediante la eliminación iterativa de ruido en lugar de la predicción secuencial del siguiente token, utilizando atención bidireccional y objetivos de entrenamiento diferentes. Existe una brecha crítica en comprender si los DLM implementan la inducción —un mecanismo canónico para el aprendizaje en contexto donde un modelo identifica un contexto repetido y copia el token subsiguiente— de la misma manera que los modelos AR. Específicamente, no está claro cómo los DLM manejan la inducción cuando el token objetivo está oculto tras un token de máscara [M] y rodeado de un contexto bidireccional, a diferencia del contexto estrictamente causal disponible para los modelos AR.

Metodología

Los autores llevaron a cabo un estudio mecánico controlado comparando transformadores AR de solo atención emparejados y DLM de máscara de absorción.

• Arquitectura del Modelo: El estudio utilizó modelos de pequeña escala (profundidades $L \in \{2, 3\}$, dimensión oculta $d_{model}=512$, 16 cabezales de atención) entrenados en secuencias de FineWeb-Edu. Ambas familias de modelos compartían arquitecturas y recetas de entrenamiento idénticas, diferenciándose solo en sus máscaras de atención (causal para AR, bidireccional para DLM) y sus objetivos de entrenamiento (predicción del siguiente token frente a difusión de máscara de absorción).
• Tarea de Evaluación: Las capacidades de inducción se midieron utilizando secuencias sintéticas de longitud 512, donde la primera mitad se repetía como la segunda mitad. La tarea consistía en predecir los tokens enmascarados en la segunda mitad basándose en la fuente no enmascarada en la primera mitad (inducción hacia adelante) y viceversa (inducción inversa).
• Técnicas de Análisis:
  • Ablaciones Causales: Se realizaron ablaciones de la media en cabezales de atención individuales para determinar su necesidad causal para la inducción.
  • Descomposición de Pesos: El estudio empleó la descomposición QK (Query-Key) y OV (Output-Value) para analizar cómo se forman los logits de atención y cómo fluye la información desde la fuente hacia el objetivo.
  • Sondeo Lineal (Linear Probing): Se entrenó un sondeo lineal en la corriente residual para detectar la presencia de información de la fracción de máscara global.
  • Experimentos de Parcheo (Patching): Se utilizaron técnicas de "parcheo" para intervenir causalmente en la representación de la tasa de máscara global para observar cambios en la entropía de la predicción.
• Variantes del Modelo: Para el análisis mecánico, los autores utilizaron DLM "plegados sin LayerNorm", convirtiendo los modelos entrenados con LayerNorm a mapas lineales para facilitar la descomposición aditiva de la corriente residual, verificando al mismo tiempo que el rendimiento permaneciera comparable con los modelos originales con LayerNorm.

Contribuciones Clave y Resultados

1. Circuito de Inducción Bidireccional

El estudio identifica que los DLM aprenden un circuito de inducción bidireccional que es simétrico en dirección. A diferencia de los modelos AR, que dependen de una única vía causal, los DLM utilizan dos vías distintas:

• Cabezales de Token Anterior y Siguiente: En la primera capa (Capa 0), cabezales específicos atienden a los vecinos inmediatos de la izquierda y la derecha del token enmascarado, escribiendo estas características de contexto local en la corriente residual.
• Cabezales de Inducción: En las capas subsiguientes (Capa 1+), los cabezales de inducción utilizan pesos QK para hacer coincidir las características escritas por los cabezales de la Capa 0 en la posición de la máscara con las características correspondientes en la posición de la fuente de la respuesta.
• Mecanismo: La descomposición QK revela que los cabezales de inducción utilizan la composición Q (query construida a partir del contexto local de la máscara) y la composición K (key construida a partir del contexto local de la fuente) para localizar la fuente coincidente. La ruta OV entonces copia el token de la fuente a la posición enmascarada.

2. Acceso al Contexto vs. Fuerza del Mecanismo

Los autores demuestran que el rendimiento de inducción superior de los DLM en comparación con los modelos AR no se debe a un mecanismo de inducción fundamentalmente más fuerte, sino al acceso al contexto bidireccional.

• Al restringirse únicamente al contexto de la izquierda (imitando la configuración de AR), los DLM rinden al mismo nivel o por debajo de los modelos AR.
• Los DLM superan a los modelos AR solo cuando tanto el vecino izquierdo como el derecho del token enmascarado son visibles.
• El puntaje de inducción en la configuración bidireccional completa se aproxima a la suma de los puntajes de las configuraciones de solo izquierda y solo derecha, confirmando que la ventaja proviene de la utilización de información de ambas direcciones.

3. Timestep Implícito mediante la Tasa de Máscara Global

El artículo proporciona evidencia causal de que los DLM computan y utilizan la fracción global de tokens enmascarados como un timestep implícito, a pesar de carecer de embeddings de timestep explícitos.

• Representación: Un sondeo lineal puede recuperar la fracción de máscara global de la corriente residual después de la primera capa de atención con alta precisión ($R^2 \approx 0.82-0.91$).
• Utilidad Causal: El parcheo de la "dirección de la tasa de máscara" desde un estado de alta corrupción hacia uno de baja corrupción aumenta causalmente la entropía de la predicción en la posición enmascarada. Esto confirma que el modelo utiliza el nivel de corrupción global para modular su incertidumbre, tratando efectivamente la tasa de máscara como un timestep.

4. Robustez y Emergencia del Circuito

• Emergencia: La inducción emerge abruptamente durante el entrenamiento (alrededor de los 8–24k pasos) y requiere al menos dos capas para funcionar, lo cual es consistente con la necesidad de una capa para escribir características y una capa subsiguiente para recuperarlas.
• Ventana Local: El mecanismo de inducción depende fuertemente de una ventana local de aproximadamente cinco tokens (el token enmascarado más dos vecinos a cada lado). El rendimiento cae significativamente si estos vecinos inmediatos están enmascarados, incluso si la tasa de corrupción global es baja.

Significancia y Reivindicaciones

Este trabajo representa un paso fundacional para comprender los mecanismos internos de los DLM, yendo más allá de la suposición de que simplemente replican los comportos de los modelos AR. Los autores afirman que:

1. Divergencia Mecánica: Los DLM implementan la inducción a través de un circuito bidireccional simétrico que difiere estructuralmente de los circuitos de inducción causales en los modelos AR, específicamente en cómo se construyen las características de query y key a partir de las posiciones enmascaradas.
2. Ventaja Contextual: La brecha de rendimiento entre los DLM y los modelos AR en el aprendizaje en contexto es atribuible a la disponibilidad de contexto bidireccional, más que a una superioridad inherente en el algoritmo de inducción en sí mismo.
3. Gestión de Estado Implícita: Los DLM derivan internamente la información temporal (timestep) a partir del estado global de la secuencia (fracción de máscara) sin embeddings explícitos, un mecanismo que está causalmente vinculado a la entropía de la predicción.

El artículo concluye que estos hallazgos proporcionan una base para comprender el aprendizaje en contexto en modelos no autorregresivos, aunque reconoce limitaciones respecto al uso de modelos pequeños de solo atención y tareas sintéticas, señalando que se requiere más investigación para determinar si estos mecanismos escalan a DLM más grandes que contienen capas MLP y tareas de lenguaje natural.