Week 8

Week 8 — Transformers: build a trainable network around attention

Key questionHow do we combine token and position representations, causal multi-head attention, the FFN, LayerNorm, and residual paths into stackable GPT blocks?

Learning objectives

Trace a Pre-Norm block through the same [you,like] / [we,like] batch of [2,2,4] tensors.
Distinguish attention's cross-position mixing from the FFN's per-position feature mixing.
Check each residual addition, two-head output projection, 4→16→4 FFN, and final [2,2,5] vocabulary logits.
Connect decoder-only causal visibility to Week 6's loss/generation interface and prepare for the tokenizer/data pipeline.

135 min estimated reading time

Week progress: 0 of 20 sections (0%)Course progress: 0 of 321 sections (0%)

Attention can read context, but it is not a complete GPT. This week, distinguish the representation carried along the identity path from the update calculated by a sublayer. Prerequisites are Week 7's causal attention and arithmetic with means, squares, and square roots.

Scroll horizontally to view all columns.

Course data table
Study unit	The question we solve
1: Add position information	The same token ID retrieves the same initial vector; a position vector additionally identifies where it appears.
2: Transform features and add a bypass	The FFN combines features within a position. A residual connection adds the sublayer's update to the incoming representation.
3: Normalize and follow the order	Calculate LayerNorm by hand, then trace the two identity paths in x+attention(LN1(x)) followed by the FFN residual update.
4: Pass representations to the output head	Run the bridge experiment, distinguish block outputs from vocabulary logits, then compare with the full two-head worked example.

Run python week08_bridge.py from the examples directory. Start with the previous chapter's single head and add components one at a time. The full block below uses separately specified fixed two-head parameters; it is not a model trained in Week 7. Encoder/decoder categories, GELU's special function, and alternative position schemes are optional reading.

Alternate reading, hand calculation, and code changes. Units may take several sessions. Existing section numbers remain for links and reference; follow the displayed learning order rather than jumping around by the older numbers.

Week 8 learning goal: assemble attention into a repeatable block

Scroll horizontally to view all columns.

Course data table
Term used below	Plain-language meaning
Block	A repeatable unit of computation
per-token / per position	Apply the same operation separately at each position
residual	Carry the incoming representation and add the current update
pre-norm	Normalize a branch's input before computing its update and adding it back
didactic weights	Teaching weights chosen for hand calculation
identity / I	The identity transformation that passes values through unchanged, not token identity

First trace input → branch update → output and the shapes. GELU's error function and matrix derivatives of residuals can wait for a second reading. Understand why each component exists before trying to memorize every formula.

Fixed teaching batch: vocabulary 0=you, 1=like, 2=AI, 3=study, 4=we, with V_vocab=5. Prompt A IDs=[[0,1]] means “you like”; B IDs=[[4,1]] means “we like.” Set B=2, T=2, C=4, n_head=2, and head_size=D=2.

ids [B,T] = [2,2]
token embeddings [B,T,C] = [2,2,4] + position embeddings [T,C] = [2,4]
residual_0 [2,2,4]
residual_0 + Attention(LN1(residual_0)) → residual_after_attention [2,2,4]
residual_after_attention + FFN(LN2(residual_after_attention)) → residual_after_ffn [2,2,4]
contextual representation [2,2,4]
final LayerNorm + LM head → logits [B,T,V_vocab] = [2,2,5]

The two like@1 positions start with the same token-plus-position vector. Attention can read different position-0 information— you@0 or we@0—so residual_after_ffn[:,1,:] can differ. This alone does not promise a sensible continuation from an untrained model.

r_0=E_{\mathrm{token}}(\mathrm{ids})+E_{\mathrm{pos}}(0{:}T),\quad r_A=r_0+\operatorname{Attention}(\operatorname{LN}_1(r_0)),\quad r_F=r_A+\operatorname{FFN}(\operatorname{LN}_2(r_A))

Knowledge check

Which operation mixes you/we into the second position, and which only transforms that position's existing features?

1. What is still missing after attention?

Scroll horizontally to view all columns.

Course data table
Needed capability	Added component	Role at like@1
Token-position information	position embeddings P	Give the same token different input representations at different locations
Cross-position communication	causal multi-head Attention	Read you or we at position 0
Nonlinear feature calculation within a row	FFN + GELU	Transform its four features after contextual information has been incorporated
Identity and gradient paths	residual + x	Add a branch update to the incoming representation
Control branch-input scales	LayerNorm	Normalize each token row before it enters the branch

\operatorname{Block}:\mathbb{R}^{B\times T\times C}\to\mathbb{R}^{B\times T\times C}

Knowledge check

Which four kinds of components do we add around attention?

2. Token embedding lookup does not encode position

Scroll horizontally to view all columns.

Course data table
token	E[token]	position 0	position 1
like	[0.60, 0.30, -0.20, 0.10]	The same lookup row	The same lookup row

Without positional input or a mask, self-attention is permutation equivariant: permuting input rows permutes corresponding output rows. GPT's fixed causal mask ties visibility to sequence indices, so arbitrary token permutations no longer preserve that property. The mask introduces order-dependent visibility but does not add an explicit learned position coordinate. Our model uses position embeddings; other designs can use mechanisms such as RoPE.

