Training objectives

What you will be able to do

See how a language model learns from plain text that nobody labelled. It uses one of two training objectives: predict the next token, or fill in a hidden token. Then you can explain why the causal mask is needed in training, and why Larkfield should not pretrain its own large model.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook builds the training examples of both objectives. It trains a small next-word model with and without the causal mask, and has the module practice. It does not download SmolLM2. About 40 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 45 minutes

What you need to know

• Cross-entropy loss as “the surprise at the right answer”, −log p, from Loss: how wrong is the network?.
• The causal mask, the padding mask, encoders and decoders, from Position and masking.
• TinyTransformer and its shapes, from Transformer blocks.

What you need on your computer (only for the local route)

• The folder ticket-router with its .venv, from Set up PyTorch: Route 2. The SmolLM2 example also needs pip install "transformers>=5,<6" and downloads about 270 MB; it is optional, because its output is on this page.

Tested with Python 3.12, PyTorch 2.14.1 and transformers 5.19.0 on the CPU (macOS).

Fill in the blank

In a language class, the teacher gives you a text with gaps: “My parcel arrived ____, please send a new lamp.” You fill each gap from the words around it. Nobody had to write an answer sheet: the original text is the answer sheet. A good student gets better at the language just by doing many of these.

The problem: labels are expensive, and language needs a lot of text

Every network in this course so far learned from labelled examples: each Larkfield ticket came with the team that Grace’s policy gives. Someone had to decide each label, and the tiny transformer still saw only 4,696 of them. In Transformer blocks, that was far too little to learn what words mean and how order works.

To learn a language, a model needs billions of sentences. Nobody can label billions of sentences by hand.

The idea: the text is its own answer key

Make the questions from the text itself. Hide part of a sentence, and ask the model to predict it. The hidden part is the right answer, so every sentence ever written becomes training data. This is self-supervised learning: supervised, because each question has a known answer, and “self” because the text supplies the answer. Two ways of hiding are common.

Next-token prediction. Read the text from the left, and at every position predict the next token. One Larkfield sentence gives one example for each token after the first:

objectives.py

from ticketnet import tokenize

sentence = "My parcel arrived broken, please send a new lamp."
tokens = tokenize(sentence)
print(len(tokens), tokens)
for i in range(1, len(tokens)):
    print(" ".join(tokens[:i]), "->", tokens[i])

The output:

11 ['my', 'parcel', 'arrived', 'broken', ',', 'please', 'send', 'a', 'new', 'lamp', '.']
my -> parcel
my parcel -> arrived
my parcel arrived -> broken
my parcel arrived broken -> ,
my parcel arrived broken , -> please
my parcel arrived broken , please -> send
my parcel arrived broken , please send -> a
my parcel arrived broken , please send a -> new
my parcel arrived broken , please send a new -> lamp
my parcel arrived broken , please send a new lamp -> .

11 tokens give 10 examples. A decoder does not run them one by one. It reads the whole sentence once, and position i predicts token i + 1, all positions at the same time. The causal mask from the last lesson makes this fair: each position sees only its own row of this list.

Masked-token prediction. Hide some tokens anywhere in the text, and predict them from both sides:

import random

random.seed(3)
hidden = sorted(random.sample(range(len(tokens)), 2))     # hide 2 of the 11 tokens
masked = ["[MASK]" if i in hidden else t for i, t in enumerate(tokens)]
print(" ".join(masked))
print("answers:", {i: tokens[i] for i in hidden})

The output:

my parcel arrived [MASK] , please send a new [MASK] .
answers: {3: 'broken', 9: 'lamp'}

[MASK] is a special token that means “a word was here”. BERT, the classic encoder, hides about 15% of the tokens of each text. Each hidden token is predicted from every word before and after it, as an encoder reads.

The two training objectives on one Larkfield sentence. Both make their questions and answers from the text itself.

Text description of the diagram

The diagram has two panels for the sentence “My parcel arrived broken, please send a new lamp.” Panel 1, next-token prediction (GPT-like): a row of the tokens “my”, “parcel”, “arrived”, “broken”, the comma, “please” and “send”. Above each one, an arrow points to its answer, the next token: “parcel”, “arrived”, “broken”, the comma, “please”, “send” and “a”. A note says that each position predicts the next token and reads only the tokens to its left. So one sentence of 11 tokens gives 10 examples. Panel 2, masked-token prediction (BERT-like): the row “my parcel arrived [MASK] , please send a new [MASK] .” with the two masks shaded. Below each mask, an arrow points to its answer, “broken” and “lamp”. A note says that about 15% of the tokens are hidden. The model predicts each one from the words on both sides, and the text holds every answer.

