Residuos de Atención

Resumen

Las conexiones residuales con PreNorm son estándar en los LLM modernos, pero acumulan todas las salidas de capa con pesos unitarios fijos. Esta agregación uniforme provoca un crecimiento no controlado del estado oculto con la profundidad, diluyendo progresivamente la contribución de cada capa. Proponemos Attention Residuals (AttnRes), que reemplaza esta acumulación fija con atención softmax sobre las salidas de capas precedentes, permitiendo que cada capa agregue selectivamente representaciones anteriores con pesos aprendidos y dependientes de la entrada. Para abordar la sobrecarga de memoria y comunicación de atender todas las salidas de capas precedentes en el entrenamiento de modelos a gran escala, introducimos Block AttnRes, que divide las capas en bloques y atiende a representaciones a nivel de bloque, reduciendo la huella de memoria mientras preserva la mayor parte de las ventajas de AttnRes completo. Combinado con comunicación de pipeline basada en caché y una estrategia de computación en dos fases, Block AttnRes se convierte en un reemplazo práctico y directo para las conexiones residuales estándar con sobrecarga mínima. Los experimentos de leyes de escalado confirman que la mejora es consistente entre tamaños de modelo, y las ablaciones validan el beneficio de la selección dependiente del contenido en profundidad. Además, integramos AttnRes en la arquitectura Kimi Linear (48B total / 3B parámetros activados) y realizamos pre-entrenamiento en 1.4T de tokens, donde AttnRes mitiga la dilución de PreNorm, produciendo magnitudes de salida y distribución de gradientes más uniformes a lo largo de la profundidad, y mejora el rendimiento downstream en todas las tareas evaluadas.

English

Residual connections with PreNorm are standard in modern LLMs, yet they accumulate all layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state growth with depth, progressively diluting each layer's contribution. We propose Attention Residuals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding layer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-dependent weights. To address the memory and communication overhead of attending over all preceding layer outputs for large-scale model training, we introduce Block AttnRes, which partitions layers into blocks and attends over block-level representations, reducing the memory footprint while preserving most of the gains of full AttnRes. Combined with cache-based pipeline communication and a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for standard residual connections with minimal overhead. Scaling law experiments confirm that the improvement is consistent across model sizes, and ablations validate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into the Kimi Linear architecture (48B total / 3B activated parameters) and pre-train on 1.4T tokens, where AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient distribution across depth, and improves downstream performance across all evaluated tasks.