Lesson 4 · Attention Residuals

You have now seen K3's answer on the sequence axis. This lesson is the depth axis, and it rests on an observation the report makes in its first sentence on the subject: the residual stream is a bottleneck of exactly the kind attention was invented to fix.

The flaw in the residual stream

A transformer passes information up through depth by accumulation. Layer 1 writes into a running state, layer 2 adds to it, layer 3 adds to that. By layer 70, that state holds contributions from 69 earlier layers, all summed together.

Residual connection: The mechanism that lets a layer add its output to a running state rather than replacing it. Introduced for image networks in 2015, and the reason deep networks train at all: gradients have a short path back to early layers.
Residual stream: The running state itself, carried from the embedding to the output. Every layer reads it and writes back to it.

Summing works, but it destroys information about origin. Layer 70 receives one vector. It cannot ask what layer 12 contributed, or weight the embedding differently from layer 69. Everything arrives pre-mixed, in fixed proportions the layer has no say over.

The report's framing is the useful one, and it is worth quoting because it tells you exactly where the idea came from: standard residuals compress all prior information into a single state h_l over depth — a bottleneck reminiscent of RNNs over time.¹

That comparison is the whole lesson. Recurrent networks had this exact problem along the sequence axis: everything read so far crushed into one hidden state, no way to reach back to a specific earlier token. Attention solved it by letting each position query all previous positions with learned, data-dependent weights. AttnRes applies that same move to depth.

    THE SAME BOTTLENECK, TWICE

    over SEQUENCE                        over DEPTH
    ─────────────────                    ──────────────

    RNN                                  Standard residual
      one hidden state carries              one residual stream carries
      everything read so far                everything computed so far
      ↓                                     ↓
      cannot reach token 12,345             cannot reach layer 12
      specifically                          specifically

    ── solved by attention ──            ── solved by AttnRes ──

    Transformer                          Attention Residuals
      each position queries all             each layer queries all
      previous positions, with              previous layers, with
      learned weights                       learned weights

    Lessons 2–3                          this lesson

One idea applied on a second axis. The report is explicit that AttnRes applies the same methodology to depth. Recognising this pattern is more valuable than memorising the mechanism, because it is how a great deal of architecture research proceeds.

How a layer queries depth

The mechanism is small. Each layer gets a learned pseudo-query, a single vector belonging to the layer rather than to any token. Every earlier layer's output serves as both key and value, with the token embedding included as the earliest source. The layer's query is compared against all of them, softmax turns the comparisons into weights, and the layer reads a weighted blend instead of a plain sum.

Pseudo-query: A learned vector, one per layer, playing the role a query plays in ordinary attention. Called "pseudo" because it does not come from a token. It encodes a fixed preference: which depths this layer wants to hear from.
softmax: Turns a list of arbitrary scores into positive weights that sum to 1, exaggerating the gaps so the largest score dominates. Every attention weight in this course, over tokens or over depth, is a softmax output.
RMSNorm: Rescales a vector so its typical element size is about 1, without shifting it. Reach for it whenever one quantity must not dominate another purely by being numerically larger.

One detail in the report matters more than its size suggests. The keys pass through RMSNorm before the comparison, which the authors say prevents layers with large-magnitude outputs from dominating the weights.¹ Without it, a layer that happens to produce big numbers would win the attention competition on volume rather than relevance.

show me the actual equation

From §2.2, eqs. 8–9. First, what serves as keys and values for layer l:

k_i = v_i = ⎧ h₁ i = 0 ← the token embedding ⎨ ⎩ f_i(h_i) 1 ≤ i ≤ l−1 ← output of layer i q_l = w_l ∈ ℝ^d ← learned, one per layer

Then the weights and the read, a softmax kernel over depth:

φ(q, k) = exp( qᵀ · RMSNorm(k) ) φ(q_l, k_i) l−1 α_{i→l} = ───────────────────── h_l = Σ α_{i→l} · v_i Σ_j φ(q_l, k_j) i=0

Compare that to ordinary attention and the only differences are what the query is (learned per layer, not projected from a token) and what the index runs over (depth, not sequence position). The RMSNorm(k) inside the kernel is the magnitude-fairness fix.

Cost, in the report's own accounting: the arithmetic is O(L²d), which is affordable because network depth is modest (L < 100). The real problem is O(Ld) memory, because every layer output must stay alive, and must cross stage boundaries under pipeline parallelism.

Why K3 does not ship the full version

Full AttnRes has a cost you can see without any math. If layer 70 attends over 69 earlier outputs, all 69 must still be in memory. Multiply by the size of each and by pipeline stages that need them shipped across devices, and the bill lands on memory and interconnect rather than arithmetic.

So K3 ships Block AttnRes. The 93 layers are partitioned into blocks of 12. Within a block, outputs are summed into one running representation, exactly as a standard residual would. Across blocks, full attention is applied over just the block-level summaries.

