The Transformer Block

Context

A transformer block is Lesson 5.3.2's attention run as several heads side by side, followed by a small two-layer network applied to every token on its own, each of the two wrapped in a layer normalization (a rescaling of each token's features to a mean of 0 and a spread of 1) and a residual connection that adds the input back. Learned position embeddings, one trained vector per position added to each token's vector, give the block the word order attention alone cannot see. An encoder is a stack of these blocks.

The pretrained encoder Section 5.4 loads is built from these same parts with trained weights, and so is every transformer that routes a customer query, tags a support ticket or ranks a search result. Knowing the block by its parts is what lets a team read such a model's configuration, its width, head count and depth, and know what each one changes.

Lesson 5.3.2 built attention and found it as blind to word order as Lesson 5.3.1's bag of embeddings. Built by hand, the block fails without an error in three places: a head split written with view alone keeps every shape and lets each token read only the first words of its query, a normalization that takes its statistics across the batch makes a token's output depend on how much padding its neighbors carry, and a deep stack without the residual connection turns every token of a query into the same vector and stops learning.

Goal

• Build multi-head attention on attention.py with view and transpose, shape the padding mask [batch, 1, 1, seq] for [batch, heads, seq, seq] scores, and merge the heads with reshape.
• Add learned position embeddings, show a reordered query's pooled vector change, and read the IndexError a sequence longer than the table raises.
• Normalize each token with nn.LayerNorm, and measure Lesson 5.1.2's BatchNorm moving a query's outputs with its batch and its padding.
• Assemble the pre-norm block and prove it equal to nn.TransformerEncoderLayer by copying its weights in and comparing with torch.allclose.
• Train a two-block encoder on Banking77 against Lesson 5.3.1's bag and the same encoder without positions in one script, and score the test queries once for the model the validation queries choose.

Prerequisites

None of this is required, and the lesson stands on its own. The first builds the same block in code, and the second is PyTorch's own page for the layer this lesson checks against.

• YouTube - "Let's build GPT: from scratch, in code, spelled out." by Andrej Karpathy, from 1:19:11 to 1:37:49. A self-attention block inserted into a small network, then multi-headed self-attention, the feed-forward layer, residual connections, and layer normalization and its relationship to batch normalization.
• Official docs - PyTorch 2.14, torch.nn.TransformerEncoderLayer. Its arguments, d_model, nhead, dim_feedforward, dropout, batch_first and norm_first with their defaults, the src_key_padding_mask its forward pass takes, and the conditions under which it runs an optimized fast path.

Table of Contents

• Multi-Head Attention
• Position Embeddings
• Layer Normalization
• The Feed-Forward Network and the Residual Connection
• The Pre-Norm Block
• A Transformer Encoder on Banking77

Theory

Every block runs in ~/dev/bootcamp, beside text.py, Lesson 5.3.1's, which reads the Banking77 files it saved in data/banking77, attention.py, Lesson 5.3.2's, and image_loop.py, Lesson 4.5.2's. Nothing here installs or downloads. The first five subtitles are one script on seeded, untrained weights, so their numbers show the mechanism and nothing about what a trained model does. The sixth saves transformer.py, which its training script, the labs and the Assignment import. Trained figures are this machine's on its default eight threads, as Section 5.2's were, and seconds are compared only inside the script that timed them. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

Multi-Head Attention

