Attention Residuals explicado: repensando la profundidad de los Transformers

Los LLM modernos son cada vez más profundos, anchos y voraces en cómputo, pero apilar más capas de Transformer no siempre se traduce en mejoras proporcionales. Una razón es que las conexiones residuales estándar agregan las salidas de capa con pesos unitarios fijos, de modo que cada capa hereda una suma uniforme de todo lo que vino antes. Esto puede diluir las representaciones tempranas, amplificar la magnitud del estado oculto y dificultar que la red reutilice de forma selectiva las características intermedias más útiles.

En lugar de tratar la profundidad como una recurrencia aditiva fija, el trabajo Attention Residuals del equipo de Kimi permite que cada capa atienda a las salidas de capas anteriores usando pesos aprendidos con softmax. En este artículo, veremos por qué la agregación residual estándar se convierte en un cuello de botella, cómo funciona Attention Residuals, por qué importa la variante por bloques y qué sugieren realmente los resultados sobre la escalabilidad de modelos de lenguaje más profundos.

Si te interesa profundizar en algunas ideas detrás de las arquitecturas Transformer modernas, te recomiendo el curso de DataCamp Transformer Models Tutorial in  PyTorch.

El problema: la señal que se ahoga

Las conexiones residuales han sido fundamentales en el deep learning durante años. En Transformers, la actualización residual estándar es:

hl=hl−1​+fl−1​(hl−1​)

Esto ayuda a que los gradientes fluyan por redes muy profundas. Pero los residuales no son solo un truco de gradiente. También definen cómo se agrega la información a lo largo de la profundidad. Si desplegamos la recurrencia, obtenemos:

hl​=h1​+∑i=1l−1​fi​(hi​)

Esto significa que el estado oculto en la capa l es básicamente el embedding y una suma con pesos uniformes de todas las salidas anteriores. En la práctica, cada contribución recibe el mismo peso.

Esto se convierte en un problema a escala. El artículo muestra que en arquitecturas PreNorm, la acumulación residual no ponderada hace que las magnitudes del estado oculto crezcan con la profundidad, aproximadamente como O(L). 

A medida que crece el flujo residual, las salidas de capas tempranas se diluyen dentro de una suma acumulada cada vez mayor. Una vez que una señal temprana se mezcla en esta acumulación, las capas más profundas no pueden recuperarla de forma selectiva. Por tanto, solo operan sobre el estado ya agregado.

Esto conduce a lo que el artículo llama el efecto de "señal que se ahoga". Una señal empírica clara de esta ineficiencia viene de estudios de poda de capas: a menudo se puede eliminar una fracción sustancial de capas en modelos entrenados con un impacto mínimo en el rendimiento. Esto sugiere que, aunque seguimos aumentando la profundidad, a los modelos les falta un mecanismo eficaz para exprimirla al máximo. En lugar de formar una cadena jerárquica de razonamiento, la red se comporta más como un relevo redundante donde las señales tempranas se diluyen progresivamente.

También hay una compensación en el pase hacia adelante. A medida que crece el flujo residual, las capas posteriores pueden necesitar producir salidas de mayor magnitud para influir de forma significativa en el estado acumulado. El artículo lo relaciona con el comportamiento PreNorm, donde las magnitudes del estado oculto aumentan monótonamente con la profundidad bajo la agregación residual estándar.

Las conexiones residuales ayudan al flujo de gradientes, y el gradiente con respecto a un estado oculto intermedio puede escribirse como:

El término identidad conserva aquí una ruta de gradiente directa. Pero los residuales siguen forzando que la ruta de agregación hacia adelante trate cada capa anterior con un peso fijo de 1,0. Esa es la limitación estructural que AttnRes intenta solucionar.

De residuales aditivos a atención a lo largo de la profundidad

El movimiento conceptual más interesante del artículo es la dualidad tiempo–profundidad. Las conexiones residuales comprimen información a lo largo de la profundidad, del mismo modo que las redes recurrentes la comprimen en el tiempo. En modelado secuencial, la atención sustituyó a la recurrencia permitiendo que cada posición accediera selectivamente a posiciones anteriores. AttnRes aplica la misma transición a la profundidad de la red. En lugar de definir el siguiente estado oculto como una suma fija de capas previas, AttnRes permite que cada capa atienda a las salidas de capas anteriores:

donde los pesos αi→l son pesos de atención con softmax sobre la profundidad y suman 1. Esos pesos se calculan como:

y el artículo utiliza:

Cada capa obtiene un vector de pseudo-consulta aprendido wl, y esa consulta atiende a claves y valores construidos a partir de las salidas de capas anteriores. Se aplica RMSNorm a las claves para que las capas con salidas naturalmente mayores no dominen el softmax solo por su escala.

El cambio clave del artículo es tratar los residuales estándar como una especie de atención lineal a lo largo de la profundidad, mientras que AttnRes lo actualiza a atención con softmax en la profundidad. En lugar de acumulación uniforme, obtenemos recuperación selectiva a través de la profundidad.

Un detalle de implementación pequeño pero importante es la inicialización. Los autores sugieren inicializar todos los vectores de pseudo-consulta a cero, lo que garantiza que los pesos de atención iniciales sean uniformes; así, el modelo comienza el entrenamiento como un promedio de peso igual en lugar de un mecanismo de atención sesgado al azar.

Figura 1: descripción general de Attention Residuals (Attention Residuals paper)

Full Attention Residuals

En Full Attention Residuals, cada capa atiende a todas las salidas de capas previas, lo que le permite agregar información de forma selectiva a lo largo de la profundidad. Esto da al modelo la máxima flexibilidad: una capa más profunda puede enfatizar a su predecesora inmediata, al embedding original o a cualquier capa anterior si ahí reside la señal más útil.

En entrenamiento estándar, la sobrecarga de memoria de Full AttnRes es menor de lo que podría parecer al principio, porque muchas salidas de capa ya se conservan para backpropagation. Pero el entrenamiento a gran escala cambia el panorama. Cuando introduces recomputación de activaciones y paralelismo en pipeline, esas salidas tempranas deben preservarse y comunicarse explícitamente para que las capas posteriores puedan atenderlas, y eso se vuelve costoso.

Así que Full AttnRes es la mejor forma de entender la idea central, pero no es la versión que la mayoría de equipos querría desplegar a escala.

Block Attention Residuals

El artículo introduce la idea de Block Attention Residuals para reducir la sobrecarga de memoria y comunicación. Las capas del modelo se dividen en bloques de forma que, dentro de cada bloque, las salidas se combinan con acumulación residual aditiva estándar, pero entre bloques, el modelo realiza atención sobre resúmenes a nivel de bloque en lugar de cada capa individual previa.

Si Bn es el conjunto de capas del bloque n, entonces la representación del bloque viene dada por:

El modelo entonces atiende al embedding b0=h1 a los resúmenes de bloques anteriores y a la suma parcial del bloque actual conforme avanza el cómputo. Esto reduce la memoria y la comunicación de O(Ld) a O(Nd), donde N es el número de bloques.

El artículo informa de que usar alrededor de ocho bloques recupera la mayor parte del beneficio de la versión completa, y la brecha de rendimiento entre Full AttnRes y Block AttnRes se estrecha a medida que aumenta la escala. Tamaños de bloque como S=2,4,8 se mantienen muy cerca de la versión completa, mientras que agrupaciones mucho más gruesas tienden a volver al comportamiento base.

La lección más útil del artículo es que no necesitas atención a lo largo de toda la profundidad y sobre cada capa para conseguir la mayor parte de la ganancia. Aunque Block AttnRes hace viable computacionalmente la atención a lo largo de la profundidad, desplegarla de forma eficiente a escala sigue requiriendo un diseño de sistemas cuidadoso. El artículo también introduce varias optimizaciones que hacen que AttnRes sea práctico para entrenamiento e inferencia en el mundo real.

Figura 2: ejemplo de comunicación por pipeline basada en caché con 4 rangos físicos y 2 etapas virtuales por rango, donde las cajas rayadas denotan el final de los bloques de AttnRes (Attention Residuals paper)

Estas son tres optimizaciones clave de sistema que convierten AttnRes en un sustituto directo de los residuales estándar:

• Cómputo en dos fases: al desacoplar la pseudo-consulta del estado oculto, podemos agrupar en lotes las consultas de todas las capas dentro de un bloque. La fase 1 calcula en paralelo la atención entre bloques, mientras que la fase 2 gestiona las dependencias secuenciales dentro del bloque usando un Online Softmax Merge. Esto reduce significativamente el I/O por capa frente a métodos multistream previos como mHC.
• Caché entre etapas: en el paralelismo por pipeline, eliminamos transferencias de datos redundantes almacenando en caché localmente los bloques recibidos previamente. Esto aporta una mejora de V× (donde V es el número de etapas virtuales, como se muestra en la figura 2) en el coste máximo por transición, permitiendo solapar completamente la comunicación con el cómputo.
• Prefill eficiente en memoria: al fragmentar (sharding) las representaciones de bloque entre dispositivos en paralelismo tensorial, la huella de memoria para secuencias de largo contexto (128K+) se reduce en órdenes de magnitud (por ejemplo, de 15 GB a 0,3 GB por dispositivo).

Resultados y análisis de rendimiento

El artículo valida AttnRes a varios niveles: leyes de escalado, ablaciones, dinámica de entrenamiento y benchmarks downstream.

Figura 3: curvas de leyes de escalado para Attention Residuals. (Attention Residuals paper)

Aun así, el resultado más destacado proviene de los experimentos de leyes de escalado. En cinco tamaños de modelo, tanto Full AttnRes como Block AttnRes logran de forma consistente una pérdida de validación menor que la línea base. Según las curvas ajustadas, Block AttnRes alcanza la misma pérdida que un modelo base entrenado con aproximadamente 1,25× más cómputo. Esto sugiere que AttnRes no solo es mejor en teoría, sino también más eficiente en cómputo.

En tareas downstream, AttnRes mejora sobre la línea base en todos los benchmarks evaluados. Algunas de las mayores ganancias se dan en tareas intensivas de razonamiento, como se muestra en la tabla 1:

Tabla 1: comparación del rendimiento de AttnRes con la línea base, ambos tras la misma receta de preentrenamiento (Attention Residuals paper)

Las mejoras son especialmente fuertes en razonamiento multietapa, matemáticas y código, lo que encaja con la hipótesis central del artículo: si las capas posteriores pueden recuperar selectivamente representaciones tempranas en lugar de heredar un agregado borroso, el razonamiento composicional debería mejorar.

En comparación con la línea base, AttnRes muestra menor pérdida de validación durante todo el entrenamiento, magnitudes de salida más acotadas a lo largo de la profundidad y magnitudes de gradiente más uniformes entre capas. La base sufre el patrón habitual de dilución PreNorm, donde las magnitudes del estado oculto crecen monótonamente con la profundidad, mientras que Block AttnRes produce un patrón más controlado y acotado gracias a la agregación selectiva en los límites de bloque.

También hay un hallazgo a nivel de arquitectura: en un barrido de cómputo y parámetros fijos, la configuración óptima para AttnRes se desplaza hacia un modelo más profundo y estrecho en comparación con la línea base. Esto suiere que AttnRes puede hacer que la profundidad adicional sea más útil que bajo la agregación residual estándar.

Conclusión

A menudo se tratan los residuales como un truco necesario de entrenamiento, es decir, un atajo de identidad que mantiene vivos los gradientes. AttnRes sugiere que esta visión es demasiado estrecha. Las rutas residuales también son el mecanismo por el que la información se enruta a lo largo de la profundidad, y la acumulación aditiva fija puede ser demasiado primitiva para modelos muy profundos.

Attention Residuals no trata tanto de parchear un problema conocido como de mejorar una decisión de diseño desatendida. El modelado secuencial evolucionó de la recurrencia a la atención porque la recurrencia fija era demasiado restrictiva. AttnRes sostiene que la profundidad podría estar lista para la misma transición.

El modelo sigue mostrando fuerte localidad, con capas que a menudo atienden más a predecesoras cercanas, pero también aprende saltos no triviales, mantiene peso persistente sobre el embedding y conserva comportamientos distintos entre las capas previas a la atención y previas al MLP. 

La conclusión práctica no es que todos los modelos deban sustituir los residuales por atención completa a lo largo de la profundidad. Más bien, replantea la agregación residual como un espacio central de diseño arquitectónico.