Interactive research explainer

Attention Residuals: Reshaping Information Aggregation Across Depth

Block Attention Residuals — Interactive Architecture Demo

Paper: Attention Residuals · Kimi Team, Moonshot AI · March 2026 · arXiv:2603.15031 · GitHub · Interactive single-page demo

KaTeX formatting requires network access; readable formula fallbacks remain available offline.

Overview 1 · The Problem 2 · Full AttnRes 3 · Block AttnRes ★ 4 · Two-Phase Compute 5 · Results Speed 0.5× 1× 2× 4×

One-Minute Overview

Core insight — time–depth duality: residual connections compress information across depth in a form analogous to how recurrent networks compress information across time. Transformers replaced fixed recurrence over sequence positions with attention; AttnRes applies an analogous transition to depth, replacing fixed residual accumulation with learned softmax attention.

Problem

PreNorm Depth Dilution

Standard residuals accumulate every layer output with a fixed coefficient of 1. Hidden-state magnitude grows with depth, early contributions become progressively diluted, and the network has no content-based mechanism for selectively retrieving earlier representations.

Approach

Attention Residuals

Each layer carries one learned pseudo-query and aggregates preceding layer outputs through softmax attention. The query is a learned parameter, while the attention weights remain input-dependent because the keys come from the token’s intermediate representations.

Scalable Form

Block AttnRes

Split L layers into N blocks. Use standard residual accumulation within each block and attention over block representations across blocks. Depth-wise storage and communication fall from O(Ld) to O(Nd).

The Problem: Depth Dilution in Standard Residual Connections

Residual connections provide an effective gradient highway, but they also prescribe how information is aggregated across depth: every layer output enters the residual stream with a fixed coefficient of 1. Unrolling the recurrence yields a uniformly weighted sum of preceding outputs. In PreNorm models, hidden-state magnitude grows with depth, progressively reducing the relative influence of any individual contribution.

h_l = h_{l-1} + f_{l-1}(h_{l-1}) \;\;\Longrightarrow\;\; h_l = h_1 + \sum_{i} f_i(h_i)

\text{teaching model:}\quad \lVert h_l \rVert \approx l+1, \qquad \text{embedding share} \approx \frac{1}{l+1}

▶ Play ⏸ Pause ⏭ Step ↺ Reset

Select Play for automatic playback or Step to advance one beat at a time. Space = play/pause, Right Arrow = step, R = reset.

Scroll horizontally to inspect the diagram.

Depth dilution in standard residual connections Eight layer cards feed a central residual trunk with a fixed coefficient of 1. Two charts show hidden-state magnitude growing with depth and the embedding share shrinking with depth, using a scalar toy model. Layers — increasing depthfixed weight × 1Residual trunk h Layer 1Attention× 1 Layer 2MLP/MoE× 1 Layer 3Attention× 1 Layer 4MLP/MoE× 1 Layer 5Attention× 1 Layer 6MLP/MoE× 1 Layer 7Attention× 1 Layer 8MLP/MoE× 1 Hidden-state magnitude12345678 Embedding share12345678 Layer index 23456789 1/21/31/41/51/61/71/81/9 Attention layerMLP/MoE layerresidual contribution (fixed × 1)magnitude (toy model)embedding share (toy model)

Teaching schematic: the bars use a scalar, equal-magnitude toy model. The exact values ‖h‖ ≈ l+1 and 1/(l+1) are not measurements reported by the paper.

Consequences: fixed unit weights cannot perform selective retrieval; residual magnitude grows with depth; early signals are progressively diluted. Prior work has observed that a non-trivial fraction of layers can be pruned with limited loss.

Core Idea: Full Attention Residuals

Replace fixed uniform residual weights with learned, content-dependent softmax attention. Each layer owns one learned pseudo-query w_l whose parameters are independent of the input. The resulting attention weights are nevertheless input-dependent because the keys are computed from each token’s actual intermediate representations. RMSNorm prevents representation scale from monopolizing the weights.

How the weights are computed

\alpha_{i \to l} = \mathrm{softmax}\!\big(w_l \cdot \mathrm{RMSNorm}(k_i)\big), \qquad h_l = \sum_{i} \alpha_{i \to l}\, v_i

v_0 = h_1 \;\;\text{(token embedding)}, \qquad v_i = f_i(h_i) \;\;\text{(layer outputs)}

One learned pseudo-query w_l per layer, initialized to zero.

▶ Play ⏸ Pause ⏭ Step ↺ Reset

Select Play for automatic playback or Step to advance one beat at a time. Space = play/pause, Right Arrow = step, R = reset.

Scroll horizontally to inspect the diagram.