E_{\mathrm{token}}(i)=E_{\mathrm{token}}(i)\ \text{at every position},\qquad \operatorname{token\_embedding}(\mathrm{ids}):[B,T,C]=[2,2,4]

Knowledge check

Why can the “like” lookup alone not tell the model that it is at position 1?

3. Position Embedding

The table below contains fixed teaching values illustrating shapes and information flow. Training learns parameter values from loss; no coordinate is manually given a semantic label such as “meaning of the first word.”

Scroll horizontally to view all columns.

didactic initial position table [T,C]=[2,4]
position t	P[t]
0	[0.05, 0.10, -0.05, 0.00]
1	[-0.10, 0.00, 0.05, 0.10]

Scroll horizontally to view all columns.

Course data table
prompt / token	E[token]	P[t]	x=E+P
A: you@0	[0.20,-0.10,0.70,0.30]	[0.05,0.10,-0.05,0.00]	[0.25,0.00,0.65,0.30]
A: like@1	[0.60,0.30,-0.20,0.10]	[-0.10,0.00,0.05,0.10]	[0.50,0.30,-0.15,0.20]
B: we@0	[-0.70,0.40,0.30,0.60]	[0.05,0.10,-0.05,0.00]	[-0.65,0.50,0.25,0.60]
B: like@1	[0.60,0.30,-0.20,0.10]	[-0.10,0.00,0.05,0.10]	[0.50,0.30,-0.15,0.20]

What happens when the same token occupies different positions?

E_{\mathrm{token}}[\text{like}]+P[0]=[0.60,0.30,-0.20,0.10]+[0.05,0.10,-0.05,0]=[0.65,0.40,-0.25,0.10]

E_{\mathrm{token}}[\text{like}]+P[1]=[0.60,0.30,-0.20,0.10]+[-0.10,0,0.05,0.10]=[0.50,0.30,-0.15,0.20]

The token embedding identifies “like”; the position embedding additionally supplies location information to the first block. The model need not reconstruct the two summands separately. Training can learn to use their combined features.

x[b,t,:]=E_{\mathrm{token}}[\mathrm{ids}[b,t]]+E_{\mathrm{pos}}[t]

Position rows [T,C]=[2,4] broadcast over the batch axis: token rows [2,2,4] plus position rows [2,4] give residual_0 [2,2,4]. The learned absolute table has only its configured positions; an out-of-range index has no row. Other position schemes, such as RoPE, address position information differently and are optional later reading.

Knowledge check

Why can P:[T,C]=[2,4] be added to [B,T,C]=[2,2,4]?

4. The block's two main computations

Attention mixes positions: query t combines values at allowed j≤t. The FFN applies the same nonlinear function separately to x[b,t,:], without directly reading x[b,j,:] for j≠t. “Communication versus local transformation” is a mnemonic for operations, not literal thought.

\operatorname{Attention}(x)_{b,t,:}=\sum_{j\le t}\alpha_{b,t,j}\,\mathrm{value\_states}_{b,j,:},\qquad \operatorname{FFN}(x)_{b,t,:}=f(x_{b,t,:})

z = LN1(x) [2,2,4]
Q, K, value_states [2,2,4] each
reshape + transpose each → [B,H,T,D] = [2,2,2,2]
scores = Q @ K.transpose(-2,-1) [B,H,T,T] = [2,2,2,2]
causal mask + Softmax over key axis [2,2,2,2]
head outputs = weights @ value_states [2,2,2,2]
transpose + concatenate heads [B,T,H·D] = [2,2,4]
output projection [B,T,C] = [2,2,4]

Knowledge check

Can the FFN at like@1 directly inspect we@0?

5. Feed-Forward Network

Think of I₄ as a four-feature pass-through matrix: ones on the main diagonal, zeros elsewhere, so xI₄=x. We use it to make the FFN arithmetic readable—not to prescribe fixed identity weights in a trained model. The FFN maps four features to sixteen intermediate features, applies a nonlinearity, then combines them into four outputs.

Attention can give A/B's like@1 different contextual information. The same FFN weights then process all four token rows separately. Different inputs can give different outputs, but positions do not have separate FFN parameters and the FFN does not directly read other rows.

\operatorname{FFN}(z)=\operatorname{GELU}(zW_1+b_1)W_2+b_2

Scroll horizontally to view all columns.

Course data table
stage	row-vector parameter / tensor shape	visible batch shape
first Linear	W1:[C,4C]=[4,16]	[2,2,4] → [2,2,16]
GELU	no learned shape change	[2,2,16] → [2,2,16]
second Linear	W2:[4C,C]=[16,4]	[2,2,16] → [2,2,4]

Change one weight to see what feature mixing means

An identity matrix passes features separately, which is useful for checking numbers but hides mixing. In a separate experiment, use z=[1,2,0,0] and change only W1[1,0] from 0 to 1. The first hidden pre-activation changes from z0 to z0+z1=3.

Scroll horizontally to view all columns.

