Resumen Técnico: Residuales de Atención Multi-Cabeza (MHAR)

1. Planteamiento del Problema

Los Transformers estándar propagan la información a través de la profundidad mediante un único flujo residual aditivo, donde cada subcapa lee solo el estado inmediatamente anterior ($h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$). Aunque las características útiles computadas en capas anteriores permanecen en la suma acumulada, estas se vuelven individualmente inabordables.

Los Residuales de Atención (Kimi, 2025) relajan esto tratando todas las salidas de las subcapas previas como una memoria, permitiendo que cada subcapa recupere una combinación aprendida y ponderada por softmax de todo el historial de profundidad. Sin embargo, este mecanismo sufre de un cuello de botella estructural: utiliza una única consulta de enrutamiento compartida ($q \in \mathbb{R}^d$) para puntuar el historial de profundidad. En consecuencia, las $d$ coordenadas de las características se ven obligadas a leer de las mismas capas en las mismas proporciones a través de una única distribución $\alpha$.

El artículo argumenta que esto crea un "compromiso forzado". A medida que el ancho del modelo aumenta, es probable que los diferentes subespacios de características desarrollen preferencias divergentes respecto a qué capas históricas son más relevantes. Una única consulta no puede satisfacer simultáneamente estas preferencias divergentes, colapsando las distintas preferencias de profundidad en una lectura única. El costo de este compromiso crece con el ancho del modelo, lo que potencialmente hace que el enrutamiento de consulta única sea perjudicial a gran escala.

2. Metodología: Residuales de Atención Multi-Cabeza (MHAR)

Los autores proponen los Residuales de Atención Multi-Cabeza (MHAR), una generalización estricta de los residuales de atención que elimina el compromiso forzado sin añadir parámetros ni un cómputo significativo.

Mecanismo Central

En lugar de una única consulta $q$, MHAR reformatea la consulta de enrutamiento en $H$ cabezas por subespacio, cada una con su propio softmax independiente sobre el historial de profundidad.

• Reformateo (Reshaping): El vector de consulta único $q \in \mathbb{R}^d$ se reformatea en un tensor de forma $(H, d/H)$.
• Enrutamiento Independiente: Para cada cabeza $h$, el modelo computa una distribución de profundidad $\alpha^{(h)}$ distinta basada en la porción del vector de origen correspondiente a esa cabeza.
• Lectura de Bloque Diagonal: La mezcla recuperada es de bloque diagonal; cada subespacio lee el historial de profundidad de forma independiente.
• Recuperación: Establecer $H=1$ recupera exactamente los residuales de atención de una sola cabeza originales.

Detalles de Implementación

• Eficiencia de Parámetros: La transformación es un reformateo, lo que añade cero parámetros sobre la línea base de una sola cabeza.
• Eficiencia de Cómputo: El exceso es insignificante, involucrando solo $H$ softmaxes pequeños sobre el eje de profundidad (longitud $\le 2L+1$) en lugar de uno grande.
• Optimización del Sistema: Reconociendo que el enrutamiento de profundidad está limitado por la memoria (requiere tráfico de $O(L^2 B T d)$), los autores proporcionan kernels de Triton fusionados. Estos kernels realizan un softmax en línea con RMSNorm fusionado, eliminando la materialización intermedia de claves y fuentes apiladas. Esto reduce el tiempo de enrutamiento en un $2\times$ respecto a torch.compile y un $6.4\times$ respecto a las implementaciones de referencia en modo eager. Esto eleva el rendimiento de entrenamiento a $0.55\text{--}0.88\times$ de la línea base.

Configuración de Hiperparámetros

El número de cabezas de enrutamiento $H$ se establece igual al número de cabezas KV del modelo. Los autores encuentran que este es un valor por defecto casi óptimo y libre de hiperparámetros que alinea la granularidad del enrutamiento con la granularidad de consumo del modelo.

3. Contribuciones Clave

