Attention-Residuals README

The project's README from github.com/MoonshotAI/Attention-Residuals. As one markdown file: README.md.

━━━━━━━━━━━━━━━━━━━━━━━━━━━
Attention Residuals
━━━━━━━━━━━━━━━━━━━━━━━━━━━

Paper  |  arXiv  |  Overview  |  Results  |  Citation

(a) Standard residuals with uniform additive accumulation. (b) Full AttnRes: each layer attends over all previous outputs. (c) Block AttnRes: layers are grouped into blocks, reducing memory from O(Ld) to O(Nd).

This is the official repository for Attention Residuals (AttnRes), a drop-in replacement for standard residual connections in Transformers that enables each layer to selectively aggregate earlier representations via learned, input-dependent attention over depth.

Overview

Standard residual connections accumulate all layer outputs with fixed unit weights. As depth grows, this uniform aggregation dilutes each layer's contribution and causes hidden-state magnitudes to grow unboundedly — a well-known problem with PreNorm.

AttnRes replaces this fixed accumulation with softmax attention over preceding layer outputs:

$$\mathbf{h}l = \sum{i=0}^{l-1} \alpha_{i \to l} \cdot \mathbf{v}_i$$

where the weights $\alpha_{i \to l}$ are computed via a single learned pseudo-query $\mathbf{w}_l \in \mathbb{R}^d$ per layer. This gives every layer selective, content-aware access to all earlier representations.

Block AttnRes

Full AttnRes is straightforward but requires O(Ld) memory at scale. Block AttnRes partitions layers into N blocks, accumulates within each block via standard residuals, and applies attention only over block-level representations. With ~8 blocks, it recovers most of Full AttnRes's gains while serving as a practical drop-in replacement with marginal overhead.

PyTorch-style pseudocode 
def block_attn_res(blocks: list[Tensor], partial_block: Tensor, proj: Linear, norm: RMSNorm) -> Tensor:
    """
    Inter-block attention: attend over block reps + partial sum.
    blocks:
        N tensors of shape [B, T, D]: completed block representations for each previous block
    partial_block:
        [B, T, D]:    intra-block partial sum (b_n^i)
    """
    V = torch.stack(blocks + [partial_block])  # [N+1, B, T, D]
    K = norm(V)
    logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)
    h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
    return h

def forward(self, blocks: list[Tensor], hidden_states: Tensor) -> tuple[list[Tensor], Tensor]:
    partial_block = hidden_states
    # apply block attnres before attn
    # blocks already include token embedding
    h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)

    # if reaches block boundary, start new block
    # block_size counts ATTN + MLP; each transformer layer has 2
    if self.layer_number % (self.block_size // 2) == 0:
        blocks.append(partial_block)
        partial_block = None

    # self-attention layer
    attn_out = self.attn(self.attn_norm(h))
    partial_block = partial_block + attn_out if partial_block is not None else attn_out

    # apply block attnres before MLP
    h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)

    # MLP layer
    mlp_out = self.mlp(self.mlp_norm(h))
    partial_block = partial_block + mlp_out

    return blocks, partial_block

Results

Scaling Laws

AttnRes consistently outperforms the baseline across all compute budgets. Block AttnRes matches the loss of a baseline trained with 1.25x more compute.

Downstream Performance (Kimi Linear 48B / 3B activated, 1.4T tokens)

Category Benchmark Baseline AttnRes
General MMLU 73.5 74.6
GPQA-Diamond 36.9 44.4
BBH 76.3 78.0
TriviaQA 69.9 71.8
Math & Code Math 53.5 57.1
HumanEval 59.1 62.2
MBPP 72.0 73.9
Chinese CMMLU 82.0 82.9
C-Eval 79.6 82.5

AttnRes improves across the board, with the largest gains on multi-step reasoning (+7.5 on GPQA-Diamond) and code generation (+3.1 on HumanEval).

Training Dynamics

AttnRes mitigates PreNorm dilution: output magnitudes remain bounded across depth and gradient norms distribute more uniformly across layers.

Citation

If you found our work useful, please cite

