Residual streams and attention over depth

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

Applied AI · advanced · Session 20

Mechanism map

  x₀ ──┬───────────────────────────────────────────┐
        │                                           │ checkpoint
   ┌────▼────┐                                       │
   │  F₀(x)  │──▶ ⊕ ──▶ x₁ = x₀ + F₀(x₀)             │
   └─────────┘                                       │
              … 12 layers of additions …             │
  x₁₂ ─┬──────────────────────────────────────────┐  │ checkpoint
   ┌───▼─────┐                                    │  │
   │ F₁₂(x)  │──▶ ⊕ ──▶ …                         │  │
   └─────────┘                                    │  │
              … 12 layers of additions …          │  │
  x₂₄ ─┬───────────────────────────────────────┐  │  │ checkpoint
       │                                       │  │  │
   ┌───▼─────────┐   q₂₄                       ▼  ▼  ▼
   │  LAYER 24   │──────▶ scores [2, 1, 0] over {x₂₄,x₁₂,x₀}
   │   query     │        softmax [0.67, 0.24, 0.09]
   └─────────────┘              │
                                ▼
                     y = 0.67·x₂₄ + 0.24·x₁₂ + 0.09·x₀
                                │
                                ▼  ⊕ residual → layer 25

The problem — Residual connection

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 — 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.

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

Why / at what price — Residual connection

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.

Check: In x_{l+1} = x_l + F_l(x_l), suppose F_l ≡ 0 for twelve layers. What is x₁₂, and what does that say about the gradient flowing back along that path?

The problem — Dilution

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 — 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.

Why / at what price — Dilution

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.

Check: The energy share of x₀ falls from 100% to 7.7% between layer 0 and layer 48. Verify that 7.7% using ‖x_l‖² ≈ 1 + 0.25·l (assumption: orthogonal contributions), then say whether x₀’s information was erased or merely diluted.

Visual support — Dilution

share of x₀ in ‖x_l‖²    (assumption: orthogonal contributions)

100 % ●
      │
 25 % ┤             ●
14.3 %┤                      ●
 7.7 %┤                               ●   present, but drowned
      └────┬────────┬─────────┬───────┬─▶ layer
           0       12        24      48

addition erases nothing: it buries

The problem — Depth states

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 — 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.

Why / at what price — Depth states

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.

Check: We keep {x₀, x₁₂, x₂₄} and not the other 45 states. On what criterion would you choose those three indices rather than {x₀, x₁, x₂}, and what exactly is lost by not keeping x₂₃?

The problem — Depth scores

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 — 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.

Why / at what price — Depth scores

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.

Check: Scores [2, 1, 0] give [0.665, 0.245, 0.090]. Recompute the softmax for [4, 1, 0]: how far does x₀ drop, and has the module thereby removed x₀ from the stream?

Visual support — Depth scores

attention (s.12): a position reads OTHER TOKENS, same layer
attention (s.20): a layer reads OTHER DEPTHS, same token

   tokens ────────▶                 depth ────────▶
   t1   t2   t3   t4                x₀    x₁₂    x₂₄
                   ▲                                ▲
               q (from t4)                  q (from layer 24)

same machinery: query, scores, softmax, mixture

The problem — Blocking

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 — Blocking

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

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.

Check: Going from 48 checkpoints (1.5 GiB) to 4 (128 MiB) divides storage by 12. Which question can layer 24 no longer ask once only x₀, x₁₂ and x₂₄ remain available?

The problem — What this does not prove

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 — What this does not prove

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 — What this does not prove

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.

Check: The 0.665 weight on x₂₄ is measured, not interpreted. Propose an observation that would contradict the reading “layer 24 carries the reasoning”, using the degenerate case [0.333, 0.333, 0.333].

Visual support — What this does not prove

what 0.665 on x₂₄ SAYS            what it does NOT say
─────────────────────────          ─────────────────────────
x₂₄’s influence inside THIS        what layer 24 actually
module, this layer,                computes
this token
                                   that the reasoning
re-measurable weight,              “lives” in layer 24
reproducible experiment
                                   test: flat weights [0.333…]
                                   → readable, zero explanation

Worked case — full trace

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%.

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

DILUTION (roughly orthogonal contributions, ‖F_l‖ = 0.5)
  ‖x_l‖² ≈ ‖x₀‖² + l × 0.25   with ‖x₀‖ = 1
  l = 0   → ‖x‖² = 1     share of x₀ = 100%
  l = 12  → ‖x‖² = 4     share of x₀ = 25%
  l = 24  → ‖x‖² = 7     share of x₀ ≈ 14.3%
  l = 48  → ‖x‖² = 13    share of x₀ ≈ 7.7%   ← present, but drowned