How many examples are hidden in Larkfield’s training tickets, with no label from a person?

from ticketnet import load_split

train_rows = load_split("data/train.csv")
n_tokens = sum(len(tokenize(row["text"])) for row in train_rows)
print(f"{n_tokens:,} tokens in {len(train_rows):,} training tickets")
print(f"{n_tokens - len(train_rows):,} next-token examples, with no label written by a person")

The output:

269,231 tokens in 4,696 training tickets
264,535 next-token examples, with no label written by a person

More than 50 times as many examples as tickets, and nobody wrote an answer for any of them. Still, as you see below, that is a tiny amount of text for learning a language.

The objective on a real model

SmolLM2-135M, which Module 7 explores, was pretrained with next-token prediction. Its training loss is the cross-entropy that you know: for each position, the surprise at the real next token, −log p. Ask it how likely each real next token of the sentence is:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "HuggingFaceTB/SmolLM2-135M"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float32).eval()

ids = tokenizer("My parcel arrived broken, please send a new lamp.", return_tensors="pt").input_ids
with torch.no_grad():
    scores = model(ids).logits                       # (1, tokens, 49152)
probs = scores[0, :-1].softmax(-1)                   # the prediction after each token
answers = ids[0, 1:]                                 # the real next tokens
p = probs[torch.arange(len(answers)), answers]
for token, chance in zip(tokenizer.convert_ids_to_tokens(answers), p.tolist()):
    print(f"{token!r:>10}  {chance:.4f}   1 in {1 / chance:,.0f}")
print("loss", round(float(-p.log().mean()), 3))

The output:

 'Ġparcel'  0.0000   1 in 215,056
'Ġarrived'  0.0036   1 in 277
 'Ġbroken'  0.0000   1 in 35,351
       ','  0.1655   1 in 6
 'Ġplease'  0.0001   1 in 9,503
   'Ġsend'  0.0486   1 in 21
      'Ġa'  0.0601   1 in 17
    'Ġnew'  0.0148   1 in 68
   'Ġlamp'  0.0000   1 in 251,563
       '.'  0.1967   1 in 5
loss 6.345

Ġ marks a space before the token. The model gives a score for each of its 49,152 tokens at each position, and the real next token is the answer. “broken” after “My parcel arrived” surprises it: 1 chance in 35,351. The comma after “broken” does not: 1 in 6. The loss, 6.345, is the average surprise. model(ids, labels=ids).loss gives the same 6.345, because Hugging Face computes this objective for you. Pretraining changes the weights, over trillions of tokens, to make this number smaller.

Why the causal mask matters in training

What happens if a model learns next-token prediction without the causal mask? Priya tries it on the tickets. NextWordModel is a small decoder: token and position vectors, 2 blocks with 4 heads, and a score for every word of the vocabulary. The only difference between the two runs is the mask:

import torch
from torch import nn
from ticketnet import Vocab, PAD, set_seed

vocab = Vocab([row["text"] for row in train_rows])
valid_rows = load_split("data/valid.csv")

def sequences(rows, length=48):
    """Each ticket's first 48 token IDs, padded with 0."""
    ids = [vocab.encode(row["text"])[:length] for row in rows]
    return torch.tensor([s + [PAD] * (length - len(s)) for s in ids])

seq_train, seq_valid = sequences(train_rows), sequences(valid_rows)

class NextWordModel(nn.Module):
    """Token + position vectors -> 2 transformer blocks -> a score for every word of the vocabulary."""

    def __init__(self, vocab_size, dim=64, length=48):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)
        self.position = nn.Embedding(length, dim)
        block = nn.TransformerEncoderLayer(dim, 4, 4 * dim, dropout=0.0, batch_first=True, norm_first=True)
        self.blocks = nn.TransformerEncoder(block, 2, enable_nested_tensor=False)
        self.out = nn.Linear(dim, vocab_size)

    def forward(self, ids, causal=True):                      # ids: (batch, tokens)
        n = ids.shape[1]
        h = self.embed(ids) + self.position(torch.arange(n))  # (batch, tokens, dim)
        mask = nn.Transformer.generate_square_subsequent_mask(n) if causal else None
        h = self.blocks(h, mask=mask, is_causal=causal)       # (batch, tokens, dim)
        return self.out(h)                                    # (batch, tokens, vocab)

loss_fn = nn.CrossEntropyLoss(ignore_index=PAD)              # padding is not an answer

