In partnership with

This is not a design choice, it is a default that nobody questioned because it works well enough. Moonshot AI's Kimi Team questioned it. The result, Attention Residuals (AttnRes, arXiv:2603.15031, MoonshotAI/Attention-Residuals, 3.3k stars), replaces that fixed accumulation with softmax attention over preceding layer outputs. Integrated into Kimi Linear (48B total / 3B activated) and pre-trained on 1.4T tokens, Block AttnRes achieves the same loss as a baseline trained with 1.25x more compute, and improves GPQA-Diamond by +7.5 points.

SnackOnAI Engineering | Senior AI Systems Researcher | Technical Deep Dive | July 27, 2026

What It Actually Does

AttnRes (Kimi Team, Moonshot AI, March 2026) replaces the fixed-weight additive residual connection with learned, input-dependent softmax attention over all preceding layer outputs. The standard residual formula is:

h_l = h_{l-1} + f(h_{l-1})   ← standard: fixed weight 1 on everything

AttnRes replaces it with:

h_l = Σ_{i=0}^{l-1} α_{i→l} · v_i   ← AttnRes: learned weights, input-dependent

where α_{i→l} are softmax weights computed from a single learned pseudo-query w_l ∈ R^d per layer. Each layer now decides, based on the current input, which earlier representations to aggregate and how much to weight them.

The practical variant, Block AttnRes, partitions layers into N blocks (default ~8 for large models), applies standard residuals within each block, and applies cross-block attention only at block boundaries. Memory footprint drops from O(Ld) (full AttnRes, quadratic in depth) to O(Nd) (block, linear in block count), while preserving most of the gains of full AttnRes.

Scope covered: the PreNorm dilution problem, the full AttnRes and Block AttnRes formulations, the depth-time duality insight, comparison to DenseFormer and mHC, scaling law results, downstream benchmarks on Kimi Linear, and the PyTorch implementation. Excluded: the Kimi Linear architecture details beyond AttnRes integration, and comparison to PostNorm variants.

The Architecture, Unpacked

Why Standard Residuals Fail at Depth

Standard residual connections with PreNorm have a known failure mode that gets worse as models get deeper. The math is straightforward: if you add every layer output with weight 1, the hidden state at layer L is a sum of L terms. As L grows, the magnitude grows as O(L), each individual layer's contribution becomes a fraction 1/L of the total, and the network progressively loses the ability to use early-layer representations effectively. The paper calls this PreNorm dilution.

The counterintuitive part: PreNorm was introduced precisely to fix gradient issues in deep networks. It does fix gradients. But it introduces magnitude growth as a side effect, and standard additive residuals do not compensate for it.

Caption: The critical distinction is that AttnRes weights sum to 1.0 by softmax construction, bounding hidden state growth. Block AttnRes reduces the number of attention sources from L (all layers) to N+1 (block reps plus partial sum), making the method practical at scale.

The Depth-Time Duality

The paper's core insight is an elegant structural analogy. Standard transformers replaced fixed recurrence over sequence positions with attention: instead of a fixed-weight LSTM update h_t = Wh_{t-1} + ..., self-attention learns dynamic weights over all positions. AttnRes applies the exact same transition to the depth dimension: instead of a fixed-weight residual accumulation, it learns dynamic weights over all preceding layers. The duality is exact: information dilution across depth is structurally identical to memory loss across a sequence, and can be solved the same way.

Comparison to Prior Work

Two prior methods address the same problem but less effectively:

DenseFormer (Pagliardini et al., 2023) gives each layer access to all previous outputs but uses fixed, input-independent scalar coefficients. It shows no gain over the baseline (loss 1.767 vs. baseline 1.766). The input-independence is the bottleneck: the aggregation cannot adapt to the content being processed.

mHC introduces input dependence through m parallel streams with learned mixing matrices, improving to 1.747. Better than DenseFormer, still not content-dependent in the same sense as softmax attention.

AttnRes uses explicit content-dependent selection via softmax: Full AttnRes achieves 1.737, Block AttnRes 1.746. The comparison is decisive: input-dependence plus explicit softmax normalization is what produces the improvement. Fixed coefficients (DenseFormer) and learned-but-not-softmax mixing (mHC) both underperform.