What. Multi-head attention runs Lesson 5.3.2's attention several times side by side, each time on a slice of every token's vector. With a model dimension (d_model, the length of every token's vector) of 64 and four heads, each token's query, key and value are cut into four pieces of 16, each head scores, softmaxes and averages with its own pieces, and the four outputs of 16 are joined back into 64 and mixed by one more linear layer, the output projection. Why. One attention gives each token one softmax, one set of weights over its query that sums to 1. Four heads give each token four such sets, each computed from its own 16 numbers, so one head can weight one word of the query while another weights a different one, and the heads cost no parameters: the four projections are 64 by 64 whatever the head count. When. In every transformer block, with a head count that divides d_model.

import torch
from torch import nn
from text import vocab, encode, MAX_LEN, train_loader
from attention import attention

torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 64, padding_idx=0)
ids, labels = next(iter(train_loader))
x = embed(ids).detach()       # [32, 48, 64]
b, s, d = x.shape
heads = 4
# 64 read as 4 by 16, then heads before positions
split = x.view(b, s, heads, d // heads).transpose(1, 2)
pad = ids == 0
mask = pad[:, None, None, :]  # True at every pad key
out, weights = attention(split, split, split, mask)
print(x.shape, split.shape)
print(mask.shape, weights.shape, out.shape)

torch.Size([32, 48, 64]) torch.Size([32, 4, 48, 16])
torch.Size([32, 1, 1, 48]) torch.Size([32, 4, 48, 48]) torch.Size([32, 4, 48, 16])

view(b, s, heads, d // heads) reads each token's 64 numbers as 4 rows of 16 without moving any data, and transpose(1, 2) swaps the positions and heads axes, so the last two dimensions of split are 48 tokens of 16 numbers. attention multiplies over the last two dimensions only and treats the leading [32, 4] as 128 separate attentions, so its weights are one 48 by 48 table per query and per head. The mask's two inserted dimensions of size 1 broadcast it over the four heads and the 48 query positions, Lesson 4.1.2's rule, so every head of a query blocks the same pad keys. The raw slices stand in for the projections here, and the class below adds them. What the heads compute:

print(weights[0, :, 0].sum(dim=-1), weights[0][:, :, pad[0]].max())
one = x[..., 16:32].unsqueeze(1)  # the second slice, as one head
alone, _ = attention(one, one, one, mask)
print(torch.allclose(alone[:, 0], out[:, 1]))
try:
    attention(split, split, split, pad.unsqueeze(1))
except RuntimeError as err:
    print(err)

tensor([1., 1., 1., 1.]) tensor(0.)
True
The size of tensor a (32) must match the size of tensor b (4) at non-singleton dimension 1

The first token has four sets of weights, one per head, each summing to 1, and no head of that query puts any weight on a pad. Head 1 is exactly attention run on numbers 16 to 31 of every token, so the heads never see each other's slices. Lesson 5.3.2's three-dimensional mask, [batch, 1, seq], is refused against these scores: broadcasting lines shapes up from the right, and its batch of 32 met the 4 heads. Joining the heads back:

try:
    out.transpose(1, 2).view(b, s, d)
except RuntimeError as err:
    print(err)
merged = out.transpose(1, 2).reshape(b, s, d)
print(merged.shape, torch.equal(merged[..., 16:32], out[:, 1]))

view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
torch.Size([32, 48, 64]) True

transpose changes how a tensor is read, not where its numbers sit, so after it the four heads' rows for one token lie far apart in memory, and view, which never copies, refuses to read them as one row of 64. reshape copies them into that order, and numbers 16 to 31 of every merged token are head 1's output, bit for bit. The whole layer, with Lesson 5.3.2's three projections for queries, keys and values and the output projection:

class MultiHead(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.q = nn.Linear(d_model, d_model)
        self.k = nn.Linear(d_model, d_model)
        self.v = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

    def split(self, h):
        b, s, d = h.shape
        return h.view(b, s, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, h, mask=None):
        b, s, d = h.shape
        q, k, v = self.split(self.q(h)), self.split(self.k(h)), self.split(self.v(h))
        out, weights = attention(q, k, v, mask)
        return self.out(out.transpose(1, 2).reshape(b, s, d))

print([sum(p.numel() for p in MultiHead(64, n).parameters()) for n in (1, 4, 8)])
torch.manual_seed(42)
mh = MultiHead(64, 4)
with torch.no_grad():
    print(mh(x, mask).shape)
try:
    MultiHead(64, 5)(x, mask)
except RuntimeError as err:
    print(err)

[16640, 16640, 16640]
torch.Size([32, 48, 64])
shape '[32, 48, 5, 12]' is invalid for input of size 98304

One head, four or eight, the layer holds 16,640 parameters, four 64 by 64 weights and four biases of 64: the heads divide the projections' outputs and add nothing. The output keeps the input's [32, 48, 64], which is what lets blocks stack. Five heads asked for 12 numbers each, d // heads rounding 12.8 down, and 5 times 12 is 60 of the 64, so view refused.

DON'T:

• DON'T merge the heads with view after transpose. It was refused, and reshape is the call that copies.
• DON'T pass a [batch, 1, seq] mask to four-dimensional scores. It was refused against [32, 4, 48, 48], and [batch, 1, 1, seq] blocks the same keys for every head and every query position.
• DON'T choose a head count that does not divide d_model. MultiHead(64, 5) failed on its first batch with a message about a shape, not about heads.

Position Embeddings

What. A position table is a second nn.Embedding, with MAX_LEN rows, one per position from 0 to 47, and the row for position i is added to the embedding of whatever token stands at i. positions(torch.arange(48)) reads all 48 rows at once, and broadcasting adds the same [48, 64] to every query of a batch. Why. Lesson 5.3.2 measured attention without positions giving a reordered query the same pooled vector, because attention treats its tokens as a set. With a position's vector added, "card" at position 1 and "card" at position 2 are two different inputs, so the same words in another order are another query. When. In every transformer that should be able to read word order, and with a table at least as long as the longest sequence it will ever be given.

def pool(h, pad):
    keep = (~pad).unsqueeze(-1).float()
    return (h * keep).sum(dim=1) / keep.sum(dim=1)

pair = torch.tensor([encode("has my card arrived"), encode("my card has arrived")])
pair_pad = pair == 0
pair_mask = pair_pad[:, None, None, :]
torch.manual_seed(42)
positions = nn.Embedding(MAX_LEN, 64)
with torch.no_grad():
    plain = pool(mh(embed(pair), pair_mask), pair_pad)
    placed = pool(mh(embed(pair) + positions(torch.arange(MAX_LEN)), pair_mask), pair_pad)
print(pair[:, :6])
print(torch.allclose(plain[0], plain[1], atol=1e-6), torch.allclose(placed[0], placed[1], atol=1e-6))
print(round((placed[0] - placed[1]).abs().max().item(), 4))

tensor([[ 64,   4,   9, 274,   0,   0],
        [  4,   9,  64, 274,   0,   0]])
True False
0.2486

pool is Lesson 5.3.1's average over the real tokens as a function, and it is how every model of this lesson turns a query's tokens into one vector. The question and the statement hold the same four ids in two orders. Through mh alone their pooled vectors agree within 1e-6, Lesson 5.3.2's result again, and with the position rows added they differ by up to 0.2486: the layer now receives different inputs for the two orders. The table's size, and its edge:

print(round(embed.weight[1:].norm(dim=1).mean().item(), 2), round(positions.weight.norm(dim=1).mean().item(), 2))
try:
    positions(torch.arange(MAX_LEN + 1))
except IndexError as err:
    print(err)

8.0 7.92
index out of range in self

nn.Embedding draws every entry of a new table from a normal distribution with mean 0 and standard deviation 1, so a row of 64 entries is about the square root of 64 long, and the position rows, 7.92 on average, are as long as the token rows, 8.0: before any training a position weighs as much in a token's vector as its word. Position 48 has no row, and the message names neither positions nor 48. encode's truncation at MAX_LEN, Lesson 5.3.1's, keeps every Banking77 query inside the table.

DON'T:

• DON'T give a learned position table a sequence longer than its rows. Position 48 raised index out of range in self, so the table's length is a hard limit on every query the model will ever read.
• DON'T read a position table as proof that a model uses order. The table makes order visible, and the sixth subtitle measures whether training uses it.

Layer Normalization

What. nn.LayerNorm(64) normalizes each token on its own: it subtracts the mean of that token's 64 features and divides by their standard deviation, then multiplies by a learned weight and adds a learned bias, 64 of each. Lesson 5.1.2's BatchNorm does the opposite: it takes each channel's mean and variance across the batch, there over 64 images and 784 positions, and here it would take each feature's over every query and every position of the batch. Why. A token's normalized vector should depend on that token. Across a batch of padded queries, BatchNorm's statistics depend on which other queries share the batch and on how many of its positions are pads. LayerNorm reads one token's own 64 numbers, and like the GroupNorm of Lesson 5.1.2's Assignment it keeps no running statistics and computes the same thing in training and evaluation mode. When. Before the attention and before the feed-forward network in every block of this lesson, and once more after the last block.

norm = nn.LayerNorm(64)
with torch.no_grad():
    y = norm(x)
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, unbiased=False, keepdim=True)
    print(torch.allclose(y, (x - mean) / torch.sqrt(var + norm.eps), atol=1e-5))
    print(round(y[0, 0].mean().abs().item(), 4), round(y[0, 0].std(unbiased=False).item(), 4), y[0, 40].abs().max().item())
print(list(norm.state_dict()), norm.weight.shape)
try:
    nn.LayerNorm(48)(x)
except RuntimeError as err:
    print(err)

True
0.0 1.0 0.0
['weight', 'bias'] torch.Size([64])
Given normalized_shape=[48], expected input with shape [*, 48], but got input of size[32, 48, 64]

dim=-1 takes each statistic over one token's 64 features, and unbiased=False divides by 64, as LayerNorm does. The first token leaves with a mean of 0 and a standard deviation of 1, and position 40, a pad whose embedding is all zeros, leaves as zeros, eps (the small number added to every variance, here 1e-5) keeping the division defined. The state holds the weight and the bias and no running statistics, and nn.LayerNorm(48) asked for a last dimension of 48. BatchNorm on the same batch:

bn = nn.BatchNorm1d(64)
with torch.no_grad():
    in_batch = bn(x.transpose(1, 2)).transpose(1, 2)
    alone = bn(x[:1].transpose(1, 2)).transpose(1, 2)
    short = bn(x[:, :20].transpose(1, 2)).transpose(1, 2)
    real = ~pad[0]
    print(round(pad.float().mean().item(), 4), round(pad[:, :20].float().mean().item(), 4), int(real.sum()))
    print(torch.equal(norm(x[:1]), y[:1]), torch.equal(norm(x[:, :20])[0], y[0, :20]))
    print(round((in_batch[0] - alone[0])[real].abs().max().item(), 4), round((in_batch[0, :20] - short[0])[real[:20]].abs().max().item(), 4))
print(len(train_loader.dataset) % 32)

0.7409 0.4094 20
True True
2.4523 1.9433
3

nn.BatchNorm1d(64) reads [batch, features, positions], so each transpose(1, 2) puts the 64 features where it reads channels and back. Of the batch's 1,536 positions 0.7409 are pads, and the first query has 20 real tokens, all inside the first 20 columns, where the pads fall to 0.4094. In training mode, which a new module starts in, BatchNorm moved that query's real tokens by up to 2.4523 between its batch and itself alone, and by up to 1.9433 when the same batch was cut to 20 columns, a change of nothing but padding. LayerNorm gave the query bit for bit the same output all three ways. 8003 % 32 is 3: every epoch of train_loader ends on a batch of three queries, whose statistics BatchNorm would take as well.

DON'T:

• DON'T normalize padded sequences with BatchNorm. The same query's real tokens moved by up to 1.9433 when only the batch's padding changed, where LayerNorm held them exactly.
• DON'T hand nn.LayerNorm the sequence length. It normalizes the last dimension, the 64 features, and nn.LayerNorm(48) was refused on [32, 48, 64].

The Feed-Forward Network and the Residual Connection

What. The feed-forward network is nn.Linear(64, 128), a ReLU and nn.Linear(128, 64), applied to every token on its own, which is why it is called position-wise: nn.Linear acts on the last dimension, so the network never mixes two tokens. The residual connection adds a part's input back to its output, h + part(h), Lesson 5.2.1's relu(path(x) + x) without the ReLU after the sum. Why. Attention's output for a token is a weighted average of value vectors, so mixing between tokens happens there and only there. The feed-forward network is where each token's vector is transformed through a nonlinearity, with a hidden layer twice as wide as the token. The residual keeps each token's own vector in the sum, which matters because averaging pulls the tokens of a query toward one another. When. Around both parts of every block: attention mixes, the feed-forward network transforms, and each one's output is added to what came in.

ffn = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 64))
print(sum(p.numel() for p in ffn.parameters()))
with torch.no_grad():
    f = ffn(x)
    print(f.shape, torch.allclose(f[:, 3], ffn(x[:, 3])), torch.allclose(f[5, 3], ffn(x[5, 3])))