def next_word_loss(model, ids, causal=True):
    scores = model(ids[:, :-1], causal=causal)                # read tokens 0 .. n-2
    targets = ids[:, 1:]                                      # the answers: tokens 1 .. n-1
    return loss_fn(scores.reshape(-1, scores.shape[-1]), targets.reshape(-1))

The answers are the same tokens, moved one place to the left: ids[:, 1:]. ignore_index=PAD tells the loss to skip the padding. nn.TransformerEncoderLayer with a causal mask works as a decoder block; the name only says that PyTorch built it for encoders first. Train both models for 3 epochs with the five steps:

def train_next_word(causal, epochs=3):
    set_seed(0)
    model = NextWordModel(len(vocab))
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    curve = []
    for epoch in range(epochs):
        model.train()
        order = torch.randperm(len(seq_train))
        for i in range(0, len(order), 32):
            opt.zero_grad()
            loss = next_word_loss(model, seq_train[order[i:i + 32]], causal=causal)
            loss.backward()
            opt.step()
            curve.append(round(loss.item(), 4))
    return model, curve

honest, honest_curve = train_next_word(causal=True)
cheater, cheater_curve = train_next_word(causal=False)
print("steps", len(honest_curve))
print("training loss, last step: with causal mask", honest_curve[-1], "| without mask", cheater_curve[-1])

The output:

steps 441
training loss, last step: with causal mask 3.4118 | without mask 0.5356

• With the causal mask
• Without a mask

Training loss of two next-word models on the tickets, the mean over the last 21 steps. Without the causal mask, the loss falls much lower, because the model reads the answers.

Text description of the chart

A line chart with the training step, 1 to 441, on the horizontal axis and the training loss, 0 to 8, on the vertical axis. Both lines start near 8, the loss of a guess among 2,864 words. With the causal mask, the loss falls to about 3.5 after 441 steps. Without a mask, it falls much faster and further, to about 0.6.

Without the mask, the training loss looks six times better. Now measure both on the validation tickets, in the situation where the model is used: writing, where the future does not exist yet.

honest.eval()
cheater.eval()
with torch.no_grad():
    print("with mask, valid loss      ", round(next_word_loss(honest, seq_valid).item(), 3))
    print("without mask, future open  ", round(next_word_loss(cheater, seq_valid, causal=False).item(), 3))
    print("without mask, future hidden", round(next_word_loss(cheater, seq_valid, causal=True).item(), 3))

def continue_text(model, text, words=8):
    ids = vocab.encode(text)
    with torch.no_grad():
        for _ in range(words):
            scores = model(torch.tensor([ids]), causal=True)     # at writing time, the future does not exist
            ids.append(int(scores[0, -1].argmax()))
    return " ".join(vocab.itos[i] for i in ids)

print("with mask:   ", continue_text(honest, "my parcel"))
print("without mask:", continue_text(cheater, "my parcel"))

The output:

with mask, valid loss       3.559
without mask, future open   0.54
without mask, future hidden 11.345

with mask:    my parcel was charged twice for the garden bench ,
without mask: my parcel my problem problem problem problem problem problem problem

continue_text() is the generation loop: predict a word, add it, predict again. Read the three losses. With the future open, the model without a mask scores 0.54, because it copies the next word from the input. When the future is hidden, its loss jumps to 11.345, worse than a blind guess among 2,864 words (ln 2,864 ≈ 7.96). It learned to copy, not to predict. The honest model, 3.559, writes a sentence that sounds like a ticket, even if the parcel and the charge do not fit together. That is what three epochs on 4,696 tickets can teach.

Masked or next token: which to use

• Next-token prediction trains a decoder, and the same loop then writes text. It is the objective of GPT-like models such as SmolLM2. Each token learns only from its left side.
• Masked-token prediction trains an encoder, which reads both sides of every word. It suits reading tasks, such as classifying tickets. It does not write text one token after another.
• Both need no labels and enormous amounts of text. That is the cost.

Tomás: “Let’s train our own LLM”

Tomás has seen the honest result of the tiny transformer. “The fix is obvious: we pretrain our own large language model on support text.” Priya puts numbers on it.

SmolLM2-135M is a small model: 30 blocks and 134,515,008 parameters. Its model card on Hugging Face states: “Pretraining tokens: 2T”, “GPUs: 64 H100”, “Precision: bfloat16”. Its larger version, SmolLM2-1.7B, was “trained on 11 trillion tokens” on 256 H100 GPUs. An H100 is a data-centre graphics card built for training.

