Attention Residuals (AttnRes / Block AttnRes) — Technical Report

TL;DR

Standard PreNorm residual connections add every layer’s output with fixed unit weight, so hidden-state magnitude grows as O(L) with depth and each layer’s relative contribution is progressively diluted. Attention Residuals (AttnRes) replaces that uniform sum with softmax attention over preceding layer outputs — each layer carries one learned d-dimensional pseudo-query and selectively aggregates earlier layers with input-dependent weights, formally completing for depth the linear-to-softmax transition that attention completed over sequence length. To stay practical under pipeline parallelism, Block AttnRes partitions the L layers into N blocks (≈8 in practice), runs full attention only over N block-level summaries, and ships with cross-stage caching plus a two-phase inference schedule. Scaling laws on five model sizes show Block AttnRes matches a baseline trained with 1.25× more compute; integrated into Kimi Linear (48B-total / 3B-activated, 1.4T tokens) it suppresses PreNorm dilution — bounded output magnitudes, more uniform per-layer gradient norms — and improves every downstream benchmark with <4% training and <2% inference overhead.

Key claims

• Under standard PreNorm residuals, unrolling the recurrence shows every layer receives the same uniformly-weighted sum of all prior layer outputs, hidden-state magnitudes grow as O(L), and a significant fraction of layers can be pruned with minimal loss — empirical evidence that early-layer information is buried by uniform accumulation [§1, §2.1].
• AttnRes generalizes the residual update to h_l = Σ_i α_{i→l} · v_i with softmax attention weights from a single learned pseudo-query w_l ∈ R^d per layer; standard residuals and prior recurrence-based variants (Highway, hyper-connections, mHC) are shown to all be instances of depth-wise linear attention, while AttnRes performs depth-wise softmax attention [§3, §6.2].
• Pseudo-queries must be initialized to zero so initial attention weights are uniform across source layers (recovering the equal-weight average at init), which is required to avoid training volatility [§5, Architecture Details].
• Block AttnRes partitions L layers into N blocks, applies full attention only over the N block-level summaries (plus an intra-block partial sum), reducing per-token memory and cross-stage communication from O(Ld) to O(Nd); N≈8 recovers most of Full AttnRes’s gains [§3.2, §5].
• Cross-stage caching under interleaved pipeline schedules cuts per-token block-representation communication from C(C-1)/2 · N_p · d down to roughly P(P-1)/2 · N_p · d + (V-1) · P² · N_p · d — a V× per-transition improvement that enables full overlap with 1F1B steady-state compute [§4.1, Eq. 7-8, Fig. 3].
• A two-phase inference schedule batches all S per-block pseudo-queries into a single Phase-1 inter-block attention against cached block representations, then merges sequential intra-block lookback via online softmax in Phase 2 — total residual-mechanism memory access per layer drops to (N/S + 5)d, well below mHC’s (8m+2)d + 2m² + 4m at typical m=4 [§4.2, Table 1].
• Scaling laws on five model sizes fit Baseline L = 1.891·C^{-0.057}, Block AttnRes L = 1.870·C^{-0.058}, Full AttnRes L = 1.865·C^{-0.057} — Block AttnRes matches a baseline trained with 1.25× more compute [§5.1, Fig. 4].
• Integrated into Kimi Linear (48B-total / 3B-activated MoE, KDA + MLA in 3:1 ratio, 1.4T tokens), AttnRes improves over the baseline on all evaluated downstream benchmarks, while training overhead is <4% with pipeline parallelism and inference latency overhead is <2% [§5, Abstract].
• Training-dynamics analysis on the 48B run shows AttnRes mitigates PreNorm dilution: output magnitudes stay bounded across depth and gradient norms distribute more uniformly layer-to-layer [§5, Abstract].
• Block AttnRes’s per-device prefill memory for block representations scales as N·(T/P)·d under sequence-sharded TP; for a 128K context at N=8 this drops from 15 GB to ~1.9 GB per device, and with 16K chunked prefill further to <0.3 GB [§4.2].

Method

The mechanism is a single d-vector per layer. At layer l, AttnRes computes softmax attention weights α_{i→l} = exp(w_l^T RMSNorm(k_i)) / Σ_j exp(w_l^T RMSNorm(k_j)) over keys k_i = f_i(h_i) for all prior layers i < l (plus the token embedding at i=0), then sets the layer input to h_l = Σ_i α_{i→l} · v_i. The pseudo-query w_l is a learned per-layer parameter decoupled from the forward pass, which is the key structural property that admits batching: every per-block query can be computed in parallel without waiting for the sequential layer outputs. RMSNorm inside the kernel function prevents layers with large-magnitude outputs from dominating the attention weights. Pseudo-queries are initialized to zero so that the initial state is exactly the standard residual recurrence.

