7 - A GPT-style language model, from scratch

In chapter 7 Emma explains the structure of a GPT-style LLM: token embeddings with position information, a stack of Transformer blocks with causal self-attention (each token only sees itself and the tokens before it), and a final projection, the LM head, that turns the last hidden state into probabilities for the next token. And she explains the training objective, "surprisingly simple in concept": predicting the next token, with the cross-entropy as the loss, backpropagated through the whole network.

In this notebook we build and train one, tiny, with PyTorch:

1. Import the libraries.
2. A corpus where the context matters.
3. From text to training examples.
4. The model: causal attention, Transformer blocks and the LM head.
5. Training: predicting the next token.
6. What the context gives: the model against a trigram.
7. Where the model looks: the attention weights.
8. Generating discharge sheets.

Step 1: Import the libraries

---

PyTorch computes the gradients of all the parameters automatically, the step the book calls backpropagating the loss "through the entire network". The pieces of the model, attention included, are written by hand.

import collections
import math
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

torch.manual_seed(0)

<torch._C.Generator at 0x110549350>

Step 2: A corpus where the context matters

---

A real LLM is pretrained on a large part of the internet. We generate 3,000 synthetic discharge sheets: a diagnosis, some general advice, the drug to take, more advice and the warning sign that should bring the patient back. The drug and the warning sign depend on the diagnosis, which appears several sentences earlier; the general advice does not depend on anything. To predict the drug well, a model has to look back beyond the last few words.

cases = {'pneumonia': ('the antibiotic', 'shortness of breath'),
         'migraine': ('ibuprofen', 'blurred vision'),
         'hypertension': ('the blood pressure tablets', 'chest pain'),
         'gastroenteritis': ('oral rehydration salts', 'blood in the stool'),
         'cystitis': ('fosfomycin', 'back pain')}
frequencies = ['twice a day', 'every eight hours', 'once a day']
general_advice = ['drink plenty of water .', 'rest at home for a few days .', 'do not drive today .', 'eat light meals .',
                  'avoid alcohol .', 'you can go back to work next week .', 'keep this sheet with you .']
DIAGNOSIS, TAKE, WARNING = ['diagnosis', ':'], ['take'], ['come', 'back', 'if', 'you', 'have']

def discharge_sheet(rng):
    """One discharge sheet as a list of words: diagnosis, some advice, the drug, more advice, the warning sign."""
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    words = DIAGNOSIS + [diagnosis, '.']
    for advice in rng.choice(general_advice, rng.integers(0, 4), replace=False):     # 0 to 3 sentences
        words += advice.split()
    words += TAKE + drug.split() + rng.choice(frequencies).split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 3), replace=False):     # 0 to 2 sentences
        words += advice.split()
    words += WARNING + warning.split() + ['.', '</s>']
    return words, diagnosis

rng = np.random.default_rng(0)
sheets = [discharge_sheet(rng) for _ in range(3000)]
vocab = sorted({w for words, _ in sheets for w in words} | {'<s>'})
token_id = {w: i for i, w in enumerate(vocab)}
stream = torch.tensor([token_id[w] for words, _ in sheets for w in ['<s>'] + words])   # all the sheets, one after another

for words, _ in sheets[:2]:
    print(" ".join(words), "\n")
print(f"Vocabulary: {len(vocab)} tokens (one per word), corpus: {len(stream):,} tokens, "
      f"longest sheet: {max(len(words) for words, _ in sheets)} tokens")

diagnosis : cystitis . rest at home for a few days . eat light meals . take fosfomycin twice a day . come back if you have back pain . </s> 

diagnosis : pneumonia . take the antibiotic once a day . keep this sheet with you . come back if you have shortness of breath . </s> 

Vocabulary: 70 tokens (one per word), corpus: 107,700 tokens, longest sheet: 59 tokens

