AttnRes

Lets understand Attention Residuals from the Kimi Team by implementing the Full Attention Residuals and Block Attention Residuals from the paper. (Forward pass) AttnRes makes each layer get a single learned query vector, scores every preceding layer's output through a softmax, and constructs its own input as a weighted combination. We build the AttnRes modules from scratch, run a forward pass with dummy data, and visualize the depth-wise attention weight matrices as heatmaps (as shown in the paper). AttnRes applies RMSNorm to the keys (layer outputs) before computing attention scores. This ensures selection is based on the direction of each output, not its magnitude. In Full AttnRes, each layer 𝑙 l attends over all preceding layer outputs 𝑣 0 , 𝑣 1 , … , 𝑣 𝑙 − 1 v 0 ​ ,v 1 ​ ,…,v l−1 ​ using a learned pseudo-query 𝑤 𝑙 ∈ 𝑅 𝑑 w l ​ ∈R d The pseudo-queries are initialized to zero so that initial weights are uniform (the paper emphasizes this is important for training

Lets understand Attention Residuals from the Kimi Team by implementing the Full Attention Residuals and Block Attention Residuals from the paper. (Forward pass) AttnRes makes each layer get a single learned query vector, scores every preceding layer's output through a softmax, and constructs its own input as a weighted combination. We build the AttnRes modules from scratch, run a forward pass with dummy data, and visualize the depth-wise attention weight matrices as heatmaps (as shown in the paper). RMSNorm AttnRes applies RMSNorm to the keys (layer outputs) before computing attention…

saved by

• Rajan Agarwal
• Alisha Arora

related reading

• Attention-Residuals/Attention_Residuals.pdf at master · MoonshotAI/Attention-Residuals · GitHubgithub.com
• A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
• The Illustrated Transformer – Jay Alammar – Visualizing machine learning one concept at a time.jalammar.github.io
• An Intuition for Attention | Jay Modyjaykmody.com
• 11.7. The Transformer Architecture — Dive into Deep Learning 1.0.3 documentationd2l.ai
• The Annotated Transformernlp.seas.harvard.edu
• Chapter 1: Transformer Interpretability - ARENAlearn.arena.education
• The Illustrated Transformer – Jay Alammar – Visualizing machine learning one concept at a time.jalammar.github.io
• The Annotated Transformernlp.seas.harvard.edu
• Inside the Transformer: The Life of a Token - Aleksa Gordićaleksagordic.com
• The Big LLM Architecture Comparisonmagazine.sebastianraschka.com
• arxiv.org/pdf/2512.24880#page=3.56arxiv.org