Full Attention Residuals teaching heatmap A lower-triangular 8 by 8 matrix of illustrative attention weights. Rows are layers 1 to 8, columns are sources h1 and f1 to f7, and each row sums to 1. Qualitative patterns reported in Figure 8 appear to the right. Lower-triangular teaching matrix — one row per layer h₁ f₁ f₂ f₃ f₄ f₅ f₆ f₇ Layer 1 Layer 2 Layer 3 Layer 4 Layer 5 Layer 6 Layer 7 Layer 8 1.00 0.300.70 0.200.650.15 0.250.150.450.15 0.100.100.150.600.05 0.100.080.120.200.450.05 0.180.070.100.120.180.300.05 0.120.060.080.100.140.200.250.05 attention weight01 Qualitative patterns reported in Figure 8① Diagonal dominance: each layer favors itsimmediate predecessor; locality remains the mainpathway.② The token embedding keeps a persistent non-trivialweight.③ Some deep layers retrieve much earlier sources,creating learned skip connections.④ Pre-Attn and Pre-MLP layers show differentmeasured specialization patterns. Matrix assembled — lower triangular; each row sums to 1.

Legend: plum cell intensity is proportional to the attention weight α; every row sums to 1. Empty dashed cells are future sources, which are never available to a layer.

Schematic teaching weights — not the numerical measurements from Figure 8.

The paper’s Figure 8 contains four measured heatmaps for a 16-head model with 16 attention layers and 16 MLP layers: Full and Block AttnRes × Pre-Attn and Pre-MLP, averaged over tokens. This simplified 8-layer matrix illustrates only the reported qualitative patterns.

Block Attention Residuals: The Practical Scalable Form

Full AttnRes must store and transmit all L preceding layer outputs. Block AttnRes divides depth into N blocks. Within a block, ordinary residual connections accumulate into a partial sum. At block completion, the partial sum becomes one cached block representation. Later layers attend over all completed block representations, plus the current partial sum when applicable.

S = \frac{L}{N} \qquad \text{teaching example: } L = 8 \text{ layers}, \; N = 4 \text{ blocks}, \; S = 2 \text{ layers per block}

\text{depth-wise state: } O(Ld) \;\longrightarrow\; O(Nd)

▶ Play ⏸ Pause ⏭ Step ↺ Reset

Select Play for automatic playback or Step to advance one beat at a time. Space = play/pause, Right Arrow = step, R = reset.

Scroll horizontally to inspect the diagram.

Block Attention Residuals teaching diagram Eight layers in four blocks of two. Within a block, outputs accumulate into a partial sum; at block completion the sum is sealed into a block-representation cache. Curved arrows show illustrative attention reads from cached block representations into later layers, with thickness proportional to alpha. Panels compare O(Ld) with O(Nd) storage and show pipeline-stage communication. token embedding h₁ = b₀ the same representation block 1 · S=2 layersL1 · Attn✓L2 · MLP✓partial sum bₙⁱ block 2 · S=2 layersstandard residualL3 · Attn✓L4 · MLP✓partial sum bₙⁱ block 3 · S=2 layersL5 · Attn✓L6 · MLP✓partial sum bₙⁱ block 4 · S=2 layersL7 · Attn✓L8 · MLP✓partial sum bₙⁱ output block-representation cache ·O(Nd)cached block representationsb₀ · token embeddingb₁ · block 1 sumb₂ · block 2 sumb₃ · block 3 sumb₄ · block 4 sum Depth-wise storage per tokenFull AttnRes layer outputsL1L2L3L4L5L6L7L8keeps all L outputs — O(Ld)Block AttnRes block representationsb₁b₂b₃b₄stores N blocks (+ b₀) — O(Nd)Teaching config: Full AttnRes keeps all 8layer outputs; Block AttnRes keeps 4 blockrepresentations (+ b₀). Pipeline-stage communicationpipeline stage 1blocks 1–2pipeline stage 2blocks 3–4block representations across stagesonly deltas cross stages Attention layerMLP/MoE layertoken embeddingpartial sum bₙⁱblock representationdirect residual / output flowshared / cached flowattention readCurved arrows = attention reads; thickness ∝ α. 1.00 0.350.65 0.300.70 0.200.250.55 0.150.250.60 0.120.180.250.45 0.100.150.200.55 0.080.120.150.250.40

Curved arrows = attention reads; thickness ∝ α. Solid arrows = direct residual or output flow; dashed arrows = shared or cached flow.

Block-count clarification: the N ≈ 8 block count comes from the scaling-law experiments and is distinct from the final 48B model's nine completed blocks (54 layers, 6 per block) described in Section 5.

All α values in this section are illustrative teaching weights.

Efficient Implementation: Two-Phase Inference

Pseudo-queries are learned parameters known in advance. All S queries in a block can therefore be packed into one matrix multiplication and applied to the shared block-representation cache. Sequential layer execution is required only for the current within-block partial sum. Online softmax merges the inter-block and intra-block results exactly.

▶ Play ⏸ Pause ⏭ Step ↺ Reset