Block AttnRes makes the scheme tractable under pipeline parallelism. Layers are partitioned into N blocks; within each block layer outputs sum into a single representation b_n = Σ_{j∈B_n} f_j(h_j), and cross-block attention is run only over the N block summaries (plus an intra-block partial sum for layers beyond the first in a block). The systems half of the paper is heavy: (1) cross-stage caching reuses previously received blocks across virtual stages under an interleaved pipeline schedule so only incremental blocks cross stage boundaries; (2) a two-phase inference algorithm batches all S=L/N per-block pseudo-queries into one inter-block attention (Phase 1) and then merges sequential intra-block updates via online softmax (Phase 2), so per-layer I/O collapses to (N/S + 5)d reads — substantially below mHC’s residual-stream I/O at typical m=4; (3) memory-efficient prefill shards block representations along the sequence dimension across TP devices and fuses the merge into the standard reduce-scatter / all-gather path.

Results

Scaling-law experiments sweep five sizes and fit power laws L = A·C^{-α}: Baseline 1.891·C^{-0.057}, Block AttnRes (N≈8) 1.870·C^{-0.058}, Full AttnRes 1.865·C^{-0.057} [§5.1, Fig. 4]. All three slopes are essentially identical; AttnRes shifts the intercept down, and Block AttnRes matches a baseline trained with 1.25× more compute. The flagship integration is into Kimi Linear (48B-total / 3B-activated MoE, KDA+MLA interleaved 3:1, 1.4T tokens). AttnRes improves every downstream task evaluated and the analysis figures show two specific PreNorm-dilution signatures fixed: per-layer output magnitude is bounded across depth (vs. the baseline’s monotone growth) and per-layer gradient norms become substantially more uniform [§5, Abstract]. Systems numbers are kept tight: <4% end-to-end training overhead under pipeline parallelism, negligible without it, and <2% inference latency overhead on typical workloads. The 128K-context prefill memory budget for block representations drops from 15 GB → 1.9 GB per device with sequence-sharded TP, and to <0.3 GB with 16K chunked prefill.

Why it’s interesting

AttnRes is best read as a direct sibling to mHC: Manifold-Constrained Hyper-Connections: both papers argue the residual stream itself — fixed scalar addition for ten years — is a load-bearing design dimension where naive accumulation actively hurts at scale, and both validate at ~30B+ MoE scale with explicit identity-mapping / gain-bound analysis. The conceptual move is different though: mHC keeps the linear depth-wise recurrence but constrains the mixing matrix to the Birkhoff polytope (preserving an identity-like spectrum); AttnRes abandons the linear recurrence entirely and substitutes softmax depth-attention, framing the linear-to-softmax transition as the same one that replaced RNNs over sequence length. The paper makes this duality precise (it shows that mHC, hyper-connections, and standard residuals are all instances of depth-wise linear attention). This is a clean point of intellectual conflict for Training stability at scale — two contemporaneous architectural fixes for PreNorm dilution that compose differently and have not yet been stacked in any published ablation. The systems work is also noteworthy in its own right: the cross-stage caching scheme and the two-phase / online-softmax inference schedule are the kind of infrastructure details that determine whether a “small architectural tweak” actually ships at scale, and they share the same “amortize across batched queries via online softmax” structure as FlashAttention’s tile-merge math (FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling). Finally, this is another Kimi Team release that ships a technical-report PDF on a public GitHub rather than via arXiv — a packaging pattern worth tracking against Open foundation-model releases (compare DeepSeek’s mHC and Moonshot’s earlier Kimi Linear / Kimi K2.5 drops).

See also

• mHC: Manifold-Constrained Hyper-Connections — contemporaneous Birkhoff-polytope residual-mixing constraint; AttnRes shows mHC is a depth-wise linear-attention variant and Full/Block AttnRes generalizes it to softmax.
• Training stability at scale — the concept page where AttnRes joins mHC, GatedNorm/Gated Attention, SSO, and QK-Norm as architecture-side stability primitives at large scale.
• A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training — outlier-driven rescaling unifies attention and residual sinks; AttnRes’s RMSNorm-inside-the-kernel and bounded-magnitude property is the same kind of explicit-rescale fix from a different angle.
• Controlled LLM Training on Spectral Sphere — optimizer-side stability fix; relevant because the AttnRes analysis (bounded output magnitudes, uniform gradient norms across depth) is exactly the property SSO claims for its weight+update spectral retraction.
• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling — the online-softmax / two-phase amortization in AttnRes inference shares structure with FlashAttention’s tile-merge accounting.
• Open foundation-model releases — Moonshot/Kimi Team continues the GitHub-tech-report packaging pattern; integrates into the Kimi Linear architecture (48B-total / 3B-activated) which is itself an open release.