[Tech Report] Attention Residuals: Breaking the Unit-Weight Bottleneck in Deep LLMs

TL;DR

Modern LLMs rely on residual connections that sum layer outputs with fixed, equal weights. Attention Residuals (AttnRes), proposed by the Kimi Team, replaces this static sum with dynamic softmax attention over depth. This allows each layer to selectively "peek" back at specific earlier representations. The result? A 48B model that outperforms baselines across the board, particularly in reasoning and coding, with a 1.25x compute efficiency gain.

The Problem: The "Dilution" of Deep Networks

In a standard Transformer, the hidden state $h_l$ is simply $h_{l-1}+f_{l-1}(h_{l-1})$. If you unroll this, every layer is just an equal-weight sum of everything that came before it.

While this "gradient highway" prevented vanishing gradients, it introduced a new pathology in the era of LLMs: PreNorm Dilution. As the model gets deeper, the magnitude of the hidden state grows $O(L)$. To have any impact, deeper layers must produce increasingly massive outputs. This makes the model's depth "ineffective"—information from early layers gets buried, and many layers become redundant.

The Insight: Time-Depth Duality

The researchers observed a beautiful symmetry: Residual connections are to depth what RNNs are to sequences.

• RNNs compress history into a single state (Time).
• Residuals compress prior layers into a single state (Depth).

Just as the Transformer replaced RNN recurrence with self-attention to allow direct access to any token, AttnRes replaces residual recurrence with depth-wise attention to allow direct access to any layer.

Methodology: How AttnRes Works

Instead of $h_l=\sum v_i$ with unit weights, AttnRes uses: $h_l=\sum {\alpha }_{iol}\cdot v_i$ where $\alpha$ is computed via softmax attention using a learned pseudo-query $w_l$ for each layer.

Scalability with Block AttnRes

Attending to every previous layer (Full AttnRes) is expensive for memory and communication during training ($O(L^2)$). To solve this, the authors introduced Block AttnRes:

1. Divide $L$ layers into $N$ blocks (e.g., 8 blocks).
2. Sum outputs within a block (intra-block).
3. Perform softmax attention across the $N$ block representations (inter-block).

Figure 1: Comparison between standard residuals and the selective aggregation of Attention Residuals.

Infrastructure Magic

To make this practical at a 48B-parameter scale, the team implemented:

• Cross-stage caching: In pipeline parallelism, instead of re-sending all prior block representations, stages cache them locally, reducing communication overhead to $O(P)$ rather than $O(C)$.
• Two-phase Inference: By batching pseudo-queries, they reduced memory I/O, keeping inference latency overhead under 2%.

Experiments & Results

The team integrated AttnRes into the Kimi Linear architecture and trained a 48B MoE model (3B activated) on 1.4T tokens.

1. Scaling Law Advantage

AttnRes consistently sits on a lower loss curve than the baseline. Block AttnRes with just 8 blocks recovers almost all the gain of Full AttnRes, providing a 1.25x compute advantage (meaning a smaller AttnRes model can match a much larger baseline).

Figure 2: Scaling curves showing AttnRes performing consistently better than standard PreNorm across compute budgets.

2. Downstream Performance

The gains were most visible in "heavy reasoning" tasks:

• GPQA-Diamond: +7.5 points
• HumanEval (Coding): +3.1 points
• Minerva Math: +3.6 points

3. Training Dynamics

AttnRes effectively "flattens" the magnitude growth. While baseline hidden states explode with depth, Block AttnRes keeps magnitudes bounded and periodic, leading to much more stable and uniform gradient distributions.

Figure 3: Learned attention patterns. Note the "diagonal dominance" (locality) but also the persistent "attention sinks" on the initial embedding layer.

A Unified View of the Residual Landscape

The paper provides a rigorous "Structured Matrix" analysis, reframing several prior works as specific instances of depth-wise mixing:

• Standard Residuals: Depth-wise all-ones matrix.
• Highway Networks: 1-semiseparable rank (dynamic gates).
• (m)HC: m-semiseparable rank (multi-stream).
• AttnRes: Dense, rank-$L$ (or rank-$N$) matrix.

Conclusion

AttnRes represents a fundamental shift in how we think about model depth. By acknowledging that not all layers need to contribute equally to the next, and providing a content-aware mechanism to select them, the Kimi Team has unlocked a more efficient way to build deep Transformers. As hardware interconnects improve, the transition from Block to Full AttnRes may become the new standard for ultra-deep architectures.

Takeaway: Depth is just another dimension to attend to. If your model is deep, don't just sum—attend.