Course data table
Quantity	Original W1[1,0]=0	After changing it to 1
First hidden value before activation	1	1+2=3
After GELU	Approximately 0.841345	Approximately 2.995950
First output after multiplication by 0.25 in the second layer	Approximately 0.210336	Approximately 0.748988

The output still has four values, but its first value now depends on two input features. That is feature mixing, not an effect of shape expansion alone. Other outputs remain unchanged here. This controlled experiment does not replace the fixed weights used in the later full-block calculation.

Attention as a whole can also be nonlinear because its weights depend on the input through Q/K and Softmax. The FFN provides an additional explicit per-position nonlinear feature transformation. The division of roles does not imply that attention is purely linear in its input.

Knowledge check

Why does the second Linear return from sixteen features to four?

6. Why use GELU?

(zW_1+b_1)W_2+b_2=z(W_1W_2)+(b_1W_2+b_2)

\operatorname{GELU}(x)=x\Phi(x)=\frac{x}{2}\left(1+\operatorname{erf}\!\left(\frac{x}{\sqrt2}\right)\right)

Scroll horizontally to view all columns.

GELU is a continuous smooth nonlinearity, not a binary switch.
Input x	Approximate GELU(x)	Intuitive comparison with ReLU
−1	−0.159	Negative input is not forced to zero
0	0	Zero still maps to zero
1	0.841	Positive input is smoothly gated
2	1.955	A large positive input is nearly unchanged

one token row z [4]
Linear W1 → hidden [16]
GELU → hidden [16]
Linear W2 → output [4]
applied independently to every [B,T] location → [2,2,4]

GELU responds smoothly; it is not a binary condition and does not mix token positions. Other architectures use different activations or gated FFNs. This block's essential point is the nonlinear step between projections.

text

Prompt A's final position after the second LayerNorm:
z = [1.2476, 0.2644, -1.5404, 0.0284]

First four hidden pre-activations under the teaching W1:
[1.2476, 0.2644, -1.5404, 0.0284]

After GELU:
[1.1152, 0.1596, -0.0952, 0.0144]

After the teaching W2=0.25I:
FFN update = [0.2788, 0.0399, -0.0238, 0.0036]

Knowledge check

What precisely changes if we remove GELU?

7. Residual Connection

Start with a scalar: y=x+F(x). If F(x)=0.1x, then y=1.1x; increasing x by 0.01 increases y by 0.011. The direct path contributes derivative 1 and the branch contributes 0.1, totaling 1.1. The identity matrix I is the multidimensional counterpart of this direct derivative. Residual paths do not guarantee stable gradients in every network.

For A, let a=Attention(LN1(residual_0)); then residual_after_attention=residual_0+a. The same rule for B adds a contextual update that can incorporate its different we@0 information. Backward adds the path contributions, echoing Week 4. This structure is helpful but is not a universal optimization guarantee.

y=x+F(x),\qquad \frac{\partial y}{\partial x}=I+\frac{\partial F(x)}{\partial x}

Scroll horizontally to view all columns.

Course data table
term	shape in this lesson
residual_0	[B,T,C]=[2,2,4]
Attention(LN1(residual_0))	[B,T,C]=[2,2,4]
residual_after_attention	[B,T,C]=[2,2,4]
FFN(LN2(residual_after_attention))	[B,T,C]=[2,2,4]
residual_after_ffn	[B,T,C]=[2,2,4]

text

Prompt A's final “like” position:
residual_0       = [ 0.5000, 0.3000, -0.1500,  0.2000]
attention_update = [ 0.9944, 0.1092, -1.4326, -0.0512]

Elementwise addition:
residual_after_attention
= [1.4944, 0.4092, -1.5826, 0.1488]

Section 10 derives that attention update from LN1, both heads, Softmax weights, and values. A residual addition is neither replacement nor concatenation: it adds the branch update to the old state element by element. The sum can change or even cancel individual old values.

Knowledge check

Which two contributions form the residual output?

8. Residual shape requirements

text

residual_0                                      [2,2,4]
Attention(LN1(residual_0))                    [2,2,4]
residual_after_attention                      [2,2,4]

FFN(LN2(residual_after_attention))            [2,2,4]
residual_after_ffn                            [2,2,4]

Causal attention mixes positions. Each head produces [B,T,D]=[2,2,2]; concatenating H=2 heads gives [2,2,4], and the output projection preserves that shape. The FFN instead mixes row features, expands temporarily to [2,2,16], and uses W2 to return to [2,2,4].

\operatorname{shape}(x)=\operatorname{shape}(F(x))=[B,T,C],\qquad H\cdot D=2\cdot2=C=4

Knowledge check

Why can FFN hidden state [2,2,16] not be directly added to [2,2,4]?

9. Layer normalization intuition

Work through four numbers to make the formula concrete

Scroll horizontally to view all columns.

Course data table
One row [1,2,3,4]	Calculation	Result
Find the center	(1+2+3+4)/4	μ=2.5
Subtract the center from each value	[1,2,3,4]−2.5	[-1.5,-0.5,0.5,1.5]
Square deviations and average	(2.25+0.25+0.25+2.25)/4	Variance=1.25
Take the square root to obtain a scale	√1.25	Approximately 1.118
Divide deviations by the scale	Temporarily ignore the tiny ε	Approximately [−1.342,−0.447,0.447,1.342]

