The Transformer and the residual stream

Understand, at a conceptual-but-precise level, what actually runs inside the model an agent is built on: scaled dot-product attention and why the 1/sqrt(d_k) factor, causal masking, multi-head attention and its QK/OV circuits, the Pre-LN block and the residual stream as a communication channel, LayerNorm vs RMSNorm, the feed-forward network, positional encodings (sinusoidal and RoPE), and the cost/memory picture (quadratic attention, FlashAttention and online softmax).

Track

Research

Module

9 · Agent cognition & reasoning

Prerequisites

None

When an agent reads a tool trace, retrieves a document, or plans a step, the raw computation underneath is a stack of Transformer blocks. Understanding that stack is the difference between treating the model as magic and treating it as a mechanism with known properties and known failure modes. This lesson opens the block piece by piece.

The core idea is simple to state. A Transformer turns a sequence of token representations into a new sequence where the representation at each position has been allowed to selectively read from other positions. The reading is not a fixed pattern like a convolution; it is computed from the content itself. That single property — content-dependent routing of information between positions — is what lets the model copy a value it saw earlier, follow a format it was shown, or attend to the one clause in a long context that matters.

Let X be the matrix whose rows are the representations of n positions. One attention head learns three projections and produces three new matrices: queries Q = X·W_Q, keys K = X·W_K, and values V = X·W_V. The query of a position is what it is looking for; the key of a position is what it offers to be found by; the value is what it hands over if it is attended to.

Query (q_i)

What position i is searching for. It is compared against every key to decide how much to read from each other position.

Key (k_j)

What position j advertises. A large dot product q_i·k_j means position i finds position j relevant.

Value (v_j)

The content position j contributes to the output when it is attended to. The output at i is a weighted average of values.

Attention weight (a_ij)

The softmax-normalized score. Because the weights are non-negative and sum to one, the output is a convex combination of values: attention mixes information but never amplifies it outside the range of the values.

Two properties follow immediately and are proven, not conjectured. First, because the scores are a bilinear form of the representations, the routing is dynamic — it depends on the input — which is exactly why attention is not a fixed-weight filter. Second, with no mask and no positional information, the operation is permutation-equivariant: it knows nothing about order by itself. That is why positional encodings exist.

The output of a head is softmax(Q·Kᵀ / sqrt(d_k) + M)·V, applied row by row, where d_k is the key dimension and M is an optional mask. Two design choices in that formula carry most of the weight.

synthetic — scaled dot-product attention (one head, schematic)
# X: (n, d)  input rows;  W_Q, W_K: (d, d_k);  W_V: (d, d_v)
Q = X @ W_Q            # (n, d_k) — what each position looks for
K = X @ W_K            # (n, d_k) — what each position offers
V = X @ W_V            # (n, d_v) — what each position hands over

S = (Q @ K.T) / sqrt(d_k)          # (n, n) raw scores, SCALED
if causal:                          # position i may not see the future
    S = where(col <= row, S, -inf)  # masked entries -> 0 after softmax
A = softmax(S, axis=-1)             # rows sum to 1 (convex weights)
O = A @ V                           # (n, d_v) output = weighted values

Why divide by sqrt(d_k)? Under an initialization argument, if query and key components are independent with unit variance, the dot product q·k has variance d_k, so its standard deviation grows like sqrt(d_k). Left unscaled, scores grow with dimension, the softmax saturates toward a one-hot distribution, and its Jacobian — diag(a) − a·aᵀ — collapses to zero. A collapsed Jacobian means no gradient flows back to W_Q and W_K, so learning stalls. Dividing by sqrt(d_k) returns the score variance to roughly one. The literature reports, and a small synthetic experiment in the source confirms, that without the scale the average attention entropy over many keys collapses as d_k grows, while with the scale it stays essentially flat. The scale fixes the starting temperature, not the final one — trained norms drift.

The causal mask makes an autoregressive language model well-defined: setting M_ij = −∞ for j > i means exp(−∞) = 0, so position i can only distribute attention over positions up to i. This is exact and composes across layers, which is what makes the key/value cache valid at inference: keys and values of already-processed positions never change when new tokens are appended.