Compare that with Larkfield. The training tickets have 269,231 tokens. 2 trillion tokens is about 7 million times more text, for the smallest model of the family.

Priya also measured SmolLM2-135M’s training speed on the Mac mini of this course. It does 5 training steps on a batch of 4 texts × 128 tokens, on two CPU cores. Each step takes about 0.8 seconds: about 600 tokens per second. At that speed, 2 trillion tokens would take about 100 years. Memory is a problem too, as Inference costs explains. The weights, the gradients and Adam’s two running averages are about 4 × 538 MB = 2.2 GB for this small model, before the activations. A 70-billion-parameter model needs about 1.1 TB for the same parts.

Pretraining is not the end, either. A pretrained model only continues text. Instruction tuning and preference training then make it an assistant; Training stages in the language models course explains these stages.

So Larkfield should not pretrain. It should start from a model that someone else pretrained, and adapt it: fine-tune it on the tickets, or prompt it. Module 7 runs such a model. The case study Fine-tune a small transformer measures what pretraining buys against the tiny transformer, with the same data.

Guided practice: build the examples in the notebook

1. Open the notebook from the box at the top of this page, and run the setup cells.
2. In section 1, make the next-token pairs of the Larkfield sentence. You see 10 lines, from my -> parcel to lamp -> ..
3. In the Your turn cell, put the number of next-token examples of the validation tickets in valid_examples: their tokens minus one for each ticket. Run the check cell. It prints Check passed.
4. In section 2, make the masked version with seed 3. You see [MASK] in place of “broken” and “lamp”.
5. In section 3, train the two next-word models for 3 epochs each. You see a lower training loss without the mask.
6. Measure both on validation with the future hidden. The model without the mask has the higher loss.

Show an example answer to step 3

valid_examples = sum(len(tokenize(row["text"])) - 1 for row in valid_rows)
print(valid_examples)

Each ticket of n tokens gives n − 1 examples: every token after the first is an answer once. The check cell compares your number without showing it.

Your turn: a longer window

Change one condition: the next-word models read each ticket’s first 96 tokens instead of 48. Before you run it, predict: what is the shape of the scores for 2 tickets? Does the old model, built for 48 positions, still work? Does the model without the mask still look better in training and worse with the future hidden?

Show the answer

seq_train, seq_valid = sequences(train_rows, length=96), sequences(valid_rows, length=96)
try:
    NextWordModel(len(vocab))(seq_valid[:2, :-1])          # the position table has 48 rows
except IndexError as error:
    print("IndexError:", error)

def train_96(causal, epochs=3):
    set_seed(0)
    model = NextWordModel(len(vocab), length=96)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for epoch in range(epochs):
        model.train()
        order = torch.randperm(len(seq_train))
        for i in range(0, len(order), 32):
            opt.zero_grad()
            loss = next_word_loss(model, seq_train[order[i:i + 32]], causal=causal)
            loss.backward()
            opt.step()
    return model.eval(), round(loss.item(), 4)

honest96, last_honest = train_96(causal=True)
cheater96, last_cheater = train_96(causal=False)
with torch.no_grad():
    print("scores", tuple(honest96(seq_valid[:2, :-1]).shape))
    print("training loss, last step: with mask", last_honest, "| without mask", last_cheater)
    print("valid, future hidden: with mask", round(next_word_loss(honest96, seq_valid).item(), 3),
          "| without mask", round(next_word_loss(cheater96, seq_valid, causal=True).item(), 3))
seq_train, seq_valid = sequences(train_rows), sequences(valid_rows)   # back to 48 tokens

IndexError: index out of range in self
scores (2, 95, 2864)
training loss, last step: with mask 3.6497 | without mask 2.6696
valid, future hidden: with mask 3.629 | without mask 5.624

The old model stops: its position table has 48 rows, and position 48 does not exist. With a table of 96 rows, the scores are (2, 95, 2864): one row of scores for each position that has a next token. The pattern stays. Without the mask, the training loss is lower (2.67 against 3.65). With the future hidden, the loss is higher (5.624 against 3.629). The gap is smaller than with 48 tokens, because in 3 epochs this model has learned to copy less well. It still learns the wrong skill.

Common mistake

Training a next-token model without the causal mask

Symptom: the training loss falls fast and far: 0.54 against 3.41 for a model with the mask. But in the generation loop, the model writes “my parcel my problem problem problem …”, and its loss with the future hidden is 11.345.

