Attention Residuals

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.

• The paper introduces Attention Residuals to replace static residual accumulation with input-dependent, softmax-based depth attention that mitigates PreNorm dilution.
• Empirical results demonstrate improved scaling behavior, balanced gradient propagation, and superior downstream performance on multi-step reasoning tasks.
• Block AttnRes partitions layers into blocks to minimize memory overhead while preserving selective depth-wise aggregation, making it scalable for large models.

Attention Residuals: Selective Depth-Wise Aggregation for Deep LLMs

Overview and Motivation

The "Attention Residuals" paper (2603.15031) examines the fundamental limitations of the standard residual connection paradigm in deep neural networks, particularly in large-scale LLMs. Classic residuals with PreNorm ensure stable gradient propagation but accumulate all preceding layer outputs with fixed, uniform weights. This undifferentiated aggregation leads to an $O(L)$ growth in hidden state magnitude and progressive attenuation of early-layer contributions, referred to as PreNorm dilution. Despite the prevalence of input-dependent mixing in sequence modeling and expert routing, most architectures retain static, additive depth-wise accumulation. The authors formalize a duality between sequential (time) recurrence and layer-wise (depth) recurrence and generalize the latter by proposing Attention Residuals (AttnRes): a mechanism that replaces static accumulation with depth-wise softmax attention.

Method: Attention Residuals (AttnRes) and Block AttnRes

Full Attention Residuals

AttnRes redefines the residual update at layer $l$ as

$$\bm{h}_l = \sum_{i=0}^{l-1} \alpha_{i \to l} \bm{v}_i,$$

where $\alpha_{i \to l}$ are softmax-normalized attention weights derived from a learned pseudo-query $\bm{w}_l$ and source-specific keys (outputs or embeddings), with

$$\alpha_{i\to l} = \frac{\exp(\bm{w}_l^\top \operatorname{RMSNorm}(\bm{k}_i))}{\sum_{j=0}^{l-1} \exp(\bm{w}_l^\top \operatorname{RMSNorm}(\bm{k}_j))}.$$

This mechanism enables selective, content-aware retrieval of prior representations, generalizing depth-wise aggregation from linear to softmax attention.

Block Attention Residuals

For scalability, Block AttnRes partitions the $L$ layers into $N$ blocks, reducing memory and communication from $O(Ld)$ to $O(Nd)$. Intra-block outputs are aggregated via summation, and inter-block attention is computed only over block-level summaries plus the token embedding. This allows the method to be deployed efficiently in large-scale settings with pipeline parallelism and activation recomputation, introducing only marginal communication and memory overhead.

System Optimizations

To address practical constraints under distributed training, the paper introduces cross-stage caching to avoid redundant block communication, and a two-phase computation schedule to amortize inter-block attention reads during inference. Sequence-sharded prefill schemes further mitigate KV cache overhead for long-context inference.

Theoretical Foundations

The paper concretizes the connection between residual architectures and attention via a unified matrix framework, showing that standard residuals and multi-stream recurrences correspond to low-rank, linear attention over the depth axis. AttnRes generalizes this with full expressive power (rank-$L$) softmax-based depth attention. The structured-matrix characterization illuminates when and why selective aggregation across layers is beneficial and clarifies the limitations of previous approaches that exclusively aggregate via compressive recurrence.

Empirical Results and Ablations

Scaling Laws and Validation Loss

Systematic scaling experiments demonstrate that AttnRes consistently dominates standard residuals and other recent generalizations (e.g., mHC(-lite), DenseFormer) across model sizes, yielding improved scaling exponents and notably lower validation loss at fixed compute. Block AttnRes with as few as 8 blocks recovers essentially all gains of the full variant, with only minimal practical overhead.

Training Dynamics

In deep Mixture-of-Experts Transformers (e.g., Kimi Linear 48B, 3B activated parameters), AttnRes produces more uniform output magnitudes and balanced gradient norms across layers, sharply limiting the PreNorm dilution observed in baselines. This results in stable, effective training even in regimes with hundreds of layers.

Downstream Performance

Block AttnRes yields consistent gains across a wide range of reasoning and language benchmarks. Particularly strong improvements are observed for multi-step logical reasoning (GPQA-Diamond, +7.5 over baseline; GSM8K, Math, and HumanEval also improved), supporting the claim that improved depth-wise access enhances compositional, non-local computation.

Ablation Studies

Ablations confirm that:

• Input-dependent, softmax-based attention is essential for optimal performance; simple input-independent or sliding-window schemes underperform.
• The scaling advantage persists across block sizes, with diminishing returns for $N > 8$.
• Multihead depth-wise aggregation and omission of RMSNorm both hurt accuracy, showing the particular effectiveness of the chosen normalization and single-vector depth attention.
• AttnRes enables effective depth-parameterized models, which favor increased depth under a fixed parameter budget—evidence of improved capacity utilization.

Analysis of Learned Depth Attention

Visualization of learned attention weights confirms intuitions about locality and global structure. Layers largely attend to their direct predecessors but reliably develop off-diagonal skip connections, and the embedding layer retains persistent weight in later blocks. The structured selectivity differentiates AttnRes from fixed recurrent or gated schemes and suggests new forms of layerwise specialization and information integration.

Implications and Future Directions

This work reframes the architecture of depth-wise information flow as an attention problem, formally unifying it with developments in sequence modeling. By achieving O($L^2$) computational cost in the depth axis, AttnRes generalizes a broad spectrum of recent residual innovations, and the proposed Block AttnRes formulation makes such techniques practical at extreme scale. Strong scaling behavior and downstream improvements indicate that deep networks can more effectively utilize additional layers when depth-wise selectivity is granted.

Block AttnRes is ready for immediate integration in any modern Transformer-based architecture due to its minimal code and system changes, and the authors anticipate finer-grained (or even full) attention over layer outputs as memory and interconnect capabilities expand.

Future research directions include input-dependent query schemes, more efficient kernelizations for depth attention, and exploration of depth-selective aggregation in other domains (e.g., vision or multi-modal networks). AttnRes opens a new orthogonal axis for architecture search beyond width and sequence mechanisms.

Conclusion

Attention Residuals introduces input-dependent, softmax-based selective aggregation over depth, strictly generalizing the standard residual paradigm. Together with a scalable blockwise variant and system optimizations, it provides superior performance across scales and tasks while imposing negligible overhead. The results substantiate the sequence-depth duality in deep architecture design and mark selective depth-wise aggregation as a new fundamental primitive for efficient model scaling and reasoning.