3,000 sheets make a corpus of 107,700 tokens with a vocabulary of only 70 words, one token per word (a real LLM uses subword tokens, as in the notebook Tokens and embeddings). The longest sheet has 59 tokens, so a context of 64 tokens always holds a whole sheet.

Step 3: From text to training examples

---

The book's example: from "The doctor diagnosed the patient with pneumonia" come the examples Input = "The", Target = "doctor"; Input = "The doctor", Target = "diagnosed"; and so on. In practice the text is cut into windows of a fixed length, here 64 tokens, and a window gives all those examples at once: the target sequence is the input sequence shifted one position to the left. The causal mask of the next step makes sure that the prediction at each position only sees the tokens up to that position.

context_length = 64

def random_windows(batch_size, generator):
    starts = torch.randint(0, len(stream) - context_length - 1, (batch_size,), generator=generator)
    x = torch.stack([stream[s:s + context_length] for s in starts])
    y = torch.stack([stream[s + 1:s + context_length + 1] for s in starts])     # the same window, one token later
    return x, y

x, y = stream[None, :context_length], stream[None, 1:context_length + 1]    # the first window of the corpus
for t in range(6):
    print(f"Input = {' '.join(vocab[i] for i in x[0, :t + 1])!r:<45} Target = {vocab[y[0, t]]!r}")

Input = '<s>'                                         Target = 'diagnosis'
Input = '<s> diagnosis'                               Target = ':'
Input = '<s> diagnosis :'                             Target = 'cystitis'
Input = '<s> diagnosis : cystitis'                    Target = '.'
Input = '<s> diagnosis : cystitis .'                  Target = 'rest'
Input = '<s> diagnosis : cystitis . rest'             Target = 'at'

Step 4: The model: causal attention, Transformer blocks and the LM head

---

The same pieces as in the notebook Attention and the Transformer, now as trainable PyTorch layers:

• Causal self-attention with 4 heads: queries, keys and values are projections of the token vectors, $\text{softmax}\big(QK^\top / \sqrt{d_k}\big)V$, and a triangular mask sets to $-\infty$ the scores of future tokens before the softmax. We keep the last attention weights to look at them in step 7.
• A Transformer block with Pre-LayerNorm, as in many modern LLMs: $x + \text{Attention}(\text{LN}(x))$, then $x + \text{FFN}(\text{LN}(x))$, with residual connections.
• Embeddings of the tokens plus learned embeddings of the positions, 2 blocks, and the LM head, $z_t = W_\text{vocab}\, h_t$, which shares its matrix with the token embeddings (weight tying).

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.qkv = nn.Linear(d_model, 3 * d_model)         # W_Q, W_K and W_V in a single matrix
        self.out = nn.Linear(d_model, d_model)             # W_O
        self.register_buffer('mask', torch.tril(torch.ones(context_length, context_length)).bool())

    def forward(self, x):
        batch, T, d = x.shape
        q, k, v = self.qkv(x).split(d, dim=2)
        q, k, v = [m.view(batch, T, self.heads, d // self.heads).transpose(1, 2) for m in (q, k, v)]
        scores = q @ k.transpose(-2, -1) / math.sqrt(d // self.heads)
        scores = scores.masked_fill(~self.mask[:T, :T], float('-inf'))   # no looking at future tokens
        self.weights = scores.softmax(dim=-1)
        return self.out((self.weights @ v).transpose(1, 2).reshape(batch, T, d))

class TransformerBlock(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attention = CausalSelfAttention(d_model, heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model))

    def forward(self, x):
        x = x + self.attention(self.norm1(x))              # Pre-LayerNorm and residual connection
        return x + self.ffn(self.norm2(x))

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model=64, heads=4, layers=2):
        super().__init__()
        self.tokens = nn.Embedding(vocab_size, d_model)
        self.positions = nn.Embedding(context_length, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, heads) for _ in range(layers)])
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.tokens.weight            # weight tying
        nn.init.normal_(self.tokens.weight, std=0.02)      # small initial vectors, as in GPT
        nn.init.normal_(self.positions.weight, std=0.02)

    def forward(self, x):
        h = self.tokens(x) + self.positions(torch.arange(x.shape[1]))
        for block in self.blocks:
            h = block(h)
        return self.lm_head(self.norm(h))                  # logits for the next token at every position