1. Arquitectura MHAR: Una modificación sin parámetros que divide la consulta de enrutamiento en cabezas independientes, eliminando el compromiso forzado del enrutamiento de profundidad de consulta única.
2. Kernels Eficientes: Kernels de Triton fusionados que hacen que el enrutamiento de profundidad sea práctico a escala, reduciendo significativamente la sobrecarga de memoria y la latencia.
3. Conversión de Preservación de Identidad: Un método para injertar MHAR en modelos preentrenados usando "residuales de atención delta" (añadiendo el término enrutado al flujo existente en lugar de reemplazarlo), lo que permite el entrenamiento intermedio sin picos de pérdida.
4. Perspectiva Mecanicista: El sondeo directo de las consultas entrenadas confirma que el desacuerdo de los subespacios crece con el ancho y es el principal motor de la degradación del enrutamiento de una sola cabeza.

4. Resultados Experimentales

Preentrenamiento (Desde Cero)

Entrenado en FineWeb-Edu a través de tres escalas (100M, 350M, 1B):

• Pérdida de Validación: MHAR mejora consistentemente la pérdida de validación sobre los Transformers estándar y los residuales de atención de una sola cabeza.
  • 100M: $-0.049$ nats frente a la línea base.
  • 350M: $-0.080$ nats frente a la línea base.
  • 1B: $-0.063$ nats frente a la línea base.
• Comportamiento de Escala: La ventaja de MHAR sobre el enrutamiento de una sola cabeza se amplía con la escala. Mientras que el enrutamiento de una sola cabeza ayuda a 100M ($-0.039$), se degrada a ser perjudicial a 350M ($+0.055$) y 1B ($+0.140$).
• Equivalencia de Cómputo: MHA logra una ganancia de equivalencia de cómputo de $1.3\text{--}1.5\times$ sobre la línea base.

Entrenamiento Intermedio (8B)

Usando un control de esquema emparejado en un modelo de 8B (Marin-8B) con un proceso de enfriamiento (annealing) en un corpus orientado a capacidades:

• Preservación de Identidad: La conversión no introduce choque de optimización; las curvas de pérdida coinciden con la línea base durante los primeros 100 pasos.
• Ganancias en Tareas de Seguimiento (Downstream): MHAR produce mejoras estadísticamente significativas en tareas de razonamiento:
  • GSM8K: $+3.2$ puntos (McNemar pareado $p=0.004$).
  • GPQA: $+3.1$ puntos (McNemar pareado $p=0.038$).
  • Las ganancias en MMLU, MATH y tareas de código no fueron estadísticamente significativas en esta configuración específica.

Estudios de Ablación

• Número de Cabezas: El número óptimo de cabezas $H$ sigue al número de cabezas KV. En regímenes de bajo entrenamiento, más cabezas siempre ayudan; en regímenes convergidos, $H=KV$ es el "punto de inflexión" de la curva.
• Ancho vs. Profundidad: El sondeo de las consultas entrenadas revela que el "desacuerdo de ancho" (preferencias de enrutamiento divergentes entre subespacios) crece monótonamente con el ancho del modelo, validando la hipótesis del compromiso forzado.
• Alineación de Cabezas: Alinear las cabezas de enrutamiento con los grupos KV de atención no ofrece un beneficio mensurable sobre el enrutamiento de subespacios arbitrario, lo que sugiere que el beneficio proviene de la multiplicidad de cabezas y no de su alineación específica con las estructuras de atención.

5. Significado y Reivindicaciones

El artículo afirma que los Residuales de Atención Multi-Cabeza resuelven una limitación arquitectónica fundamental en cómo los Transformers acceden al historial de profundidad. Al permitir que diferentes subespacios de características atiendan a diferentes capas de forma independiente, MHAR escala efectivamente donde el enrutamiento de consulta única falla.

Los autores enfatizan que:

• La mejora no se debe a la adición de parámetros o cómputo, sino a la eliminación de un cuello de botella estructural.
• El método es práctico a escala gracias a los kernels fusionados que mitigan la naturaleza limitada por memoria del enrutamiento de profundidad.
• El enfoque es versátil, funcionando desde cero y mediante un entrenamiento intermedio de preservación de identidad.
• La necesidad de un enrutamiento multi-cabeza es una consecuencia directa del ancho del modelo; a medida que los modelos crecen, el "compromiso forzado" de una única consulta se convierte en una carga, haciendo que MHAR sea un componente crítico para escalar arquitecturas conscientes de la profundidad.