Lesson 45: The Transformer Architecture¶

Lesson 44 built one attention operation. A Transformer (Vaswani et al., 2017) wraps that operation into a reusable block and stacks many of them. This lesson fills in the three pieces attention alone is missing: multiple attention heads (so a layer can track several kinds of relationships at once), positional encoding (so order isn't invisible to a mechanism that otherwise treats a sequence as an unordered set), and the encoder block (attention plus a per-position feedforward network, wired together with residual connections and normalization).

In [1]:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

Multi-head attention¶

A single attention operation computes one similarity pattern between every pair of positions. Multi-head attention splits the model dimension into several smaller chunks (heads), runs attention independently within each, and concatenates the results — letting different heads specialize in different kinds of relationships (e.g. one head tracking nearby positions, another tracking a specific long-range dependency) instead of averaging everything into one pattern.

In [2]:

D, H = 16, 4  # model dimension, number of heads
mha_torch = nn.MultiheadAttention(D, H, batch_first=True)

def multihead_attention(x, mha):
    B, T, _ = x.shape
    d_head = D // H
    Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
    bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
    Q = (x @ Wq.T + bq).view(B, T, H, d_head).transpose(1, 2)
    K = (x @ Wk.T + bk).view(B, T, H, d_head).transpose(1, 2)
    V = (x @ Wv.T + bv).view(B, T, H, d_head).transpose(1, 2)
    scores = Q @ K.transpose(-2, -1) / np.sqrt(d_head)
    weights = F.softmax(scores, dim=-1)
    out = (weights @ V).transpose(1, 2).reshape(B, T, D)
    return out @ mha.out_proj.weight.T + mha.out_proj.bias

torch.manual_seed(0)
x = torch.randn(2, 5, D)
out_manual = multihead_attention(x, mha_torch)
out_torch, _ = mha_torch(x, x, x, need_weights=False)

print(f'model dim = {D}, heads = {H}, dim per head = {D // H}')
print(f'max abs diff vs nn.MultiheadAttention: {(out_manual - out_torch).abs().max().item():.2e}')

model dim = 16, heads = 4, dim per head = 4
max abs diff vs nn.MultiheadAttention: 6.33e-08

Positional encoding¶

Attention computes similarity between content vectors — nothing about the formula from Lesson 44 refers to where in the sequence a position sits. Permute the input sequence, and attention's output permutes along with it, identically: it is fundamentally a set operation, blind to order. Positional encoding fixes this by adding a unique, deterministic pattern to each position before attention runs, so position becomes part of what gets compared. The original Transformer paper's choice is a fixed (not learned) sinusoid at multiple frequencies:

$$PE_{(pos, 2i)} = \sin(pos / 10000^{2i/D}), \qquad PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/D})$$

In [3]:

def positional_encoding(T, D):
    pos = torch.arange(T).unsqueeze(1).float()
    i = torch.arange(D).unsqueeze(0).float()
    angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
    angles = pos * angle_rates
    pe = torch.zeros(T, D)
    pe[:, 0::2] = torch.sin(angles[:, 0::2])
    pe[:, 1::2] = torch.cos(angles[:, 1::2])
    return pe

pe = positional_encoding(50, 32)
plt.figure(figsize=(6, 4))
plt.imshow(pe.numpy().T, cmap='RdBu', aspect='auto')
plt.xlabel('position'); plt.ylabel('encoding dimension')
plt.title('Sinusoidal positional encoding')
plt.colorbar(fraction=0.046)
plt.show()

No description has been provided for this image

Does this actually matter? A task that requires knowing order¶

Build a task that's impossible to solve from content alone: a fixed vector A and a fixed vector B appear at two random positions in a noisy sequence, and the label is simply "does A appear before B?" A model that mean-pools attention output over the sequence can only tell which two tokens are present, never their order, unless position is somehow injected.

In [4]:

T = 8
vec_a, vec_b = torch.randn(D), torch.randn(D)

def make_order_dataset(rng, n):
    X, y = [], []
    for _ in range(n):
        pos_a, pos_b = rng.choice(T, size=2, replace=False)
        seq = torch.randn(T, D) * 0.1
        seq[pos_a] = vec_a
        seq[pos_b] = vec_b
        X.append(seq)
        y.append(1.0 if pos_a < pos_b else 0.0)
    return torch.stack(X), torch.tensor(y, dtype=torch.float32)

rng = np.random.default_rng(3)
X_train, y_train = make_order_dataset(rng, 400)
X_test, y_test = make_order_dataset(rng, 150)

class TinyAttnClassifier(nn.Module):
    def __init__(self, use_pos_enc):
        super().__init__()
        self.use_pos_enc = use_pos_enc
        self.mha = nn.MultiheadAttention(D, 4, batch_first=True)
        self.fc = nn.Sequential(nn.Linear(D, 16), nn.ReLU(), nn.Linear(16, 1))
        if use_pos_enc:
            self.register_buffer('pe', positional_encoding(T, D))

    def forward(self, x):
        if self.use_pos_enc:
            x = x + self.pe
        attn_out, _ = self.mha(x, x, x, need_weights=False)
        return self.fc(attn_out.mean(dim=1)).squeeze(-1)

