advanced · Session 20

Residual streams and attention over depth

Understand residual dilution, retrieval of earlier states, and depth-checkpoint trade-offs.

120 min6 mechanismscausal lab

What you will be able to do

Open the lab

Study method

Study this session as a causal investigation. Before every formula or interaction, write what you expect to change and what must remain fixed. During the calculation, retain units, shapes, and intermediate values so an error can be located without restarting at random. After the result, translate the number or state into one sentence about system behavior. Always finish with a counterexample or boundary value. This discipline separates understanding a mechanism from merely recognizing its vocabulary and makes the lab reproducible by another learner. Record the evidence that changed your initial prediction.

Build the mechanism step by step

1. Residual connection

The problem: Stacking 48 raw transformations (x → F(x)) multiplies 48 Jacobians: the gradient vanishes or explodes, and one bad layer destroys everything before it. “Pure” deep networks simply fail to learn.

The idea: Addition changes everything: x_{l+1} = x_l + F_l(x_l). Each layer ADDS its transformation to the stream instead of replacing it; if F ≡ 0, x passes through unchanged — the identity path guarantees a direct gradient to every layer.

x_{l+1}=x_l+F_l(x_l)

Why / at what price: The identity path makes deep training possible — the invention that unlocked hundred-plus-layer networks. The price: nothing is ever removed from the stream; everything added stays, and accumulates.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

2. Dilution

The problem: 48 layers each add their contribution. What becomes of the information laid down by x₀? Common intuition: “overwritten”. If that were true, all retrieval would be impossible — you must compute before concluding.

The idea: Under the assumption of roughly orthogonal contributions (‖F_l‖ = 0.5), the energy follows ‖x_l‖² ≈ 1 + 0.25·l: x₀’s share goes from 100% to 25% (l = 12), 14.3% (l = 24), 7.7% (l = 48). Diluted, never erased — addition is conservative.

Why / at what price: The precise diagnosis — drowned, not destroyed — is what makes the rest possible. The price: the orthogonality assumption does all the computational work; the 7.7% is a toy-model order of magnitude, not a measurement of the real model.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

3. Depth states

The problem: If layer 40 needs the pre-transformation representation, all it has is x₃₉ — the sum of everything. Going back is impossible: the stream keeps no snapshots.

The idea: The solution: photograph along the way. Keep a few chosen x_j — {x₀, x₁₂, x₂₄} — as retrieval points along depth, the informational analogue of training checkpoints.

Why / at what price: Explicit return points, addressable later. The price: each snapshot is expensive — 32 MiB per checkpoint here (4,096 tokens × d_model 4096 × BF16) — and the choice of indices is a bet made at design time.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

4. Depth scores

The problem: Three snapshots available: which one, and when? A fixed choice — always x₁₂? — would be blind to content: the right snapshot depends on the token and on the question the layer is asking.

The idea: Session 12’s mechanism, rotated 90°: a query from the current layer scores the retained states — scores [2, 1, 0] over {x₂₄, x₁₂, x₀} → softmax [0.665, 0.245, 0.090] — then mixes the values. Retrieval is chosen, no longer endured.

Why / at what price: Attention over J = 3 states costs little (scores [3], weights [3]). The price: the module can degenerate — flat scores [1,1,1] give a [0.333…] average that redoes, at extra cost, what the residual already gave for free. A choosing mechanism is only worth having if it chooses.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

5. Blocking

The problem: Storing all 48 states = 1.5 GiB for a single 4,096-token sequence — unacceptable. But every deleted snapshot is a question later layers can no longer ask. Where do you cut?

The idea: Group: one checkpoint per block of 12 layers → 4 states, 128 MiB (÷12). Granularity becomes a slider — 48, 12, 4, or 1 states: from “nearly redundant” to “no choice left”.

Why / at what price: Blocking makes the mechanism affordable. The price: retrieval turns coarse — between x₁₂ and x₂₄ nothing is addressable any more — and checkpoint placement has no canonical answer, only a per-task trade-off.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

6. What this does not prove

The problem: The module shows 0.665 on x₂₄ — the immediate temptation: “layer 24 carries the reasoning”. The weights are readable, the story is seductive… and the conclusion does not follow.

The idea: The bounded reading: 0.665 measures x₂₄’s influence in THIS module, at THIS layer, for THIS token. The residual path carries x₂₄ regardless, and perfectly readable weights can accompany a module that explains nothing — the [0.333…] case proves it.

