Lesson 11 of 20Assembling the Transformer

Residuals, Normalization & the Feed-Forward Layer

With position injected, attention has everything it needs to mix information across tokens. Two more pieces turn that into an actual transformer block, the repeating unit stacked to build every modern LLM: a way to keep gradients flowing through deep stacks, and a layer that lets each token "think" on its own.

First: residual connections. Recall from the backpropagation lesson that gradients are products of local derivatives chained together, stack enough layers and that product can shrink toward zero (or blow up), making deep networks hard to train. The fix is almost embarrassingly simple: instead of replacing a vector with f(x), compute x + f(x). The +x gives gradients a direct shortcut path backward through every layer, no matter how deep the stack gets.

Two more pieces round out a block:

Layer norm rescales a vector to have zero mean and unit variance before the next sublayer, keeping activations in a well-behaved range as they pass through many stacked blocks.
A feed-forward network, a small two-layer network (expand, ReLU, contract) applied independently to each token's vector. Where attention mixes information across tokens, the feed-forward layer lets each token "think" about what it just gathered. This sublayer alone holds a large fraction of a transformer's total parameters.

Step one token through a full transformer block

Token embedding

[1.40, 0.30]

+ Positional encoding

Self-attention output

Add residual

withPos + attention output

Layer norm

Feed-forward output

Add residual

norm output + feed-forward output

Layer norm (block output)

× N blocks, stacked — same structure, different learned weights each time

That's a complete transformer block: embed → add position → attend → residual → norm → feed-forward → residual → norm. A real model stacks dozens of these, each with its own learned weights, and every stack refines the same token vectors a little further before a final layer turns them into next-token predictions.

Here's the entire block in plain Python, reusing the exact same sentence and attention math from the self-attention lesson, continuing with the token "it":

Python

import math

def matmul(A, B):
    n, m, p = len(A), len(A[0]), len(B[0])
    return [[sum(A[i][k] * B[k][j] for k in range(m)) for j in range(p)] for i in range(n)]

def transpose(A):
    return [list(row) for row in zip(*A)]

def softmax_rows(M):
    result = []
    for row in M:
        biggest = max(row)
        exps = [math.exp(v - biggest) for v in row]
        total = sum(exps)
        result.append([e / total for e in exps])
    return result

tokens = ["The", "cat", "sat", "on", "the", "mat", "because", "it", "was", "tired"]
embeddings = [
    [0.1, 0.1], [1.5, 0.2], [0.3, 1.2], [-0.2, 0.4], [0.1, 0.1],
    [-0.3, -1.0], [0.6, -0.6], [1.4, 0.3], [0.2, -0.3], [1.0, -0.8],
]
W_Q = [[1, 0.2], [0.2, 1]]
W_K = [[1, -0.2], [-0.2, 1]]
W_V = [[0.5, 0.5], [-0.5, 0.5]]

Q = matmul(embeddings, transpose(W_Q))
K = matmul(embeddings, transpose(W_K))
V = matmul(embeddings, transpose(W_V))
d_k = len(K[0])
scores = [[s / math.sqrt(d_k) for s in row] for row in matmul(Q, transpose(K))]
attention_weights = softmax_rows(scores)
attention_outputs = matmul(attention_weights, V)

def positional_encoding_2d(pos):
    pe0 = math.sin(pos)
    pe1 = math.cos(pos)
    return [pe0 * 0.3, pe1 * 0.3]

def layer_norm(v):
    mean = sum(v) / len(v)
    variance = sum((x - mean) ** 2 for x in v) / len(v)
    std = math.sqrt(variance + 1e-5)
    return [(x - mean) / std for x in v]

FF_W1 = [[0.6, -0.4], [0.3, 0.7], [-0.5, 0.2], [0.1, -0.6]]
FF_B1 = [0.1, -0.1, 0.05, 0]
FF_W2 = [[0.5, -0.3, 0.2, 0.4], [-0.2, 0.6, 0.1, -0.5]]
FF_B2 = [0, 0.1]

def relu(x):
    return max(0, x)

def feed_forward(v):
    hidden = [relu(row[0] * v[0] + row[1] * v[1] + b) for row, b in zip(FF_W1, FF_B1)]
    return [b + sum(w * h for w, h in zip(row, hidden)) for row, b in zip(FF_W2, FF_B2)]

