Source-linked AI summary
Attention Residuals
Kimi Team, Guangyu Chen, Yu Zhang, Jianlin Su, Weixin Xu, Siyuan Pan, Yaoyu Wang, Yucheng Wang, Guanduo Chen, Bohong Yin, Yutian Chen, Junjie Yan, Ming Wei, Y. Zhang, Fanqing Meng, Chao Hong, Xiaotong Xie, Shaowei Liu, Enzhe Lu, Yunpeng Tai, Yanru Chen, Xin Men, Haiqing Guo, Y. Charles, Haoyu Lu, Lin Sui, Jinguo Zhu, Zaida Zhou, Weiran He, Weixiao Huang, Xinran Xu, Yuzhi Wang, Guokun Lai, Yulun Du, Yuxin Wu, Zhilin Yang, Xinyu Zhou
TL;DR
PreNorm residuals uniformly accumulate layer outputs, causing hidden-state growth and dilution across depth. The paper replaces this accumulation with learned softmax attention and introduces a blockwise variant for scalable training, finding consistent gains across evaluations.
Problem
PreNorm residuals use fixed unit weights to aggregate preceding layer outputs, causing hidden-state magnitudes to grow as O(L) and progressively diluting each layer’s contribution.
Method
AttnRes uses learned, input-dependent softmax attention over preceding layer outputs, while Block AttnRes attends over block-level representations to reduce large-scale memory and communication costs.
Results
AttnRes consistently outperforms the baseline across compute budgets, and downstream performance improves across all evaluated tasks in the Kimi Linear model.
Takeaways & Limitations
Block AttnRes preserves most of Full AttnRes’s gains with about 8 blocks, making selective depth-wise aggregation a practical residual replacement at scale.
Takeaways & Limitations
Full AttnRes requires retaining and communicating layer outputs, causing memory and communication overhead that grows as O(Ld) under activation recomputation and pipeline parallelism.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper identifies fixed, uniform depth-wise residual accumulation as a source of hidden-state growth and diluted layer contributions, then proposes learned attention over preceding representations. Block AttnRes makes this approach practical at scale, while experiments show consistent improvements over standard residuals.
- Standard residuals uniformly aggregate all preceding layer outputs, leaving no mechanism to selectively emphasize or suppress depth-wise contributions.
- PreNorm accumulation grows hidden-state magnitudes as O(L), progressively diluting each layer’s relative contribution and burying early-layer information.
- AttnRes replaces fixed accumulation with softmax attention using learned pseudo-queries, enabling selective, content-aware retrieval across depth.
- Block AttnRes partitions layers into blocks and attends over block-level representations to reduce memory and communication overhead for large-scale training.
- AttnRes consistently outperforms the baseline across compute budgets, while Block AttnRes matches the loss of a baseline trained with 1.25× more compute.
- The paper combines cross-stage caching and two-phase computation with Block AttnRes, reporting marginal training overhead and less than 2% inference latency overhead.
2 Motivation
This section explains that residual learning provides a direct identity path for information and gradients but uniformly compresses earlier layer outputs into a single state. The paper motivates selective depth-wise access because this compression prevents layer-specific retrieval and contributes to output growth.
- Residual learning lets gradients bypass transformations through identity mappings, supporting stable training at depth.
- Expanding the residual recurrence shows that each hidden state sums the token embedding with all preceding layer outputs.
- Fixed unit coefficients treat every layer contribution uniformly, unlike learned gates that adapt the interpolation between transformation and identity paths.
- Because each layer accesses only one compressed state, residual and gated approaches cannot selectively retrieve individual earlier outputs.
3 Attention Residuals: A Unified View of Time and Depth
Attention Residuals replace fixed depth-wise accumulation with learned attention, while Block AttnRes groups layers into block representations to reduce scaling overhead. The block design preserves selective cross-depth access while reducing stored and communicated representations.
- Attention Residuals: Attention Residuals apply softmax attention over preceding layer outputs, using learned layer-specific queries to select depth-wise representations.The query is a learned vector per layer, and RMSNorm prevents large-magnitude outputs from dominating attention weights.
- Full Attention Residuals: Full AttnRes attends over all preceding outputs with O(L^2d) arithmetic and O(Ld) memory, with memory overlapping backpropagation activations in vanilla training.Under activation recomputation and pipeline parallelism, retaining and transmitting all outputs makes memory and communication overhead grow as O(Ld).
- Block Attention Residuals: Block AttnRes partitions layers into N blocks, sums outputs within each block, and applies attention across block-level representations and the token embedding.Within a block, later layers also attend to the partial sum of preceding outputs, while the final output aggregates all block representations.
- Block Attention Residuals: Block AttnRes reduces memory and communication from O(Ld) to O(Nd), while computation falls from O(L^2) to O(N^2).N=L recovers Full AttnRes, N=1 reduces to standard residuals, and empirically N≈8 recovers most benefits across model scales.
- Block Attention Residuals: Blockwise computation uses partial intra-block sums and parallel inter-block results, merging them with online softmax while preserving exact equivalence.The block count also bounds KV-cache size and defines dispatch granularity for the blockwise optimization.
4 Infrastructure Design
Block AttnRes addresses the memory, communication, and inference costs of depth-wise attention through block compression, cross-stage caching, and two-phase computation. These optimizations reduce redundant transfers and accesses while keeping memory and latency overhead low.
- Pipeline communication: Block AttnRes propagates block-level representations across pipeline stages, avoiding the full layer-output history required by Full AttnRes.Layers are partitioned into blocks, reducing stored representations from L to N.
- Pipeline communication: Cross-stage caching reduces peak per-transition communication from O(C) to O(P), a V × improvement during interleaved pipeline execution.Previously received blocks remain local, so later virtual stages transmit only incremental blocks; the backward pass uses the same scheme.
- Pipeline communication: With P=4 and V=2, caching eliminates 6 redundant block transmissions in the second virtual stage.Block boundaries need not align with physical stage boundaries, so only transitions completing new blocks transmit them.
- Two-phase computation: The two-phase schedule batches inter-block queries and processes intra-block lookback sequentially, reducing Full AttnRes per-layer I/O from O(Ld) to O((S+N)d).Phase 1 reuses cached block representations, while Phase 2 merges sequential results with online softmax.
- Two-phase computation: Block AttnRes further reduces stored representations from L to N by compressing each block into one vector, while amortized access costs are measured across N layers.The schedule keeps Phase 2 close to standard residual I/O and allows partial overlap of Phase 1 with first-layer computation.
5 Experiments
Experiments show that AttnRes improves scaling, training dynamics, downstream performance, and depth-wise information flow relative to standard residual accumulation, while Block AttnRes preserves most benefits with lower overhead.
- 5.1 Scaling Laws: All variants have similar scaling slopes, but Full and Block AttnRes consistently achieve lower loss than the Baseline across the compute range.The fitted curves are Baseline L = 1.891 × C^-0.057, Block AttnRes L = 1.870 × C^-0.058, and Full AttnRes L = 1.865 × C^-0.057.
- 5.2 Main Results: Block AttnRes produces consistently lower validation loss, bounded periodic output magnitudes, and more uniform gradient magnitudes across depth than the Baseline.Selective aggregation resets accumulation at block boundaries, while competitive softmax weights regulate source contributions.
- 5.2 Main Results: Block AttnRes matches or outperforms the Baseline on all evaluated benchmarks, including gains of +7.5 on GPQA-Diamond, +3.6 on Minerva Math, and +3.1 on HumanEval.Knowledge-oriented benchmarks also improve, including MMLU (+1.1) and TriviaQA (+1.9).
- 5.3 Ablation Study: Full AttnRes achieves loss 1.737 and Block AttnRes 1.746, outperforming DenseFormer and mHC by using content-dependent softmax selection over depth.DenseFormer reaches 1.767 and mHC reaches 1.747 under the reported comparison.
- 5.3 Ablation Study: Block AttnRes retains strong performance with coarse grouping: block sizes S=2, 4, and 8 remain near loss 1.746, whereas S=16 and 32 approach the Baseline.With S=4, memory overhead remains constant per layer.
6 Discussions
The paper frames residual variants as depth-wise aggregation mechanisms and uses structured mixing matrices to compare their expressivity and source-selection patterns. AttnRes generalizes these mechanisms from recurrence or linear attention to depth-wise softmax attention, while Block AttnRes interpolates between standard residuals and Full AttnRes.
- Depth-wise recurrence: Residual connections propagate information across depth through an additive recurrence analogous to RNN state updates.The hidden state acts as the recurrent state, and each layer transformation acts like one update step.
- Structured matrices: Depth mixing matrices represent how each layer weights outputs from earlier layers, providing a common lens for comparing residual variants.The matrix may be fixed, learned, or input-dependent, and may be low-rank or dense.
- Structured matrices: Standard residuals use an all-ones lower-triangular mixing matrix, whereas Highway uses input-dependent scalar-gated weights that sum to one.Highway remains 1-semiseparable and can be interpreted as softmax-free stick-breaking attention over depth.
- Structured matrices: Multi-stream methods increase effective depth-mixing rank through parallel states and learned transitions, while mHC constrains transitions to stabilize cumulative products.The effective rank is tied to the number of streams, and mHC uses doubly stochastic transition matrices.
- AttnRes variants: Full AttnRes uses dense, input-dependent attention over preceding layer outputs, while Block AttnRes shares source weights within blocks and interpolates between N and N + S effective rank.Standard residual corresponds to N=1, and Full AttnRes to N=L.
- Practicality: The structured view identifies depth-wise attention sinks and motivates blockwise designs that reduce per-layer memory I/O from O(Ld) to O((S+N)d).Blockwise computation exploits pseudo-queries that are independent of sequential layer outputs.
7 Related Work
The related work contrasts normalization-based residual variants, multi-state recurrences, and cross-layer connectivity methods. AttnRes is positioned as direct, input-dependent cross-layer attention over depth rather than another refinement of single-state recurrence.
- Normalization and depth stability: PreNorm preserves an identity gradient path but allows hidden-state magnitudes to grow as O(L), while PostNorm bounds magnitudes but can cause gradient vanishing.The discussion presents normalization placement as a trade-off between magnitude control and gradient propagation.
- Multi-state recurrence: Multi-state recurrence methods widen the state with parallel streams or matrix states to address the inability to selectively retrieve individual earlier-layer contributions.Examples include Hyper-Connections, mHC, DDL, and SiameseNorm.
- Cross-layer connectivity: Cross-layer connectivity methods range from static or learned scalar weights to input-dependent aggregation, giving layers direct access to earlier representations.Examples include DenseNet, ELMo, DenseFormer, ANCRe, and MUDDFormer.
Conclusion
The conclusion presents AttnRes as a replacement for uniform residual accumulation with learned, input-dependent depth-wise attention, and Block AttnRes as its scalable approximation. The reported gains persist across scales, while blockwise implementation remains practical with limited overhead.
- Contribution: AttnRes replaces fixed, uniform residual accumulation with learned, input-dependent attention over depth.The method is motivated by the sequence-depth duality and validated through ablations and scaling-law experiments.
- Scalable implementation: Block AttnRes partitions layers into blocks and attends over block-level representations because Full AttnRes requires cross-layer memory that grows as O(Ld).About 8 blocks recover most of Full AttnRes’s gains, according to the conclusion.
- Scalable implementation: Cross-stage caching and a two-phase computation strategy make Block AttnRes practical at scale with marginal training overhead and minimal inference overhead.The conclusion describes this as a drop-in design for large-scale training constraints.
A Contributions
The listed contributors are ordered by contribution significance, with project leadership authors appearing last. The section provides the author roster rather than substantive technical contribution descriptions.
- Contribution ordering: The paper states that authors are listed in order of the significance of their contributions.This ordering convention is explicitly described in the section material.
- Author roster: Guangyu Chen, Yu Zhang, Jianlin Su, Weixin Xu, Siyuan Pan, Yaoyu Wang, and Yucheng Wang are listed first.These names appear in the first author group.
- Author roster: Guanduo Chen, Bohong Yin, Yutian Chen, Junjie Yan, Ming Wei, and Y. Zhang are also listed among the contributors.The roster continues with additional contributors across subsequent author groups.
- Author roster: Fanqing Meng, Chao Hong, Xiaotong Xie, Shaowei Liu, Enzhe Lu, and Yunpeng Tai are listed among the contributors.These names form another author group in the supplied roster.
- Author roster: Yanru Chen, Xin Men, Haiqing Guo, and Y. Charles are listed among the contributors.These names appear in the next author group.
- Author roster: Haoyu Lu, Lin Sui, Jinguo Zhu, Zaida Zhou, Weiran He, Weixiao Huang, Xinran Xu, and Yuzhi Wang are listed among the contributors.This group contains eight listed authors.
- Author roster: Guokun Lai, Yulun Du, and Yuxin Wu are listed among the contributors.They appear in the final supplied author group.
B Optimized Inference I/O for Full Attention Residuals
Because the pseudo-query is fixed before execution, inter-block accesses can be batched across layers in a two-phase schedule, while block partitioning remains an inference scheduling device.
- The two-phase schedule batches inter-block accesses across layers because each layer’s pseudo-query is input- and hidden-state-independent.This reduces total I/O below the naïve implementation’s depth-linear bound.
- Block partitioning is used only for inference scheduling and does not alter the model architecture or replace per-layer sources with block summaries.
- Phase 1 jointly computes inter-block attention for all layers in a block, while Phase 2 processes intra-block dependencies sequentially.
Phase 1: Batched Inter-block Attention
Phase 1 reuses preceding blocks’ key–value pairs across all queries in the current block, then handles intra-block dependencies separately and sequentially.
- The preceding (n−1)S key–value pairs are read once from HBM and reused across all S queries in block n.The inter-block read cost is 2(n−1)Sd, counting both keys and values.
- Summed across blocks, inter-block reads equal dL(N−1), using SN=L.
- Phase 1 writes one d-dimensional output per layer after covering all sources before the current block.
- Within each block, layer t reads t−1 intra-block key–value pairs sequentially at cost 2(t−1)d.
- Phase 2 writes one output per layer after processing intra-block dependencies.
Total Amortized I/O per Layer
Amortizing both phases yields per-layer I/O that depends on block size and block count rather than total depth, with linear scaling in S+N.
- Read per layer = (S+N−2)d and write per layer = 2d after dividing total I/O by L.
- Batching reduces per-layer I/O from O(L) to O(S+N).
- The schedule assigns most traffic to inter-block attention while keeping sequential computation local within each block.