16576
torch.Size([32, 48, 64]) True True

The network holds 16,576 parameters, a 64 by 128 weight and 128 biases in, a 128 by 64 weight and 64 biases out. Position 3 of all 32 queries sent alone, and query 5's position 3 alone, give what the whole batch gave there: each token goes through on its own. What the residual keeps, measured on the first query's 20 real tokens:

def spread(h, keep):
    real = h[keep]
    gap = (real - real.mean(dim=0)).norm(dim=-1).mean()
    return round((gap / real.norm(dim=-1).mean()).item(), 4)

with torch.no_grad():
    mixed = mh(norm(x), mask)
    print(spread(x[0], real), spread(mixed[0], real), spread((x + mixed)[0], real))
    print(round(x[0][real].norm(dim=-1).mean().item(), 3), round(mixed[0][real].norm(dim=-1).mean().item(), 3))

0.9611 0.2288 0.9495
8.257 0.984

spread is how far a query's real tokens sit from their own average, over their average length: 0 means one vector repeated. The first query's 20 embeddings start at 0.9611. One attention layer brings them to 0.2288, because every output is an average of the same 20 values with weights that differ a little from token to token. With the input added back, 0.9495: each token keeps its own vector, 8.257 long on average, and attention's output, 0.984 long, is a correction to it.