LayerNorm uses the population variance here: divide by four, not the sample-variance denominator three. A small positive ε is inside the square root to avoid division by zero for a constant row. Learned γ and β then scale and shift the normalized values; their usual initial values are 1 and 0.

This is a separate arithmetic exercise, not a replacement for the fixed token rows. LayerNorm computes the mean and scale over each position's own C features, not across batch examples or all positions in a column.

For the illustrative row [1,2,3,4], μ=2.5. Subtract μ, divide by sqrt(variance+ε), then apply γ and β. A's you@0 and like@1 and B's two rows each use their own statistics.

\mu_{b,t}=\frac{1}{C}\sum_{c=1}^{C}x_{b,t,c},\quad \sigma^2_{b,t}=\frac{1}{C}\sum_{c=1}^{C}(x_{b,t,c}-\mu_{b,t})^2,\quad \operatorname{LN}(x)_{b,t,c}=\gamma_c\frac{x_{b,t,c}-\mu_{b,t}}{\sqrt{\sigma^2_{b,t}+\varepsilon}}+\beta_c

text

Start with gamma=[1,1,1,1] and beta=[0,0,0,0], ignoring the tiny epsilon effect in the displayed approximation:

x         = [1, 2, 3, 4]
mean      = (1+2+3+4) / 4 = 2.5
deviation = [-1.5, -0.5, 0.5, 1.5]
variance  = (2.25+0.25+0.25+2.25) / 4 = 1.25
std       = sqrt(1.25) ≈ 1.118

normalized
= deviation / std
≈ [-1.342, -0.447, 0.447, 1.342]

Scroll horizontally to view all columns.

Course data table
normalization	statistics are computed over	train / eval behavior
LayerNorm here	each x[b,t,:] row over its C=4 features	same per-row rule for one generation prompt or a batch
BatchNorm (typical)	per channel over batch and often spatial/time examples	keeps running statistics for evaluation

BatchNorm is useful in other settings. This LayerNorm design is convenient for autoregressive text models because a row's normalization does not depend on other batch examples or sequence length. Input and output are [2,2,4]; γ and β are [4].

Why normalize? First observe a scale difference

Compare x=[1,2,3,4] with 100x+1000=[1100,1200,1300,1400]. Their relative order is the same but their scale and offset differ. LayerNorm lets a branch primarily see deviations relative to its row's mean and spread.

Scroll horizontally to view all columns.

Course data table
Quantity	x	100x+1000
Mean	2.5	1250
Population variance	1.25	12500
After normalization (approximately)	[−1.342,−0.447,0.447,1.342]	[−1.342,−0.447,0.447,1.342]

Epsilon prevents the normalized results from being exactly identical. Normalization is not information-preserving for every task: overall scale or offset may matter. In Pre-Norm, the identity path still carries raw x, while gamma/beta adjust the normalized branch. Scale control does not guarantee faster training or decreasing loss at every step.

Knowledge check

Which values determine the LayerNorm mean of x[1,0,:]?

10. Pre-Norm Transformer Block

residual_0 [2,2,4] → LN1 → normalized_for_attention [2,2,4]
causal MHA → attention_update [2,2,4]
residual_0 + attention_update → residual_after_attention [2,2,4]
LN2 → normalized_for_ffn [2,2,4]
FFN 4→16→4 → ffn_update [2,2,4]
residual_after_attention + ffn_update → residual_after_ffn [2,2,4]
Final LayerNorm → h [2,2,4]

r_A=r_0+\operatorname{Attention}(\operatorname{LN}_1(r_0)),\qquad r_F=r_A+\operatorname{FFN}(\operatorname{LN}_2(r_A))

The full teaching attention uses simple projections: Head 1 selects normalized features 0–1, Head 2 selects features 2–3, and within each head Q=K=V. W_O is identity. It still calculates QKᵀ, scaling, masking, Softmax, and weighted values; only the projections are simplified.

q^{(1)}=k^{(1)}=v^{(1)}=[z_0,z_1],\qquad q^{(2)}=k^{(2)}=v^{(2)}=[z_2,z_3],\qquad W_O=I_4

Scroll horizontally to view all columns.

The final like@1 LN1 row is identical in both prompts; position 0 differs.
LN1 Row	Prompt A	Prompt B
Position 0	[-0.2156,-1.2939,1.5095,0.0000]	[-1.6731,0.6591,0.1521,0.8619]
Position 1: like	[1.2206,0.3715,-1.5390,-0.0531]	[1.2206,0.3715,-1.5390,-0.0531]

text

Prompt A, Head 1: scores for the final query
q = [1.2206, 0.3715]
k_you = [-0.2156, -1.2939]
k_like = [1.2206, 0.3715]

score_you
= (1.2206×-0.2156 + 0.3715×-1.2939) / sqrt(2)
≈ -0.5260

score_like
= (1.2206×1.2206 + 0.3715×0.3715) / sqrt(2)
≈ 1.1511

Softmax([-0.5260, 1.1511])
≈ [0.1575, 0.8425]