Multi-head attention runs h heads in parallel, each in a dimension d/h, concatenates their outputs and projects. A useful rewrite shows the concatenate-then-project is algebraically a sum over heads, and each head factors into two near-independent circuits: a QK circuit that decides WHERE a position reads from, and an OV circuit that decides WHAT is moved from the attended position and where it is written. Running h small heads costs the same in parameters and FLOPs as one head of full width; the gain is h distinct attention patterns per position, each operating in a low-rank subspace.

A modern decoder block is two sub-layers wrapped in residual connections. With the Pre-LN arrangement, each block computes u = h + MHA(Norm(h)) and then h_next = u + FFN(Norm(u)). Every sub-layer reads a normalized copy of the current state and writes its result back by addition.

Unrolling that recurrence reveals the central picture: the final representation is the initial embedding PLUS the sum of the contributions of every head and every feed-forward network across all layers. This invites reading the residual stream as a communication channel with no processing of its own. Each component reads from it through a linear projection and writes to it by adding. Because heads read and write in low-rank subspaces, two heads can communicate through different subspaces without interfering, and a head in an early layer can feed a head in a later layer (composition through the query, key, or value).

LayerNorm

Normalizes each vector by subtracting its mean and dividing by its standard deviation, then applies a learned gain and bias. Geometrically it projects onto the hyperplane orthogonal to the all-ones vector and rescales.

RMSNorm

Drops the mean-centering and normalizes only by the root-mean-square. Its authors hypothesize that re-centering invariance is dispensable and report comparable quality at lower runtime. Both forms are scale-invariant, and that invariance is what stabilizes training.

Feed-forward network (FFN)

Applied independently per position: W2·σ(W1·x). It can be read as a key–value memory — each row of W1 is a key matched against the input, each column of W2 a value added to the stream when the key fires. Attention moves information between positions; the FFN transforms it within a position.

Because bare attention is order-blind, position must be injected. The original approach adds a fixed sinusoidal vector to each embedding. It has a clean property — for a fixed shift k, the encoding of position p+k is a linear function of the encoding of position p — but adding position to content mixes the two inside the q·k dot product, producing cross terms.

Rotary positional encoding (RoPE) takes a different route: it does not add position, it rotates the query and the key by an angle proportional to their position, just before the dot product. The proven consequence (established, not conjectured) is that the resulting score depends on the two positions ONLY through their difference n − m, and is invariant to shifting all positions by a constant. Rotation also preserves vector norms. Three design facts follow: positional information enters multiplicatively in the QK circuit and never contaminates the values or the FFN; high-frequency pairs distinguish nearby positions while low-frequency pairs encode long distances; and the relative property alone does NOT guarantee generalization beyond training lengths, because unseen distances drive low-frequency rotations to angles the model never observed.

synthetic — RoPE relative-position property (schematic check)
# R(p) rotates a vector by an angle proportional to position p.
# PROVEN identity:  (R(m) q) . (R(n) k)  ==  q . R(n - m) k
# => the score depends on positions only through (n - m).

score_mn = dot(rope(q, pos=m), rope(k, pos=n))
score_shift = dot(rope(q, pos=m + c), rope(k, pos=n + c))
assert isclose(score_mn, score_shift)   # invariant to a common shift c
# Caveat (UNKNOWN in general): holds for seen relative distances;
# extrapolation to distances > training length is not guaranteed.

The projections cost on the order of n·d² operations per layer; the score matrix Q·Kᵀ and its multiplication by V cost on the order of n²·d. The attention matrix itself is n×n per head if materialized — this quadratic term is the structural cost of letting every position see every other. At autoregressive inference with a cache, generating one token costs on the order of n·d in attention because only one new row of scores is computed, so generating n tokens is on the order of n²·d overall, and the cache occupies two tensors (keys and values) per layer and position.

The insight behind FlashAttention is that on current accelerators the bottleneck of standard attention is not arithmetic but data movement between the large, slow high-bandwidth memory and the small, fast on-chip memory. The standard implementation writes the full n×n score matrix out and reads it back several times. The fix computes attention in blocks without ever materializing the full matrix, using an online-softmax recurrence that keeps a running maximum, a running normalizer, and an unnormalized accumulator, correcting the accumulator as each new block arrives.

