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) — это встраиваемая замена стандартных residual-связей в Transformers: каждый слой вместо простого накопления суммой использует обучаемое, зависящее от входа внимание по выходам всех предыдущих слоёв.
Это обученная модель, которую можно скачать, или техника, которую я реализую сам?Это архитектурная техника, описанная формулами и «псевдокодом в стиле PyTorch», а не предобученный чекпоинт; вы встраиваете её в residual-связи своего Transformer.
Как применить это в собственной модели?В README нет команды установки или запускаемого фрагмента кода — только формулы и отсылка к «псевдокоду в стиле PyTorch»; вы реализуете Full AttnRes или более дешёвый вариант Block AttnRes внутри своего Transformer по описанию из статьи.
Нужен ли какой-то конкретный фреймворк глубокого обучения?Да, эталонная реализация описана в терминах PyTorch, поэтому для её применения понадобится кодовая база Transformer на PyTorch.
Масштабируется ли Full AttnRes на большие модели без дополнительных издержек?Нет — Full AttnRes требует O(Ld) памяти по мере роста глубины, что при больших масштабах становится непрактичным; Block AttnRes решает это, группируя слои примерно в 8 блоков и выполняя внимание только по представлениям уровня блоков, немного жертвуя качеством ради памяти O(Nd).
Кому это полезно?Она рассчитана на исследователей и инженеров, которые строят или обучают LLM на базе Transformer и хотят альтернативу стандартным residual-связям, а не на конечных пользователей готового приложения.
Насколько AttnRes улучшает показатели по сравнению со стандартными residual-связями?В тестах законов масштабирования он стабильно снижает loss при любом бюджете вычислений, причём Block AttnRes догоняет базовую модель, обученную с 1.25-кратным объёмом вычислений; на модели Kimi Linear 48B он даёт до +7.5 балла на GPQA-Diamond и +3.1 на HumanEval.
Почему стандартные residual-связи становятся проблемой в глубоких Transformers?Стандартные residual-связи суммируют выходы всех слоёв с фиксированным единичным весом, из-за чего с ростом глубины вклад каждого слоя размывается, а величины скрытых состояний растут неограниченно — это известная проблема PreNorm, которую призван исправить AttnRes.
Что такое Kimi Linear и почему на ней тестируют AttnRes?Kimi Linear — модель Moonshot AI с 48 млрд параметров (3 млрд активных), обученная на 1.4 трлн токенов; она служит испытательным стендом, чтобы показать выигрыши 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