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
What is this?AttnRes (Attention Residuals) is a drop-in replacement for standard residual connections in Transformers, letting each layer use learned, input-dependent attention over all previous layer outputs instead of simple additive accumulation.
Is this a trained model I can download, or a technique I implement myself?It is an architectural technique described with equations and 'PyTorch-style pseudocode', not a pretrained checkpoint; you integrate it into your own Transformer's residual connections.
How do I actually use this in my own model?The README gives no install command or runnable code snippet, only the formulas and a reference to 'PyTorch-style pseudocode'; you implement Full AttnRes or the cheaper Block AttnRes variant inside your Transformer following the paper's description.
Do I need a specific deep learning framework?Yes, the reference implementation is described in PyTorch terms, so you'd need a PyTorch-based Transformer codebase to apply it.
Does Full AttnRes scale to large models without extra cost?No — Full AttnRes needs O(Ld) memory as depth grows, which gets impractical at scale; Block AttnRes fixes this by grouping layers into about 8 blocks and attending only over block-level representations, trading a little performance for O(Nd) memory.
Who is this useful for?It's aimed at researchers and engineers building or training Transformer-based LLMs who want an alternative to standard residual connections, not end users of a finished application.
How much does AttnRes improve performance over standard residuals?It consistently lowers loss at every compute budget in scaling-law tests, with Block AttnRes matching a baseline trained with 1.25x more compute; on a Kimi Linear 48B model it adds up to +7.5 points on GPQA-Diamond and +3.1 on HumanEval.
Why do standard residual connections become a problem in deep Transformers?Standard residuals sum all layer outputs with fixed unit weights, which dilutes each layer's contribution as depth increases and makes hidden-state magnitudes grow unboundedly — a known PreNorm issue AttnRes is built to fix.
What is Kimi Linear, and why is it used to test AttnRes?Kimi Linear is Moonshot AI's 48B-parameter (3B activated) model trained on 1.4T tokens, used as the testbed to show AttnRes's downstream gains across benchmarks like MMLU, Math, HumanEval, and 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