Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

Resumen Técnico: Atención Exclusiva de Autoatención más allá de AdamW

Planteamiento del Problema
Los mecanismos de autoatención estándar en los Transformers permiten que la salida para un token en la posición $i$ incluya una mezcla ponderada de vectores de valor, incluyendo el propio vector de valor del token ($v_i$). Aunque es expresivo, este diseño puede crear un "atajo representacional" donde la salida de la atención permanece fuertemente alineada con $v_i$, aun cuando las vías residuales ya preservan la información del token local. Este fenómeno, denominado "sesgo de similitud de atención", fue abordado por Zhai [2] mediante la Autoatención Exclusiva (XSA, por sus siglas en inglés), la cual elimina explícitamente la componente de la salida de la atención paralela al propio vector de valor del token.

Sin embargo, el estudio original de XSA evaluó el método exclusivamente dentro de un stack de entrenamiento de NanoGPT utilizando el optimizador AdamW. Los autores del trabajo original identificaron explícitamente la compatibilidad de XSA con optimizadores alternativos, como Muon, como una pregunta abierta. Esta es una brecha crítica porque la investigación moderna de modelos de lenguaje a pequeña escala (por ejemplo, el stack nanochat) utiliza cada vez más estrategias de optimización híbridas (Muon para parámetros de matrices, AdamW para embeddings) y características arquitectónicas como embeddings de valor y codificación de posición rotatoria (RoPE). No está claro si las restricciones geométricas de XSA permanecen estables o beneficiosas cuando se transfieren a estos regímenes de entrenamiento modernos basados en Muon.

Metodología
El estudio realiza una reproducción controlada y un experimento de transferencia utilizando el stack de entrenamiento nanochat. La metodología involucra:

1. Implementación: Se integra una forma parametrizada de XSA en nanochat. La operación se define como $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$, donde $y_i$ es la salida de atención estándar, $v_i$ es el vector de valor del propio token, y $\alpha \in [0, 1]$ controla la fuerza de la proyección. Con $\alpha=1$, la componente auto-alineada se elimina por completo.
2. Validación de Corrección: La implementación se valida numéricamente para asegurar que, con $\alpha=0$, la ruta XSA sea equivalente a la atención estándar. Ejecuciones de entrenamiento cortas verifican que la línea base y las rutas con $\alpha=0$ exhiben un comportamiento de pérdida inicial coincidente.
3. Configuración Experimental: Se evalúan dos profundidades de modelo bajo condiciones de entrenamiento emparejadas (mismo orden de datos, familia de optimizadores, programa de tasa de aprendizaje y longitud de contexto):
   • d12: 12 capas, dimensión oculta 768, 12 cabezales.
   • d24: 24 capas, dimensión oculta 1024, 16 cabezales.
   • Optimizador: Los parámetros de matriz se optimizan con Muon; el resto de los parámetros usan AdamW.
   • Datos: ClimbMix-400B.
4. Métricas: El estudio rastrea la pérdida de entrenamiento, los bits por byte (bpb) de validación, las puntuaciones agregadas de downstream (Base CORE y ChatCORE) y los diagnósticos de atención (alineación de coseno entre la salida y el valor propio, comportamiento diagonal).

Contribuciones Clave

• Reproducción en Stack Muon: La primera integración de XSA en un stack de entrenamiento basado en Muon, abordando directamente la pregunta abierta sobre la compatibilidad de optimizadores dejada por el estudio original.
• Implementación Parametrizada: Una implementación flexible que permite la atención estándar, la XSA completa y la proyección parcial mediante un parámetro de fuerza $\alpha$, junto con controles de selección de capa.
• Corrección e Instrumentación: Validación del límite $\alpha=0$ y corrección de artefactos de registro (convenciones de índice de paso y suavizado) que anteriormente distorsionaban las comparaciones tempranas de entrenamiento.
• Evaluación de Dos Profundidades: Un análisis comparativo en las profundidades d12 y d24 utilizando un conjunto exhaustivo de métricas.
• Análisis Mecanístico y de Sistemas: Diagnósticos de auto-alineación de atención y una evaluación de la sobrecarga de tiempo de ejecución (runtime overhead).

Resultos
Los resultados empíricos indican que XSA es estable bajo el stack basado en Muon, pero produce beneficios de downstream mixtos dependiendo de la profundidad del modelo y la métrica específica:

• Estabilidad de Optimización: Las trayectorias de pérdida de entrenamiento para los modelos XSA y de línea base permanecen estrechamente alineadas en ambas profundidades. No hay evidencia de divergencia, NaNs o inestabilidad, confirmando que XSA no interrumpe la geometría de actualización de Muon.
• Rendimiento de Validación (Bits-per-Byte): XSA mejora consistentemente el bpb de validación mantenida:
  • d12: Reducción de 0.8484 a 0.8457 (mejora relativa del 0.32%).
  • d24: Reducción de 0.7193 a 0.7171 (mejora relativa del 0.31%).
• Métricas de Downstream (CORE):
  • d24: XSA mejora tanto Base CORE (0.2593 a 0.2606, +0.50%) como ChatCORE (0.3638 a 0.3682, +1.21%).
  • d12: Aunque el bpb mejora, Base CORE disminuye de 0.1602 a 0.1508 (-5.87%).
• Diagnósticos Mecanísticos: Los diagnósticos de atención muestran una reducción en la auto-alineación y una reducción en el comportamiento diagonal bajo XSA, consistente con el mecanismo pretendido de promover la mezcla contextual.
• Sobrecarga de Tiempo de Ejecución: La sobrecarga se caracteriza como modesta. Aunque una estimación inicial sugirió un ~5%, las mediciones repetidas de rendimiento (throughput) situaron la diferencia cerca de la variabilidad del sistema en ejecuciones cortas, lo que sugiere que no hay una penalización prohibitiva.

Significancia y Reivindicaciones
El artículo concluye que la Atención Exclusiva de Autoatención es un sesgo inductivo arquitectónico compatible con los stacks de entrenamiento basados en Muon, pero su utilidad no es universal. La principal significancia radica en demostrar que la restricción representacional de XSA se transfiere limpiamente a un régimen de optimización moderno sin desestabilizar el entrenamiento.

Sin embargo, el autor enfatiza que las ganancias de downstream son dependientes de la escala y la métrica. La mejora en el bpb de validación en ambas profundidades sugiere un beneficio representacional, pero la degradación de Base CORE en d12 indica que los beneficios de eliminar la auto-alineación pueden depender de la profundidad del modelo, la capacidad de las vías residuales/de valor para preservar la información local, o la tarea de evaluación específica. El estudio no afirma que XSA sea un reemplazo universalmente superior para la atención estándar, sino más bien una restricción estable y de baja disrupción que ofrece beneficios en contextos específicos (notablemente en d24 en este estudio).

El trabajo sirve como una verificación de robustez más que como una reproducción exacta bit por bit, resaltando las complejas interacciones entre la modificación arquitectónica, la geometría del optimizador y las métricas de evaluación en la investigación moderna de modelos de lenguaje.