The Year in AI Papers
The year/Moonshot AI

Paper 2603.15031

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, et al.

Published
Mar 2026
Research lab
Moonshot AI
Citations
42
GitHub
3.5K stars

01 In brief

Summary

The paper proposes Attention Residuals (AttnRes), replacing fixed unit-weight residual connections in LLMs with learned, input-dependent softmax attention over preceding layer outputs.

This addresses PreNorm dilution, where hidden-state magnitudes grow with depth, diluting layer contributions.

A scalable variant, Block AttnRes, partitions layers into blocks, attending over block-level representations to reduce memory and communication from O(Ld) to O(Nd).

Infrastructure optimizations include cache-based pipeline communication and a two-phase inference strategy, making Block AttnRes a drop-in replacement with minimal overhead.

Scaling law experiments show consistent improvements, with Block AttnRes matching a baseline trained with 1.25x more compute.

Integrated into the Kimi Linear architecture (48B total/3B activated) and pre-trained on 1.4T tokens, AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient distributions, and improves performance on all evaluated downstream tasks.

Ablations validate the benefit of content-dependent depth-wise selection, with about 8 blocks recovering most of the gains of full AttnRes.

The method generalizes prior residual variants, which are shown to be depth-wise linear attention, while AttnRes performs depth-wise softmax attention.

Training overhead is marginal, and inference latency overhead is less than 2% on typical workloads.

The paper also includes an architecture sweep showing AttnRes favors deeper, narrower models under fixed compute, and analyzes learned attention patterns, revealing preserved locality and layer specialization.

The work is open-sourced at https://github.com/MoonshotAI/Attention-Residuals.

02 From the paper

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.