1. 1Split keys and values into blocks and process them one block at a time, keeping per-query running statistics (max m, normalizer l, accumulator O).
2. 2For a new block, compute its scores, find the block maximum, and update the running maximum m.
3. 3Rescale the stored normalizer and accumulator by exp(old_max − new_max) so everything is expressed relative to the current maximum.
4. 4Add the new block's contribution to l and O, then move on; at the end divide O by l to get the exact output.
5. 5Because the result is algebraically identical to full softmax, FlashAttention is EXACT — it is not an approximation; it changes memory traffic (O(n²) down to O(n) extra memory), not the function.

Three consequences matter directly for anyone building an agent on top of these models. First, in a retrieval-augmented system, everything retrieved enters the SAME softmax as the instruction. Attention does not, by construction, distinguish "instruction" from "data" — that separation exists only to the extent that training induced it. This is the mechanistic root of indirect prompt injection: text placed in the context competes for attention with the system prompt. Provenance of retrieved content must therefore be enforced outside the model.

Second, the quadratic cost and attention dilution mean that stuffing more context is neither free nor neutral: as more positions carry similar scores, the weight on the truly relevant one drops unless its score advantage grows. Selecting and ordering what you retrieve is part of the design, not an afterthought.

Third, following a format shown earlier, or repeating a tool-call pattern, has a plausible mechanism in induction-style copying — the same mechanism that would also copy an UNWANTED pattern injected into the context. Engineering corollary, consistent with this lab's thesis: any conclusion that depends on the model having copied a datum correctly from context (an invoice number, an account id, a path) must be re-verified by deterministic code before it has effects.

• ▸Attention dilution: softmax spreads a total mass of one, so as the number of similar-scoring positions grows, the weight on the relevant one falls — long contexts can bury the needle.
• ▸Attention sinks: because a head cannot "look nowhere", trained heads may dump mass on content-free positions (often the first tokens) when there is nothing useful to copy; this is a measured behavior, not a bug you can assume away.
• ▸Positional extrapolation failure: relative encodings do not guarantee correct behavior beyond the training length; effective context must be MEASURED on your own task, not assumed from the nominal window.
• ▸The mechanistic circuit reading (QK/OV, residual stream) is precise in attention-only toy models and loses sharpness once MLPs and normalization are involved; specific-subspace claims need causal evidence.
• ▸The math of what each component computes is PROVEN; what a trained model actually does is largely UNKNOWN and an active research area — do not over-read an interpretation as a guarantee.

Expected direction: Without the scale, score variance grows with d_k, so the softmax saturates toward one-hot distributions; its Jacobian diag(a) − a·aᵀ collapses toward zero, so almost no gradient flows back to W_Q and W_K and those projections stop learning where to attend. A smaller learning rate does not address the cause — the gradient magnitude through a saturated softmax is already near zero regardless of step size — it just slows everything else down. The scale controls the starting temperature of the attention distribution so gradients stay informative at initialization; it is a conditioning fix, not a step-size issue.

Why is the output of an attention head always a convex combination of the value vectors, and what does that rule out?

Because the attention weights are non-negative and sum to one (they come from a softmax), the output lies within the convex hull of the values. It rules out amplification: attention can mix and move information between positions but cannot produce a value outside the range of the inputs it attends to.

What exactly does the "residual stream as a communication channel" reading claim, and where does it stop being exact?

It claims the final representation is the initial embedding plus the additive contributions of every head and FFN, each reading via a linear projection and writing by addition, so components communicate through low-rank subspaces. It is exact in the algebra except for normalization, which is non-linear; and asserting that a particular subspace carries a particular variable is a causal claim needing empirical evidence, not an algebraic consequence.

FlashAttention is often described as "faster attention". What does it actually change and not change?

It changes memory traffic and peak memory (avoiding materializing the n×n matrix, reducing extra memory from O(n²) to O(n)) by using an online-softmax recurrence computed in blocks. It does not change the output (it is exact, not an approximation) and does not reduce the FLOPs, which stay on the order of n²·d. Its wall-clock speedup is hardware- and length-dependent.

Why does the attention mechanism make indirect prompt injection a structural risk rather than a mere bug?

Because retrieved or tool-returned text enters the same softmax as the system instruction; attention has no built-in notion of "instruction" versus "data". Any separation is only as strong as what training induced, so injected text competes for attention on equal footing. The defense — enforcing provenance and re-verifying outside the model — must live outside the attention computation.