Masked attention and the full transformer

Meera had a number she did not like: three hours. That was Crumb’s average time to answer a support email during the Diwali rush. “Most of the replies are the same ten sentences,” she said. “Your cake is out for delivery. Sorry it’s late. Your refund is on its way. Can a computer not draft those?”

Arjun saw his chance. He had attention working, with positions and several heads. He would train it to write replies, one character at a time, on Crumb’s past support replies. He started training before lunch.

By 12:30 he was grinning. The loss (the number training tries to shrink; lower means fewer wrong guesses) had dropped from 3.6 to 0.003 in 300 steps. Practically perfect. He typed "Your " as a prompt and asked the model to continue.

'Your      d   ou               6  c \n  b    s'

Spaces, a stray d, a lonely 6. Priya looked at his screen for about ten seconds. “When it learns to predict character five,” she asked, “can it see character five?”

The model that copied the answer

Here is what had happened. To train a model to write, you show it text and ask it to predict each next token (a small unit of text; here, a single character). From “Y”, predict “o”. From “Yo”, predict “u”. From “You”, predict “r”.

Training does all of those predictions in one pass, in parallel, by feeding the whole sentence in at once. That is what makes transformers fast to train. But Arjun’s attention layer let every position look at every other position, including the ones after it. So when position 3 was asked “what comes after You?”, it could attend to position 4 and read off “r”.

The model had not learned to write. It had learned to copy from one seat to the right. Loss collapsed because the answers were visible. Then, at generation time, there was nothing to the right to copy from, because the future had not been written yet. The model had never practised guessing, so it produced noise.

“So my genius model,” said Arjun, “is the student who memorised where the answer key was printed.”

Priya’s exam paper

“Remember school exams where the answers were printed upside down at the bottom of the page?” Priya asked. “What did honest students do?”

“Folded the paper so they could only see the questions up to where they were.”

“That fold is a mask. During training, each position may look at itself and everything before it, never after. Then the only way to get the next character right is to actually learn how Crumb’s replies go.”

“But then why not train one position at a time, like it writes?”

“Because that would be slow. One 50-character reply gives you 49 next-character lessons. With a mask you get all 49 in a single pass, and none of them can cheat. The mask is what lets you train in parallel and honestly.”

Following “Your cake is late” through a mask

Let us trace the mask on four word-level tokens: Your, cake, is, late. Attention first computes a score for every pair (query row, key column). To keep the arithmetic easy, suppose every score happens to be 0:

scores	Your	cake	is	late
Your	0	0	0	0
cake	0	0	0	0
is	0	0	0	0
late	0	0	0	0

Without a mask, softmax (which turns each row into positive weights adding to 1) gives every row 0.25, 0.25, 0.25, 0.25. Row cake puts a quarter of its attention on is, which is the very word it is supposed to predict. Cheating.

With a mask, we add $-\infty$ to every score above the diagonal before the softmax. Since $e^{-\infty }=0$, those positions get exactly zero weight:

weights	Your	cake	is	late
Your	1.00	0	0	0
cake	0.50	0.50	0	0
is	0.33	0.33	0.33	0
late	0.25	0.25	0.25	0.25

Row by row: Your can only see itself. cake splits its attention between Your and itself. is sees three words. late sees all four. Every row still sums to 1, and no row sees the future. With real scores the numbers would be uneven, but the zeros above the diagonal would be identical.

Now the training signal. The model makes four predictions in one pass, one per row:

Position	Can see	Must predict
1	Your	cake
2	Your cake	is
3	Your cake is	late
4	Your cake is late	(whatever comes next, e.g. “.”)

Four honest lessons from one sentence.

Opening the hood: the masked attention formula

The mask is a matrix $M$ added to the scores:

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}+M\right)V,M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}$$

$Q$, $K$ and $V$ are the query, key and value matrices from the self-attention lesson, $d_k$ is the key length, and $i$ and $j$ are the row (the token asking) and column (the token being looked at). Adding 0 changes nothing; adding $-\infty$ erases. This is called masked attention or causal attention, because each token can only be influenced by what came before it, as causes come before effects.

You will also meet padding masks. When you batch emails of different lengths, short ones get filler tokens at the end so the batch forms a neat rectangle. A padding mask stops real tokens attending to the filler. Same trick, different shape.

