arXiv. arXiv:2603.15031 (Submitted 16 March 2026, v1)
Abstract
Residual connections with PreNorm are standard in modern LLMs, yet they accumulate all layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state growth with depth, progressively diluting each layer's contribution. We propose Attention Residuals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding layer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-dependent weights. To address the memory and communication overhead of attending over all preceding layer outputs for large-scale model training, we introduce Block AttnRes, which partitions layers into blocks and attends over block-level representations, reducing the memory footprint while preserving most of the gains of full AttnRes. Combined with cache-based pipeline communication and a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for standard residual connections with minimal overhead. Scaling law experiments confirm that the improvement is consistent across model sizes, and ablations validate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into the Kimi Linear architecture (48B total / 3B activated parameters) and pre-train on 1.4T tokens, where AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient distribution across depth, and improves downstream performance across all evaluated tasks.
Summary
The paper treats residual connections as a depth-wise mixing rule, not only as a gradient path. Under PreNorm, every layer output is added with coefficient one. Hidden-state magnitude can grow with depth, so a new layer contributes a smaller fraction of the residual stream. AttnRes replaces that unit sum with softmax weights over earlier layer outputs. Block AttnRes keeps the same idea at coarser granularity so that pipeline training remains feasible. The reported evidence is a scaling-law comparison plus a 48B / 3B Kimi Linear run on 1.4T tokens. The results support improved information flow under that recipe. They do not establish a universal gain for every dense model or post-training setting.
Detailed reading
Depth-wise mixing as a sequence problem
A PreNorm step may be written h_l = h_(l-1) + f_l(h_(l-1)). Unrolling yields the embedding plus every previous layer output. The identity path is useful for optimization. It provides no mechanism for preferring one earlier representation over another.
We observe a formal duality between depth-wise accumulation and the sequential recurrence in RNNs.
The comparison is structural. Recurrence compresses history into a single state along time; residuals compress history into a single state along depth. Sequence modeling replaced recurrence with attention. The paper applies the same replacement to depth: h_l = Σ α_(i→l) · v_i, with softmax weights from a learned per-layer pseudo-query w_l. Sources are RMS-normalized before scoring so that a large residual magnitude does not dominate by scale alone.
Indeed, standard residual connections and prior recurrence-based variants can all be shown to perform depth-wise linear attention; AttnRes generalizes them to depth-wise softmax attention, completing for depth the same linear-to-softmax transition that proved transformative over sequences.
Full AttnRes is therefore not a second copy of token self-attention. It is a routing mechanism over the network’s own layer outputs. Softmax introduces competition: increasing one source’s weight decreases the others. Ordinary residuals never do this.
Why Block AttnRes is required at scale
In ordinary training, layer outputs are already retained for backpropagation, so Full AttnRes adds limited extra memory. Under activation recomputation and pipeline parallelism, those tensors must be stored and shipped. Memory and communication then grow as O(L d). Block AttnRes partitions layers into blocks, reduces each block with standard residuals, and attends only over the block-level summaries.
This brings both memory and communication down to O(N d), and together with infrastructure optimizations (§4), Block AttnRes serves as a drop-in replacement for standard residual connections with marginal training cost and negligible inference latency overhead.
Intra-block mixing remains additive. Inter-block mixing uses attention over block summaries. Fine-grained access to every layer is exchanged for a cheaper communication graph. The paper reports that, with cache-based pipeline communication and a two-phase inference scheme, typical inference latency overhead is less than 2%.
Experimental evidence
Scaling law experiments confirm that AttnRes consistently outperforms the baseline across compute budgets, with Block AttnRes matching the loss of a baseline trained with 1.25× more compute.
The scaling study trains a PreNorm baseline, Full AttnRes, and Block AttnRes at five activated-parameter sizes with an 8,192-token context. At 5.6 PFLOP/s-days, Block AttnRes reaches validation loss 1.692 versus 1.714 for the fitted baseline, which the paper interprets as a 1.25× compute advantage.
Analysis of the resulting training dynamics reveals that AttnRes mitigates PreNorm dilution, with output magnitudes remaining bounded across depth and gradient norms distributing more uniformly across layers. On downstream benchmarks, our final model improves over the baseline across all evaluated tasks.
The large run uses Kimi Linear: 48B total parameters, 3B activated, 27 Transformer blocks, and 1.4T pretraining and mid-training tokens. Reported dynamics show periodic, bounded output magnitudes rather than monotonic growth, and more even gradient norms across depth. Downstream gains are described as consistent, with larger lifts on GPQA-Diamond, mathematics, and code.
Evidence boundary
These results support improved information flow under the reported architecture and recipe. They do not establish that every dense model, optimizer, or post-training setting will obtain the same gain. The work is a proposing-team preprint; independent replications are not yet available. A lower pretraining loss is not, by itself, better instruction following or lower serving cost. A large depth-wise weight indicates that a layer was used, not that the route admits a human-readable causal explanation. Block AttnRes is the form of the proposal that remains compatible with distributed training.