Resumen Técnico: Residuales de Atención con Roles Desacoplados

Declaración del Problema

En las arquitecturas Transformer estándar, el flujo residual actúa como una memoria fija donde cada capa hereda el estado inmediatamente anterior mediante una suma serial. Las arquitecturas recientes de enrutamiento de profundidad, como los Residuales de Atención por Bloques (Block AttnRes), mejoran esto al permitir que las subclayas recuperen una mezcla dependiente del contenido de representaciones anteriores, en lugar de solo el predecesor inmediato. Esto transforma la profundidad de un orden computacional fijo en un eje de recuperación aprendido.

Sin embargo, Block AttnRes introduce una restricción de diseño específica: utiliza una distribución de profundidad única para construir la entrada compartida para las tres proyecciones de atención (Consultas, Claves y Valores). Esto acopla dos roles funcionalmente distintos:

1. Geometría de Coincidencia: Determinada por las Consultas (Queries) y las Claves (Keys), que deciden dónde mirar (qué posiciones de tokens coinciden).
2. Recuperación de Contenido: Determinada por los Valores (Values), que deciden qué contenido se recupera.

El artículo postula que forzar a que la representación óptima para decidir "dónde mirar" sea también la mejor para decidir "qué leer" es un acoplamiento innecesario. La cuestión central es si el desacoplamiento de estos roles mejora el preentrenamiento de modelos de lenguaje cuando todos los demás factores se mantienen constantes.

Metodología: Residuales de Atención con Roles Desacoplados (RD-AttnRes)

El autor introduce los Residuales de Atención con Roles Desacoplados (RD-AttionRes), una relajación arquitectónica mínima de Block AttnRes.

• Arquitectura: RD-AttnRes conserva el conjunto de fuentes de la arquitectura matriz, la topología de bloques, la normalización, las rutas de MLP y los mecanismos de escritura (writeback). La única modificación ocurre en la lectura de profundidad del lado de la atención.
• Mecanismo de Desacoplamiento:
  • Consultas y Claves (Q/K): Comparten una única consulta de ruta aprendida ($q^{QK}_\ell$) para computar una mezcla ($h^{QK}_\ell$) a partir de la memoria de profundidad residual. Esto preserva una geometría de coincidencia coherente.
  • Valores (V): Aprenden una consulta de ruta independiente ($q^V_\ell$) para computar una mezcla separada ($h^V_\ell$) sobre el mismo conjunto exacto de fuentes residuales.
• Eficiencia de Parámetros: La modificación añade solo un vector de dimensión $d$ (ancho del modelo) por capa. Para las escalas probadas (120M y 343M de parámetros), esto representa un aumento insignificante de parámetros (<0.01%).
• Recuperación Exacta: Si la consulta de ruta de los valores se vincula a la consulta Q/K ($q^V_\ell = q^{QK}_\ell$), RD-AttnRes se reduce exactamente a la arquitectura matriz Block AttnRes. Esto asegura que la propuesta sea una extensión estricta y auditable en lugar de un sistema con múltiples cambios simultáneos.
• Costo Computacional: El método añade un trabajo de $O(BTK_\ell d)$ por capa de atención (donde $B$ es el tamaño del lote, $T$ es la longitud de la secuencia y $K_\ell$ es el número de fuentes), pero no introduce nuevos términos de atención token-a-token de $O(T^2d)$.

Diseño Experimental

El estudio emplea un protocolo experimental emparejado riguroso para aislar el efecto del desacoplamiento de roles:

• Conjuntos de Datos: FineWeb-Edu (muestra de 10BT), con una división determinista para entrenamiento, validación y prueba.
• Modelos: Dos escalas de Transformer Pre-Norm de solo decodificador:
  • S: 120M de parámetros (16 capas, ancho 576).
  • M: 343M de parámetros (24 capas, ancho 832).