Building a full block: what attention cannot do alone

Masked attention moves information between tokens. A transformer block wraps it with three more pieces. Each exists because something broke without it.

The feed-forward network: thinking time

Attention is a team meeting: everyone shares notes. But meetings do not do the work. After attention, each token passes, on its own, through a small feed-forward network: two linear layers with an activation function between them (an activation is a simple bend like GELU that lets the network learn non-straight-line patterns, as you saw in Module 28).

The middle layer is usually four times wider than the model. In the original transformer, each token’s 512 numbers expanded to 2,048 and shrank back to 512. That takes $2\times 512\times 2048\approx 2.1$ million weights per block, about twice attention’s $4\times 512\times 512\approx 1.05$ million. So roughly two thirds of a block’s weights live in the feed-forward network. Researchers believe much of a model’s stored factual knowledge sits here: attention gathers context, and the feed-forward layer looks up what that context implies.

Residual connections: annotate in pencil

Arjun’s first deep version, eight blocks stacked, would not train: the loss sat stuck at 3.5. The fix was the residual connection. Instead of replacing a token’s vector with a sub-layer’s output, you add the output to it:

$$x_{\text{new}}=x+\text{SubLayer}(x)$$

Priya’s analogy: a recipe card passed down a line of bakers, each annotating in pencil rather than rewriting it from scratch. The original stays readable, and each baker only has to add a small correction. If a token’s vector is [1.0, 2.0] and the attention layer contributes [0.1, -0.2], the result is [1.1, 1.8]: the original, gently adjusted.

There is a training reason too. When training sends its correction signals (gradients) backwards through the stack, the “+ x” path is a highway they can travel without being shrunk at every layer. Without it, signals reaching the early layers fade to nothing, and those layers stop learning. That was Arjun’s stuck loss.

Layer normalisation: the thermostat

As vectors pass through many blocks, their numbers can drift: some grow huge, some shrink. Layer normalisation (layer norm) rescales each token’s vector so its numbers have an average of 0 and a spread of 1, then applies two learned adjustments, a scale $\gamma$ and a shift $\beta$:

$$\text{LN}(x)=\gamma \cdot \frac{x-\mu }{\sigma }+\beta$$

where $\mu$ is the average of the vector’s numbers and $\sigma$ their standard deviation (spread). Worked through: [2, 4, 6, 8] has average 5 and standard deviation about 2.236, so it becomes about [-1.34, -0.45, 0.45, 1.34]. Same pattern, steady size. It is a thermostat that keeps every block working at a comfortable temperature.

The full block, and the full model

Putting it together, one block of a modern GPT-style model:

flowchart TD
    A["Input vectors"] --> B["Layer norm"]
    B --> C["Masked multi-head attention"]
    C --> D["Add (residual)"]
    A --> D
    D --> E["Layer norm"]
    E --> F["Feed-forward network"]
    F --> G["Add (residual)"]
    D --> G
    G --> H["Output vectors (same shape)"]

In maths, with $x$ as the block’s input:

$$h=x+\text{MaskedMHA}(\text{LN}(x)),\text{out}=h+\text{FFN}(\text{LN}(h))$$

A block’s output has exactly the same shape as its input, so blocks stack like trays on an oven rack. The original transformer used 6. GPT-2 small used 12. Large modern models use dozens.

The whole model, from text to next-token guess:

1. Tokenise the text and look up each token’s embedding (its learned list of numbers).
2. Add positions so order is visible.
3. Run the blocks. Each one lets tokens share information (masked attention) and think it over (feed-forward).
4. Final layer norm, then a linear layer that turns each token’s vector into one score per entry in the vocabulary. These raw scores are called logits.
5. Softmax turns the last position’s logits into next-token probabilities, and one token is picked.
6. Append it and repeat. That is generation.

The original 2017 transformer, built for translation, had two stacks: an encoder that read the whole German sentence with no mask, and a decoder that wrote the English with masked attention. Each decoder block had one extra sub-layer, cross-attention, where the decoder’s queries looked at the encoder’s keys and values. What you have built is the decoder half, minus cross-attention. Why many famous models keep only one half is the next lesson.

Building a tiny GPT in PyTorch

Step 1: the mask on its own. Two lines, and you can see the zeros.

import torch