Scroll horizontally to view all columns.

Each head normalizes over the two visible positions for the final query. Different position-0 rows produce different attention weights.
Final Query	Scaled Scores Head 1	Weights Head 1	Scaled Scores Head 2	Weights Head 2
Prompt A	[-0.5260,1.1511]	[0.1575,0.8425]	[-1.6427,1.6768]	[0.0349,0.9651]
Prompt B	[-1.2709,1.1511]	[0.0815,0.9185]	[-0.1979,1.6768]	[0.1330,0.8670]

text

Prompt A's values are the LN1 features selected by each head:

Head 1 output
= 0.1575×[-0.2156,-1.2939] + 0.8425×[1.2206,0.3715]
≈ [0.9944, 0.1092]

Head 2 output
= 0.0349×[1.5095,0.0000] + 0.9651×[-1.5390,-0.0531]
≈ [-1.4326, -0.0512]

Concatenate both heads; teaching W_O=I_4:
attention_update ≈ [0.9944,0.1092,-1.4326,-0.0512]

text

Prompt A final row
attention_update         = [ 0.9944, 0.1092, -1.4326, -0.0512]
residual_0               = [ 0.5000, 0.3000, -0.1500,  0.2000]
residual_after_attention = [ 1.4944, 0.4092, -1.5826,  0.1488]
LN2                      = [ 1.2476, 0.2644, -1.5404,  0.0284]
FFN update               = [ 0.2788, 0.0399, -0.0238,  0.0036]
residual_after_ffn       = [ 1.7732, 0.4492, -1.6064,  0.1524]
final LayerNorm h        = [ 1.3128, 0.2134, -1.4933, -0.0330]

Prompt B final row
attention_update         = [ 0.9847, 0.3949, -1.3141,  0.0686]
residual_0               = [ 0.5000, 0.3000, -0.1500,  0.2000]
residual_after_attention = [ 1.4847, 0.6949, -1.4641,  0.2686]
LN2                      = [ 1.1475, 0.4158, -1.5843,  0.0209]
FFN update               = [ 0.2508, 0.0687, -0.0224,  0.0027]
residual_after_ffn       = [ 1.7356, 0.7637, -1.4865,  0.2713]
final LayerNorm h        = [ 1.2100, 0.3787, -1.5462, -0.0425]

The prompts start with the same final residual_0 row. Their final-position states first differ after attention reads different position-0 information; those differences then pass through the residual sum, LN2, FFN, and final LayerNorm. Pre-Norm supplies row-normalized branch inputs while its direct paths carry residual states without this branch's normalization.

Knowledge check

What does the direct path carry in r_A=r_0+Attention(LN1(r_0))?

11. A complete Transformer block

python

import math

import torch
import torch.nn.functional as F


torch.set_printoptions(precision=4, sci_mode=False)

ids = torch.tensor([
    [0, 1],  # you like
    [4, 1],  # we like
])

token_table = torch.tensor([
    [ 0.20, -0.10,  0.70,  0.30],  # you
    [ 0.60,  0.30, -0.20,  0.10],  # like
    [-0.40,  0.80,  0.50, -0.30],  # AI
    [ 0.10,  0.20,  0.90,  0.40],  # study
    [-0.70,  0.40,  0.30,  0.60],  # we
])

position_table = torch.tensor([
    [ 0.05, 0.10, -0.05, 0.00],
    [-0.10, 0.00,  0.05, 0.10],
])

B, T = ids.shape
C, num_heads = 4, 2
head_size = C // num_heads

residual_0 = F.embedding(ids, token_table) + position_table[:T]

# Teaching LayerNorm: gamma=1, beta=0, epsilon=1e-5.
normalized_for_attention = F.layer_norm(
    residual_0,
    normalized_shape=(C,),
    eps=1e-5,
)

# Teaching Q/K/V projections:
# Head 1 selects channels 0:2; Head 2 selects channels 2:4.
heads = normalized_for_attention.view(B, T, num_heads, head_size)
heads = heads.transpose(1, 2)  # [B,H,T,D]
q = heads
k = heads
value_states = heads

scores = (q @ k.transpose(-2, -1)) / math.sqrt(head_size)
causal_mask = torch.tril(torch.ones(T, T, dtype=torch.bool))
masked_scores = scores.masked_fill(~causal_mask, float("-inf"))
attention_probs = F.softmax(masked_scores, dim=-1)

head_outputs = attention_probs @ value_states
attention_update = head_outputs.transpose(1, 2).contiguous().view(B, T, C)
# Teaching W_O is I_4, so the projection leaves attention_update unchanged.

residual_after_attention = residual_0 + attention_update
normalized_for_ffn = F.layer_norm(
    residual_after_attention,
    normalized_shape=(C,),
    eps=1e-5,
)

# Sparse teaching FFN:
# W1 copies four channels into the first four of sixteen hidden units.
hidden = torch.zeros(B, T, 4 * C)
hidden[..., :C] = normalized_for_ffn
hidden = F.gelu(hidden, approximate="none")

# W2 selects those four units and multiplies them by 0.25.
ffn_update = 0.25 * hidden[..., :C]
residual_after_ffn = residual_after_attention + ffn_update