Switch between the three modes below and drag the layer slider. Watch the source count in each.

The shape to take away: a layer in the middle of the network reads from a handful of sources rather than one, and rather than seventy. Overhead falls from O(Ld) to O(Nd), and the report reports that N ≈ 8 recovers most of the benefit across model scales.¹

K3's actual partition, and a wrinkle

93 layers into 8 blocks of 12 does not divide evenly. Seven full blocks take 84 layers, leaving a partial final block of 9, so 8 blocks in all. The report then counts 9 when it adds the embedding as one more source. Two different nines, easy to conflate: nine layers in the last block, and nine total sources once the embedding joins the eight blocks. This is the kind of untidy detail that tells you a number came from a real training run, and it is worth checking whenever you read a config table.

What it costs, honestly

Raschka's read of the report puts numbers on the overhead: roughly 4% in training cost and 2% in inference cost, for gains he describes as consistent but modest in validation loss and downstream performance.² A few percent for a small, reliable gain is an unglamorous trade, which is worth stating plainly.

That is a genuinely unglamorous return, and saying so is the honest read. It also explains why the block form exists at all: at full-AttnRes memory cost the trade would likely not be worth making, and the block approximation is what brings the price down to where a few percent buys a real if small gain.

Worth knowing for the deployment question

Block AttnRes was designed with inference in mind, not just training. Because a block's representation is computed once at its boundary and then shared by every layer after it, the inference-time state stays bounded. The report notes this lets the parallel inter-block results merge with the sequential intra-block partial sums using online softmax, significantly reducing inference time cost. K3 also ships dedicated kernels for Block AttnRes, and §5.4.2 describes a two-phase schedule that keeps it off the critical path.¹

A contrast worth carrying

K3 is not the only 2026 frontier model to decide the residual stream needed work, and the comparison sharpens what AttnRes actually is. DeepSeek V4 attacked the same bottleneck by making the residual path wider, using manifold-constrained hyper-connections. K3 instead kept the width and made the path selective across layers.²

    TWO ANSWERS TO ONE BOTTLENECK

    DeepSeek V4 — widen the pipe        Kimi K3 — choose what to read
    ┌───────────────────────────┐       ┌───────────────────────────┐
    │  more channels carrying   │       │  same channels, but each  │
    │  more information         │       │  layer weights its own    │
    │  forward                  │       │  sources across depth     │
    └───────────────────────────┘       └───────────────────────────┘
    cost: width, so parameters          cost: keeping earlier outputs
          and memory everywhere               reachable (bounded by N)

Two research groups, one diagnosis, two treatments. Neither is settled as correct. Noticing that two labs independently identified the residual stream as the limiting factor is the more durable observation.

Check yourself

What problem with standard residual connections does AttnRes address?

Earlier layers arrive pre-summed, in fixed proportions
Gradients vanish before reaching the earliest layers
The residual stream is too narrow to carry enough
Position information degrades as it passes up depth

Summation destroys origin. Layer 70 gets one vector and cannot weight layer 12 differently from layer 69. AttnRes gives each layer a learned query so it can choose, which is the same fix attention made for sequence.

Why does K3 use Block AttnRes rather than the full form?

The full form must keep every layer output alive in memory
The arithmetic of the full form is too costly at 93 layers
The full form gives worse validation loss than the blocked variant
The full form cannot use learned pseudo-queries at this model scale

Compute is not the constraint; memory and cross-stage communication are. Blocking cuts overhead from O(Ld) to O(Nd), and N ≈ 8 keeps most of the benefit. It also bounds inference-time state, which matters for serving.

Why does AttnRes apply RMSNorm to the keys before computing weights?

So large-magnitude layers cannot dominate on volume
To stop the softmax overflowing in low precision
To encode the relative depth order of source layers
To cut the memory needed to keep outputs alive

The report's stated reason is fairness in the attention competition: without normalisation a layer wins weight by producing big numbers rather than relevant ones. Small detail, but it is what makes the learned weights mean something.

Where we are

    ✅ SEQUENCE AXIS      KDA, chunkwise, g_min, Gated MLA, NoPE   L2–L3
    ✅ DEPTH AXIS         Block AttnRes, N = 8, pseudo-queries     L4
    ▢  WIDTH AXIS         LatentMoE, SiTU-GLU, Quantile Balancing  L6–L7
    ▢  INPUT              MoonViT-V2 native vision                 L8
    ▢  OPTIMISATION       Per-Head Muon, scaling law               L9
    ▢  SYSTEM             serving, cache, quantization             L11
    ▢  BEHAVIOUR          post-training, effort, limitations       L12

Course progress. Two of the three information-flow axes are done. Ticks mark what you can now explain from memory.

Two of the three axes are done, which is the natural point to stop moving forward. Lesson 5 is a review, interleaving everything from lessons 1 to 4. It will feel harder than these lessons did, and that is the intent: retrieving something you half-remember is what makes it stick, whereas rereading it only feels like learning.