The Code, Annotated

Snippet One: The Block AttnRes Forward Pass

This is the complete implementation from the repo. It is fewer than 20 lines of actual PyTorch logic.

# Source: https://github.com/MoonshotAI/Attention-Residuals
# Design intent: apply softmax attention over block-level representations
# before each attention sublayer and each MLP sublayer.
# The key choices: (1) attend before the sublayer, not after,
# (2) use RMSNorm on keys but not values (preserves magnitude information in V),
# (3) applied at BOTH attn and MLP positions per transformer layer.

def block_attn_res(
    blocks: list[Tensor],     # N completed block representations [B, T, D] each
    partial_block: Tensor,    # current intra-block partial sum [B, T, D]
    proj: Linear,             # per-sublayer learned pseudo-query w_l ∈ R^d
    norm: RMSNorm             # RMSNorm for key computation only
) -> Tensor:
    """
    Inter-block attention: h_l = Σ α_i · v_i where α_i = softmax(w_l · RMSNorm(v_i))
    This is the depth-wise analog of self-attention over sequence positions.
    """
    V = torch.stack(blocks + [partial_block])    # [N+1, B, T, D]
    K = norm(V)                                  # normalize keys: stable dot products
    # ← THIS is the trick: a single per-sublayer scalar query w_l (d-dim vector)
    #   gives each sublayer its own learned preference for which depth to attend to.
    #   This is strictly cheaper than a full query matrix: O(d) params per sublayer
    #   instead of O(d²), yet still gives per-token, per-position depth selection.
    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

def forward(
    self,
    blocks: list[Tensor],    # accumulated block representations from previous blocks
    hidden_states: Tensor    # current intra-block running sum
) -> tuple[list[Tensor], Tensor]:
    partial_block = hidden_states

    # ← Apply AttnRes BEFORE attention (not after): the depth aggregation
    #   provides the input to the attention sublayer, not its output.
    #   This ordering means the attention sublayer can refine the depth-aggregated
    #   representation, not the other way around.
    h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)

    # At block boundary: seal the current block and start a new partial sum.
    # block_size counts ATTN + MLP submodules; each transformer layer has 2.
    if self.layer_number % (self.block_size // 2) == 0:
        blocks.append(partial_block)   # ← partial_block is the block rep
        partial_block = None           # ← reset for next block

    # Standard self-attention sublayer, operating on the depth-aggregated h
    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 too: both sublayers benefit from depth access.
    #   This doubles the number of AttnRes calls per layer (attn + MLP),
    #   but both calls share the same blocks list (constant memory).
    h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)

    mlp_out = self.mlp(self.mlp_norm(h))
    partial_block = partial_block + mlp_out

    return blocks, partial_block   # blocks grows only at block boundaries

Caption: The pseudo-query proj.weight.squeeze() is the architectural minimalism that makes AttnRes cheap: one d-dimensional vector per sublayer instead of a full query matrix. The softmax(0) over the N+1 source axis is what enforces the sum-to-1 property that bounds hidden state magnitudes.

Snippet Two: Standard Residual vs. AttnRes Side-by-Side

# BEFORE: Standard residual connection (what every LLM uses today)
def standard_residual_forward(hidden_states: Tensor, sublayer, norm) -> Tensor:
    # h_l = h_{l-1} + f(norm(h_{l-1}))
    # Weight on h_{l-1}: always 1.0, always, regardless of content.
    # At layer 100 in a 128-layer model, this single layer's contribution
    # is now 1/100 of the total hidden state. Early layers are drowned out.
    return hidden_states + sublayer(norm(hidden_states))

# AFTER: AttnRes (content-aware, bounded magnitude)
def attnres_forward(
    blocks: list[Tensor],
    partial_block: Tensor,
    sublayer, norm, attn_res_proj, attn_res_norm
) -> Tensor:
    # h_l = Σ α_i · v_i where Σ α_i = 1.0 (softmax)
    # → hidden state magnitude is a weighted average of sources, not a growing sum
    # → each source (block rep or partial sum) can receive 0 to 100% of the weight
    # → the weight depends on the current token's content (via the pseudo-query)
    # ← THIS is the key difference: the network decides which layers matter
    #   for each specific token at each specific position.
    h = block_attn_res(blocks, partial_block, attn_res_proj, attn_res_norm)
    return partial_block + sublayer(norm(h))

