Attention Residual Connections

original paper from kimi - https://arxiv.org/pdf/2603.15031

attention with sequences was such a success that researchers are trying to find out everywhere they can replace dumb aggregations, transformers over RNNs were a banger and now Kimi has released Attention Residuals to replace your standard residual connections.

the problem is simple - residuals define how information aggregates in depth but this depth-wise aggregation is not the smartest, future layers can't selectively emphasize or suppress individual layer contributions, because in standard residuals the contributions from the layers all have fixed unit weighting. this unweighted accumulation causes hidden-state magnitudes to grow with depth diluting each layers contribution, early-layer information is buried and cannot be selectively retrieved, information lost through aggregation cannot be selectively recovered in deeper layers.

the solution came from the same place it did for RNNs - attention. for RNNs the problem was tokens couldn't selectively access previous tokens, the fix was attention over sequences. here the problem is layers can't selectively access previous layers, so the fix is attention over depth.

AttnRes replaces fixed unit weights with learned, input-dependent softmax attention weights over all previous layer outputs - letting each layer selectively retrieve exactly what it needs from any earlier layer.

the same trick that beat RNNs is now being applied one level deeper. attention is everywhere.

to fully understand this new concept i read the paper and was able to implement qwen 3.5b but use this new attention residuals connections instead of it's standard residual connection, full implementation can be found here [link]

the complete code can be found here https://github.com/sijirama/waffle-house/blob/main/Qwen3_5_From_Scratch/index.py

Full and Block attention residuals

Kimi released two versions of the attention residual connection — full attention residuals and block attention residuals. both versions replace the standard fixed-weight residual addition with learned attention weights, the difference is just in how granular that attention is.

Full attention residuals

in full attention residuals, every layer gets to attend over every single previous layer output with a learned pseudo-query vector per layer. the pseudo-query is the clever part — instead of computing the query from the current input (which would force sequential computation), each layer has a fixed learned vector that acts as its query. this independence means attention weights for any group of layers can be computed in parallel without waiting for their sequential outputs, which is a meaningful efficiency win.

the problem shows up at scale. in large scale distributed training where the model is split across hundreds of GPUs using pipeline parallelism, the dominant cost is not local computation but cross-stage communication. every layer output needs to be transmitted across GPU racks, and with full attention residuals that means transmitting all L layer outputs at every pipeline stage transition.

Block Attention Residuals

block attention residuals solve this by grouping layers into N blocks. within each block, standard fixed-weight residual additions accumulate normally into a single block summary. across blocks, full learned attention is applied over those N block summaries, instead of transmitting all layer outputs you only transmit block summaries.

empirically N=8 blocks recovers most of the gains of full attention residuals, and the paper shows block AttnRes matches the loss of a baseline trained with 1.25x more compute, as a drop-in replacement with less than 2% inference latency overhead.

The migration — standard residuals → Block AttnRes in Qwen 3.5

to fully understand attention residuals i replaced the standard residuals in Qwen 3.5 with block attention residuals. i only tested for a forward pass but this was genuinely fun to implement and the architecture changes are minimal enough that training would be a straightforward next step.

things changed in the codebase

• each layer now needs its own learnable pseudo-query, one for before the attention sublayer and one for before the MLP sublayer. these are simple 1D parameter tensors initialized to zero, which the paper specifies ensures uniform attention weights at the start of training and prevents instability.

self.attn_res_norm = RMSNorm(cfg["emb_dim"], eps=cfg.get("rms_norm_eps", 1e-6))
self.mlp_res_norm = RMSNorm(cfg["emb_dim"], eps=cfg.get("rms_norm_eps", 1e-6))

self.attn_res_proj = nn.Parameter(torch.zeros(cfg["emb_dim"]))
self.mlp_res_proj = nn.Parameter(torch.zeros(cfg["emb_dim"]))

two norms were also added. one per sublayer, to normalize the block summary keys before attention, preventing large magnitude blocks from dominating the attention weights.

• the growing list of block summaries lives in Qwen3_5Model and gets initialized fresh every forward pass. each transformer block receives the current blocks list, potentially appends to it, and returns it alongside the new hidden state so the next block inherits the full history.

  python

  def forward(self, in_idx):
    x = self.tok_emb(in_idx)
    blocks: List[Tensor] = []
    
    for block in self.trf_blocks:
        blocks, x = block(x, mask, self.cos, self.sin, blocks)
    
    x = self.final_norm(x)
    logits = self.out_head(x.to(self.cfg["dtype"]))
    return logits

• the block forward pass is where the actual change happens. instead of the standard x = x + sublayer(x), each sublayer now receives a selectively aggregated input from all previous block summaries before running.

    def forward(self, x, mask, cos, sin, blocks:List[Tensor]):
        
        partial_block = x

        h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)

        # if reaches block boundary, start new block
        if self.layer_idx % (self.block_size // 2) == 0:
            blocks.append(partial_block)
            partial_block = None

        h = self.norm1(h)

        if self.layer_type == "full_attention":
            attn_out = self.token_mixer(h, mask, cos, sin)
        else:
            attn_out = self.token_mixer(h)

        partial_block = partial_block + attn_out if partial_block is not None else attn_out

        h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)

        mlp_out = self.norm2(h)
        mlp_out = self.ff(mlp_out)

        partial_block = partial_block + mlp_out

        return blocks , partial_block

• at block boundaries (every block_size // 2 layers), the current partial_block - the accumulated sum of that block's layer outputs, gets frozen and appended to blocks as a completed block summary. within a block, standard residual additions still apply. across blocks, attention is what is used.
• this is the core of the entire implementation - copied directly from the kimi paper

def block_attn_res(blocks: list[Tensor], partial_block: Tensor, proj: nn.Parameter, norm: RMSNorm):
    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.to(K.dtype), K)
    h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
    return h

• V is all previous block summaries plus the current partial block stacked together, every information source this layer can attend over. K is V normalized for stability. the first einsum computes a dot product between the pseudo-query proj and every key, producing one attention score per block per token. softmax across the blocks dimension gives weights that sum to 1 per token. the second einsum computes the weighted sum, each block's full representation scaled by its attention weight, summed into the new hidden state.

the complete code can be found here https://github.com/sijirama/waffle-house/blob/main/Qwen3_5_From_Scratch/index.py

genuinely one of the most fun implementations i've done. replacing residuals in qwen 3.5 felt like surgery.

Why It Matters

the paper validates AttnRes across 5 model sizes and the improvement is consistent across all of them, block AttnRes matches the loss of a baseline trained with 1.25x more compute, as a drop-in replacement with less than 2% inference latency overhead.

the point is simple. standard residuals solved vanishing gradients and nobody fought the uniform weighting. AttnRes fixes it, and does so cheaply enough to be actually practical. the same technique that made transformers better than RNNs, replace fixed aggregation with learned attention now applies deeper in the architecture.