Attention Residuals: Fixing a Decade-Old Bottleneck in Deep Networks

Abstract

The humble residual connection introduced by He et al. [He et al., 2016] for image recognition and canonised by the Transformer [Vaswani et al., 2017] has remained architecturally frozen for nearly a decade. Every modern LLM from GPT-4 to LLaMA to DeepSeek uses the same equation: $$h_l = h_{l-1} + f_{l-1}(h_{l-1})$$ In March 2026, the Kimi team released Attention Residuals [Kimi Team, 2026], proposing a principled replacement: replace the fixed additive accumulation with learned, input-dependent, softmax attention over all preceding layer outputs. The analogy is elegant, the Transformer solved the sequential bottleneck of RNNs; AttnRes applies the same cure to the depth dimension. This article traces the story from the original Transformer's residual connections, through the identified pathology, to the proposed fix, its math, its implementation, and its empirical results.

1. Background: The Transformer and Its Residuals

1.1. The Transformer Architecture (2017)

The 2017 paper "Attention Is All You Need" [Vaswani et al., 2017] was a watershed moment in NLP. By dispensing entirely with recurrence and convolutions and relying solely on self-attention, Vaswani et al. achieved state-of-the-art BLEU scores on WMT 2014 English-German (28.4 BLEU) and English-French (41.8 BLEU) while being dramatically more parallelisable than RNN-based models.

The core insight was to replace the sequential hidden state of RNNs where every token $h_t$ depends on $h_{t-1}$ preventing parallelism with direct pairwise attention between all positions:

Multi-Head Attention extends this by projecting queries, keys, and values into $h$ subspaces, computing attention in each, and concatenating:

The encoder stacks $N{=}6$ such layers, and each sub-layer (attention or FFN) is wrapped with a residual connection followed by layer normalisation:

1.2. Why Residual Connections Matter?

Residual connections, due to He et al. [He et al., 2016], were introduced for image classification to allow gradients to bypass transformations. The update rule is simply:

Unrolling this recurrence yields:

The gradient with respect to any intermediate hidden state $\mathbf{h}_l$ is:

The identity term $I$ is always preserved, providing a direct gradient highway from the loss to any layer regardless of depth, the key enabler of training networks with hundreds of layers.

2. The Problem: Limitations of Fixed Residual Connections

2.1. The Three Failure Modes of PreNorm Residuals

Modern LLMs universally adopt PreNorm [Xiong et al., 2020] (normalise before the sublayer, not after), because PostNorm causes gradient vanishing at depth. But PreNorm introduces a different problem. The standard update is:

Since Norm scales inputs to bounded magnitude but residual accumulation is unbounded, the hidden state magnitude grows as $\|\mathbf{h}_l\| = O(L)$. Each layer's relative contribution therefore shrinks as depth increases. The Kimi team [Kimi Team, 2026] identify three concrete failure modes:

1. No selective access. Different layer types (attention vs. MLP) receive the same aggregated state, even though they may benefit from different weightings of earlier representations.

2. Irreversible loss. Information lost through aggregation cannot be selectively recovered by deeper layers. Empirically, a significant fraction of layers can be pruned with minimal performance loss [Gromov et al., 2025], which suggests that many layers are redundant.

3. Output growth. Deeper layers must learn ever-larger outputs to remain influential over the accumulated residual, which can destabilise training and exacerbate magnitude disparity.

2.2. The Depth–Sequence Duality

Formally, residual connections perform depth-wise linear attention: when $\phi(q, k) = q^\top k$ (linear kernel), depth-wise attention collapses to $\mathbf{h}_l = \sum_i \phi(w_l, v_i)\, v_i = \sum_i v_i$, the standard additive residual. Existing refinements like Highway Networks [Srivastava et al., 2015] and Hyper-Connections [Zhu et al., 2025] all remain within this linear (or gated-linear) attention regime. AttnRes completes the transition to softmax attention over depth.

Table 1: Taxonomy of residual variants. AttnRes is the first to combine input-dependent weights with selective access to all individual earlier layer outputs.

3. Attention Residuals: A Unified View of Time and Depth: The Proposed Fix

The limitations of fixed residual connections are reminiscent of similar bottlenecks in sequence modeling — suggesting we seek similar solutions for the depth dimension.

3.1. The Depth–Time Duality

Like RNNs over time, residual connections compress all prior information into a single state $\mathbf{h}_l$ over depth. The Transformer improved upon RNNs by replacing recurrence with attention, allowing each position to selectively access all previous positions with data-dependent weights. AttnRes proposes the same methodology for depth:

where $\alpha_{i \to l}$ are layer-specific attention weights satisfying $\sum_{i=0}^{l-1} \alpha_{i \to l} = 1$. Unlike sequence length (which can reach millions of tokens), network depth is typically modest ($L < 1000$), making $O(L^2)$ attention over depth computationally feasible.

