Attention Residuals
Learned Depth-Wise Softmax Attention for Residual Connections
Kimi Team (Moonshot AI)
Read the Paper on arXivPaper Overview
Attention Residuals (AttnRes) proposes replacing the fixed, uniform residual accumulation used in all modern LLMs with learned softmax attention over preceding layer outputs — enabling each layer to selectively retrieve and emphasize information from any earlier layer based on the current input.
Standard residual connections with PreNorm accumulate all layer outputs with fixed unit weights: h_l = h_1 + f_1(h_1) + f_2(h_2) + ... + f_{l-1}(h_{l-1}). This uniform aggregation causes PreNorm dilution — hidden-state magnitudes grow as O(L) with depth, progressively diluting each layer's relative contribution. Information buried in early layers cannot be selectively recovered.
Key contributions:
- Full AttnRes: Each layer computes softmax attention over all preceding layer outputs using a single learned pseudo-query vector per layer — replacing uniform accumulation with content-dependent, depth-wise selection
- Block AttnRes: A practical variant that partitions L layers into N blocks, applying attention at the block level — reducing memory from O(Ld) to O(Nd) while preserving most of the gain
- Infrastructure for scale: Cross-stage caching and a two-phase inference strategy make Block AttnRes a drop-in replacement with < 4% training overhead and < 2% inference latency overhead
- Scaling validation: AttnRes consistently outperforms the baseline across model sizes, matching a baseline trained with 1.25× more compute
The method is integrated into the Kimi Linear architecture (48B total / 3B activated parameters) and pre-trained on 1.4 trillion tokens, where it yields bounded output magnitudes, uniform gradient distribution, and consistent improvements across all evaluated benchmarks.
Chapter Roadmap
Click any topic to jump in
Pipeline Summary
Attention Residuals replace the standard residual sum with learned attention over all previous layers.
PreNorm Dilution
Each sublayer contributes only O(1/√L) to the output direction, capping useful transformer depth.
Full Attention Residuals
Softmax attention over depth gives each layer a learned mix of all predecessor features (Eqs. 1–3).
Block Attention Residuals
Partitioning into blocks drops O(L²) depth attention to O(LB + (L/B)²), the √L-optimal trade-off.
Infrastructure for Scale
Two-phase inference and block activation caching shrink KV memory, enabling 8-GPU trillion-param serving.
Scaling Laws & Benchmarks
Steeper compute-optimal slope delivers +2 MMLU and +4 GSM8K at equal FLOPs — depth is truly usable again.
The complete Attention Residuals framework — from the PreNorm dilution problem through Full and Block AttnRes to practical deployment at scale on Kimi Linear 48B.
Key Points
Replaces fixed unit-weight residual accumulation with learned softmax attention over depth
Full AttnRes: each layer attends to ALL previous outputs via a single learned pseudo-query w_l per layer
Block AttnRes: partition L layers into N≈8 blocks, intra-block standard residuals, inter-block full attention
Memory reduced from O(Ld) to O(Nd) — practical for pipeline parallelism at scale
Cross-stage caching eliminates redundant communication; two-phase inference keeps latency overhead < 2%
Scaling laws: Block AttnRes matches baseline trained with 1.25× more compute across all model sizes
Integrated into Kimi Linear 48B (3B active): GPQA-Diamond +7.5, Math +3.6, HumanEval +3.1, MMLU +1.1
Training dynamics: bounded output magnitudes and more uniform gradient distribution across depth
Attention Residuals replace the vanilla residual with a learned attention over all previous layer outputs: . This turns depth-wise gradient flow from a sum-of-many-small-signals into a learned path, fixing PreNorm's signal-dilution problem.