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
related reading
- The Annotated Transformernlp.seas.harvard.edu
- transformer_attention.pdfarxiv.org
- The Annotated Transformernlp.seas.harvard.edu
- Bridging the Attention Gap: Complete Replacement Models for Complete Circuit Tracinginterp.open-moss.com
- DeltaNet Explained (Part I) | Songlin Yangsustcsonglin.github.io
- Chapter 1: Transformer Interpretability - ARENAlearn.arena.education
- Chapter 1: Transformer Interpretability - ARENAlearn.arena.education
- Linear Attention Fundamentals | Hailey Schoelkopfhaileyschoelkopf.github.io
- An Intuition for Attention | Jay Modyjaykmody.com
- ali (@waterloo_intern) on Xx.com
- Kimi Linear: An Expressive, Efficient Attention Architecturearxiv.org
- 2502.11089arxiv.org