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

saved by

• Rajan Agarwal
• 6829