Chapter 17 Assembling a Complete Transformer Block

The attention output now has the same width as its input, and a causal mask prevents future access. Could we simply repeat that operation several times and call it a Transformer? Consider what is still missing. Without position information, an unrestricted attention calculation treats a reordered list as the same collection with reordered outputs. A mixture can also replace a useful input representation. And mixing between positions is not the same job as computing new nonlinear combinations within a position.

A complete Transformer block addresses these issues with several cooperating operations. We will first see what each operation repairs, then follow one position through a complete, explicitly specified block. The block uses the projections and two-head convention from Chapters 15–16 and Chapter 7’s distinction between a forward calculation and learning. All weights below are supplied teaching values.

17.1 Give equal tokens a way to occupy different positions

Example 17.1 — Reordering without position information. Take Chapter 15’s unrestricted attention and move input row 3 first, followed by rows 1 and 2. Because all positions share the same projection matrices, Q, K, and V reorder together. The output becomes old row 3, old row 1, old row 2. No numerical computation can infer “this row used to be last” from an index it was never given.

This property is called permutation equivariance: apply a reordering to the input, and the output follows the same reordering. It is not invariance, which would mean an unchanged output. The claim assumes all positions are visible, or that any visibility relation is reordered along with its positions. Holding a causal triangle fixed while moving tokens changes which tokens can interact; the same claim does not apply. Causal visibility supplies ordering constraints, but is not a promise that all useful positional relationships will be represented well.

A direct way to supply position is to give each slot another C-coordinate vector and add it to the token representation. Keep the Chapter 15 token matrix E, now explicitly naming it E to distinguish it from the block input. Supply this position table P:

$$E=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix},\qquad P=\begin{bmatrix}0&0\\0&1\\1&0\end{bmatrix},\qquad X=E+P=\begin{bmatrix}1&0\\0&2\\2&1\end{bmatrix}.$$

Each matrix has shape [3,2]. The second position becomes [0+0,1+1]=[0,2]. Addition preserves the two-coordinate width; concatenation would double it. A repeated token can now have different input vectors at different slots. The model still has to learn to use that distinction; these chosen coordinates do not encode a universal notion of distance.

P can be a learned table or a fixed rule. Our Mini GPT will use a learned finite table. A table with 32 rows has defined indices 0–31; it does not provide a trained row 32 merely because the rest of the code accepts another token. Even a position rule defined for larger indices does not by itself prove effective long-context behavior. Modern alternatives belong in Chapter 25.

17.2 Keep a direct route for the incoming representation

Suppose an attention sublayer computes an update U with the same shape as X. Instead of replacing X, add the update: $R=X+U$. This is a residual connection. At one position, [1,0]+[0,2]=[1,2]. The original two numbers have a direct additive route to R, while the sublayer can contribute a correction. If the sublayer returns zero, that route preserves X exactly.

The two arrays must refer to the same examples, positions, and coordinate width. Matching dimensions are necessary but not sufficient: adding another sample’s vectors would still be wrong. Broadcasting a single [1,2] row across all three positions is mathematically allowed by many libraries, but it is not the promised one-update-per-position residual. Our residual function rejects unequal shapes. The output projection of multi-head attention restores the required width before this addition.

A residual path does not guarantee that every old fact survives or that gradients never vanish. An update can cancel part of its input, and later operations can discard distinctions. It gives the computation an additive path, not a reliability certificate.

17.3 Normalize within one position, not across the sentence

The magnitudes entering successive transformations can differ. A normalization step offers a consistent within-row scale for the next sublayer, with learned adjustments. Start with row [1,3]: its mean is 2, its deviations are [-1,1], and the mean squared deviation is 1. Dividing by the square root of that variance gives [-1,1]. For an equal-valued row [10,10], however, the variance is zero, so we must specify how to avoid division by zero.

Layer normalization, or LayerNorm, adds a small positive epsilon to the variance before the square root, then permits a learned scale and offset for each coordinate. For a row x with C coordinates, our exact rule is:

$$\mu=\frac{1}{C}\sum_{c=1}^{C}x_c,\qquad \sigma^2=\frac{1}{C}\sum_{c=1}^{C}(x_c-\mu)^2,\qquad \operatorname{LN}(x)_c=\gamma_c\frac{x_c-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_c.$$

The subscript c selects a coordinate, not a position. We divide by C because we are defining the average squared deviation of these C values, not estimating a population variance from a sample with a C−1 correction. Gamma and beta each have C learned entries, shared across positions. Set gamma to all ones and beta to all zeros for this example, with epsilon=0.00001. Then [1,3] becomes approximately [-0.999995,0.999995]; [10,10] becomes [0,0]. With another beta, that constant row would become beta.

For [B,T,C], we normalize only the last C axis, separately for every example and position. Mixing future positions into these statistics could leak information even if attention were correctly masked. This LayerNorm computes statistics from the current row during both training and model inference; it does not keep a sentence-wide running average. Epsilon and the learned scales and offsets mean the final row need not have exactly unit variance or zero mean. LayerNorm formula and axes.

17.4 Compute new combinations inside each position

Attention communicates between positions. We also want to transform the coordinates at a single position nonlinearly. Chapter 7 already used a linear transform, a ReLU that replaces negative entries with zero, and another transform. Apply the same small network separately to every position, using shared weights. This is a position-wise feed-forward network, or FFN.

For a complete two-to-three-to-two example, give it row [-1,1], zero biases, and these matrices:

$$W_1=\begin{bmatrix}1&-1&0\\0&1&1\end{bmatrix},\qquad W_2=\begin{bmatrix}1&0\\0&1\\1&-1\end{bmatrix}.$$

The first multiplication gives [-1,2,1]: for instance its middle coordinate is $(-1)(-1)+1\times1=2$. ReLU changes that row to [0,2,1]. The second multiplication gives [1,1]: the last coordinate is $0\times0+2\times1+1\times(-1)=1$. Width expands from 2 to 3, then returns to 2 so a residual can be added. In general:

$$\operatorname{FFN}(x)=\operatorname{ReLU}(xW_1+b_1)W_2+b_2.$$

For model width C and hidden width F, W_1 is [C,F], b_1 [F], W_2 [F,C], and b_2 [C]. Each position uses the same matrices, but its own input values. Without the nonlinear operation, two such linear transformations would collapse into a single affine transformation: one matrix multiplication plus a bias. With ReLU, different coordinates can be removed for different inputs. This FFN has no operation that reads another position; it transforms whatever information the current row already contains.

17.5 Follow one complete pre-norm block

Example 17.2 — Every operation on the first position. Take X=E+P above. Use the two one-coordinate causal heads and W_O from Example 16.3, the FFN just specified, both LayerNorms with epsilon=0.00001, unit gamma and zero beta, and no dropout. Put each normalization before its sublayer. This is a pre-norm block:

$$N_1=\operatorname{LN}_1(X),\quad U=\operatorname{MHA}(N_1),\quad R=X+U,\quad N_2=\operatorname{LN}_2(R),\quad Y=R+\operatorname{FFN}(N_2).$$

The two LN subscripts distinguish separate parameter sets, not two different normalization axes. For X’s first row [1,0], the mean is 0.5, variance 0.25, and normalized row is [a,-a], where $a=0.5/\sqrt{0.25001}\approx0.999980$. At the first position, causality leaves only its own key. Each head’s sole allowed weight is exactly 1; their outputs are a and −a. Concatenation followed by W_O gives [0,2a], so the first residual is [1,2a].

That residual’s mean is $(1+2a)/2$. Let $u=(2a-1)/(2\sqrt{((2a-1)/2)^2+0.00001})\approx0.999980$. Its normalized row is [-u,u]. The FFN first produces [-u,2u,u], keeps [0,2u,u] after ReLU, and returns [u,u]. The final first row is therefore [1+u,2a+u], approximately [1.999980,2.999940]. Both residual additions and both normalizations have participated; this is not the earlier stand-alone attention output.

Table 17.1 — Complete measured scalar row trace, rounded to six decimals. The independent scalar route and NumPy route agree within absolute tolerance 0.000000000001.

Stage	First position	Whole-array shape
Token plus position, X	[1.000000,0.000000]	[3,2]
First norm, N_1	[0.999980,-0.999980]	[3,2]
Attention output, U	[0.000000,1.999960]	[3,2]
First residual, R	[1.000000,1.999960]	[3,2]
Second norm, N_2	[-0.999980,0.999980]	[3,2]
FFN after ReLU	[0.000000,1.999960,0.999980]	[3,3]
FFN output	[0.999980,0.999980]	[3,2]
Final block output, Y	[1.999980,2.999940]	[3,2]

The full output, including positions that consult several keys, is:

$$Y\approx\begin{bmatrix}1.999980&2.999940\\0.999912&1.476742\\2.999964&3.746396\end{bmatrix}.$$

Before the block

1. Stored model numbersToken table [V, C]; position table [T_max, C]
2. An operationLook up rows, then add → X [B, T, C]
3. A list of numbersExample X₁ = [1, 0]

Position is supplied here; it is not a new token appended to the sequence.

First residual path

1. An operationLN₁ over C → N₁ [B, T, C]
2. Stored model numbersQ/K/V projection weights [C, C]
3. An operationSplit → [B, H, T, d]; QKᵀ/√d → [B, H, T, T]
4. An operationMask → row Softmax → mix V [B, H, T, d]
5. An operationJoin → W_O → U [B, T, C]
6. A list of numbersADD X + U → R; R₁ = [1, 1.999960]

X goes directly to ADD as well as through the normalized attention branch.

Second residual path

1. An operationLN₂ over C → N₂ [B, T, C]
2. Stored model numbersW₁ [C, F], b₁ [F]; W₂ [F, C], b₂ [C]
3. An operationN₂W₁+b₁ → ReLU [B, T, F] → W₂,b₂ [B, T, C]
4. A list of numbersADD R + FFN → Y; Y₁ = [1.999980, 2.999940]

R also goes directly to the second ADD. Mini GPT applies a final norm after this block.

Figure 17.1 — Each sublayer returns to the residual width. Original teaching diagram, CC BY-SA 4.0. Values and rules are declared in the chapter and shared data.

Figure 17.1 shows the two residual routes separately. Position addition occurs before the block. Attention temporarily introduces head and source-position axes, then returns to [B,T,C]. The FFN expands the final coordinate axis, not the position axis. These distinctions let the diagram map to small_block in code/part-iv/mechanisms.py: norm1, attention, residual1, norm2, hidden, ffn, and output are named trace entries, not hidden framework stages. reference.py repeats the calculation using scalar loops independently of NumPy’s vectorized operations.

norm1 = layer_norm(x)
attention = two_heads(norm1)["output"]
r = residual(x, attention)
norm2 = layer_norm(r)
hidden = np.maximum(norm2 @ w1, 0)
y = residual(r, hidden @ w2)
assert y.shape == x.shape

This excerpt has the example’s zero biases. The complete function supplies the declared matrices and checks the residual shapes. The recorded trace is in data/part-iv/numpy-run.json. Later, a batch simply applies the position-wise operations independently to each row and preserves its own attention relation; it does not turn normalization or the FFN into cross-example mixing.

17.6 A different order is a different block

The original Transformer uses normalization after each sublayer’s residual addition, a post-norm order. Written with the same names for the operation types:

$$R_{\text{post}}=\operatorname{LN}_1\!\left(X+\operatorname{MHA}(X)\right),\qquad Y_{\text{post}}=\operatorname{LN}_2\!\left(R_{\text{post}}+\operatorname{FFN}(R_{\text{post}})\right).$$

Post-norm feeds X itself into attention; pre-norm feeds LN(X). Post-norm also normalizes the second sum before returning it. Moving the same boxes on a diagram therefore changes numerical results. They are not interchangeable spellings. The ordering also affects optimization, but no one small forward example establishes which trains better across tasks. Original post-norm order, pre/post comparison, Table 1.

Our Mini GPT chooses pre-norm and adds a final LayerNorm after its block, before the vocabulary head. That final norm is outside this block’s two residual paths. The original Transformer includes dropout; the bounded teaching project sets dropout to zero. A decoder-only block here has causal self-attention and an FFN; an encoder-decoder model additionally needs a decoder cross-attention sublayer to read the encoded source. “Block” is a compositional pattern, not a claim that every Transformer has identical internals.

17.7 What we can now locate

A token representation now receives a position vector, exchanges allowed information, preserves additive paths, normalizes its own coordinates, and undergoes a shared nonlinear transformation. Matching shapes permit those operations to connect, while axis meanings determine whether they connect correctly. None of these additions turns an untrained forward pass into a learned language model or proves reliable long-distance clue use.

The next architectural question is which positions should be visible and which outputs should be trained. The same general building blocks can support reading a whole input, continuing a prefix, or producing a new sequence conditioned on another one. Chapter 18 compares those choices before we build our small generator.

17.8 Exercises

1. E contains the same row [2,1] at two positions. Their position rows are [0,0] and [0,1]. Calculate both block inputs. Does their difference prove that the network learned order?
2. Normalize [2,4] using this chapter’s epsilon, gamma, and beta. Then normalize [20,20]. Explain the denominator’s coordinate count and the role of epsilon.
3. A tensor is [2,5,16]. A proposed LayerNorm averages across the five positions; an FFN maps 16 coordinates to 64 then back to 16. Identify which operation violates this block contract and whether a causal mask fixes it.
4. For the supplied W_1 and W_2, calculate the FFN result for [1,-1]. Compare it with the result for [-1,1]; identify the entries removed by ReLU.
5. An attention output [2,5,16] is added to a tensor [1,5,16]. A library broadcasts successfully. Why should this block’s residual assertion reject it?
6. A checkpoint uses pre-norm; someone loads its matrices into a post-norm implementation because every shape matches. Explain two changed calculations and what must be checked besides shapes.

17.9 Reference answers

1. Inputs become [2,1] and [2,2]. The model has been given a numerical distinction, not demonstrated learned use of it. It could ignore the differing coordinate or remove the distinction later. A controlled task with order changes would supply behavioral evidence.
2. Mean 3 and variance 1 give [-1,1]/sqrt(1.00001), approximately [-0.999995,0.999995]. For [20,20], both deviations are zero, so the result is [0,0]. Divide the squared deviations by two because there are two coordinates in this row. Epsilon makes the denominator positive even when variance is zero.
3. Averaging across positions violates coordinate-only LayerNorm and may allow a future token to change an earlier row. Attention masking cannot block that separate path. The FFN’s [2,5,16] → [2,5,64] → [2,5,16] shapes are appropriate if the same weights are applied independently to every position.
4. [1,-1] @ W_1 gives [1,-2,-1]; ReLU leaves [1,0,0]; multiplying W_2 gives [1,0]. For [-1,1], the first stage is [-1,2,1], then [0,2,1], then [1,1]. The negative entries differ, so the composed function is not a single linear map.
5. The second tensor has one example, not two separately computed updates. Broadcasting might accidentally reuse one sample’s representation in another sample. Equal full shapes are part of the residual contract; the successful library operation does not prove the intended example alignment.
6. Attention receives LN(X) in pre-norm but X in post-norm; the post-norm output also normalizes the final residual sum. Check operation order, parameter identities, epsilon, position policy, masking, and a known forward trace. A compatible collection of tensor sizes is insufficient to reproduce the checkpoint’s function.

17.10 References

• Vaswani et al. (2017), Attention Is All You Need. Sections 3.1, 3.3, and 3.5 locate residual/post-norm operations, position-wise FFNs, and position addition in the original architecture.
• Ba, Kiros, and Hinton (2016), Layer Normalization. Section 3 and Eq. 5 explain within-example statistics and learned gains/offsets; our epsilon and row arithmetic are explicit implementation choices.
• Xiong et al. (2020), On Layer Normalization in the Transformer Architecture. Table 1 gives the two operation orders and the final norm for a pre-norm stack. This chapter does not reproduce its optimization experiments.
• PyTorch, LayerNorm, 2.7. Specifies the final-axis statistics, variance denominator, and affine parameters checked against the explicit implementation.