model = TinyGPT(len(vocab))
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

Parameters: 108,672

108,672 parameters. A GPT-3-sized model has more than a million times as many, with the same pieces: more dimensions, more heads and many more blocks.

Step 5: Training: predicting the next token

---

1,500 steps with batches of 32 windows. At every position of every window the model predicts the next token, and the loss is the average cross-entropy of all those predictions. loss.backward() computes the gradients of every parameter, embeddings, attention projections, feed-forward layers, normalizations and the LM head, and the optimizer, AdamW (a variant of gradient descent), updates them all together.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
generator = torch.Generator().manual_seed(0)
losses = []
for step in range(1, 1501):
    x, y = random_windows(32, generator)
    logits = model(x)
    loss = nn.functional.cross_entropy(logits.reshape(-1, len(vocab)), y.reshape(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
model.eval()
print(f"Loss at the first step: {losses[0]:.3f} (ln {len(vocab)} = {math.log(len(vocab)):.3f}), "
      f"mean of the last 100 steps: {np.mean(losses[-100:]):.3f}")

plt.figure(figsize=(8, 3.5))
plt.plot(losses, lw=0.8)
plt.xlabel('step'); plt.ylabel('cross-entropy'); plt.title('Training by next-token prediction'); plt.grid(True); plt.show()

Loss at the first step: 4.262 (ln 70 = 4.248), mean of the last 100 steps: 0.308

At the first step the loss is 4.262, almost exactly $\ln 70$: the untrained model gives about the same probability to each of the 70 tokens. At the end it is around 0.308. It cannot reach zero, because part of each sheet is random, which diagnosis, which pieces of advice and in which order, and no model can predict that. What it can predict almost exactly are the words that follow from the context.

Step 6: What the context gives: the model against a trigram

---

We generate 500 new sheets and, for each one, give the model the text up to the word before the drug, and then up to the word before the warning sign, and let it complete the phrase greedily, always taking the most probable token, until the full stop. We do the same with a trigram model, the n-gram of chapter 8 with a context of two words, counted on the same 3,000 sheets.

def next_token_probs(words):
    with torch.no_grad():
        return model(torch.tensor([[token_id[w] for w in words]]))[0, -1].softmax(dim=-1).numpy()

trigram_counts = collections.defaultdict(collections.Counter)
for words, _ in sheets:
    padded = ['<s>', '<s>'] + words
    for i in range(2, len(padded)):
        trigram_counts[tuple(padded[i - 2:i])][padded[i]] += 1

def complete(prefix, predict):
    words = list(prefix)
    while words[-1] != '.':
        words.append(predict(words))
    return " ".join(words[len(prefix):-1])

gpt_next = lambda words: vocab[next_token_probs(words).argmax()]
trigram_next = lambda words: trigram_counts[tuple((['<s>', '<s>'] + words)[-2:])].most_common(1)[0][0]

test_rng = np.random.default_rng(1000)
scores = {'GPT': [0, 0], 'trigram': [0, 0]}
for _ in range(500):
    words, diagnosis = discharge_sheet(test_rng)
    drug, warning = cases[diagnosis]
    drug_prefix = ['<s>'] + words[:words.index(TAKE[-1]) + 1]
    warning_prefix = ['<s>'] + words[:len(words) - len(warning.split()) - 2]
    for name, predict in [('GPT', gpt_next), ('trigram', trigram_next)]:
        scores[name][0] += complete(drug_prefix, predict).startswith(drug)
        scores[name][1] += complete(warning_prefix, predict) == warning
for name, (drug_ok, warning_ok) in scores.items():
    print(f"{name:<8} right drug: {drug_ok / 500:.1%}, right warning sign: {warning_ok / 500:.1%}")

example_rng = np.random.default_rng(7)
while True:                                          # an example with at least 15 tokens between diagnosis and drug
    example, _ = discharge_sheet(example_rng)
    if example.index(TAKE[-1]) - len(DIAGNOSIS) >= 15:
        break
prefix = ['<s>'] + example[:example.index(TAKE[-1]) + 1]
p = next_token_probs(prefix)
print(f"\n{' '.join(prefix[1:])} ...")
print("GPT, next token:", ", ".join(f"{vocab[i]} {p[i]:.3f}" for i in np.argsort(-p)[:3]))
counts = trigram_counts[tuple(prefix[-2:])]
print("Trigram, next token:", ", ".join(f"{w} {c / sum(counts.values()):.3f}" for w, c in counts.most_common(3)))

GPT      right drug: 100.0%, right warning sign: 100.0%
trigram  right drug: 22.2%, right warning sign: 22.2%

diagnosis : hypertension . do not drive today . rest at home for a few days . keep this sheet with you . take ...
GPT, next token: the 0.999, plenty 0.001, fosfomycin 0.000
Trigram, next token: the 0.410, fosfomycin 0.201, ibuprofen 0.198

The GPT model completes the right drug and the right warning sign in 100.0% of the 500 new sheets. The trigram gets 22.2% of each: it only sees the last two words, the same for every diagnosis, so it always completes the same phrase and is right only when the diagnosis happens to be the one that goes with it. In the example, with three sentences of advice between the diagnosis and the drug, the model gives 0.999 to "the" (of "the blood pressure tablets"), while the trigram spreads its probability over the beginnings of all the drugs.

Step 7: Where the model looks: the attention weights

---

To predict the drug, the model must find the diagnosis. We pass the example sheet up to the word before the drug and show, for each of the 4 heads of the 2 blocks, the attention weights from that last position to every previous token.

with torch.no_grad():
    model(torch.tensor([[token_id[w] for w in prefix]]))
diagnosis_position = prefix.index(DIAGNOSIS[-1]) + 1
rows, labels = [], []
for b, block in enumerate(model.blocks, start=1):
    weights = block.attention.weights[0, :, -1].numpy()             # (heads, positions): from the last token
    for head, w in enumerate(weights, start=1):
        rows.append(w); labels.append(f'block {b}, head {head}')
        print(f"block {b}, head {head}: weight on the diagnosis {w[diagnosis_position]:.3f}, "
              f"largest on {prefix[w.argmax()]!r} ({w.max():.3f})")

plt.figure(figsize=(14, 4))
plt.imshow(np.array(rows), cmap='Blues', vmin=0, vmax=1, aspect='auto')
plt.xticks(range(len(prefix)), prefix, rotation=90); plt.yticks(range(len(labels)), labels)
plt.gca().get_xticklabels()[diagnosis_position].set_color('red')
plt.colorbar(label='attention weight'); plt.title(f'Attention from {prefix[-1]!r} to the previous tokens (the diagnosis in red)')
plt.tight_layout(); plt.show()

block 1, head 1: weight on the diagnosis 0.037, largest on 'take' (0.380)
block 1, head 2: weight on the diagnosis 0.819, largest on 'hypertension' (0.819)
block 1, head 3: weight on the diagnosis 0.006, largest on '.' (0.148)
block 1, head 4: weight on the diagnosis 0.043, largest on 'take' (0.174)
block 2, head 1: weight on the diagnosis 0.015, largest on '.' (0.268)
block 2, head 2: weight on the diagnosis 0.035, largest on ':' (0.367)
block 2, head 3: weight on the diagnosis 0.035, largest on 'take' (0.139)
block 2, head 4: weight on the diagnosis 0.008, largest on '.' (0.197)

Nobody told the model where the diagnosis is. Even so, in the first block, head 2 puts 0.819 of its attention on "hypertension", 21 tokens back, and almost nothing on the advice in between. The other heads and the second block look at nearby tokens, "take", the full stops, the colon. The training loss alone, predicting the next token, has made one head specialize in finding the diagnosis.

Step 8: Generating discharge sheets

---

To generate, the model reads the text so far, turns the logits of the last position into probabilities with a softmax at temperature $T$, samples the next token, appends it and repeats, one token at a time, until </s>. We generate 200 sheets with $T = 1$ and check whether each one is coherent: the drug and the warning sign are the ones that go with its diagnosis.

def generate(sample_rng, temperature=1.0):
    words = ['<s>']
    while words[-1] != '</s>' and len(words) < context_length:
        with torch.no_grad():
            logits = model(torch.tensor([[token_id[w] for w in words]]))[0, -1].numpy() / temperature
        p = np.exp(logits - logits.max())
        words.append(vocab[sample_rng.choice(len(vocab), p=p / p.sum())])
    return words[1:]

def coherent(words):
    text = " ".join(words)
    return any(text.startswith(" ".join(DIAGNOSIS + [d])) and " ".join(TAKE + drug.split()) in text
               and " ".join(WARNING + warning.split() + ['.', '</s>']) in text for d, (drug, warning) in cases.items())

sample_rng = np.random.default_rng(42)
generated = [generate(sample_rng) for _ in range(200)]
for words in generated[:3]:
    print(" ".join(words), "\n")
print(f"Coherent sheets: {sum(coherent(w) for w in generated)} of {len(generated)}")
print("First incoherent one:", " ".join(next(w for w in generated if not coherent(w))))
print(f"Sheets identical to one in the corpus: {sum(tuple(w) in {tuple(s) for s, _ in sheets} for w in generated)} of {len(generated)}")

diagnosis : pneumonia . do not drive today . keep this sheet with you . eat light meals . take the antibiotic twice a day . come back if you have shortness of breath . </s> 

diagnosis : gastroenteritis . keep this sheet with you . eat light meals . drink plenty of water . take oral rehydration salts every eight hours . keep this sheet with you . come back if you have blood in the stool . </s> 

diagnosis : pneumonia . avoid alcohol . rest at home for a few days . take the antibiotic every eight hours . come back if you have shortness of breath . </s> 

Coherent sheets: 191 of 200
First incoherent one: diagnosis diagnosis : hypertension . you can go back to work next week . eat light meals . take the blood pressure tablets twice a day . drink plenty of water . keep this sheet with you . come back if you have chest pain . </s>

Sheets identical to one in the corpus: 88 of 200

191 of the 200 generated sheets are coherent. The advice changes from sheet to sheet, because it is sampled, and 88 of the 200 happen to be identical to a sheet of the corpus: with so few possible combinations, the model often lands on one it has seen, but most of the sheets are new. The first incoherent one repeats the word "diagnosis": a sampled token with a small probability that the model then has to continue as best it can. At the scale of this toy, it is the same phenomenon the book calls hallucination: the model generates what is probable, not what is true.

Conclusion

---

• A GPT-style model is a stack of Transformer blocks with causal attention between token embeddings and an LM head that gives the probabilities of the next token.
• It is trained with a single objective, predicting the next token, and a single loss, the cross-entropy, whose gradient reaches every parameter of the network at once.
• Attention lets it use information from far back in the text, which an n-gram cannot do: that is the difference between predicting the right drug almost always and guessing.
• Generating text is repeating the prediction one token at a time, sampling from the probabilities.
• Our model has about a hundred thousand parameters and 3,000 very regular sheets; the LLMs of the book have billions of parameters and a large part of the internet, but the objective and the architecture are the same. The notebook Pretraining and fine-tuning takes the next step of the book: adapting a pretrained model to a new task.