Las conexiones residuales son una de las partes menos cuestionadas del diseño moderno de transformadores. En las arquitecturas PreNorm, cada capa vuelve a agregar su salida a un estado oculto en ejecución, lo que mantiene estable la optimización y permite entrenar modelos profundos. Los investigadores de Moonshot AI sostienen que este mecanismo estándar también introduce un problema estructural: todas las salidas de las capas anteriores se acumulan con pesos unitarios fijos, lo que hace que la magnitud del estado oculto crezca con la profundidad y debilite progresivamente la contribución de cualquier capa individual.
El equipo de investigación propone los Residuales de Atención (AttnRes) como reemplazo directo de la acumulación residual estándar. En lugar de obligar a cada capa a consumir el mismo flujo residual mezclado uniformemente, AttnRes permite que cada capa agregue representaciones anteriores utilizando la atención softmax sobre la profundidad. La entrada a la capa (l) es una suma ponderada de la incorporación del token y las salidas de la capa anterior, donde las ponderaciones se calculan sobre las posiciones de profundidad anteriores en lugar de sobre las posiciones de secuencia. La idea central es simple: si la atención mejoró el modelado de secuencia al reemplazar la recurrencia fija en el tiempo, se puede aplicar una idea similar a la dimensión de profundidad de una red.
Por qué los residuos estándar se convierten en un cuello de botella
El equipo de investigación identificó tres problemas con la acumulación residual estándar. En primer lugar, no hay acceso selectivo: todas las capas reciben el mismo estado agregado, aunque las capas de atención y las capas de retroalimentación o MoE pueden beneficiarse de diferentes combinaciones de información anterior. En segundo lugar, existe una pérdida irreversible: una vez que la información se combina en un único flujo residual, las capas posteriores no pueden recuperar selectivamente representaciones anteriores específicas. En tercer lugar, está el crecimiento de la producción: las capas más profundas tienden a producir mayores producciones para seguir siendo influyentes dentro de un estado acumulado en constante crecimiento, lo que puede desestabilizar la formación.
Este es el marco principal del equipo de investigación: los residuos estándar se comportan como una recurrencia comprimida sobre capas. AttnRes reemplaza esa recurrencia fija con atención explícita sobre los resultados de la capa anterior.
Full AttnRes: atención sobre todas las capas anteriores
En Full AttnRes, cada capa calcula los pesos de atención sobre todas las fuentes de profundidad anteriores. El diseño predeterminado no utiliza una consulta condicionada por entrada. En cambio, cada capa tiene un vector de pseudoconsulta específico de la capa aprendida wl ∈ Rd, mientras que las claves y los valores provienen de la incrustación del token y de las salidas de la capa anterior después de RMSNorm. El paso RMSNorm es importante porque evita que las salidas de capas de gran magnitud dominen los pesos de atención en profundidad.
Full AttnRes es sencillo, pero aumenta el costo. Por token, requiere aritmética O (L2 d) y memoria (O (Ld)) para almacenar las salidas de la capa. En el entrenamiento estándar, esta memoria se superpone en gran medida con las activaciones que ya se necesitan para la retropropagación, pero bajo el recálculo de la activación y el paralelismo de canalización, la sobrecarga se vuelve más significativa porque esas salidas anteriores deben permanecer disponibles y es posible que deban transmitirse a través de etapas.
Block AttnRes: una variante práctica para modelos grandes
Para que el método sea utilizable a escala, el equipo de investigación de Moonshot AI presenta Block AttnRes. En lugar de revisar cada salida de capa anterior, el modelo divide las capas en N bloques. Dentro de cada bloque, las salidas se acumulan en una única representación de bloque y la atención se aplica solo a esas representaciones a nivel de bloque más la incrustación del token. Esto reduce la sobrecarga de memoria y comunicación de O(Ld) a O(Nd).
El equipo de investigación describe la comunicación de canalización basada en caché y una estrategia de cálculo de dos fases que hacen que Block AttnRes sea práctico en la inferencia y el entrenamiento distribuido. Esto da como resultado una sobrecarga de capacitación de menos del 4 % bajo el paralelismo de canalización, mientras que el repositorio reporta una sobrecarga de latencia de inferencia de menos del 2 % en cargas de trabajo típicas.
Resultados de escala
El equipo de investigación evalúa cinco tamaños de modelos y compara tres variantes en cada tamaño: una línea de base PreNorm, Full AttnRes y Block AttnRes con aproximadamente ocho bloques. Todas las variantes dentro de cada grupo de tamaño comparten los mismos hiperparámetros elegidos en la línea de base, lo que, según el equipo de investigación, hace que la comparación sea conservadora. Las leyes de escala ajustadas se expresan como:
Línea de base: L = 1,891 x C-0,057
Res. atención del bloque: L = 1,870 x C-0,058
Res. atención completa: L = 1,865 x C-0,057
La implicación práctica es que AttnRes logra una menor pérdida de validación en todo el rango de cálculo probado, y Block AttnRes iguala la pérdida de una línea de base entrenada con aproximadamente 1,25 veces más cálculo.
Integración en Kimi Linear
Moonshot AI también integra AttnRes en Kimi Linear, su arquitectura MoE con 48 mil millones de parámetros totales y 3 mil millones de parámetros activados, y lo entrena previamente en 1,4 mil millones de tokens. Según el artículo de investigación, AttnRes mitiga la dilución PreNorm manteniendo las magnitudes de salida más limitadas en profundidad y distribuyendo los gradientes de manera más uniforme entre las capas. Otro detalle de la implementación es que todos los vectores de pseudoconsulta se inicializan a cero, por lo que los pesos de atención iniciales son uniformes en todas las capas de origen, lo que reduce efectivamente AttnRes a un promedio de igual peso al inicio del entrenamiento y evita la inestabilidad temprana.
En la evaluación posterior, los avances reportados son consistentes en todas las tareas enumeradas. Informa mejoras de 73,5 a 74,6 en MMLU, de 36,9 a 44,4 en GPQA-Diamond, de 76,3 a 78,0 en BBH, de 53,5 a 57,1 en Matemáticas, de 59,1 a 62,2 en HumanEval, de 72,0 a 73,9 en MBPP, de 82,0 a 82,9 en CMMLU y 79,6 a 82,5 en C-Eval.
Conclusiones clave
Attention Residuals reemplaza la acumulación residual fija con atención softmax sobre las capas anteriores. El diseño predeterminado de AttnRes utiliza una pseudoconsulta específica de la capa aprendida, no una consulta condicionada por entrada. Block AttnRes hace que el método sea práctico al reducir la memoria profunda y la comunicación de O(Ld) a O(Nd). El equipo de investigación de Moonshot informa una pérdida de escala menor que la línea base PreNorm, con Block AttnRes igualando aproximadamente 1,25 veces más cálculo de línea base. En Kimi Linear, AttnRes mejora los resultados en los puntos de referencia de razonamiento, codificación y evaluación con una sobrecarga limitada.
Consulte Paper and Repo. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.