Select Play for automatic playback or Step to advance one beat at a time. Space = play/pause, Right Arrow = step, R = reset.

Scroll horizontally to inspect the diagram.

Two-phase inference for Block AttnRes Left panel: Phase 1 packs four pseudo-queries into one query matrix, reads the shared key-value cache once, and produces four outputs with maximum and log-sum-exp statistics. Right panel: Phase 2 distributes the Phase 1 outputs over a dashed bus to four sequential within-block layers, each merged with the current partial sum by an exact online-softmax merge. Phase 1: Batched Inter-Block Attention block representations= shared K,V cache b₀ b₁ b₂ b₃ shared K,V readcache read once S = 4 pseudo-queries learned parameters, known in advance wₗ wₗ₊₁ wₗ₊₂ wₗ₊₃ Q matrix [wₗ … wₗ₊₃] softmax(Q · RMSNorm(K)) out 1out 2out 3out 4S Phase 1 outputs — one per layer statistics:max and LSE Phase 2: Sequential Intra-BlockAttention + Online Softmax Merge ⊕ = exact online-softmax merge Phase 1 outputs ⊕within-block layer 1hₗ ⊕within-block layer 2hₗ₊₁ ⊕within-block layer 3hₗ₊₂ ⊕within-block layer 4hₗ₊₃ partial sum bₙⁱ— sequentiallyupdated

Paper’s typical configuration: L=128, N=8, S=16 → residual-mechanism I/O = (N/S + 5)d = 5.5d per token per layer; standard residuals = 3d; mHC with m=4 = 34d. The compact diagram uses four pseudo-queries and four cached representations only as a teaching configuration; it is not the configuration used to derive 5.5d.

End-to-end inference-latency overhead is below 2% on typical workloads.

Results: Bounded Magnitudes, More Uniform Gradients, Consistent Gains

Two qualitative trends summarize the mechanism’s effect: Block AttnRes keeps output magnitudes bounded across depth, and its gradient magnitudes are substantially more uniform across depth than the baseline’s.

▶ Play ⏸ Pause ⏭ Step ↺ Reset

Select Play for automatic playback or Step to advance one beat at a time. Space = play/pause, Right Arrow = step, R = reset.

Scroll horizontally to inspect the diagram.

Qualitative magnitude and gradient trends Two handcrafted qualitative charts over 27 Transformer-block positions. Left: baseline output magnitude grows monotonically while Block AttnRes stays bounded and periodic. Right: baseline gradient magnitude concentrates in shallow layers while Block AttnRes is substantially more uniform across depth. Output magnitude per Transformer block Gradient magnitude per Transformer block 127Transformer block index (1–27) 127Transformer block index (1–27) Baseline: monotonic growth with depth Block AttnRes: attention resets accumulation atblock boundaries, producing a bounded periodicpattern Baseline: gradients concentrate in the shallowestlayers Block AttnRes: substantially more uniform acrossdepth

Qualitative schematic based on the trends reported in Figure 5; not the paper’s original numerical series.

1.25×

Baseline training compute required at equal validation loss

< 2%

End-to-end inference-latency overhead on typical workloads

N ≈ 8

Approximate block count recovering most Full AttnRes gains in scaling-law experiments

+7.5

GPQA-Diamond improvement; HumanEval +3.1 and MMLU +1.1

Benchmark	Baseline	AttnRes	Change
MMLU	73.5	74.6	+1.1
MMLU-Pro	52.2	52.2	0.0 (tie)
GPQA-Diamond	36.9	44.4	+7.5
BBH	76.3	78.0	+1.7
ARC-Challenge	64.6	65.7	+1.1
HellaSwag	83.2	83.4	+0.2
TriviaQA	69.9	71.8	+1.9
GSM8K	81.7	82.4	+0.7
MGSM	64.9	66.1	+1.2
Math	53.5	57.1	+3.6
CMath	84.7	85.1	+0.4
HumanEval	59.1	62.2	+3.1
MBPP	72.0	73.9	+1.9
CMMLU	82.0	82.9	+0.9
C-Eval	79.6	82.5	+2.9

Benchmark comparison reported for the final model (all 15 benchmarks)

All 15 reported benchmarks matched or improved over the baseline; MMLU-Pro was tied.

The final Kimi Linear configuration used 48B total and 3B activated parameters, 27 Transformer blocks or 54 layers, six layers per AttnRes block, nine completed blocks plus the token embedding, ten depth-wise sources in total, and approximately 1.4T training tokens.

Based on Attention Residuals by the Kimi Team, Moonshot AI, arXiv:2603.15031. Displayed teaching weights and explicitly labeled schematic curves are educational illustrations rather than original paper measurements.

arXiv paper · Official repository

Interactive single-page demo · Keyboard: Space = play/pause, Right Arrow = step, R = reset.