# Memory comparison:
# Standard residual: O(Ld) — grows with depth L
# Full AttnRes:      O(Ld) — same, attends over all L outputs (L keys)
# Block AttnRes:     O(Nd) — N blocks, not L layers; N << L at scale
#                            e.g. 54-layer model with block_size=6: N=9 blocks
#                            memory for sources: 9 vectors vs 54 vectors

Caption: The standard residual's weight of 1.0 is the problem. AttnRes replaces that with softmax weights that sum to 1.0, bounding magnitude. Block AttnRes reduces the number of sources from 54 (all layers) to 9 (block reps), trading some selectivity for memory efficiency.

It In Action: End-to-End Scaling Law Experiment

Setup: Five model sizes, three variants each (PreNorm baseline, Full AttnRes, Block AttnRes with ~8 blocks), trained with identical hyperparameters per size group. 8192-token context, cosine LR schedule.

The scaling law fits:

Baseline:      L = 1.891 × C^{-0.057}
Block AttnRes: L = 1.870 × C^{-0.058}   ← lower intercept, same exponent
Full AttnRes:  L = 1.865 × C^{-0.057}   ← lowest intercept

Interpretation:
- The exponent α is essentially the same (~0.057-0.058): AttnRes does not
  change the scaling rate, it shifts the loss curve down uniformly.
- Lower intercept A means the same compute budget reaches a lower loss level.
- Block AttnRes intercept 1.870 vs baseline 1.891: matches baseline loss at
  C_{baseline} × 1.25. In practice: run the same training, get a 25% compute discount.

Comparison at equal compute (fixed validation loss target 1.746):

Baseline:      needs compute C
Full AttnRes:  needs compute C × 0.78   (22% less compute for same loss)
Block AttnRes: needs compute C × 0.80   (20% less compute for same loss)
DenseFormer:   needs compute C × 1.00   (no improvement — fixed coefficients)
mHC:           needs compute C × 0.92   (8% savings — better than baseline)

Downstream benchmarks on Kimi Linear (48B total / 3B activated, 1.4T tokens):

Benchmark

Baseline

AttnRes

Delta

MMLU

73.5

74.6

+1.1

GPQA-Diamond

36.9

44.4

+7.5

BBH

76.3

78.0

+1.7

TriviaQA

69.9

71.8

+1.9

Math

53.5

57.1

+3.6

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

Pattern: The largest gains appear on multi-step reasoning tasks (GPQA-Diamond +7.5, Math +3.6) and code generation (HumanEval +3.1). The smallest gains appear on memorization-heavy tasks (MMLU +1.1, TriviaQA +1.9). This is consistent with the theoretical prediction: AttnRes enables richer depth-wise integration that helps tasks requiring compositional reasoning more than tasks that primarily require recall.

Training dynamics (measured, not claimed):

Standard baseline: hidden-state magnitude grows with depth (Figure 5 in paper)
AttnRes:          magnitude remains bounded across all 54 layers
Gradient norms:   more uniform across depth with AttnRes vs. pronounced spike pattern in baseline

Why This Design Works (and What It Trades Away)

The depth-time duality is the correct frame. Transformers use attention to decide which tokens matter for computing each output. AttnRes uses attention to decide which layers matter for computing each hidden state. The mechanism is identical: softmax over a set of vectors with learned queries. The axis changes from sequence length to depth. This is not a metaphor. The math is structurally the same.

The softmax normalization is load-bearing. DenseFormer proves this: fixed, input-independent cross-layer access with learned scalar coefficients gains nothing. The gain comes specifically from (1) input-dependent weighting and (2) softmax normalization that ensures weights sum to 1.0. Both conditions must hold. mHC has input dependence but no softmax normalization, and achieves smaller gains than AttnRes.

What this trades away. Two costs: parameter count and compute overhead.