final_hidden = F.layer_norm(
    residual_after_ffn,
    normalized_shape=(C,),
    eps=1e-5,
)

# Reuse Week 6's five candidate scoring rules.
W_out = torch.tensor([
    [1.0, 0.0,  0.0, 0.0],
    [0.0, 1.0,  2.0, 2.0],
    [0.0, 0.0,  1.0, 1.0],
    [0.0, 3.0, -1.0, 0.0],
    [1.0, 2.0,  0.0, 0.0],
])
bias = torch.tensor([-0.2, 0.1, 0.0, 0.0, 0.0])

logits = final_hidden @ W_out.T + bias
vocabulary_probs = F.softmax(logits, dim=-1)

print("final attention probabilities:", attention_probs[:, :, -1, :])
print("attention update:", attention_update[:, -1, :])
print("residual after attention:", residual_after_attention[:, -1, :])
print("FFN update:", ffn_update[:, -1, :])
print("final hidden:", final_hidden[:, -1, :])
print("final logits:", logits[:, -1, :])
print("final vocabulary probabilities:", vocabulary_probs[:, -1, :])

python

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalSelfAttention(nn.Module):
    def __init__(self, model_dim: int, num_heads: int, dropout: float):
        super().__init__()
        if model_dim % num_heads != 0:
            raise ValueError("model_dim must be evenly divisible by num_heads")
        self.num_heads = num_heads
        self.head_size = model_dim // num_heads
        self.qkv = nn.Linear(model_dim, 3 * model_dim, bias=False)
        self.output = nn.Linear(model_dim, model_dim, bias=False)
        self.attention_dropout = nn.Dropout(dropout)
        self.output_dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, time_steps, model_dim = x.shape
        q, k, value_states = self.qkv(x).chunk(3, dim=-1)
        q = q.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)
        k = k.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)
        value_states = value_states.view(
            batch_size, time_steps, self.num_heads, self.head_size
        ).transpose(1, 2)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_size)
        causal = torch.ones(time_steps, time_steps, device=x.device, dtype=torch.bool).tril()
        scores = scores.masked_fill(~causal, float("-inf"))

        attention_probs = F.softmax(scores, dim=-1)
        dropped_probs = self.attention_dropout(attention_probs)
        heads = dropped_probs @ value_states

        merged = heads.transpose(1, 2).contiguous().view(batch_size, time_steps, model_dim)
        return self.output_dropout(self.output(merged))


class TransformerBlock(nn.Module):
    def __init__(self, model_dim: int, num_heads: int, dropout: float):
        super().__init__()
        self.ln1 = nn.LayerNorm(model_dim)
        self.attention = CausalSelfAttention(model_dim, num_heads, dropout)
        self.ln2 = nn.LayerNorm(model_dim)
        self.ffn = nn.Sequential(
            nn.Linear(model_dim, 4 * model_dim),
            nn.GELU(),
            nn.Linear(4 * model_dim, model_dim),
            nn.Dropout(dropout),
        )

    def forward(self, residual: torch.Tensor) -> torch.Tensor:
        attention_update = self.attention(self.ln1(residual))
        residual_after_attention = residual + attention_update

        ffn_update = self.ffn(self.ln2(residual_after_attention))
        residual_after_ffn = residual_after_attention + ffn_update
        return residual_after_ffn


sample = torch.randn(2, 2, 4)
block = TransformerBlock(model_dim=4, num_heads=2, dropout=0.1)
block.eval()
output = block(sample)
print(output.shape)  # torch.Size([2, 2, 4])

z = LN1(x) [2,2,4]
Q, K, value_states [2,2,4] each
reshape + transpose each → [B,H,T,D] = [2,2,2,2]
scores = Q @ K.transpose(-2,-1) [B,H,T,T] = [2,2,2,2]
causal mask + Softmax over key axis [2,2,2,2]
head outputs = weights @ value_states [2,2,2,2]
transpose + concatenate heads [B,T,H·D] = [2,2,4]
output projection [B,T,C] = [2,2,4]

ln1/ln2 normalize each token row without changing [2,2,4]. Attention mixes positions and returns [2,2,4]; the FFN applies 4→16→4 within each row. Both residual additions combine matching [2,2,4] tensors; the second result is the block output. eval() disables Dropout. In this fixed CPU example with unchanged parameters, repeated evaluation gives the same outputs; this is not a guarantee of bitwise determinism for every backend or sampling procedure.

Knowledge check

Which two lines perform residual addition, and what shapes do they return?

12. Stack multiple blocks

x^(0) = token + position embeddings [2,2,4]
Block^(0)(x^(0)) → x^(1) contextual states [2,2,4]
Block^(1)(x^(1)) → x^(2) richer contextual states [2,2,4]
Final LayerNorm(x^(2)) → h [2,2,4]

x^{(\ell+1)}=\operatorname{Block}^{(\ell)}(x^{(\ell)}),\qquad \ell=0,\ldots,L-1

Each block has its own LayerNorm, Q/K/V/output projections, and FFN parameters unless weight sharing is explicitly designed. Training dropout does not change the shapes. n_layer counts blocks, not heads.