3.2. Full Attention Residuals

The attention weights are computed via a kernel function $\phi: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}_{\geq 0}$, adopting $\phi(\mathbf{q}, \mathbf{k}) = \exp(\mathbf{q}^\top \text{RMSNorm}(\mathbf{k}))$ with softmax normalisation over depth:

For each layer $l$, the query and key–value pairs are defined as:

where $\mathbf{w}_l \in \mathbb{R}^d$ is a layer-specific learnable pseudo-query vector. The RMSNorm inside $\phi$ prevents layers with large-magnitude outputs from dominating the attention weights. The input to layer $l$ is then:

Full AttnRes requires $O(L^2 d)$ arithmetic and $O(Ld)$ memory to store layer outputs. Since depth is far smaller than sequence length, the arithmetic cost is modest. However, at scale — where activation recomputation and pipeline parallelism are employed — both memory and communication overhead grow as $O(Ld)$, motivating the Block variant below.

3.3. Block Attention Residuals

Block AttnRes partitions the $L$ layers into $N$ blocks: within each block, layer outputs are reduced to a single representation via summation; across blocks, full attention is applied over only $N$ block-level representations. This reduces both memory and communication from $O(Ld)$ to $O(Nd)$.

Intra-Block Accumulation. The $L$ layers are divided into $N$ blocks of $S = L/N$ layers each. Each block is summarised by summing its layer outputs:

Inter-Block Attention. For the $i$-th layer in block $n$, the value matrix attends over all completed block representations, with $\mathbf{b}_0 = \mathbf{h}_1$ ensuring the token embedding is always included as a source:

Efficiency. Since each layer now attends over $N$ block representations rather than $L$ individual outputs, memory reduces from $O(L)$ to $O(N)$ and computation from $O(L^2)$ to $O(N^2)$. The block count $N$ interpolates between two extremes: $N = L$ recovers Full AttnRes, while $N = 1$ reduces to standard residual connections. Empirically, $N \approx 8$ recovers most of the benefit across model scales, requiring only eight stored hidden states per token.

# ── Inter-block attention ─────────────────────────────────────────
            def block_attn_res(blocks: list[Tensor], partial_block: Tensor,
                            proj: Linear, norm: RMSNorm) -> Tensor:
                """
                Inter-block attention: attend over block reps + partial sum.
                blocks:        N tensors of shape [B, T, D] — completed block representations
                partial_block: [B, T, D]  — intra-block partial sum (b_n^i)
                """
                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.weight.squeeze(), K)
                h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
                return h


            # ── Single-layer forward pass ──────────────────────────────────────
            def forward(self, blocks: list[Tensor],
                        hidden_states: Tensor) -> tuple[list[Tensor], Tensor]:
                partial_block = hidden_states

                # Apply AttnRes before self-attention (blocks include token embedding)
                h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)

                # If block boundary reached, start a new block
                # block_size counts ATTN + MLP; each transformer layer has 2
                if self.layer_number % (self.block_size // 2) == 0:
                    blocks.append(partial_block)
                    partial_block = None

                # Self-attention sub-layer
                attn_out = self.attn(self.attn_norm(h))
                partial_block = partial_block + attn_out if partial_block is not None else attn_out

                # Apply AttnRes before MLP
                h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)

                # MLP sub-layer
                mlp_out = self.mlp(self.mlp_norm(h))
                partial_block = partial_block + mlp_out

                return blocks, partial_block

Figure 2: PyTorch-style pseudo code for Block Attention Residuals. block_attn_res computes softmax attention over block representations using a learned pseudo-query $\mathbf{w}_l$; forward is a single-layer pass that maintains partial_block ($\mathbf{b}_n^i$, intra-block residual) and blocks ($[\mathbf{b}_0, \ldots, \mathbf{b}_{n-1}]$, inter-block history).

4. Infrastructure Design

Block AttnRes introduces additional system challenges compared to standard residual connections. For large-scale model training, block representations must be propagated across pipeline stages. Two key challenges arise: repeated access to accumulated block representations increases inference latency, and long-context prefilling amplifies the memory cost of caching block representations. These are addressed via cross-stage caching in training and a two-phase computation strategy in inference.

4.1. Training

For small-scale training, AttnRes adds negligible overhead — activations are already saved for backpropagation. Under large-scale distributed training, pipeline parallelism is the primary challenge. Consider an interleaved pipeline schedule with $P$ physical stages and $V$ virtual stages per physical stage, where each physical stage produces on average $N_p$ block representations of dimension $d$ per token. With $C = PV$ total chunks, naïvely transmitting all accumulated blocks at every transition incurs:

Cross-stage caching eliminates this redundancy by caching blocks locally across virtual stages. Each transition conveys only the ${\sim}PN_p$ incremental blocks accumulated since the previous virtual stage, reducing total communication to:

This reduces peak per-transition cost from $O(C)$ to $O(P)$ — a $V\times$ improvement — enabling full overlap with computation during steady-state 1F1B. The measured end-to-end training overhead under pipeline parallelism is less than 4%.

4.2. Inference: Two-Phase Computation

During inference, block representations serve as a shared KV cache reused across layers. A naïve implementation requires $O(L \cdot N)$ memory accesses. Since pseudo-query vectors $\mathbf{w}_l$ are decoupled from the forward computation, all $S = L/N$ queries within a block can be batched into a single matrix multiplication — amortising memory access from $S$ reads to 1. The strategy proceeds in two phases:

• Phase 1 — Parallel inter-block attention: Batch all $S$ queries against the cached block representations simultaneously, returning outputs and softmax statistics (max and log-sum-exp). Reduces reads from $S$ times to just once per block.

• Phase 2 — Sequential intra-block attention: Compute attention sequentially using the evolving partial sum $\mathbf{b}_n^{i-1}$, then merge with Phase 1 outputs via online softmax. Admits kernel fusion with surrounding operations, further reducing I/O overhead.

4.3. Memory Access Cost Comparison

With the two-phase design, total per-layer memory access cost remains only $\left(\frac{N}{S} + 3\right)d$ reads and $2d$ writes — substantially lower than prior residual generalisations. The end-to-end inference latency overhead is less than 2% on typical workloads.

Table 1: Memory access cost per token per layer. Typical values: $L{=}128$, $N{=}8$, $S{=}L/N{=}16$, $m{=}4$. Block AttnRes achieves the lowest I/O cost of any method, at just 5.5d — only marginally above the 3d baseline of standard residuals.

5. Experiments

5.1. Architecture Details

The model backbone is Kimi Linear [Zhang et al., 2025], a Mixture-of-Experts (MoE) Transformer following the DeepSeek-V3 [DeepSeek-AI, 2025] design. It interleaves Kimi Delta Attention (KDA) and Multi-Head Latent Attention (MLA) layers in a 3:1 ratio, each followed by a MoE feed-forward layer. The only modification is the addition of AttnRes to the residual connections — all other components remain unchanged. AttnRes introduces only one RMSNorm and one pseudo-query vector $\mathbf{w}_l \in \mathbb{R}^d$ per layer — a negligible fraction of total parameter count. Crucially, all pseudo-query vectors are initialised to zero, ensuring uniform attention weights at the start of training.

5.2. Scaling Laws

Five model sizes are swept (194M to 528M activated parameters), each trained with three variants: a PreNorm baseline, Full AttnRes, and Block AttnRes ($N \approx 8$). Following standard practice, power-law curves are fitted as:

where $\mathcal{L}$ is validation loss and $C$ is compute in PFLOP/s-days. Results:

• Baseline: $\mathcal{L} = 1.891 \times C^{-0.057}$

• Block AttnRes: $\mathcal{L} = 1.870 \times C^{-0.058}$

• Full AttnRes: $\mathcal{L} = 1.865 \times C^{-0.057}$

All variants exhibit a similar scaling slope $\alpha$, but AttnRes consistently achieves a lower intercept $A$ — meaning it is uniformly better at every compute budget. The key headline: Block AttnRes matches the loss of the baseline trained with $1.25\times$ more compute. The gap between Full and Block AttnRes narrows with scale, shrinking to just 0.001 at the largest model size.

Table 2: Baseline vs Block AttnRes ($N=8$) vs Full AttnRes vs mHC(-lite). Bold denotes best per-row result. $L_b = L/2$ denotes the number of Transformer blocks. All models trained with context length 8192.

5.3. Main Results

The largest model is the full Kimi Linear 48B configuration: 27 Transformer blocks (54 layers) with 8 out of 256 routed experts plus 1 shared expert, yielding 48B total / 3B activated parameters. Block AttnRes applies with 6 layers per block, producing 9 blocks plus the token embedding — 10 depth-wise sources in total.

Training recipe:

• Pre-trained on 1.4T tokens following the Kimi Linear recipe

• Optimiser: Muon [Liu et al., 2025] with WSD learning rate schedule

• Global batch size: 8M tokens; context window: 4096 tokens

• Stage 1: WSD pre-training on 1T tokens → Stage 2: mid-training on ~400B high-quality tokens → Stage 3: context extension to 32K tokens

5.4. Training Dynamics

Three diagnostic signals reveal why AttnRes works:

• Validation loss: AttnRes achieves consistently lower loss throughout training, with the gap widening during the decay phase and resulting in a notably lower final loss.