Parameter count: each sublayer adds one d-dimensional pseudo-query vector (attn_res_proj) and one RMSNorm (attn_res_norm). For a 48B model with 54 layers and 2 sublayers per layer, that is 108 additional parameter pairs, each of size d (typically 4096-8192). Total overhead: negligible at model scale.

Compute overhead: each sublayer executes one additional attention operation over N+1 sources. For Block AttnRes with 8 blocks: 9 matrix-vector products plus a softmax, applied 108 times per forward pass. The paper reports this as "minimal overhead" without publishing exact percentage. The 1.25x compute advantage at the same parameter count means the overhead is well below 20% in practice, and the loss improvement more than compensates.

The deeper tradeoff: Block AttnRes requires explicit block boundary management in the training infrastructure. A block-aware training loop must track which block each layer belongs to, accumulate block representations at boundaries, and pass the block list through the forward stack. This is not a standard transformer architecture, and retrofitting it onto existing codebases requires understanding the block state management, not just swapping a function call.

Technical Moats

The pseudo-query design is the efficiency insight that competitors will miss. Full attention over depth would require a query matrix Q ∈ R^{d×d} per sublayer to compute q_l = W_Q h_l, then cross-layer key-query products. AttnRes uses a single learned vector w_l ∈ R^d as the query. This is the outer product approximation: it forces the attention to use a rank-1 "direction of interest" per sublayer rather than an arbitrary linear map. This costs fewer parameters and less compute than full attention, while still achieving input-dependent depth weighting. It is the minimum design that captures the benefit.

The two-phase computation strategy enables pipeline parallelism. Training large models across pipeline-parallel stages creates a communication challenge: Block AttnRes needs block representations from previous stages. The paper describes a cache-based pipeline communication approach where completed block representations are cached and forwarded. This is a non-trivial distributed systems design that makes Block AttnRes compatible with pipeline-parallel training at scale. Teams attempting to reproduce AttnRes at frontier scale will need to implement this cache forwarding, which requires understanding the pipeline parallelism infrastructure at a low level.

Uniform gradient distribution is the training stability moat. The paper's Figure 5 shows that standard PreNorm baselines have non-uniform gradient magnitude across depth: some layers receive much larger gradients than others, creating effective learning rate imbalance. AttnRes produces substantially more uniform gradient norms across depth. This is not just a cosmetic improvement: uniform gradient flow means more of the model's parameters are updating meaningfully per step, improving sample efficiency throughout training. At frontier scale, this training efficiency compound across billions of tokens.

Contrarian Insights

Insight One: The 1.25x compute advantage claim requires careful interpretation, and some teams will misapply it. The scaling law experiment shows Block AttnRes matches a baseline trained with 1.25x more compute at the same parameter count. This does not mean you can run AttnRes for 80% of the planned compute and get an equivalent model. It means that at the same training compute, the AttnRes model achieves a loss level that the baseline would only reach at 1.25x compute. The practical implication is a validation loss improvement for fixed compute budget, not a training compute reduction. Teams that read "1.25x compute advantage" as "we can stop training 20% earlier" will be wrong. The correct reading is "we get a better model for the same training budget."

Insight Two: The routing collapse result from independent analysis suggests AttnRes's gains may be more fragile than the paper implies. A subsequent paper (arXiv:2606.22325) analyzed an open 0.6B AttnRes model trained on Qwen3 and found that depth routing concentrates heavily: the top source receives 0.643 of the routing weight against a uniform baseline of 0.245. The concentration piles onto two hubs: the current block (recency, 0.42) and the token embeddings (source 0, 0.26). This means AttnRes is not learning rich, diverse depth retrieval in practice. It is primarily learning "attend to the most recent block and the token embedding." That is still better than fixed uniform weights, but it calls into question whether the mechanism is doing what the paper claims, or whether it is finding a simpler solution that achieves similar loss without genuine depth-wise selective retrieval.

Surprising Takeaway

AttnRes shifts the optimal architecture toward deeper, narrower models. One of the paper's ablations shows that with AttnRes, the loss-optimal architecture for a fixed parameter budget changes: deeper models (more layers) become preferable relative to wider models (larger hidden dimension). This is the opposite of the typical wisdom that very deep narrow models are harder to train due to gradient issues. AttnRes mitigates the gradient non-uniformity that makes depth problematic, so the architecture search over depth vs. width finds a different optimum. Teams designing new model architectures on AttnRes should expect to find more depth, less width compared to standard residual architectures at the same parameter count.