• Protocolo: Cinco semillas emparejadas por escala. Tanto Block AttnRes como RD-AttnRes se entrenan con la misma inicialización, orden de datos, esquemas de optimización y un presupuesto de 2.0B de tokens.
• Controles: Controles direccionales a 0.5B de tokens:
  • Tied (Vinculado): Fuerza a la ruta de valores a coincidir con la ruta Q/K (recuperando el modelo matriz).
  • Average (Promedio): Ejecuta dos rutas pero promedia sus mezclas (probando si la ejecución dual por sí sola ayuda).
  • Frozen-V (V Congelado): Mantiene la ruta de valores uniforme (probando si una ruta fija ayuda).
• Diagnósticos: Se mide la divergencia de Jensen–Shannon (JS) entre las distribuciones de profundidad de Q/K y V para verificar si las rutas aprenden estrategias distintas.

Resultados Clave

1. Mejora en la Pérdida de Validación: RD-AttnRes mejora la pérdida de log-verosimilitud negativa (NLL) de validación en todos los 10 pares de comparaciones (5 semillas × 2 escalas).
   • Escala 120M: Reducción media de NLL de 0.0301, lo que corresponde a una reducción de la perplejidad del 2.97%.
   • Escala 343M: Reducción media de NLL de 0.0247, lo que corresponde a una reducción de la perplejidad del 2.43%.
2. Persistencia de la Ganancia: La mejora es visible en todos los hitos de evaluación (0.25B, 0.5B, 1.0B y 2.0B de tokens), aunque la magnitud absoluta de la brecha se reduce ligeramente a medida que avanza el entrenamiento.
3. Validación de Controles:
   • Los controles Tied y Average se comportan de manera similar al Block AttnRes original, indicando que la ganancia no se debe al recuento de parámetros o a la duplicación de la computación.
   • El control Frozen-V tiene un desempeño deficiente, confirmando que la ruta de valores debe ser aprendida, no fija.
4. Divergencia de Rutas: Las distribuciones aprendidas de Q/K y V no colapsan a la misma solución. La divergencia JS promedio entre capas es aproximadamente 0.218 (120M) y 0.215 (343M), confirmando que el modelo aprende lecturas de profundidad distintas para la coincidencia frente a la recuperación de contenido.

Significancia y Reivindicaciones

El artículo realiza una reivindicación estrecha y acotada en lugar de una universal. El autor no afirma que RD-AttnRes sea superior a todas las arquitecturas de enrutamiento residual o que establezca un comportamiento de escalado universal.

En su lugar, los resultados respaldan la conclusión específica de que dentro del régimen de entrenamiento probado (FineWeb-Edu, 2B de tokens, escalas de modelos específicas), la lectura de profundidad QKV compartida en Block AttnRes es una restricción consequential. La evidencia sugiere que la coincidencia (Q/K) y la recuperación de contenido (V) se benefician de diferentes lecturas de profundidad.

Limitaciones Reconocidas:

• El estudio está limitado a un corpus web en inglés y modelos de menos de mil millones de parámetros.
• La brecha de rendimiento absoluto se reduce en horizontes de entrenamiento más largos; se requiere trabajo futuro para determinar si la brecha persiste o crece a los 4B–10B de tokens.
• Existe una advertencia de sistemas: la implementación actual incurre en una penalización de rendimiento (throughput) (~9–10%) y un mayor uso de memoria en comparación con el Block AttnRes de referencia, aunque los kernels fusionados podrían mitigar esto.
• Las comparaciones con otras arquitecturas recientes (por ejemplo, MUDDFormer, Delta AttnRes) son conceptuales más que baselines empíricos emparejados.

Conclusión:
El artículo establece que cuando la profundidad se convierte en un eje de recuperación aprendido, el intercambio debe seguir la función descendente (rol) en lugar de la conveniencia de implementación. Al desacoplar la ruta de valores con un gasto mínimo de parámetros, RD-AttnRes logra mejoras consistentes en la pérdida de preentrenamiento, validando la hipótesis de que la geometría de coincidencia y la recuperación de contenido son tareas distintas que requieren representaciones de profundidad distintas.