1 episodes

Paper: Attention Residuals

Attention Residuals (Kimi Team, arXiv 2603.15031, Mar 2026) replace the fixed unit-weight residual sum with a softmax over the embedding and all earlier layer outputs, using one pseudo-query per layer and RMS-normalized keys; Block AttnRes bounds memory. The module derives the zero-query equivalence to the PreNorm residual, probes the stream-growth claim, implements a ResidualStrategy in nlpcore, and draws depth-attention maps.

Progress — loading…

Objectives

Language

  • Explain why deep layers may need selective access to earlier representations (token identity, intermediate features)

Math & data structures

  • Softmax over depth; RMSNorm scale invariance; induction proof of equivalence at initialization
  • Random-walk vs aligned growth of sums; memory O(L d) vs O(N d)

Code

  • NumPy full and block AttnRes with depth maps; ResidualStrategy 'attnres' in torch
  • Stream-growth probe with seeds; key-normalization break-it

Episodes

  1. Depth-wise attention over previous layer outputs (arXiv 2603.15031).

Exercises

ExerciseTierLadderTime
Attention Residuals from scratch: depth softmax, reduction to the standard residual, Block AttnRes browserPredict · Scratch · Harden · Compose60 min
Probe the claim: how the residual stream grows with depth, and break AttnRes's key normalization browserPredict · Break it · Probe a claim50 min
Port: AttnRes as a PyTorch ResidualStrategy, and a standard vs AttnRes vs Block AttnRes GPT ablation gpuPort · Compose · Probe a claim90 min

Labs

Library contributions

Files in nlpcore this module teaches you to write.

Sources

  1. Kimi Team, Attention Residuals. arXiv 2603.15031 (v1 16 Mar 2026). · arXiv:2603.15031
  2. Pagliardini, M., Mohtashami, A., Fleuret, F. & Jaggi, M. DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging. arXiv 2402.02622 (Feb 2024). · arXiv:2402.02622
  3. Assignment note: AttnRes is used in Kimi K3 (see Capstone: A Kimi-K3-Style Model, Small, arXiv 2607.24653).
  4. Probe the claim: how the residual stream grows with depth, and break AttnRes's key normalization and notebook (this curriculum): synthetic stream-growth probe and key-norm break.
  5. He, K., Zhang, X., Ren, S. & Sun, J. Deep Residual Learning for Image Recognition. arXiv 1512.03385 (2015). · arXiv:1512.03385