T = 4
scores = torch.zeros(T, T)                                   # the all-zero scores from above
allowed = torch.tril(torch.ones(T, T, dtype=torch.bool))     # lower triangle = may look
scores = scores.masked_fill(~allowed, float("-inf"))         # block the future
print(torch.softmax(scores, dim=-1).round(decimals=2))

tensor([[1.0000, 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000, 0.0000],
        [0.3300, 0.3300, 0.3300, 0.0000],
        [0.2500, 0.2500, 0.2500, 0.2500]])

Exactly the table you traced. torch.tril keeps the lower triangle; ~ flips True and False; masked_fill writes $-\infty$ wherever the mask says “blocked”.

Step 2: one block. This uses PyTorch’s built-in nn.MultiheadAttention, which does the head-splitting from last lesson for you.

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

class Block(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(                                  # expand 4x, then shrink
            nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model))

    def forward(self, x, causal=True):
        T = x.size(1)
        future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)  # True = blocked
        h = self.ln1(x)
        a = self.attn(h, h, h, attn_mask=future if causal else None, need_weights=False)[0]
        x = x + a                       # residual 1
        x = x + self.ffn(self.ln2(x))   # residual 2
        return x

Watch the mask convention: nn.MultiheadAttention wants True for blocked positions, the opposite of the tril mask in Step 1, so here we use triu (upper triangle) with diagonal=1. Mixing these up is a classic bug. The causal switch exists only so you can reproduce Arjun’s mistake in Step 4.

Step 3: the whole model.

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model=64, n_heads=4, n_layers=4, max_len=128):
        super().__init__()
        self.tok = nn.Embedding(vocab_size, d_model)    # token embeddings
        self.pos = nn.Embedding(max_len, d_model)       # learned positions, like GPT-2
        self.blocks = nn.ModuleList(Block(d_model, n_heads) for _ in range(n_layers))
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)      # vector -> one logit per token

    def forward(self, ids, causal=True):                # ids: (batch, T)
        x = self.tok(ids) + self.pos(torch.arange(ids.size(1)))
        for block in self.blocks:
            x = block(x, causal)
        return self.head(self.ln_f(x))                  # logits: (batch, T, vocab)

Step 4: train it on Crumb replies, with and without the mask. Character-level tokens keep this self-contained: each distinct character is one token.

replies = [
    "Your eggless chocolate cake is out for delivery.",
    "Your plum cake is out for delivery and will reach you soon.",
    "Sorry, your delivery is late. Your cake will reach you by 6 pm.",
    "Our eggless cakes are baked in a separate kitchen.",
    "Your refund for the late delivery is on its way.",
    "Sorry about the delay. Your refund is on its way.",
]
text = "\n".join(replies) + "\n"
chars = sorted(set(text))                       # our vocabulary: every distinct character
stoi = {c: i for i, c in enumerate(chars)}      # character -> id
data = torch.tensor([stoi[c] for c in text])

def batch(n=16, T=32):
    """Pick n random 32-character windows; targets are the same windows shifted by one."""
    starts = torch.randint(0, len(data) - T - 1, (n,))
    x = torch.stack([data[s:s + T] for s in starts])
    y = torch.stack([data[s + 1:s + T + 1] for s in starts])
    return x, y