Each Pre-Norm branch receives a normalized input, while the identity paths carry the residual stream. This architecture applies a final LayerNorm after the stack to prepare the LM head's input. It is not a third residual branch.

Knowledge check

What input does Block 2 consume?

13. From Transformer output to vocabulary logits

After final LayerNorm, each vocabulary candidate has its own four-feature scoring vector. Its logit uses all four coordinates of the current h row in a dot product, plus its bias. The four features do not separately correspond to five scores.

z_v=w_v\cdot h+b_v=\sum_{c=1}^{C}w_{v,c}h_c+b_v

Scroll horizontally to view all columns.

The same teaching output head processes every position in both prompts.
Candidate v	w_v reused from Week 6	b_v
you	[1,0,0,0]	−0.2
like	[0,1,2,2]	0.1
AI	[0,0,1,1]	0
study	[0,3,−1,0]	0
we	[1,2,0,0]	0

text

Prompt A final hidden
h_A = [1.3128, 0.2134, -1.4933, -0.0330]

For example, the logit for study:
z_study = [0,3,-1,0] · h_A + 0
        = 3(0.2134) - (-1.4933)
        ≈ 2.1336

All logits [you,like,AI,study,we]
z_A = [1.1128, -2.7390, -1.5262, 2.1336, 1.7397]

Prompt B final hidden
h_B = [1.2100, 0.3787, -1.5462, -0.0425]
z_B = [1.0100, -2.6987, -1.5887, 2.6821, 1.9674]

Displayed values are rounded; the code calculates with unrounded intermediate values.

W_{\mathrm{out}}\in\mathbb{R}^{V_{\mathrm{vocab}}\times C}=[5,4],\qquad \mathrm{logits}=H W_{\mathrm{out}}^{\mathsf T}+b

Scroll horizontally to view all columns.

Different contexts can produce different distributions even if argmax happens to select the same token.
Prompt final row	Vocabulary probabilities [you,like,AI,study,we]	Highest-scoring candidate
A	[0.1742,0.0037,0.0124,0.4835,0.3261]	study
B	[0.1108,0.0027,0.0082,0.5897,0.2886]	study

Scroll horizontally to view all columns.

Course data table
quantity	shape
h	[B,T,C]=[2,2,4]
W_out	[V_vocab,C]=[5,4]
b_vocab	[V_vocab]=[5]
logits	[B,T,V_vocab]=[2,2,5]

During training, the LM head scores every position to form B×T next-token tasks. A generation step uses only logits[:,-1,:]. Logits are raw scores; vocabulary Softmax converts them into probabilities.

Knowledge check

Why does [2,2,4] become [2,2,5]?

14. Shapes through a complete GPT forward pass

text

ids A/B                                      [2,2]
token embeddings                             [2,2,4]
position embeddings                          [2,4] (broadcast across B)
residual_0 = token + position                 [2,2,4]
one or more Transformer blocks               [2,2,4]
final LayerNorm                              [2,2,4]
LM head                                      [2,2,5]

training: logits.reshape(B*T,V_vocab)        [4,5]
          targets.reshape(B*T)                [4]
          cross entropy                       scalar

generation: logits[:, -1, :]                 [2,5]
            sampled/argmax next_id            [2,1]
            appended ids                       [2,3]

python