• Output magnitude: The baseline suffers from the PreNorm dilution problem — hidden-state magnitudes grow monotonically with depth, compelling deeper layers to learn increasingly large outputs. Block AttnRes confines this growth within each block, as selective aggregation at block boundaries resets the accumulation, yielding a bounded periodic pattern.

• Gradient magnitude: The baseline's fixed unit weights provide no mechanism for regulating gradient flow, causing disproportionately large gradients in the earliest layers. AttnRes's learnable softmax weights introduce competition for probability mass, resulting in a substantially more uniform gradient distribution across depth.

5.5. Downstream Benchmark Performance

Block AttnRes matches or outperforms the baseline on every evaluated benchmark. Gains are most pronounced on multi-step reasoning and code generation tasks, consistent with the hypothesis that improved depth-wise information flow benefits compositional tasks.

Table 3: Performance comparison of AttnRes with the baseline after the same pre-training recipe. AttnRes wins on every benchmark. The largest gains are on multi-step reasoning (GPQA-Diamond: +7.5) and math/code tasks.

5.6. Ablation Study

Ablations on the 16-layer model validate key design choices. All models share identical hyperparameters and compute budget.

Table 4: Ablation study on 16-layer model. Key takeaways below.

6. Discussions

6.1. Sequence–Depth Duality

Residual connections propagate information over depth exactly as RNNs propagate over time. Test-Time Training (TTT) formalises the sequence side: each recurrent step is gradient descent on a self-supervised loss:

When $f$ is linear, this reduces to vanilla linear attention $\mathbf{S}_t = \mathbf{S}_{t-1} + \mathbf{k}_t\mathbf{v}_t^\top$. The standard residual exhibits the same additive form along depth. Methods like Highway Networks, mHC, and DDL all refine the recurrent update while remaining within the recurrence paradigm. AttnRes goes a step further — replacing depth-wise recurrence with direct cross-layer attention, just as Transformers replaced temporal recurrence with self-attention.

6.2. Residual Connections as Structured Matrices

Every residual variant can be expressed as $\mathbf{h}_l = \sum_{i=0}^{l-1} M_{i \to l}\,\mathbf{v}_i$ for a lower-triangular depth mixing matrix $\mathbf{M} \in \mathbb{R}^{L \times L}$. The key insight is that existing residual variants correspond to depth-wise linear attention, while AttnRes performs depth-wise softmax attention:

• Standard residual: $\mathbf{M}$ is all-ones lower-triangular — rank-1 semiseparable.

• Highway: $\mathbf{M}$ is 1-semiseparable with input-dependent gates — still linear attention.

• mHC: $\mathbf{M}$ is $m$-semiseparable — linear attention with matrix-valued states.

• Full AttnRes: $\mathbf{M}$ is dense, rank-$L$ — depth-wise softmax attention, the most expressive.

The structured-matrix perspective also reveals depth-wise attention sinks in AttnRes — certain layers consistently attract high weight regardless of input, mirroring the sequence-wise attention sinks in standard Transformers.

7. Conclusion

References

1. He, K., Zhang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. CVPR, 2016. arXiv:1512.03385
2. Vaswani, A., et al. Attention is All You Need. NeurIPS, 2017. arXiv:1706.03762
3. Xiong, R., et al. On Layer Normalization in the Transformer Architecture. 2020. arXiv:2002.04745
4. Kimi Team. Attention Residuals. 2026. arXiv:2603.15031
5. Gromov, A., et al. The Unreasonable Ineffectiveness of the Deeper Layers. 2025. arXiv:2403.17887
6. Srivastava, R. K., Greff, K., and Schmidhuber, J. Highway Networks. 2015. arXiv:1505.00387
7. Pagliardini, M., et al. DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging. 2024. arXiv:2402.02622
8. Xie, Z., et al. mHC: Manifold-Constrained Hyper-Connections. 2026. arXiv:2512.24880
9. Liu, J., et al. Muon is Scalable for LLM Training. 2025. arXiv:2502.16982
10. Zhang, Y., et al. Kimi Linear: An Expressive, Efficient Attention Architecture. 2025. arXiv:2510.26692
11. DeepSeek-AI. DeepSeek-V3 Technical Report. 2025. arXiv:2412.19437
12. Bachlechner, T., et al. ReZero is All You Need: Fast Convergence at Large Depth. 2020. arXiv:2003.04887

Cite This Article

If you found this article useful, please cite it as:

BibTeX

@article{vikram2026attnres,
            author  = {Vikram},
            title   = {Attention Residuals: Fixing a Decade-Old Bottleneck in Deep Networks},
            year    = {2026},
            month   = {March},
            url     = {https://vikrampande7.github.io/blog/AttnRes.html},
            note    = {Blog post}
            }