Loading...
Learned Depth-Wise Softmax Attention for Residual Connections
Kimi Team (Moonshot AI)
Read the Paper on arXivAttention 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:
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.
Click any topic to jump in
Attention Residuals replace the standard residual sum with learned attention over all previous layers.
Each sublayer contributes only O(1/√L) to the output direction, capping useful transformer depth.
Softmax attention over depth gives each layer a learned mix of all predecessor features (Eqs. 1–3).
Partitioning into blocks drops O(L²) depth attention to O(LB + (L/B)²), the √L-optimal trade-off.
Two-phase inference and block activation caching shrink KV memory, enabling 8-GPU trillion-param serving.
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.
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.