def train_eval(use_pos_enc, seed, epochs=300, lr=0.01):
    torch.manual_seed(seed)
    model = TinyAttnClassifier(use_pos_enc)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        opt.zero_grad()
        loss = F.binary_cross_entropy_with_logits(model(X_train), y_train)
        loss.backward()
        opt.step()
    with torch.no_grad():
        return ((model(X_test) > 0).float() == y_test).float().mean().item()

no_pe_accs = [train_eval(False, seed) for seed in range(5)]
pe_accs = [train_eval(True, seed) for seed in range(5)]

print(f'without positional encoding: mean test acc = {np.mean(no_pe_accs):.1%}  (+/- {np.std(no_pe_accs):.1%})')
print(f'with positional encoding:    mean test acc = {np.mean(pe_accs):.1%}  (+/- {np.std(pe_accs):.1%})')

without positional encoding: mean test acc = 44.9%  (+/- 3.0%)
with positional encoding:    mean test acc = 100.0%  (+/- 0.0%)

Without positional encoding, the model is stuck at chance — it can only ever report which two tokens showed up, never their relative order, no matter how long it trains. Adding the fixed sinusoid gives every position a distinct signature the model can key off, and the task becomes trivial.

The encoder block¶

A single Transformer layer is multi-head attention, a small per-position feedforward network, and two residual connections (Lesson 35) with layer normalization (Lesson 35's batch norm, but normalizing across the feature dimension for each individual token instead of across the batch):

In [5]:

class EncoderBlock(nn.Module):
    def __init__(self, D, H, FF):
        super().__init__()
        self.mha = nn.MultiheadAttention(D, H, batch_first=True)
        self.ln1 = nn.LayerNorm(D)
        self.ff = nn.Sequential(nn.Linear(D, FF), nn.ReLU(), nn.Linear(FF, D))
        self.ln2 = nn.LayerNorm(D)

    def forward(self, x):
        attn_out, _ = self.mha(x, x, x, need_weights=False)
        x = self.ln1(x + attn_out)   # residual + norm around attention
        ff_out = self.ff(x)
        x = self.ln2(x + ff_out)     # residual + norm around the feedforward network
        return x

FF = 32
block = EncoderBlock(D, H, FF)
layer_torch = nn.TransformerEncoderLayer(d_model=D, nhead=H, dim_feedforward=FF,
                                          batch_first=True, dropout=0.0)
layer_torch.eval()

# copy torch's weights into our block so the two are directly comparable
block.mha.load_state_dict(layer_torch.self_attn.state_dict())
block.ln1.load_state_dict(layer_torch.norm1.state_dict())
block.ln2.load_state_dict(layer_torch.norm2.state_dict())
block.ff[0].load_state_dict({'weight': layer_torch.linear1.weight, 'bias': layer_torch.linear1.bias})
block.ff[2].load_state_dict({'weight': layer_torch.linear2.weight, 'bias': layer_torch.linear2.bias})

x2 = torch.randn(2, 6, D)
out_block = block(x2)
out_layer = layer_torch(x2)
print(f'encoder block max abs diff vs nn.TransformerEncoderLayer: {(out_block - out_layer).abs().max().item():.2e}')
print(f'output shape unchanged from input: {tuple(out_block.shape)} == {tuple(x2.shape)}')

encoder block max abs diff vs nn.TransformerEncoderLayer: 0.00e+00
output shape unchanged from input: (2, 6, 16) == (2, 6, 16)

The block's output has exactly the same shape as its input — same trick as Lesson 35's residual block, and for the same reason: it means blocks can be stacked arbitrarily deep, each one refining the same sequence of vectors, without any reshaping between them. A full Transformer encoder is just N copies of this block stacked in sequence (nn.TransformerEncoder in PyTorch); a decoder adds a second attention step per block that attends to the encoder's output, plus a causal mask (Lesson 44) on its own self-attention. Vision Transformers (Lesson 46) reuse the encoder side almost unchanged — the only real difference is what gets fed in as the initial sequence of vectors.

Exercise¶

Change H (heads) from 4 to 1, keeping D=16 fixed, and rerun the multi-head attention validation. With a single head, is there still a meaningful difference from Lesson 44's single-head attention function?
In the order-detection task, change T (sequence length) from 8 to 32. Does the with-positional-encoding model's accuracy hold up, or does the longer sequence make the task harder in a way positional encoding alone doesn't fix?
EncoderBlock above uses "post-norm" (LayerNorm applied after the residual add, matching the original 2017 paper). Many modern Transformers use "pre-norm" instead: x = x + self.mha(self.ln1(x)). Implement pre-norm and compare the two at greater depth (stack 10 blocks) — does one train more stably, echoing Lesson 35's vanishing-gradient story?

← Lesson 44: The Attention Mechanism

Lesson 46: Vision Transformers (ViT) →