DON'T:

• DON'T expect the feed-forward network to carry information between words. Token 3 alone gave what the batch gave at position 3, so only attention lets a word reach another.
• DON'T stack attention layers without adding their input back. One layer took the first query's spread from 0.9611 to 0.2288, and Lab Exercise 2 measures what six do.

The Pre-Norm Block

What. The pre-norm block normalizes before each part and adds the part's output to the unnormalized stream: h = h + attn(norm1(h)), then h = h + ffn(norm2(h)). PyTorch ships the same block as nn.TransformerEncoderLayer, which normalizes after each sum by default and before each part with norm_first=True. Why. A block built by hand can run, keep every shape and compute the wrong thing, which Lab Exercise 1 shows, and copying its weights into the library's layer and comparing outputs proves it number for number. The pre-norm order keeps the stream between blocks a plain sum of the input and every part's output, the residual connection unbroken from the first block to the last. When. Every time a transformer block is written by hand, before it trains.

class Block(nn.Module):
    def __init__(self, d_model, heads, d_ff):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = MultiHead(d_model, heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model))

    def forward(self, h, mask=None):
        h = h + self.attn(self.norm1(h), mask)
        return h + self.ffn(self.norm2(h))

torch.manual_seed(42)
block = Block(64, 4, 128)
reference = nn.TransformerEncoderLayer(64, 4, 128, dropout=0.0, batch_first=True, norm_first=True)
print(sum(p.numel() for p in block.parameters()), sum(p.numel() for p in reference.parameters()))
print([(name, tuple(p.shape)) for name, p in reference.named_parameters()][:4])

33472 33472
[('self_attn.in_proj_weight', (192, 64)), ('self_attn.in_proj_bias', (192,)), ('self_attn.out_proj.weight', (64, 64)), ('self_attn.out_proj.bias', (64,))]

Both hold 33,472 parameters, the attention's 16,640, the feed-forward network's 16,576 and the two LayerNorms' 128 each. dropout=0.0 turns off the dropout (Lesson 4.4.2's) the layer applies by default at 0.1, which Block does not have, and batch_first=True makes it read [batch, seq, features]. The layer keeps the three input projections as one (192, 64) weight, in_proj_weight, and the rest under out_proj, linear1, linear2, norm1 and norm2. The weights copied across, key by key:

mine = block.state_dict()
theirs = {
    "self_attn.in_proj_weight": torch.cat([mine["attn.q.weight"], mine["attn.k.weight"], mine["attn.v.weight"]]),
    "self_attn.in_proj_bias": torch.cat([mine["attn.q.bias"], mine["attn.k.bias"], mine["attn.v.bias"]]),
    "self_attn.out_proj.weight": mine["attn.out.weight"],
    "self_attn.out_proj.bias": mine["attn.out.bias"],
    "linear1.weight": mine["ffn.0.weight"], "linear1.bias": mine["ffn.0.bias"],
    "linear2.weight": mine["ffn.2.weight"], "linear2.bias": mine["ffn.2.bias"],
    "norm1.weight": mine["norm1.weight"], "norm1.bias": mine["norm1.bias"],
    "norm2.weight": mine["norm2.weight"], "norm2.bias": mine["norm2.bias"],
}
print(reference.load_state_dict(theirs))
block.eval()
reference.eval()
with torch.no_grad():
    ours = block(x, mask)
    print(torch.allclose(ours, reference(x, src_key_padding_mask=pad), atol=1e-5))
    print(torch.allclose(ours[~pad], reference(x)[~pad], atol=1e-5))
    post = nn.TransformerEncoderLayer(64, 4, 128, dropout=0.0, batch_first=True)
    post.load_state_dict(theirs)
    post.eval()
    print(torch.allclose(ours, post(x, src_key_padding_mask=pad), atol=1e-5))

<All keys matched successfully>
True
False
False

torch.cat stacks the query, key and value weights into one (192, 64) in that order, every key matched, and the two agree within 1e-5 at every position of all 32 queries. src_key_padding_mask is the [batch, seq] tensor ids == 0, True at every pad key, the same convention as attention.py's mask, and the layer broadcasts it over heads and queries itself. Under no_grad in evaluation mode the layer takes an optimized fast path, and the outputs still agree. Without the mask every token attends to the pads, and the real tokens differ. post, the default norm_first=False, normalizes after each sum, norm(h + attn(h)), a different function of the same weights.

DON'T:

• DON'T check a block against the reference without the padding mask on both sides. Without src_key_padding_mask the real tokens of the same batch disagreed.
• DON'T take nn.TransformerEncoderLayer's defaults for the block of this lesson. norm_first=False gave another output from the same weights, and the default dropout of 0.1 would add noise in training mode that Block does not have.

A Transformer Encoder on Banking77

What. The encoder embeds the token ids, adds the position rows, runs two blocks with the padding mask, normalizes the result once more, pools the real tokens and maps the pooled 64 numbers to 77 intent logits. It trains with train_epoch and accuracy from image_loop.py, which take any model, exactly as Lesson 5.3.1 trained its bag: order and the weights reseeded with 42 before each model, then ten epochs of Adam at 0.01. Beside it, in the same script and by the same recipe, train that bag and the same encoder without the position table. Why. The block matched PyTorch's, so what is left is what it buys on this data. The encoder without positions is order-blind, by Lesson 5.3.2's result and pool, so it isolates what the blocks' mixing buys, and the encoder with positions adds what order buys. One recipe for all three makes their accuracies comparable, and the bag's run reproduces Lesson 5.3.1's, which checks that the recipe is the same. When. Whenever a new architecture is proposed for a task a simpler model already does, the simpler model trained the same way being the number to beat.

