Attention Residuals

2026-03-16

AI Summary

This work addresses the issue of uncontrolled growth in hidden states and diluted per-layer contributions caused by fixed-weight residual connections in deep large language models. To overcome this, the authors propose Attention Residuals (AttnRes), which, for the first time, integrates input-dependent softmax attention into residual connections to enable content-aware, dynamic aggregation of information across layers. They further introduce a block-level AttnRes design with a two-stage computation strategy and cache-based pipeline communication, enhancing scalability while controlling computational overhead. Evaluated on the Kimi Linear architecture (48B/3B activated parameters) pretrained on 1.4T tokens, AttnRes substantially mitigates representation dilution under PreNorm and consistently improves performance across diverse downstream tasks, demonstrating both effectiveness and robustness at scale.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• residual connections
• hidden-state growth
• layer contribution dilution
• PreNorm
• depth scaling

Innovation

• Attention Residuals
• Block AttnRes
• PreNorm dilution
• content-dependent aggregation
• scalable residual connection