query = tokens.index("it")
embed = embeddings[query]
with_pos = [embed[i] + positional_encoding_2d(query)[i] for i in range(2)]
attn_out = attention_outputs[query]
after_residual_1 = [with_pos[i] + attn_out[i] for i in range(2)]
after_norm_1 = layer_norm(after_residual_1)
ffn_out = feed_forward(after_norm_1)
after_residual_2 = [after_norm_1[i] + ffn_out[i] for i in range(2)]
after_norm_2 = layer_norm(after_residual_2)

stages = [
    ("embedding", embed),
    ("+ positional encoding", with_pos),
    ("self-attention output", attn_out),
    ("+ residual", after_residual_1),
    ("layer norm", after_norm_1),
    ("feed-forward output", ffn_out),
    ("+ residual", after_residual_2),
    ("layer norm (block output)", after_norm_2),
]
for name, v in stages:
    print(f"{name:<28} [{v[0]:.3f}, {v[1]:.3f}]")

import math

def matmul(A, B):
    n, m, p = len(A), len(A[0]), len(B[0])
    return [[sum(A[i][k] * B[k][j] for k in range(m)) for j in range(p)] for i in range(n)]

def transpose(A):
    return [list(row) for row in zip(*A)]

def softmax_rows(M):
    result = []
    for row in M:
        biggest = max(row)
        exps = [math.exp(v - biggest) for v in row]
        total = sum(exps)
        result.append([e / total for e in exps])
    return result

tokens = ["The", "cat", "sat", "on", "the", "mat", "because", "it", "was", "tired"]
embeddings = [
    [0.1, 0.1], [1.5, 0.2], [0.3, 1.2], [-0.2, 0.4], [0.1, 0.1],
    [-0.3, -1.0], [0.6, -0.6], [1.4, 0.3], [0.2, -0.3], [1.0, -0.8],
]
W_Q = [[1, 0.2], [0.2, 1]]
W_K = [[1, -0.2], [-0.2, 1]]
W_V = [[0.5, 0.5], [-0.5, 0.5]]

Q = matmul(embeddings, transpose(W_Q))
K = matmul(embeddings, transpose(W_K))
V = matmul(embeddings, transpose(W_V))
d_k = len(K[0])
scores = [[s / math.sqrt(d_k) for s in row] for row in matmul(Q, transpose(K))]
attention_weights = softmax_rows(scores)
attention_outputs = matmul(attention_weights, V)

def positional_encoding_2d(pos):
    pe0 = math.sin(pos)
    pe1 = math.cos(pos)
    return [pe0 * 0.3, pe1 * 0.3]

def layer_norm(v):
    mean = sum(v) / len(v)
    variance = sum((x - mean) ** 2 for x in v) / len(v)
    std = math.sqrt(variance + 1e-5)
    return [(x - mean) / std for x in v]

FF_W1 = [[0.6, -0.4], [0.3, 0.7], [-0.5, 0.2], [0.1, -0.6]]
FF_B1 = [0.1, -0.1, 0.05, 0]
FF_W2 = [[0.5, -0.3, 0.2, 0.4], [-0.2, 0.6, 0.1, -0.5]]
FF_B2 = [0, 0.1]

def relu(x):
    return max(0, x)

def feed_forward(v):
    hidden = [relu(row[0] * v[0] + row[1] * v[1] + b) for row, b in zip(FF_W1, FF_B1)]
    return [b + sum(w * h for w, h in zip(row, hidden)) for row, b in zip(FF_W2, FF_B2)]

query = tokens.index("it")
embed = embeddings[query]
with_pos = [embed[i] + positional_encoding_2d(query)[i] for i in range(2)]
attn_out = attention_outputs[query]
after_residual_1 = [with_pos[i] + attn_out[i] for i in range(2)]
after_norm_1 = layer_norm(after_residual_1)
ffn_out = feed_forward(after_norm_1)
after_residual_2 = [after_norm_1[i] + ffn_out[i] for i in range(2)]
after_norm_2 = layer_norm(after_residual_2)

stages = [
    ("embedding", embed),
    ("+ positional encoding", with_pos),
    ("self-attention output", attn_out),
    ("+ residual", after_residual_1),
    ("layer norm", after_norm_1),
    ("feed-forward output", ffn_out),
    ("+ residual", after_residual_2),
    ("layer norm (block output)", after_norm_2),
]
for name, v in stages:
    print(f"{name:<28} [{v[0]:.3f}, {v[1]:.3f}]")

One thing this block glossed over: the attention step above only used a single set of Q/K/V projections. Real transformers run several in parallel, next.

← Positional Encoding: Absolute and Rotary Multi-Head Attention →