Cause: without the mask, position i can read token i + 1, which is its own answer. The model learns to copy that answer from the input. When it writes, there is no next token to copy, so it has learned nothing useful.

Fix: use the causal mask in every next-token model, in training and in generation. Measure a decoder the way it is used: with the future hidden, and with a few generated samples, not only the training loss.

What it means for a real AI system: a training loss that looks too good often means that the model can see its answer. This is the same leakage as a test set that leaks into training, in Train, validation and test sets in the machine learning course.

See it end to endTrain a tiny language modelCase study

The tiny language model is trained with next-token prediction, with characters as tokens, on six plays.

Knowledge check

Check your understanding

Why do next-token and masked-token prediction need no human labels? Give the training examples that the sentence 'parcel arrived late' gives for next-token prediction.

Show the answer

Both objectives hide part of a real text and ask the model to predict it. The hidden part is the right answer, so the text is its own answer key, and nobody has to write a label. This is self-supervised learning. The three tokens “parcel arrived late” give two next-token examples: “parcel” → “arrived”, and “parcel arrived” → “late”. For masked-token prediction, one could hide “arrived” and ask for it from “parcel [MASK] late”.

Check your understanding

Tomás says: SmolLM2-135M fits in 2.2 GB for training, so Larkfield can pretrain one on a laptop. What does he miss?

Show the answer

Memory is only one cost; the amount of text and the compute time are the larger ones. The model card states 2 trillion pretraining tokens on 64 H100 GPUs. Larkfield’s training tickets have 269,231 tokens, about 7 million times less text. On the course’s Mac mini, SmolLM2-135M trains at about 600 tokens per second, so 2 trillion tokens would take about 100 years. The 2.2 GB also counts only weights, gradients and Adam’s states, without activations. Larger models need far more: about 1.1 TB for a 70-billion-parameter model. Larkfield should adapt a pretrained model instead.

Module practice

In the module practice, you train a tiny transformer classifier and annotate its input and output shapes. First follow the worked example, then do it on your own.

Worked example: a smaller transformer’s shapes

Priya builds a smaller TinyTransformer: 32 numbers per token, 2 heads, 1 block. She sends 4 validation tickets through it, step by step, and prints the shape after each step.

from ticketnet import TinyTransformer, token_ids, labels, train, count_parameters, TEAMS

set_seed(0)
small = TinyTransformer(len(vocab), dim=32, heads=2, blocks=1)
small.eval()
ids = token_ids(valid_rows[:4], vocab)
with torch.no_grad():
    h = small.embed(ids) + small.position(torch.arange(ids.shape[1]))
    print("ids               ", tuple(ids.shape))
    print("token + position  ", tuple(h.shape))
    h = small.blocks(h, src_key_padding_mask=ids == PAD)
    print("after 1 block     ", tuple(h.shape))
    keep = (ids != PAD).unsqueeze(-1).float()
    mean = (small.norm(h) * keep).sum(1) / keep.sum(1)
    print("average per ticket", tuple(mean.shape))
    print("scores            ", tuple(small.out(mean).shape))
print("parameters", count_parameters(small))

The output:

ids                (4, 93)
token + position   (4, 93, 32)
after 1 block      (4, 93, 32)
average per ticket (4, 32)
scores             (4, 5)
parameters 108421

Her annotation, one line for each shape:

• (4, 93): 4 tickets, padded to the longest of them, 93 token IDs.
• (4, 93, 32): each token ID becomes 32 numbers, its token vector plus its position vector.
• (4, 93, 32): the block keeps the shape; the padding mask hides the padding.
• (4, 32): the average over each ticket’s real tokens removes the token dimension.
• (4, 5): one score for each of the 5 teams.

Your task

Train small on the training tickets for 5 epochs. Use train(), seed 0 and learning rate 0.001. Then send the first 16 validation tickets through the trained model. Before you run it, write down the shape of the input IDs and of the scores, and annotate each one.

Show the model answer

ids_train, ids_valid = token_ids(train_rows, vocab), token_ids(valid_rows, vocab)
y_train, y_valid = labels(train_rows), labels(valid_rows)
set_seed(0)
small = TinyTransformer(len(vocab), dim=32, heads=2, blocks=1)
history = train(small, ids_train, y_train, ids_valid, y_valid, epochs=5, lr=1e-3, log=None)
for row in history:
    print(row)

small.eval()
batch16 = token_ids(valid_rows[:16], vocab)
with torch.no_grad():
    scores = small(batch16)