@misc{chen2026attnres,
  title         = {Attention Residuals},
  author        = {Kimi Team  and Chen, Guangyu  and Zhang, Yu  and Su, Jianlin  and Xu, Weixin  and Pan, Siyuan  and Wang, Yaoyu  and Wang, Yucheng  and Chen, Guanduo  and Yin, Bohong  and Chen, Yutian  and Yan, Junjie  and Wei, Ming  and Zhang, Y.  and Meng, Fanqing  and Hong, Chao  and Xie, Xiaotong  and Liu, Shaowei  and Lu, Enzhe  and Tai, Yunpeng  and Chen, Yanru  and Men, Xin  and Guo, Haiqing  and Charles, Y.  and Lu, Haoyu  and Sui, Lin  and Zhu, Jinguo  and Zhou, Zaida  and He, Weiran  and Huang, Weixiao  and Xu, Xinran  and Wang, Yuzhi  and Lai, Guokun  and Du, Yulun  and Wu, Yuxin  and Yang, Zhilin  and Zhou, Xinyu},
  year          = {2026},
  archiveprefix = {arXiv},
  eprint        = {2603.15031},
  primaryclass  = {cs.CL}
}

Attention-Residuals (MoonshotAI/Attention-Residuals)

Searching Database…

3.5K ·GitHub
Attention-Residuals
这是什么？AttnRes（Attention Residuals，注意力残差）是 Transformer 中标准残差连接的即插即用替代方案，让每一层用可学习的、依赖输入的注意力关注之前所有层的输出，而不是简单地把各层输出相加。
这是可以下载的现成模型，还是要我自己动手实现的技术？它是一种用公式和“PyTorch 风格伪代码”描述的架构技术，不是预训练权重；你需要把它集成到自己 Transformer 的残差连接中。
我到底该怎么在自己的模型里用它？README 没有给出安装命令或可运行的代码片段，只有公式以及对“PyTorch 风格伪代码”的引用；你需要按照论文的描述，在自己的 Transformer 中实现 Full AttnRes，或实现更省资源的 Block AttnRes 变体。
我需要特定的深度学习框架吗？需要。参考实现是按 PyTorch 的术语描述的，所以你需要一个基于 PyTorch 的 Transformer 代码库才能应用它。
Full AttnRes 能毫无额外代价地扩展到大模型吗？不能——随着深度增长，Full AttnRes 需要 O(Ld) 的内存，规模一大就不现实；Block AttnRes 解决了这个问题：它把层分成约 8 个块，只在块级表示上做注意力，用少量性能损失换来 O(Nd) 的内存。
这对谁有用？它面向构建或训练基于 Transformer 的大语言模型、想为标准残差连接寻找替代方案的研究人员和工程师，而不是成品应用的最终用户。
相比标准残差，AttnRes 能带来多大性能提升？在缩放定律测试中，它在每个算力预算下都能稳定降低损失，其中 Block AttnRes 相当于用多 1.25 倍算力训练出的基线；在 Kimi Linear 48B 模型上，它在 GPQA-Diamond 上最多提升 7.5 分，在 HumanEval 上提升 3.1 分。
为什么标准残差连接在深层 Transformer 中会成为问题？标准残差用固定的单位权重把所有层输出相加，随着深度增加，每一层的贡献被稀释，隐藏状态的幅值会无界增长——这正是 AttnRes 要解决的著名 PreNorm 问题。
Kimi Linear 是什么？为什么用它来测试 AttnRes？Kimi Linear 是月之暗面（Moonshot AI）的 48B 参数（激活 3B）模型，训练用了 1.4T token，被用作测试平台，以展示 AttnRes 在 MMLU、Math、HumanEval 和 C-Eval 等基准上的下游收益。
What are similar projects?MSA, bdh, Engram, llama3-from-scratch, jeeves, train-llm-from-scratch, LittleBit, minGPT
Which topics is it in?attention, model-training, deep-neural-networks, arxiv, foundation-models, reasoning, attention-mechanism, paper, transformers, architecture, neural-network, academic-research
Creator
MoonshotAI
Size
1.7 MB
Stars
3.5K
Updated
CommitMar 17
Age
7mo
Commits · 90 Days