def train(causal, steps=300):
    torch.manual_seed(0)
    model = TinyGPT(len(chars))
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for step in range(steps + 1):
        x, y = batch()
        logits = model(x, causal)
        loss = F.cross_entropy(logits.reshape(-1, len(chars)), y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 100 == 0:
            print(f"causal={causal} step {step:3d} loss {loss.item():.3f}")
    return model

@torch.no_grad()
def generate(model, prompt, n=40, causal=True):
    ids = torch.tensor([[stoi[c] for c in prompt]])
    for _ in range(n):
        logits = model(ids[:, -128:], causal)
        next_id = logits[0, -1].argmax()            # pick the most likely next character
        ids = torch.cat([ids, next_id.view(1, 1)], dim=1)
    return "".join(chars[i] for i in ids[0])

cheat = train(causal=False)
print(repr(generate(cheat, "Your ", causal=False)))
honest = train(causal=True)
print(repr(generate(honest, "Your ", causal=True)))

The key line is y = data[s + 1 : s + T + 1]: the targets are the inputs shifted one place left, so position $t$ is graded on character $t+1$. You should see something like this (exact numbers vary by machine):

causal=False step   0 loss 3.585
causal=False step 100 loss 0.185
causal=False step 200 loss 0.013
causal=False step 300 loss 0.003
'Your      d   ou               6  c \n  b    s'
causal=True step   0 loss 3.569
causal=True step 100 loss 0.252
causal=True step 200 loss 0.141
causal=True step 300 loss 0.129
'Your plum cake is out for delivery. waterery '

There is the whole lesson in eight lines. Without the mask: near-perfect loss, garbage writing. With the mask: a higher loss, and real Crumb sentences, until it runs out of things it memorised and invents “waterery”. Six replies is a tiny dataset, so the model mostly memorises; with thousands it would start to generalise.

Where it breaks

Flipped masks. Some libraries want True = allowed, others True = blocked. Both run without error. The only defence is a test, like the one in “Try it yourself” below.

Loss that looks too good. A text model whose training loss drops near zero on real data in minutes is almost always leaking the answer: a missing mask, or targets not shifted. Treat a suspiciously good loss as a bug report.

Training on perfect prefixes. During training, every prefix the model sees is real Crumb text. During generation, it sees its own output, mistakes included. One odd character (“waterery”) and it is in territory it never practised. This gap, sometimes called exposure bias, is one reason models can drift into nonsense in long outputs.

Deep stacks without residuals or norms. Arjun’s eight-block model without residuals sat stuck at a loss of 3.5. If a deep model will not learn at all, check those two pieces first.

Try it yourself

Prove the mask works. With an untrained TinyGPT(vocab_size=1000), take a batch of token ids, change only the last token, and compare the logits at every position.

model = TinyGPT(vocab_size=1000).eval()
ids = torch.randint(0, 1000, (1, 10))
ids2 = ids.clone()
ids2[0, -1] = (ids2[0, -1] + 1) % 1000      # change only the final token
with torch.no_grad():
    diff = (model(ids) - model(ids2)).abs().sum(dim=-1)
print(diff)

Which positions changed? Now run it again with causal=False passed to both model(...) calls.

Show the answer

With the mask, only the last position changes; every earlier position shows a difference of 0, because no earlier token may look at the last one. With causal=False, every position changes, because every token can see the future. That is precisely Arjun’s cheating bug, made visible.

Common mistakes

• Training a text generator without a causal mask. Loss falls beautifully, then generation is garbage. Always run the “change the last token” test.
• Forgetting to shift the targets. If labels equal inputs instead of inputs shifted by one, the model learns to output the token it is looking at. Same symptom: perfect loss, useless model.
• Thinking attention does all the work. In most models the feed-forward layers hold around two thirds of each block’s weights.
• Removing residuals to “simplify”. Without them, deep stacks often refuse to train.
• Assuming a low training loss means a good writer. Arjun’s honest model memorised six replies. Always look at what it generates, and evaluate on text it has not seen (Module 23).

Recap

• Training predicts every next token in parallel; without a mask, each position can copy the answer from the next one.
• Masked (causal) attention adds $-\infty$ to future scores so they get zero weight after softmax.
• A transformer block is masked attention plus a feed-forward network, each wrapped in layer norm and a residual connection.
• Residuals keep deep stacks trainable; layer norm keeps numbers at a steady size; the feed-forward layers hold most of the weights.
• Blocks keep their input shape, so you stack many; a final linear layer turns vectors into logits, and softmax turns logits into next-token probabilities.

On Friday, Arjun showed Meera the honest model’s output. “It writes ‘Your plum cake is out for delivery’. A week ago it wrote spaces and a 6.”

“What changed?”

“Before, it was the student who peeked at the answer key. It got perfect marks in practice and failed the real exam. Now it can only look backwards, so the only way to score is to learn how our replies actually go. And it’s built from a stack of identical layers: each one lets the words compare notes, then gives each word a moment to think, without throwing away what it already knew.”

“Can we use it for Diwali?”

“Not this one. It has read six emails. The real thing is a model pretrained on a huge chunk of the internet.” He paused. “Which raises a question Priya wants to answer next: why do some famous models only read, like the ones behind our search box, and others only write?”

Go deeper

• Attention Is All You Need, section 3.1 and figure 1.
• The Illustrated Transformer
• PyTorch tutorials