print("ids", tuple(batch16.shape), "-> scores", tuple(scores.shape))
print("predicted:", [TEAMS[i] for i in scores.argmax(1)[:4].tolist()])
print("real:     ", [row["team"] for row in valid_rows[:4]])

{'epoch': 1, 'train_loss': 1.4929, 'train_acc': 0.3844, 'valid_loss': 1.4961, 'valid_acc': 0.355}
{'epoch': 2, 'train_loss': 1.0709, 'train_acc': 0.5901, 'valid_loss': 1.0662, 'valid_acc': 0.6083}
{'epoch': 3, 'train_loss': 0.6611, 'train_acc': 0.7621, 'valid_loss': 0.6765, 'valid_acc': 0.7417}
{'epoch': 4, 'train_loss': 0.4474, 'train_acc': 0.8441, 'valid_loss': 0.4822, 'valid_acc': 0.81}
{'epoch': 5, 'train_loss': 0.322, 'train_acc': 0.8963, 'valid_loss': 0.358, 'valid_acc': 0.8617}
ids (16, 120) -> scores (16, 5)
predicted: ['payment', 'delivery', 'delivery', 'delivery']
real:      ['payment', 'delivery', 'delivery', 'delivery']

Shapes: (16, 120) is 16 tickets padded to 120 token IDs, because one of these 16 tickets reaches the 120-token limit. Inside, the shapes are (16, 120, 32) after the embeddings and after the block, and (16, 32) after the average. The scores are (16, 5): one score for each team.

Result: after 5 epochs, the small model is right on 86.2% of the validation tickets, and its loss still falls. It has 108,421 parameters, fewer than the tiny transformer’s 291,397, and it is still learning. Like the larger model, it starts with no knowledge of language.

Module check

Check your understanding

Retrieval: what does a transformer block contain, and what are the two masks in a transformer for?

Show the answer

A block contains multi-head attention and a feed-forward layer for each token. Each of the two has a residual connection (add the part’s output to its input) and a layer normalization (rescale each token’s numbers). The shape in equals the shape out, so blocks stack. The padding mask hides padding tokens, so that a ticket’s result does not depend on the length of its batch. The causal mask hides every token after the current one, so that a model that predicts the next token cannot read it.

Check your understanding

Predict: you set the position table of the trained tiny transformer to zero, and send a ticket and the same ticket with its words shuffled. What happens to the scores, and why?

Show the answer

The two sets of scores are exactly the same: in the lesson, both were −3.571, −2.937, 10.644, −2.797 and 0.005. Attention compares every query with every key, whatever their places are. The feed-forward layer, the norms and the average treat every token alike. Without position vectors, nothing in the model depends on the order. Shuffling the words only shuffles the token outputs, and their average is the same.

Check your understanding

Completion check: explain why a causal model must not inspect future tokens during training.

Show the answer

A causal model learns next-token prediction: at each position, the answer is the next token of the same text. If a position could read the tokens after it, it could copy its answer from the input instead of predicting it. The training loss would look very good, but the model would learn to copy, not to predict. When the model writes text, the future tokens do not exist yet, so there is nothing to copy, and it fails. In the lesson, the model trained without the mask reached a training loss of 0.54, against 3.41 with the mask. With the future hidden, its validation loss was 11.345, worse than a blind guess, and it wrote “my parcel my problem problem …”. The causal mask makes training match generation: each position sees only what it will see when it writes.

Recap

In this module, you assembled a transformer and saw how it is trained:

• A block is multi-head attention and a feed-forward layer, each with a residual connection and a layer normalization; the shape stays the same, so blocks stack.
• Position embeddings give attention the word order. The padding mask hides padding; the causal mask hides the future. Encoders read everything; decoders read only the past.
• Trained from scratch on 4,696 tickets, the tiny transformer is Larkfield’s weakest router: 91.2% on test, against 99.2% for the linear model.
• Next-token and masked-token prediction turn plain text into training examples with no labels. Without the causal mask, a model learns to copy.
• Pretraining at the scale of SmolLM2-135M takes 2 trillion tokens on 64 data-centre GPUs: about 100 years on this course’s Mac. Adapt a pretrained model instead.

Reset: in the notebook, select Runtime › Disconnect and delete runtime in Colab, then run all cells from the top. On your computer, run python get_data.py --reset for the data and ticketnet.py, or delete .venv and make it again as in Reset the setup. Then train the tiny transformer again with seed 0, as in Transformer blocks: Train it.

Next, in Module 7, you load a model that someone else pretrained, and read its 30 blocks with what you learned here.

Next: Training versus inference