The final LayerNorm is there because in the pre-norm order nothing normalizes the stream that leaves the last block. Save the pieces as transformer.py:

import torch
from torch import nn
from attention import attention
from text import MAX_LEN

class MultiHead(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.q = nn.Linear(d_model, d_model)
        self.k = nn.Linear(d_model, d_model)
        self.v = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

    def split(self, h):
        b, s, d = h.shape
        return h.view(b, s, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, h, mask=None):
        b, s, d = h.shape
        q, k, v = self.split(self.q(h)), self.split(self.k(h)), self.split(self.v(h))
        out, weights = attention(q, k, v, mask)
        return self.out(out.transpose(1, 2).reshape(b, s, d))

class Block(nn.Module):
    def __init__(self, d_model, heads, d_ff):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = MultiHead(d_model, heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model))

    def forward(self, h, mask=None):
        h = h + self.attn(self.norm1(h), mask)
        return h + self.ffn(self.norm2(h))

def pool(h, pad):
    keep = (~pad).unsqueeze(-1).float()
    return (h * keep).sum(dim=1) / keep.sum(dim=1)

class Encoder(nn.Module):
    def __init__(self, vocab_size, classes, d_model=64, heads=4, d_ff=128, depth=2, positions=True):
        super().__init__()
        self.tokens = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.positions = nn.Embedding(MAX_LEN, d_model) if positions else None
        self.blocks = nn.ModuleList([Block(d_model, heads, d_ff) for _ in range(depth)])
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, classes)

    def forward(self, ids):
        pad = ids == 0
        h = self.tokens(ids)
        if self.positions is not None:
            h = h + self.positions(torch.arange(ids.shape[1]))
        for block in self.blocks:
            h = block(h, pad[:, None, None, :])
        return self.head(pool(self.norm(h), pad))

class Bag(nn.Module):
    def __init__(self, vocab_size, classes, d_model=64):
        super().__init__()
        self.tokens = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.head = nn.Linear(d_model, classes)

    def forward(self, ids):
        return self.head(pool(self.tokens(ids), ids == 0))

Bag is Lesson 5.3.1's bag of embeddings with its average written as pool, built in the same order, the table and then the head, so seed 42 gives it the same weights. nn.ModuleList is a list whose modules' parameters the encoder registers, which a plain Python list would not. The three models, trained in one script:

import time
import torch
from torch.utils.data import DataLoader
from text import vocab, intents, order, train_set, train_loader, val_set, val_loader, test_loader
from image_loop import accuracy, train_epoch
from transformer import Encoder, Bag

models = {}
for kind in ("bag", "no positions", "encoder"):
    order.manual_seed(42)
    torch.manual_seed(42)
    if kind == "bag":
        model = Bag(len(vocab), len(intents))
    else:
        model = Encoder(len(vocab), len(intents), positions=kind == "encoder")
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    seconds = 0.0
    for epoch in range(1, 11):
        start = time.perf_counter()
        loss = train_epoch(model, optimizer, train_loader)
        seconds += time.perf_counter() - start
        if epoch in (1, 5, 10):
            print(kind, epoch, f"{loss:.4f}", f"{accuracy(model, val_loader):.4f}")
    models[kind] = model
    print(kind, sum(p.numel() for p in model.parameters()), f"{seconds / 10:.2f}")

bag 1 2.3178 0.7565
bag 5 0.1572 0.8665
bag 10 0.0353 0.8760
bag 145933 0.21
no positions 1 2.1410 0.6325
no positions 5 0.3697 0.8030
no positions 10 0.2926 0.7880
no positions 213005 2.26
encoder 1 2.3685 0.5995
encoder 5 0.2952 0.8070
encoder 10 0.1536 0.8295
encoder 216077 2.33

The bag's three rows are Lesson 5.3.1's to the digit, so the reseeding reproduced its run, and kind == "encoder" turns the position table on for the third model alone. After ten epochs the bag scores 0.8760 on the validation queries, the encoder 0.8295 and the encoder without positions 0.7880. The bag, with the fewest parameters, 145,933 against 216,077, leads both transformers, and the position table's 3,072 parameters lift the encoder by 4.15 points. The no-position encoder's training loss ended at 0.2926 and the encoder's at 0.1536, where the bag's reached 0.0353. The last figure on each summary line is the seconds per epoch, and the bag's epoch costs about a tenth of an encoder's, because it runs no attention. How each model fits its own training queries, and what each makes of the validation queries with their words shuffled:

def shuffled(ids, generator):
    out = ids.clone()
    for row in out:
        n = int((row != 0).sum())
        row[:n] = row[:n][torch.randperm(n, generator=generator)]
    return out

ids, labels = val_set.tensors
mixed = shuffled(ids, torch.Generator().manual_seed(42))
train_eval = DataLoader(train_set, batch_size=250)
for kind, model in models.items():
    model.eval()
    with torch.no_grad():
        on_mixed = (model(mixed).argmax(dim=1) == labels).float().mean().item()
    print(kind, f"{accuracy(model, train_eval):.4f}", f"{accuracy(model, val_loader):.4f}", f"{on_mixed:.4f}")

bag 0.9961 0.8760 0.8760
no positions 0.9239 0.7880 0.7880
encoder 0.9629 0.8295 0.8125

