Residual connections have been broken this whole time

I’m late to this one. Kimi’s Attention Residuals paper came out in late February and I only read it properly last week. It kept nagging at me, because it pokes at something so basic that I’d stopped seeing it: how a transformer passes information from one layer to the next.

The short version is that every modern LLM uses residual connections that treat all previous layers equally. Layer 1’s output and layer 50’s output get blended together with the same weight, no matter what. The model has no way to prioritise. Early computations get buried under everything that came after, and later layers have to shout louder and louder to be heard. You end up with wasted capacity, less stable training and weaker multi-step reasoning. The Kimi team finally did something about it.

The problem nobody talks about

A standard residual connection works like this. At each layer you take the previous hidden state, run it through the layer’s transformation (attention or FFN), and add the result back:

h_l = h_{l-1} + f_{l-1}(h_{l-1})

It’s simple and elegant, and it’s what made training deep networks possible in the first place (He et al., 2015). But unroll that recurrence and look at what layer 32 actually receives:

h_32 = h_1 + f_1(h_1) + f_2(h_2) + ... + f_31(h_31)

Every layer’s output is added with a weight of exactly 1. The embedding, layer 1, layer 17, layer 31, all treated identically. Layer 32 has no way to say “I care a lot about what layer 4 worked out, and layer 12 is irrelevant to me right now.”

That sounds academic until you look at what happens in practice.

With PreNorm, which has been the default since GPT-2, the magnitude of the hidden state grows as O(L) with depth. The accumulated residual stream keeps getting bigger, so each new layer’s contribution is a smaller share of the whole. Fifty layers in, early information is buried under the sheer size of everything stacked on top of it.

The paper puts it well: “early-layer information is buried and cannot be selectively retrieved.” It also cites work showing that a fair fraction of layers in deep networks can be pruned with little loss. Those layers aren’t doing much. They add noise to a residual stream that’s already bloated.

Three ways it goes wrong

The paper names three structural problems with fixed, unit-weight residuals.

The first is that there’s no selective access. Attention layers and MLP layers receive the same accumulated state, even though they do very different jobs. One routes information between positions; the other transforms features. Both get the same compressed history and neither can pick out the earlier computation it needs.

The second is that the loss is irreversible. Once something is folded into the residual stream it can’t be pulled back out. If layer 8 computed something useful and layers 9 to 15 piled on top of it, layer 16 can’t dig it back up.

The third is output growth. Later layers have to produce larger and larger outputs just to stay visible against the growing residual, which destabilises training. You can measure it: output magnitudes rise steadily with depth in standard PreNorm transformers.

The fix: attention over depth

The idea comes from a duality that feels obvious once someone has pointed it out. RNNs had exactly this problem along the sequence. They squeezed all prior timesteps into one hidden state, and each new step couldn’t reach back selectively to earlier ones. The Transformer fixed that by replacing the recurrence with attention over every previous position.

Attention Residuals (AttnRes) does the same thing over depth instead of sequence length. Rather than adding every previous layer’s output with weight 1, each layer computes learned softmax weights over all prior outputs:

h_l = sum(alpha_{i->l} * v_i)  for i in [0, l-1]

where alpha weights sum to 1 via softmax

Each layer gets one learned pseudo-query vector w_l, just d dimensions. It takes dot-product attention against RMSNorm’d keys from all previous layer outputs. Because of the softmax, the weights are bounded and sum to 1, so output magnitudes stay under control however deep the network goes.

The parameter cost is tiny: one RMSNorm and one d-dimensional query per layer. In a 48B model it’s a rounding error.

The practical version: Block AttnRes

Full AttnRes sounds great until you try to train a 48B model with pipeline parallelism. Every layer’s output has to be kept alive for every later layer and shipped across pipeline stages. Memory grows as O(Ld), and the communication overhead makes it impractical.

Block AttnRes is the engineering compromise. Group the L layers into N blocks (they use about 8). Inside a block, layers sum as usual. Across blocks, layers attend over the N block-level representations with softmax. Memory drops from O(Ld) to O(Nd).

It’s a neat design. Because the pseudo-queries don’t depend on the forward computation, all the queries in a block can be batched into a single matrix multiply. Add cross-stage caching, so blocks already received don’t get sent again, and a two-phase computation that amortises memory reads, and Block AttnRes becomes a practical drop-in replacement.

The overhead is less than 4% in training under pipeline parallelism, and less than 2% in inference latency.

Does it work?

On every metric they tested, yes.

They swept five model sizes from 194M to 528M activated parameters, and Block AttnRes beat the PreNorm baseline across the whole compute range. At 5.6 PFLOP/s-days it matched the loss of a baseline trained with 1.25x more compute. Full AttnRes does slightly better again, but the difference shrinks to 0.001 at the largest scale.

Then they put Block AttnRes into Kimi Linear (48B total, 3B activated, MoE) and pre-trained on 1.4T tokens. The before and after:

• GPQA-Diamond: 36.9 to 44.4 (+7.5 points)
• Minerva Math: 53.5 to 57.1 (+3.6 points)
• HumanEval: 59.1 to 62.2 (+3.1 points)
• MMLU: 73.5 to 74.6 (+1.1 points)
• BBH: 76.3 to 78.0 (+1.7 points)

The biggest gains are on multi-step reasoning, GPQA-Diamond and maths, which is what you’d hope. Those are the tasks where a later layer most needs to pull out something an earlier layer worked out, rather than settle for a diluted average of everything.

Why it matters if you don’t train models

Most people reading this don’t pre-train foundation models. I don’t either. I still think it matters, for a few reasons.

Start with cost. If Block AttnRes matches a baseline trained with 1.25x more compute, then a lab using it can train an equivalent model with about 20% less compute. Frontier training runs cost tens of millions of dollars, so that’s real money, and some of it should eventually show up in API prices and inference costs for everything built on top.

Then look at where the gains land. They aren’t on trivia. They’re on multi-step reasoning (+7.5 on GPQA-Diamond) and code (+3.1 on HumanEval), which is the work agents and coding assistants actually do: plan several steps, hold a chain of thought across a complicated function, build a conclusion on intermediate results. Letting later layers reach back to earlier steps is precisely what that work needs.

And it says something about where architecture research is going. “Just make it bigger” is running out of road. I suspect a lot of the next round of gains will come from revisiting parts of the transformer everyone assumed were fine. Residual connections were one of them. If a change this small, one learned vector per layer, gives consistent gains across every benchmark, there are probably more sitting in the plumbing.

The training dynamics are the real story

The benchmark numbers are convincing. The training plots are where you can see what’s actually been fixed.

In the baseline, per-block output magnitudes climb steadily with depth and reach about 14 at the final block. That’s PreNorm dilution in action: deeper layers have to produce bigger outputs to be heard over the accumulated residual. With Block AttnRes the magnitudes stay bounded, in a periodic pattern that resets at each block boundary. The network is no longer in an arms race with its own residual stream.

The gradients tell the same story from the other side. The baseline has very large gradients in the earliest layers and almost nothing in the later ones, because early layers affect everything downstream while late ones barely register. AttnRes spreads the gradient much more evenly, so every layer gets a useful learning signal.

Validation loss sits below the baseline throughout training, and the gap widens during the decay phase. That suggests the benefit builds as training goes on.

The structured matrix view

My favourite theoretical part of the paper frames every residual variant as a depth-mixing matrix M. Standard residuals give an all-ones lower-triangular M. Highway networks give a 1-semiseparable M. AttnRes gives a dense, rank-L matrix whose weights depend on the input.

Seen that way, standard residuals are doing depth-wise linear attention with fixed weights, and AttnRes is doing depth-wise softmax attention with learned, input-dependent weights. It’s the same move Transformers made along the sequence. Highway networks and mHC sit somewhere in between.

You can push the analogy further. Test-Time Training treats residual connections as gradient descent steps, one update per layer. Under that lens, AttnRes swaps gradient descent (adding things up) for attention (choosing what to retrieve), which mirrors the historical shift from RNNs to Transformers.

What the paper doesn’t settle

I have a few reservations.

The baseline uses hyperparameters tuned for standard residuals. AttnRes might do better still with its own tuning. Equally, the baseline gets home advantage, so if anything the reported gains are probably conservative.

The 48B results come from a single training recipe, Kimi Linear’s. We don’t yet know how AttnRes behaves with other architectures, recipes or data mixes. The five-size scaling sweep is encouraging, but there’s always a question of whether it generalises beyond one lab’s setup.

Eight blocks is presented as the practical choice, but the ablation (Figure 6) shows that more, smaller blocks get close to Full AttnRes at the cost of more memory. As hardware and cross-stage bandwidth improve, smaller blocks or even full AttnRes may become practical. It’s a moving target.

And the paper doesn’t compare against some concurrent work on cross-layer access, such as ANCRe and MUDDFormer, though it does compare against DenseFormer and mHC and beats both.

Worth reading in full

Residual connections are the plumbing of modern LLMs. Everyone uses them and nobody thinks about them. The Kimi team showed the plumbing has a real, measurable flaw, and that replacing the fixed sum with softmax attention over depth gives consistent improvements from small models up to 48B, for under 2% extra inference latency.

What stays with me is the symmetry between depth and sequence. Both mean aggregating information from a growing history, and both do better when the model can choose what to look at. Transformers made that change for sequences in 2017. AttnRes makes it for depth.

If you train models, read the paper. Block AttnRes is a practical drop-in replacement, the code is on GitHub (MoonshotAI/Attention-Residuals), and the engineering that makes it scale is well documented. The rest of the transformer learned input-dependent weights years ago. The skip connections were overdue.