TL;DR For Engineers

  • AttnRes (arXiv:2603.15031, MoonshotAI/Attention-Residuals, 3.3k stars) replaces the fixed-weight additive residual connection with softmax attention over preceding layer outputs: h_l = Σ α_{i→l} · v_i where α are learned, input-dependent, and sum to 1.0 by softmax construction.

  • Block AttnRes partitions layers into ~8 blocks, applies standard residuals within blocks, and uses cross-block attention at boundaries only. Memory: O(Nd) instead of O(Ld). Applied twice per transformer layer (before attention sublayer and before MLP sublayer).

  • Scaling law results: Block AttnRes matches a baseline trained with 1.25x more compute. Baseline: L = 1.891 × C^{-0.057}. Block AttnRes: L = 1.870 × C^{-0.058}. DenseFormer (fixed cross-layer coefficients): no improvement. The softmax normalization and input-dependence are both required.

  • Kimi Linear (48B/3B activated, 1.4T tokens): GPQA-Diamond +7.5, Math +3.6, HumanEval +3.1. All 15 benchmarks matched or improved over baseline (MMLU tied). Largest gains on multi-step reasoning and code.

  • The mechanism: one d-dimensional pseudo-query per sublayer (not a full query matrix). The rank-1 approximation is why AttnRes is computationally cheap relative to full cross-layer attention.

Explain It Like I'm New

When a modern AI model processes text, it does so in layers, like a factory assembly line. Each station (layer) takes the work from the previous station, does something to it, and passes it forward. The standard way to do this adds all previous work together with equal weight: station 47's contribution is treated identically to station 1's contribution when computing station 48's output.

This seems reasonable until you think about it carefully. A deep model might have 128 layers. By the time you reach layer 128, you are adding 127 previous contributions together with equal weight 1. The total signal is enormous, and any single layer's individual contribution is now just 1/127th of the total. Early layers get progressively "drowned out" as the model gets deeper. The paper calls this PreNorm dilution.

Attention Residuals solves this by replacing the fixed equal-weight summation with a learned, input-dependent choice. Instead of always giving every previous layer weight 1, each layer asks: given what I am currently processing, which earlier representations are most relevant, and how much should each contribute? The answer is computed by a small learned query that produces soft weights summing to 1.

Think of it as the difference between a researcher who reads all their notes with equal attention (current approach) versus one who highlights the most relevant notes for each specific question they are trying to answer (AttnRes).

Because the weights always sum to 1 (via softmax), the total signal magnitude stays bounded regardless of depth, which also improves how gradients flow backward during training, making the learning process more stable and efficient at depth.

See It In Action

  • Attention Residuals Interactive Architecture Demo Source: Kimi / anonymous implementation | https://7inif7p6jcz2y.kimi.page/ A single-page interactive visualization of the AttnRes mechanism with step-by-step animation of the depth-wise attention computation. The best non-paper resource for building intuition about how block representations are accumulated and how weights concentrate in practice.

  • AttnRes GitHub Repository (MoonshotAI/Attention-Residuals) Source: Moonshot AI | https://github.com/MoonshotAI/Attention-Residuals Contains the complete PyTorch pseudocode, scaling law result figures, benchmark table, and training dynamics visualizations. The README's three-panel overview figure (standard residual vs. Full AttnRes vs. Block AttnRes) is the most efficient single visual for understanding the design.

  • Attention Residuals Paper (arXiv:2603.15031) Source: Kimi Team, Moonshot AI | https://arxiv.org/abs/2603.15031 Full mathematical treatment including the DenseFormer/mHC comparison (Table 1), the scaling law fitted curves (Figure 4), the ablation over cross-layer granularity, and the Kimi Linear integration details.

  • "Attention Residuals: Teaching transformers to choose which layers matter" (Vizuara Newsletter) Source: Vizuara | https://www.vizuaranewsletter.com/p/attention-residuals-teaching-transformers The best independent analysis, correctly identifying the depth-time duality as the core insight and providing the clearest explanation of why DenseFormer failed (input-independence) while AttnRes succeeds.