shuffled reorders the real tokens of every validation query by a seeded torch.randperm and leaves the pads in place, and val_set.tensors is the TensorDataset's queries and labels. Each line is a model's accuracy on its training queries, on the validation queries, and on the validation queries shuffled. The bag and the encoder without positions score exactly what they scored on the real order, the bag because an average has no order and the encoder by Lesson 5.3.2's result and pool. The encoder with positions falls from 0.8295 to 0.8125 on the shuffled queries, so it reads order, and shuffled it still scores above the encoder without positions. At Lesson 5.3.1's recipe both transformers fit their own training queries less closely than the bag, 0.9239 and 0.9629 against 0.9961, and validate lower: ten epochs at 0.01 trained the average of word vectors further than the attention stack. The validation queries choose, and the test queries score the choice once, Lesson 5.2.3's rule:

scores = {kind: accuracy(model, val_loader) for kind, model in models.items()}
chosen = max(scores, key=scores.get)
print(chosen, f"{scores[chosen]:.4f}", f"{accuracy(models[chosen], test_loader):.4f}")

bag 0.8760 0.8808

max(scores, key=scores.get) returns the key whose value is largest, so the validation queries choose the bag, and the 3,080 test queries score it once at 0.8808, a little above its validation 0.8760. That is the figure to hold against the encoder Lesson 5.4.1 loads, built from these same parts and pretrained.

DON'T:

• DON'T expect a transformer to beat a bag of embeddings because it can read order. Trained by the bag's own recipe on these 8,003 queries, the encoder with positions scored 0.8295 and the bag 0.8760.
• DON'T read a gain from position embeddings as order read. Shuffling cost the encoder 0.0170 of its 0.0415 lead over the encoder without positions, and shuffled it still led.
• DON'T compare models trained on different batch orders. order continues its stream from model to model, so each model reseeds it first, and all three read the same batches.
• DON'T score the test queries for all three and keep the best. That makes the test queries the choosing set, and the choice here was made on the validation queries before the test queries were read.

Lab Exercises

Lab Exercise 1: A Head Split Without the Transpose