DEPTH MIXTURE, checkpoints {x₂₄, x₁₂, x₀}, scores [2, 1, 0]
  exp = [7.389, 2.718, 1.000]   sum = 11.107
  softmax = [0.665, 0.245, 0.090]                   ✅ sums to 1.000
  y = 0.665·x₂₄ + 0.245·x₁₂ + 0.090·x₀
  x₀ weighs 9% after mixing versus a 14.3% share of residual energy:
  retrieval is now chosen rather than merely endured.

DEGENERATE CASE: scores [1.0, 1.0, 1.0]
  softmax = [0.333, 0.333, 0.333]                   ❌
  y = mean of the three checkpoints. The module retrieves nothing:
  it redoes, more expensively, what the residual already did for free.

STORAGE COST (d_model = 4096, 4,096 tokens, BF16)
  1 checkpoint = 4096 × 4096 × 2 = 33,554,432 B = 32 MiB
  all 48 layers stored → 48 × 32 MiB = 1,536 MiB = 1.5 GiB   ❌
  blocks of 12 → 4 checkpoints → 128 MiB                     ✅ ÷12

SHAPE CHECK: q ∈ R^{4096}, stacked states [J=3, 4096]
  → scores [3] → weights [3] → y ∈ R^{4096}. Moving to J=4 changes
  [3]→[4], never the dimension of y.

Depth-checkpoint granularity (48 layers, 32 MiB each)

Spacing Storage Retrieval quality
Every layer (48) 1.5 GiB finest, scores nearly redundant between neighbours
Every 4 (12 states) 384 MiB fine, still separates block start from block end
Every 12 (4 states) 128 MiB coarse: one point per processing phase
x₀ only (1 state) 32 MiB no choice left: the mixture degenerates to a long residual

Causal lab

Predict → change one variable → run → explain the delta

/interactives/curriculum/moe-attnres.html?lang=en

Common errors

“After 48 layers, the information in x₀ has been overwritten.”

Nothing is overwritten: addition is conservative and x₄₈ still contains x₀ in full. What changes is the signal-to-noise ratio — 7.7% of energy versus 100% at the start. That is exactly why a depth mixture can lift it back to 9% explicit weight: it does not restore lost information, it re-selects it.

“A 0.665 weight on x₂₄ proves layer 24 does the work.”

That weight describes x₂₄’s influence inside this particular retrieval module, at this layer, for this token. It says nothing about what layer 24 computes, nor about the residual path that carries x₂₄ regardless. The [0.333, 0.333, 0.333] case makes it plain: perfectly readable weights can belong to a module that explains nothing at all.

Boundary, evidence, and sources

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.

  • Owner-supplied bilingual course packet, Chapter 15.
  • He et al., “Deep Residual Learning for Image Recognition”, CVPR (2016).
  • Elhage et al., “A Mathematical Framework for Transformer Circuits” (residual-stream view), Anthropic (2021).
  • Course source packet supplied by the owner; named-product details remain source-reported until primary verification.

Transfer challenge

Your checkpoint budget is 256 MiB for 48 layers (32 MiB each).

  1. Choose a placement — how many states, which indices — and pre-register what your choice makes impossible to retrieve.
  2. Incident: layer 30 turns out to need the state of layer 18. Does your placement cope? Quantify the gap with the dilution toy model.
  3. State the measurement that would make you change placements.

Synthesis and exit ticket

  • Residual connection
  • Dilution
  • Depth states
  • Depth scores
  • Blocking
  • What this does not prove

Ticket: mechanism · trace · observation · boundary · evidence · next experiment

Instructor notes: Frame the problem before naming the mechanism. Collect an initial prediction and retain it for the exit ticket.

Instructor notes: Connect every step to the next with a causal verb. Flag any merely decorative arrow.

Instructor notes: Open with the failure story: “before 2015, past roughly twenty layers, networks got WORSE in training and in test”. Collect hypotheses; write “gradient” on the board the moment it appears — the beat makes it concrete with F ≡ 0.

Instructor notes: Have the formula written, then set F_l ≡ 0 on the board. Watching x₁₂ = x₀ emerge conveys the identity path faster than any talk about vanishing gradients.

Instructor notes: Answer: x₁₂ = x₀ — the stream passes through unchanged, and the gradient flows back along the identity path unattenuated (Jacobian = I). Expected wrong answer: “x₁₂ = 0”, confusing F ≡ 0 with x ≡ 0.

Instructor notes: Vote before any computation: “after 48 layers, how much of x₀ remains — 0%, 8%, 50%?”. The majority vote for “zero” is the beat’s raw material; post it and return to it after the 7.7%.