Community Conversation

  • Kimi Team (Moonshot AI) via arXiv:2603.15031 https://arxiv.org/abs/2603.15031 The paper's discussion section is unusually candid about the computational overhead question: the authors acknowledge that full AttnRes is impractical at scale and that Block AttnRes is the deployable variant, which is why all downstream results use Block AttnRes rather than Full AttnRes.

  • Meta-Attention paper (arXiv:2605.28384) on AttnRes composability https://arxiv.org/pdf/2605.28384 Section 2.5 describes AttnRes as operating on "the depth dimension" while their own method operates on "the mechanism dimension," and argues the two are orthogonal and jointly differentiable. This suggests AttnRes can compose with other attention variants without interference.

  • "All Routes Lead to Collapse" paper (arXiv:2606.22325) on routing concentration https://arxiv.org/pdf/2606.22325 Section 5.4 analyzes a 0.6B AttnRes model and finds depth routing concentrates onto two hubs (current block, weight 0.42; token embedding, weight 0.26) rather than distributing evenly across depth. This is the most important independent validation of AttnRes, and also the most important challenge to the "selective depth retrieval" narrative.

  • HuggingFace papers discussion (arXiv:2603.15031) https://huggingface.co/papers/2603.15031 3,360 collections additions since March 2026. The comment threads include researchers noting the similarity to DenseNet (dense connections in CNNs) and the connection to Highway Networks, providing useful historical context for the depth-reuse idea.

The Residual Connection Is Not a Sacred Cow

Residual connections have been in transformers since they inherited them from ResNets in 2017. In nine years, the research community has collectively treated them as fixed infrastructure. AttnRes is the first approach that achieves consistent gains across model sizes with an explicit benchmark comparison against prior alternatives, and does so while shipping a practical deployable variant (Block AttnRes) rather than just a theoretical proposal.

The results are clean: 1.25x compute advantage on scaling laws, improvements across all 15 downstream benchmarks, largest gains on tasks that require multi-step reasoning. The mechanism is simple: replace fixed weight 1 with learned softmax weights. The implementation is 20 lines of PyTorch.

The open question from independent analysis is whether those learned weights are actually doing rich selective depth retrieval, or whether they are finding a simpler solution that concentrates on recency and embeddings. The distinction matters for understanding when AttnRes will help most. The gains are real either way.

References

Attention Residuals (AttnRes, Kimi Team, Moonshot AI, arXiv:2603.15031) replaces the fixed-weight additive residual connection with learned, input-dependent softmax attention over preceding layer outputs, bounding hidden state growth and distributing gradient flow more uniformly across depth. Block AttnRes (the practical variant with ~8 blocks) achieves a 1.25x compute advantage on scaling laws and improves Kimi Linear (48B/3B activated, 1.4T tokens) by up to +7.5 points on GPQA-Diamond, with consistent gains across all 15 evaluated benchmarks. The core insight is a direct depth-time duality: the same mechanism that replaced fixed RNN recurrence with self-attention over sequence positions can replace fixed residual accumulation with learned attention over depth.

Sponsored Ad

If you enjoy practical AI insights, check out SnackOnAI and support the newsletter by subscribing, sharing, and exploring our sponsored ad — it helps us keep building and delivering value 🚀

Don’t Let Tax Season Cost You Year-Round

That pit in your stomach is trying to tell you something: Waiting until spring is costing you peace of mind.

When tax season feels like a crisis, it’s usually because the right financial information isn’t organized ahead of time. Deductions, education expenses, and important documents all become a last-minute scramble.

Listen to your gut. You can start preparing now.

BELAY’s experienced tax prep professionals help you stay organized year-round, so tax season becomes simpler, less stressful, and actually manageable.

Start with BELAY’s free Personal Tax Prep Checklist and take the first step toward a smoother tax season.

Don’t spend another spring stressing over paperwork. Get help now and leave the pit in your stomach behind for good.

Download the Free Checklist

Recommended for you

© 2026 Snack On AI.
beehiivPowered by beehiiv