In heads.py, subclass MultiHead from transformer.py as Unsplit, whose split returns h.view(b, self.heads, s, d // self.heads) with no transpose, the shape attention needs reached in one call. Seed 42, build an nn.Embedding(len(vocab), 64, padding_idx=0), a MultiHead(64, 4), and an Unsplit(64, 4) loaded with the MultiHead's weights.

Encode "why was my card payment declined" and "why was my card payment approved" as one batch. For each layer print its output's shape, whether it matches nn.MultiheadAttention(64, 4, batch_first=True) given the same weights and key_padding_mask, and how far each of the six real tokens' outputs moves between the two queries. Then, for the Unsplit layer alone, replace each word of the first query in turn with "transfer" and print how far the real tokens' outputs move. Say in a comment what the broken split lets a token read, and why no shape caught it.

Lab Exercise 2: Six Blocks Without the Shortcut

In depth.py, subclass Block as Plain, whose forward returns ffn(norm2(attn(norm1(h)))), the same parts with no residual connection. Build each encoder with a function that seeds 42, reseeds order, builds Encoder(len(vocab), len(intents), depth=depth) and, for the broken version, replaces model.blocks with an nn.ModuleList of Plain(64, 4, 128) blocks.

At depth 6, with and without the shortcut, print the spread of the first training batch's first query after each block and the gradient norm one batch's loss leaves on the first and last blocks' query projections. Then train depths 2 and 6, with and without the shortcut, for three epochs each at Adam's 0.001, and print the last epoch's training loss, the validation accuracy and how many different intents the model calls across the validation queries. Say in a comment what the shortcut changes at each depth, and what the broken deep encoder predicts.

import torch
from torch import nn
from text import vocab, encode
from transformer import MultiHead

class Unsplit(MultiHead):
    def split(self, h):
        b, s, d = h.shape
        return h.view(b, self.heads, s, d // self.heads)

torch.manual_seed(42)
embed = nn.Embedding(len(vocab), 64, padding_idx=0)
right = MultiHead(64, 4)
wrong = Unsplit(64, 4)
wrong.load_state_dict(right.state_dict())
reference = nn.MultiheadAttention(64, 4, batch_first=True)
with torch.no_grad():
    reference.in_proj_weight.copy_(torch.cat([right.q.weight, right.k.weight, right.v.weight]))
    reference.in_proj_bias.copy_(torch.cat([right.q.bias, right.k.bias, right.v.bias]))
    reference.out_proj.weight.copy_(right.out.weight)
    reference.out_proj.bias.copy_(right.out.bias)
pair = torch.tensor([encode("why was my card payment declined"), encode("why was my card payment approved")])
pad = pair == 0
with torch.no_grad():
    h = embed(pair)
    expected, _ = reference(h, h, h, key_padding_mask=pad)
    for name, layer in (("right", right), ("wrong", wrong)):
        out = layer(h, pad[:, None, None, :])
        moved = [round((out[0, t] - out[1, t]).abs().max().item(), 4) for t in range(6)]
        print(name, out.shape, torch.allclose(out, expected, atol=1e-5), moved)

right torch.Size([2, 48, 64]) True [0.1428, 0.2042, 0.1638, 0.1811, 0.1864, 0.2321]
wrong torch.Size([2, 48, 64]) False [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

base = torch.tensor([encode("why was my card payment declined")])
base_mask = (base == 0)[:, None, None, :]
with torch.no_grad():
    before = wrong(embed(base), base_mask)
    for position in range(6):
        other = base.clone()
        other[0, position] = vocab["transfer"]
        after = wrong(embed(other), base_mask)
        print(position, round((before[0, :6] - after[0, :6]).abs().max().item(), 4))
# view(b, heads, s, 16) regroups the numbers in memory order: head 0's
# rows 0 to 5 hold words 0 and 1, the other heads' rows 0 to 5 hold
# pads, and the mask keeps rows 0 to 5 as keys. Every real token reads
# words 0 and 1 alone. Every shape was right, so nothing was raised.

0 0.2538
1 0.1303
2 0.0
3 0.0
4 0.0
5 0.0

import math
import torch
from torch import nn
from text import vocab, intents, order, train_loader, val_loader
from image_loop import loss_fn, accuracy, train_epoch
from transformer import Block, Encoder

class Plain(Block):
    def forward(self, h, mask=None):
        h = self.attn(self.norm1(h), mask)
        return self.ffn(self.norm2(h))

def build(depth, shortcut):
    torch.manual_seed(42)
    order.manual_seed(42)
    model = Encoder(len(vocab), len(intents), depth=depth)
    if not shortcut:
        model.blocks = nn.ModuleList([Plain(64, 4, 128) for _ in range(depth)])
    return model

def spreads(model, ids):
    pad = ids == 0
    keep = ~pad[0]
    out = []
    with torch.no_grad():
        h = model.tokens(ids) + model.positions(torch.arange(ids.shape[1]))
        for block in model.blocks:
            h = block(h, pad[:, None, None, :])
            real = h[0][keep]
            gap = (real - real.mean(dim=0)).norm(dim=-1).mean()
            out.append(round((gap / real.norm(dim=-1).mean()).item(), 4))
    return out

ids, labels = next(iter(train_loader))
for shortcut in (True, False):
    model = build(6, shortcut)
    print(shortcut, spreads(model, ids))
    loss_fn(model(ids), labels).backward()
    print(f"{model.blocks[0].attn.q.weight.grad.norm().item():.2e}", f"{model.blocks[-1].attn.q.weight.grad.norm().item():.2e}")

True [0.9736, 0.9645, 0.952, 0.9512, 0.9491, 0.9351]
1.30e-02 1.20e-02
False [0.2177, 0.0015, 0.0, 0.0, 0.0, 0.0]
7.50e-02 1.79e-08

for depth in (2, 6):
    for shortcut in (True, False):
        model = build(depth, shortcut)
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
        for epoch in range(3):
            loss = train_epoch(model, optimizer, train_loader)
        acc = accuracy(model, val_loader)
        model.eval()
        with torch.no_grad():
            calls = torch.cat([model(x).argmax(dim=1) for x, y in val_loader]).unique()
        print(depth, shortcut, f"{loss:.4f}", f"{acc:.4f}", len(calls))
print(round(math.log(len(intents)), 4), intents[calls[0]])
# At depth 2 the shortcut is worth 7.45 points after three epochs, and
# the plain stack still learns. At depth 6 the plain stack learns
# nothing: its loss sits at ln 77 and it calls one intent for every
# query, while the same six blocks with shortcuts score best of the four.

2 True 1.1886 0.6690 77
2 False 1.4814 0.5945 75
6 True 1.0667 0.6910 77
6 False 4.3297 0.0170 1
4.3438 card_payment_fee_charged

Theory Assessment

Q1: Four Dimensions of Weights

After view(b, s, 4, 16).transpose(1, 2), attention returned weights shaped [32, 4, 48, 48], and MultiHead held 16,640 parameters at one, four or eight heads. Say what each dimension indexes, why the mask must be [32, 1, 1, 48], why the heads merge with reshape and not view, and why heads cost no parameters.

Hint: Write which two dimensions q @ k.transpose(-2, -1) multiplies over and which it treats as a batch, then ask where the heads' numbers come from.

Q2: Whose Statistics

BatchNorm moved the first query's real tokens by up to 2.4523 between its batch and itself alone, and by up to 1.9433 when the batch was cut to 20 columns, while LayerNorm gave the same output all three ways. Say what each takes its statistics over, why padded queries are where the difference shows, what the epoch's last batch of three adds, and why LayerNorm needs no running statistics.

Hint: For one real token's normalized vector, list every number that went into the mean it was shifted by.

Q3: A Match at Every Position

With its weights copied into nn.TransformerEncoderLayer(64, 4, 128, dropout=0.0, batch_first=True, norm_first=True), the hand-built block agreed with it within 1e-5 at every position, disagreed on the real tokens without src_key_padding_mask, and disagreed with norm_first at its default. Say what in_proj_weight holds and in which order, what src_key_padding_mask marks, why dropout=0.0 is part of the check, and what the third result says about the two orders.

Hint: Ask what an agreement within 1e-5 at every position proves that two equal parameter counts do not.

Q4: A Bag With Extra Steps

On the shuffled validation queries the bag and the encoder without positions scored exactly what they scored on the real ones, 0.8760 and 0.7880, and a teammate concludes that the encoder without positions is a bag of embeddings with extra steps. Say what the two models share, what the encoder computes that the bag cannot, and what the encoder with positions, 0.8295 on the real order and 0.8125 shuffled, says about how much of its lead is order.

Hint: Change one word of a query and ask which of the other words' vectors move in each model, then separate a model's sensitivity to order from what a component bought.

import torch
from text import vocab, intents, encode
from transformer import Encoder

torch.manual_seed(42)
model = Encoder(len(vocab), len(intents), positions=False)
pair = torch.tensor([encode("my card was declined"), encode("my card was stolen")])
pad = pair == 0
with torch.no_grad():
    h = model.tokens(pair)
    print((h[0, :3] - h[1, :3]).abs().max().item())
    for block in model.blocks:
        h = block(h, pad[:, None, None, :])
    print(round((h[0, :3] - h[1, :3]).abs().max().item(), 4))

0.0
0.4193

Assignment

Assignment: Positions Without a Table

The second subtitle's learned table has 48 rows and refused position 48, and encode cuts every query at MAX_LEN, so the tail of every longer query never reaches the model. Write fixed_positions.py, which replaces the table with positions computed from a formula, trains the encoder with them, and serves the long validation queries whole.

One construction this lesson did not use. The sinusoidal position encoding of the original transformer gives position p a vector whose even entries 2i are sin(p / 10000 ** (2i / d)) and whose odd entries 2i + 1 are the cos of the same angle, d being the model dimension. Nothing in it trains, and it is defined for every position.

Write sinusoid(length, d), returning the [length, d] table, and print the lengths of its first three rows. Subclass Encoder as Fixed, built with positions=False, whose forward adds sinusoid(ids.shape[1], 64) to the token embeddings, train it by the sixth subtitle's recipe and print epochs 1, 5 and 10. Then collect the validation queries longer than 48 tokens from text.py's val frame with tokenize, encode each whole with vocab.get(t, 1), pad them to the longest, and print how many there are, the longest length, and how many Fixed calls right cut to 48 tokens and whole. Say in a comment how the fixed table compares with the learned one, what serving the whole queries changed, and what the learned table would do with them.

Hints:

• torch.arange(0, d, 2) / d gives every 2i / d at once, and a [length, 1] column of positions divided by 10000 ** (...) broadcasts into the [length, d / 2] table of angles, whose sines fill table[:, 0::2] and cosines table[:, 1::2].
• The learned table's answer to a long query is the second subtitle's IndexError, so the comparison on whole queries runs on Fixed alone.

import torch
from text import vocab, intents, order, tokenize, val, train_loader, val_loader
from image_loop import accuracy, train_epoch
from transformer import Encoder, pool

def sinusoid(length, d):
    pos = torch.arange(length).unsqueeze(1)
    angle = pos / 10000 ** (torch.arange(0, d, 2) / d)
    table = torch.zeros(length, d)
    table[:, 0::2] = torch.sin(angle)
    table[:, 1::2] = torch.cos(angle)
    return table

class Fixed(Encoder):
    def forward(self, ids):
        pad = ids == 0
        h = self.tokens(ids) + sinusoid(ids.shape[1], self.tokens.embedding_dim)
        for block in self.blocks:
            h = block(h, pad[:, None, None, :])
        return self.head(pool(self.norm(h), pad))

print(sinusoid(48, 64).norm(dim=1)[:3])
order.manual_seed(42)
torch.manual_seed(42)
model = Fixed(len(vocab), len(intents), positions=False)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(1, 11):
    loss = train_epoch(model, optimizer, train_loader)
    if epoch in (1, 5, 10):
        print(epoch, f"{loss:.4f}", f"{accuracy(model, val_loader):.4f}")

tensor([5.6569, 5.6569, 5.6569])
1 2.2588 0.5955
5 0.3787 0.7915
10 0.2018 0.8165

lengths = [len(tokenize(text)) for text in val["text"]]
long = [i for i, n in enumerate(lengths) if n > 48]
rows = [[vocab.get(t, 1) for t in tokenize(val["text"].iloc[i])] for i in long]
width = max(len(r) for r in rows)
whole = torch.tensor([r + [0] * (width - len(r)) for r in rows])
y = torch.tensor([intents.index(val["category"].iloc[i]) for i in long])
model.eval()
with torch.no_grad():
    cut_right = (model(whole[:, :48]).argmax(dim=1) == y).sum().item()
    whole_right = (model(whole).argmax(dim=1) == y).sum().item()
print(len(long), width, cut_right, whole_right)
# The fixed table scores 0.8165 against the learned table's 0.8295, with
# nothing to train. It reads the 17 long queries whole, up to 84 tokens,
# and calls the same 13 right as when they were cut to 48: on these
# queries the tail added nothing. The learned table raises IndexError.

17 84 13 13

Additional Resources

Optional, and nothing later depends on either. One shows the original order of the block's parts, the other the normalization's own contract.

• Book chapter - Dive into Deep Learning, "The Transformer Architecture" by Zhang, Lipton, Li and Smola, under CC BY-SA 4.0. The position-wise feed-forward network and the residual connection with layer normalization built as classes, in the original order that normalizes after each sum, then the encoder and the decoder assembled from them.
• Official docs - PyTorch 2.14, torch.nn.LayerNorm. The formula, the statistics taken over the last dimensions named by normalized_shape, the biased variance, the per-element weight and bias, and the note that the layer uses statistics computed from its input in both training and evaluation mode.

Summary

The transformer block turned out to be four small parts held together by one sum. Attention mixes the tokens of a query, the feed-forward network transforms each token alone, LayerNorm keeps each token on one scale using nothing but its own numbers, and the residual stream carries every token's own vector from the first block to the last while each part adds a correction to it. The block's definition fit in two lines of forward, and what proved them right was a match with PyTorch's layer at every position, not that they ran.

The failures were all silent, and each was a question of what a number is allowed to depend on. A split without its transpose kept every shape and read a query by its first two words, a normalization across the batch let a token's output depend on its neighbors' padding, and a deep stack without the shortcut averaged a query into one vector by its third block. In each case the code ran, and only a measurement of what the output did and did not depend on found the defect.

And Banking77 gave a plain answer to what the machinery is worth on 8,003 labeled queries. Trained exactly as Lesson 5.3.1 trained its bag, the bag validated above both transformers and fitted its training queries more closely too, and the position table lifted the encoder by four points, of which shuffling the words took back fewer than two. The same blocks with weights already trained on text are another matter, and that is where Section 5.4 starts.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

q = h.view(b, s, heads, d // heads).transpose(1, 2)
out, w = attention(q, k, v, pad[:, None, None, :])
out = out.transpose(1, 2).reshape(b, s, d)  # then out_proj

h = tokens(ids) + positions(torch.arange(ids.shape[1]))
for block in blocks:
    # mixes tokens
    h = h + block.attn(block.norm1(h), mask)
    h = h + block.ffn(block.norm2(h))  # one token at a time
logits = head(pool(norm(h), ids == 0))