Transformers, ELI5 · Part 3 / 10

One block: mix, then think

A block makes two moves — attention lets words look at each other, the FFN lets each think alone — and hands back the same shape it got.

Each word is now a row of numbers — a [3 × 4] grid for “the cat sat” (rows = words, columns = meaning). A block enriches that grid with two moves.

Move 1 — Attention: words look at each other

Every word still sits alone. But sat only means something next to cat. lets each word take a weighted mix of the others — relevant words get big weights, the rest near zero. It’s the only step where words share information.

How does it pick the weights? Each word sends out a (“what am I looking for?”) and every word offers a key; dot them (the Chapter 2 move) for a score, softmax into weights, then blend everyone’s values. And it’s so a word can only look back — never at the future it’s trying to guess.

Classic puzzle: “The chicken didn’t cross the road because it was too tired.” What is it?

Try it: tired → it leans on chicken; wide → road. Same sentence, context decides.

Many heads at once

One pass catches one kind of relationship. So a block runs several in parallel, each tuned to something different.

Try it: Head 1 looks at the previous word, Head 2 at the nouns. The block blends them all.

see this in PyTorch

Each token builds a query, key, and value, scores how much it should attend to every earlier token, and mixes their values into its new representation.

# Single-head self-attention: every token decides how much to "look at" the others.
import torch
import torch.nn as nn
import torch.nn.functional as F

seq, d_model = 4, 8        # 4 tokens, each a vector of size 8
x = torch.randn(seq, d_model)   # x: (seq, d_model) -- our input tokens

# Each token asks a Question, offers a Key (label), and carries a Value (content).
to_q = nn.Linear(d_model, d_model)   # learns to build queries
to_k = nn.Linear(d_model, d_model)   # learns to build keys
to_v = nn.Linear(d_model, d_model)   # learns to build values

q = to_q(x)   # q: (seq, d_model) -- what each token is looking for
k = to_k(x)   # k: (seq, d_model) -- what each token offers to be matched on
v = to_v(x)   # v: (seq, d_model) -- the info each token will share

d_k = q.size(-1)                          # size of each query/key vector
scores = q @ k.mT / d_k ** 0.5            # scores: (seq, seq) -- how much each token looks at each other; scale keeps numbers calm

# Causal mask: a token may only look at itself and tokens before it, never the future.
mask = torch.triu(torch.ones(seq, seq, dtype=torch.bool), diagonal=1)  # True = future
scores = scores.masked_fill(mask, float('-inf'))   # block the future with -inf

weights = F.softmax(scores, dim=-1)   # weights: (seq, seq) -- each row sums to 1 (attention %)
out = weights @ v                     # out: (seq, d_model) -- each token = weighted blend of values

Move 2 — The feed-forward network: think alone

After mixing, each word thinks by itself in the (next chapter). The headline: attention shares, the FFN thinks.

The shape never changes

A block takes a [3 × 4] grid and hands back a [3 × 4] grid — same shape out as in. So you can feed it into another block, and another.

Try it: click the stages — attention rewrites the sat row, the FFN widens and folds it back.

see this in PyTorch

A transformer block stacks two simple moves, each added back onto the input: first the tokens look at each other (attention), then each token thinks on its own (the feed-forward net).

# A transformer block: each token gets refined by (1) looking at other tokens,
# then (2) thinking on its own. We keep both as submodules and just wire them up.
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, attention, ffn):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)  # steadies values before attention
        self.norm2 = nn.LayerNorm(d_model)  # steadies values before the ffn
        self.attention = attention          # tokens look at each other
        self.ffn = ffn                      # each token thinks for itself

    def forward(self, x):                   # x: (batch, seq, d_model)
        # pre-norm: normalize FIRST, then add the result back onto x (residual).
        # the "+ x" is a shortcut so the original info is never lost.
        x = x + self.attention(self.norm1(x))  # step 1: mix info across tokens
        x = x + self.ffn(self.norm2(x))         # step 2: refine each token alone
        return x                                # x: (batch, seq, d_model), same shape in = out

Next: inside the FFN.

Sources · 3

Full bibliography →