Instructor notes: Have them guess x₀’s share at layer 48 before computing it. The room almost always says “zero”; the 7.7% that follows plants the dilution-versus-erasure distinction for the rest of the session.

Instructor notes: Answer: ‖x₄₈‖² ≈ 1 + 48 × 0.25 = 13, x₀’s share = 1/13 ≈ 7.7% — diluted, not erased: addition keeps everything, it is the signal-to-noise ratio that drops. Recall the figure holds under the orthogonality assumption (toy model). Expected wrong answer: concluding “erased” from the small number.

Instructor notes: Confront the curve with the opening vote (“zero”): the 7.7% point is low but not zero, and that is the whole difference — you can re-select what is buried, not what is erased.

Instructor notes: Hook analogy: a document with no version history. “How do you get back to the original draft?” — you cannot, unless you saved copies along the way. That is exactly what the beat installs.

Instructor notes: Have each person pick three checkpoint indices out of 48 and justify them. The disagreements in the room are the argument: there is no canonical placement, only a trade-off.

Instructor notes: Answer: a coverage criterion — {x₀, x₁₂, x₂₄} samples distinct processing phases; {x₀, x₁, x₂} stores three nearly identical states. Without x₂₃, fine addressing of the end of block 12-24 disappears: “just before x₂₄” becomes an unaskable question. Expected wrong answer: “the earliest layers matter most”.

Instructor notes: Recall session 12 with one question: “the three roles q, k, v — who remembers?”. The tokens → depths rotation takes one sentence if the base is there; otherwise, ninety seconds of recall before the beat.

Instructor notes: Have the softmax of [2, 1, 0] computed by hand, then swap in [1, 1, 1] without warning. The flip to 0.333 should be experienced as a failure, not read as a counterexample.

Instructor notes: Answer: softmax([4, 1, 0]) ≈ [0.936, 0.047, 0.017] — x₀ falls from 0.090 to 0.017. But no: x₀ remains fully in the residual stream; the module merely stopped re-selecting it. Expected wrong answer: “x₀ is removed” — the module/stream confusion of common error 1.

Instructor notes: Have someone name, for the depth version, what plays the roles of keys and values (the retained states, both times). The symmetry with session 12 must be stated by a learner, not by the slide.

Instructor notes: Budget in hand: “1.5 GiB of snapshots for ONE sequence — who signs off?”. Nobody signs. Blocking then arrives as an accounting necessity, not an implementation trick.

Instructor notes: Pose it as an engineer with a budget: “1.5 GiB of checkpoints, or 128 MiB?” Make them justify the choice by the target task, not by an aesthetic preference for granularity.

Instructor notes: Answer: any question aimed inside a block — for instance “the state just before layer 20”: between x₁₂ and x₂₄ nothing is addressable any more. Expected wrong answer: answering in MiB (the cost) instead of in lost questions (the granularity).

Instructor notes: Play the over-claimer: present the 0.665 as “proof that layer 24 thinks”. Let the room dismantle the argument itself; the beat then supplies the vocabulary of the bounded reading.

Instructor notes: Close by asking outright: “what experiment would refute your reading of the 0.665?” Reject any answer that does not propose a measurement; it is the only guard against decorative interpretability.

Instructor notes: Expected answer: a measurement like “on a controlled set, the weights stay ≈ [0.333…] AND performance does not move when the module is removed” — readable weights with no causal effect refute the reading. Demand a measurement; reject any narrative rephrasing.

Instructor notes: Have a third column added aloud: “which experiment would decide?” (module ablation, probe on x₂₄, task comparison). The table is only complete with that column.

Instructor notes: Walk line by line. Locate an inconsistency at the first faulty step, not only on the final line.

Instructor notes: Have learners fill the final row before revealing it: that trade-off is what decides in production.

Instructor notes: Retain initial and final values. Do not allow simultaneous changes that make the delta impossible to attribute.

Instructor notes: For each claim, have the group produce the smallest counterexample before giving the correction.

Instructor notes: Separate verifiable mechanism, reported implementation choice, and experimental result. Evidence precision must match claim precision.

Instructor notes: Expected: 8 states at most (256/32). Uniform placement every 6 — x₀, x₆, …, x₄₂ —: layer 30 reads x₁₈ exactly, the incident is absorbed. With 4 states every 12, the nearest is x₁₂ or x₂₄: 6 layers away, ≈ 1.5 added energy units in the toy model (0.25 × 6) between the wanted state and the available one. Switch measurement: a reconstruction probe on the layers actually requested, not an intuition. Misconception to harvest: “denser at the start” — dilution mostly makes LATE phases hard to address. Ten minutes, pairs.

Instructor notes: Rebuild the chain without looking at the slides, then fill the ticket in at most six lines. Compare with the opening prediction and name what actually changed.