class MiniGPT(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        model_dim: int,
        num_heads: int,
        num_layers: int,
        dropout: float,
    ):
        super().__init__()
        self.context_length = context_length
        self.token_embedding = nn.Embedding(vocab_size, model_dim)
        self.position_embedding = nn.Embedding(context_length, model_dim)
        self.blocks = nn.ModuleList([
            TransformerBlock(model_dim, num_heads, dropout)
            for _ in range(num_layers)
        ])
        self.final_norm = nn.LayerNorm(model_dim)
        self.lm_head = nn.Linear(model_dim, vocab_size)

    def forward(
        self,
        token_ids: torch.Tensor,
        targets: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        B, T = token_ids.shape
        if T > self.context_length:
            raise ValueError("Sequence length exceeds context_length")

        positions = torch.arange(T, device=token_ids.device)
        residual = (
            self.token_embedding(token_ids)
            + self.position_embedding(positions)
        )

        for block in self.blocks:
            residual = block(residual)

        final_hidden = self.final_norm(residual)
        logits = self.lm_head(final_hidden)

        if targets is None:
            return logits, None

        loss = F.cross_entropy(
            logits.reshape(B * T, logits.size(-1)),
            targets.reshape(B * T),
        )
        return logits, loss


model = MiniGPT(
    vocab_size=5,
    context_length=8,
    model_dim=4,
    num_heads=2,
    num_layers=2,
    dropout=0.1,
)

inputs = torch.tensor([
    [0, 1],  # you like
    [4, 1],  # we like
])
targets = torch.tensor([
    [1, 2],  # like AI
    [1, 0],  # like you; same target as Week 7's Prompt B
])

logits, loss = model(inputs, targets)
print(logits.shape)  # torch.Size([2, 2, 5])
print(loss.shape)    # torch.Size([])

\mathcal{L}=\operatorname{CrossEntropy}(\operatorname{reshape}(\mathrm{logits},[BT,V_{\mathrm{vocab}}]),\operatorname{reshape}(\mathrm{targets},[BT]))

T=2 is the illustration's window. Generation appends one [B,1] token-ID column to the stored sequence, then crops an overlong model input to the configured context size before the next forward pass.

Knowledge check

For B=2, T=2, V_vocab=5, which shapes enter cross-entropy?

15. Encoders, decoders, and decoder-only models

Ask what input the task permits: whole-sentence classification can use the complete sentence; left-to-right continuation can use only its existing prefix; translation can let the generating side read the full source and the generated target prefix. These visibility patterns explain the architecture labels. A decoder block is not the tokenizer.decode function.

Scroll horizontally to view all columns.

Course data table
family	permitted context	typical role / interface
encoder-only	Tokens can read bidirectional context in the input sequence	representation / understanding tasks
encoder-decoder	The encoder reads the source bidirectionally; the decoder reads its target prefix causally and cross-attends to the source	input sequence → generated output sequence
decoder-only GPT	Position t reads only j≤t in the same prompt	[B,T,V_vocab] next-token logits；last-position generation

\mathrm{decoder\text{-}only}:\qquad P(x_{1:T})=\prod_{t=1}^{T}P(x_t\mid x_{<t})

For [you,like], a decoder-only model's second position can read positions 0 and 1, not the future. Scores have [B,H,T,T]=[2,2,2,2]; a lower-triangular [T,T] mask can broadcast over batch and heads. “Decoder-only” is an architecture label, not a claim that the model cannot build prompt representations or merely decodes IDs.

Knowledge check

Which family matches Mini GPT, and why?

16. Where does a Transformer's capability come from?

data / tokenizer produce IDs
embeddings + positions create [B,T,C]
causal Attention exchanges allowed token information
FFN + GELU transform local channels; residual/norm support depth
LM head exposes vocabulary logits
loss + backprop + optimizer update parameters

\theta\leftarrow\theta-\eta\nabla_{\theta}\mathcal{L}

θ includes token/position tables, attention projections, FFN, LayerNorm, and LM-head parameters. In A/B, attention lets you/we affect like@1; position vectors supply location information; the FFN and successive blocks transform states. Data, loss, and optimization determine what parameter values are learned.

Knowledge check

Why is attention necessary but not sufficient for the architecture taught here?

17. Concepts to recognize but not study in detail yet

Scroll horizontally to view all columns.

Course data table
advanced concept	What it changes—and what it does not
RoPE / relative positions	Changes how position enters Q/K; does not remove the need to account for order
weight tying	Shares input-embedding and LM-head weights; preserves the logits interface
Flash Attention	Changes implementation and resource use; preserves the intended attention calculation up to numerical differences
KV cache	Reuses prior keys and values during generation; preserves causal visibility
padding / masking details	Handles unequal sequence lengths within a batch; serves a different purpose from causal masking
gated FFNs / other norms	Changes the local feature branch; its output still needs the [B,T,C] interface

\mathrm{cache}_t=\{K_{\le t},V_{\le t}\}

A KV-cache update stores the new position's keys and values at each layer rather than recomputing all earlier ones. Attention must still use only the permitted prefix. RoPE and learned absolute position tables are different position strategies; this course does not require them together.

Knowledge check

Does a KV cache change which prior positions causal attention permits?

18. Seven essential Week 8 ideas

Token embeddings represent identity and position embeddings supply location; their sum is still [B,T,C].
Causal multi-head attention is the operation that mixes token positions within this basic block.
With H=2 and D=2, head tensors are [B,H,T,D], concatenation gives [B,T,4], and the output projection returns [B,T,C].
The FFN applies the same nonlinear [4]→[16]→[4] transformation independently to each row.
Each residual addition combines matching [B,T,C] tensors and provides a direct identity/gradient path.
Pre-Norm performs normalization → branch computation → residual addition twice. LayerNorm uses one token row's features, not other batch rows.
Blocks stack as [B,T,C]→[B,T,C]. The LM head returns [B,T,V_vocab]; training uses all positions, while each generation step uses the last.

\operatorname{Concat}(\mathrm{head}_1,\mathrm{head}_2)\in\mathbb{R}^{B\times T\times(H D)}=\mathbb{R}^{2\times2\times4}

Knowledge check

Classify attention, the FFN, LayerNorm, residual addition, and the LM head by their jobs.

19. Week 8 → Week 9

raw text
tokenizer-selected tokens
vocabulary IDs [B,T]
token + position embeddings [B,T,C]
Transformer
logits [B,T,V_vocab]

\mathrm{text}\xrightarrow{\mathrm{tokenizer}}\mathrm{ids}\in\{0,\ldots,V_{\mathrm{vocab}}-1\}^{B\times T}

The [you,like] and [we,like] examples assume a tokenizer mapping these words into one vocabulary. Week 9 compares segmentation choices and distinguishes training a tokenizer from encoding text. A saved model and its tokenizer must agree on ID meanings; even an in-range ID can select the wrong token if the mapping differs.

Knowledge check

What must a saved model and tokenizer agree on?