Why / at what price: Interpretation discipline keeps a mechanism from being sold as an explanation. The price: giving up seductive narratives — every reading must come with an experiment that could refute it, or it is decorative interpretability.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

Development from the course source

Chapter 15 — Residual dilution and Attention Residuals

15.1 Residual streams

Goal: understand how deep Transformers preserve and update information. A residual connection adds a layer’s change to the current representation: x_next = x + F(x).

Intuition: Each editor keeps the current draft and adds a revision instead of rewriting from a blank page.

Step by step

  • Start x₀=[1,2].

  • Layer change F(x₀)=[0.5,−0.2].

  • Residual result x₁=[1.5,1.8].

  • A later layer receives both the old path and the new change.

Worked example: In a very deep network, repeatedly mixing everything into one stream may make some earlier signals difficult to recover. This is called residual dilution in the supplied narrative.

Why it matters: Residual connections help optimization, but they do not guarantee every useful earlier feature stays equally accessible.

Quick check: If x=[3,1] and F(x)=[−1,2], what is x+F(x)? Answer: [2,3].

15.2 Attention over depth

Goal: let a layer choose among earlier representations, not only accept the latest accumulated stream. Attention Residuals, abbreviated AttnRes, use attention-like weighting over depth.

Intuition: A researcher compares several saved drafts and chooses the most relevant passages for the next revision.

Step by step

  • Keep selected earlier layer outputs or block summaries.

  • Compute relevance scores for the current layer.

  • Normalize the scores into weights.

  • Form a weighted mixture and use it in the new computation.

Worked example: three depth states are h₁=[1,0], h₂=[0,2], h₃=[3,1]. Weights [0.2,0.3,0.5] give 0.2h₁+0.3h₂+0.5h₃ = [1.7,1.1].

Why it matters: Attention over depth costs storage and computation. Blockwise AttnRes reduces this cost by keeping or attending to coarser checkpoints rather than every layer output.

Quick check: What do AttnRes weights sum to after softmax? Answer: 1.

15.3 Evidence boundary

Goal: avoid treating an architectural sketch as a released specification. Exact checkpoint spacing, equations and layer placement must be verified from the authoritative Kimi K3 source.

Intuition: A building concept drawing explains the idea but is not the final engineering blueprint.

Step by step

  • Mechanism-level lesson: learned access to earlier depth states may fight dilution.

  • Product-level details: exact block sizes and reported gains require primary evidence.

Worked example: This distinction keeps the course useful even if future releases revise names or numbers.

Why it matters: Scientific reading means tracking what is derived, what is measured and what is merely reported.

Quick check: Which source should settle an exact implementation dispute? Answer: official code or an authoritative technical report for the exact model version.

Complete worked case

Depth scores [2,1,0] over {x₂₄,x₁₂,x₀} give softmax ≈ [0.67,0.24,0.09]. The mixture favors x₂₄, the most recent checkpoint, while the other two together retain 33%.

Reading method: write the data, state every object shape, perform one transformation, and interpret the result before continuing.

Validity boundary

AttnRes is presented as a course mechanism and source-reported case study; exact Kimi K3 details require a primary source.

Evidence status: Mixed: established mechanisms + source-reported Kimi K3-style choices.

Quick checks

1. What does Residual connection?

Addition changes everything: x_{l+1} = x_l + F_l(x_l). Each layer ADDS its transformation to the stream instead of replacing it; if F ≡ 0, x passes through unchanged — the identity path guarantees a direct gradient to every layer.

2. What does Dilution?

Under the assumption of roughly orthogonal contributions (‖F_l‖ = 0.5), the energy follows ‖x_l‖² ≈ 1 + 0.25·l: x₀’s share goes from 100% to 25% (l = 12), 14.3% (l = 24), 7.7% (l = 48). Diluted, never erased — addition is conservative.

3. What does Depth states?

The solution: photograph along the way. Keep a few chosen x_j — {x₀, x₁₂, x₂₄} — as retrieval points along depth, the informational analogue of training checkpoints.

4. What does Depth scores?

Session 12’s mechanism, rotated 90°: a query from the current layer scores the retained states — scores [2, 1, 0] over {x₂₄, x₁₂, x₀} → softmax [0.665, 0.245, 0.090] — then mixes the values. Retrieval is chosen, no longer endured.

Sources and evidence boundary

Scope: Mixed: established mechanisms + source-reported Kimi K3-style choices.