7 - Pretraining and fine-tuning, in miniature

In chapter 7 Emma shows the team an iceberg: a huge mass under the water labelled pre-training and a small tip labelled fine-tuning. Almost no one trains an LLM from scratch for a specific task. First the model is pretrained on a huge amount of text, predicting the next token, and it comes out with a general knowledge of the language and of the world but optimized for nothing beyond completing text. Then it is fine-tuned for a task, with much less data and compute; one of the types in the book is instruction tuning, teaching the model to answer questions or follow instructions.

In this notebook we reproduce the two phases with a tiny GPT, like the one in the notebook A GPT-style language model:

1. Import the libraries.
2. The pretraining corpus: sheets and notes about 16 diagnoses.
3. Pretraining.
4. What the pretrained model knows, and what it does not do.
5. The fine-tuning data: questions about 10 of the 16 diagnoses.
6. Fine-tuning, with and without pretraining text.
7. The same questions, trained from scratch.

Step 1: Import the libraries

---

import copy
import math
import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0)

<torch._C.Generator at 0x10bf1e730>

Step 2: The pretraining corpus: sheets and notes about 16 diagnoses

---

The pretraining text is the kind of text a hospital produces, with no questions and no answers: 4,000 discharge sheets, like those of the GPT notebook, and 4,000 short clinical notes that state the same facts in other words ("reflux is treated with omeprazole", "patients with gout should watch for..."). There are 16 diagnoses, each with its drug and its warning sign. As in a real LLM, the vocabulary is fixed before pretraining and already contains the words of the questions of step 5, although the pretraining text never uses some of them.

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'), 'asthma': ('the inhaler', 'blue lips'), 'gout': ('colchicine', 'a hot swollen joint'),
         'heart failure': ('furosemide', 'swollen legs'), 'hypothyroidism': ('levothyroxine', 'palpitations'),
         'diabetes': ('metformin', 'extreme thirst'), 'reflux': ('omeprazole', 'black stools'), 'anemia': ('iron tablets', 'fainting'),
         'shingles': ('aciclovir', 'a rash near the eye'), 'thrombosis': ('the anticoagulant', 'sudden breathlessness'),
         'conjunctivitis': ('the eye drops', 'loss of vision'), 'otitis': ('amoxicillin', 'a stiff neck')}
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 .']
note_templates = ['{dx} is treated with {drug} .', 'patients with {dx} should watch for {warning} .',
                  '{drug} is used for {dx} .', 'in {dx} the warning sign is {warning} .']
drug_questions = ['what should a patient with {dx} take ?', 'which drug is used for {dx} ?', 'what is the treatment for {dx} ?']
warning_questions = ['when should a patient with {dx} come back ?', 'which warning sign matters in {dx} ?']
TAKE, QUESTION, ANSWER = 'take', ['question', ':'], ['answer', ':']

def discharge_sheet(rng):
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    words = ['diagnosis', ':'] + diagnosis.split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 4), replace=False):
        words += advice.split()
    words += [TAKE] + drug.split() + rng.choice(frequencies).split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 3), replace=False):
        words += advice.split()
    return words + ['come', 'back', 'if', 'you', 'have'] + warning.split() + ['.', '</s>'], diagnosis

def clinical_note(rng):
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    template = note_templates[rng.integers(len(note_templates))]
    return template.format(dx=diagnosis, drug=drug, warning=warning).split() + ['</s>']

def question_answer(diagnosis, template):
    drug, warning = cases[diagnosis]
    answer = drug if template in drug_questions else warning
    return QUESTION + template.format(dx=diagnosis).split() + ANSWER + answer.split() + ['.', '</s>'], answer

rng = np.random.default_rng(0)
documents = [discharge_sheet(rng)[0] for _ in range(4000)] + [clinical_note(rng) for _ in range(4000)]
documents = [documents[i] for i in np.random.default_rng(1).permutation(len(documents))]    # sheets and notes mixed

all_questions = [question_answer(d, t)[0] for d in cases for t in drug_questions + warning_questions]
vocab = sorted({w for doc in documents + all_questions for w in doc} | {'<s>'})
token_id = {w: i for i, w in enumerate(vocab)}
def to_stream(docs):
    return torch.tensor([token_id[w] for doc in docs for w in ['<s>'] + doc])
pretraining_stream = to_stream(documents)

for doc in documents[:4]:
    print(" ".join(doc))
print(f"\nDiagnoses: {len(cases)}, vocabulary: {len(vocab)} tokens, pretraining corpus: {len(pretraining_stream):,} tokens")

patients with asthma should watch for blue lips . </s>
diagnosis : hypothyroidism . you can go back to work next week . keep this sheet with you . avoid alcohol . take levothyroxine twice a day . eat light meals . keep this sheet with you . come back if you have palpitations . </s>
diagnosis : hypertension . keep this sheet with you . rest at home for a few days . take the blood pressure tablets twice a day . avoid alcohol . drink plenty of water . come back if you have chest pain . </s>
diagnosis : pneumonia . take the antibiotic every eight hours . do not drive today . come back if you have shortness of breath . </s>

Diagnoses: 16, vocabulary: 131 tokens, pretraining corpus: 182,218 tokens

Step 3: Pretraining

---

The model of the GPT notebook, with the same objective: predicting the next token in windows of 64 tokens. train takes batches from one stream of text or, if we give it a second one, half of each batch from each; we will need that in step 6. Pretraining is the long phase: 3,000 steps.

context_length = 64

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.qkv, self.out = nn.Linear(d_model, 3 * d_model), nn.Linear(d_model, d_model)
        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 = [m.view(batch, T, self.heads, d // self.heads).transpose(1, 2) for m in self.qkv(x).split(d, dim=2)]
        scores = (q @ k.transpose(-2, -1) / math.sqrt(d // self.heads)).masked_fill(~self.mask[:T, :T], float('-inf'))
        return self.out((scores.softmax(dim=-1) @ 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))
        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, self.positions = nn.Embedding(vocab_size, d_model), 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
        nn.init.normal_(self.tokens.weight, std=0.02)
        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))

def windows(stream, n, generator):
    starts = torch.randint(0, len(stream) - context_length - 1, (n,), generator=generator)
    return (torch.stack([stream[s:s + context_length] for s in starts]),
            torch.stack([stream[s + 1:s + context_length + 1] for s in starts]))

def train(model, stream, steps, lr, mix_with=None, seed=0):
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    generator = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        if mix_with is None:
            x, y = windows(stream, 32, generator)
        else:                                               # half of the batch from each stream
            (x1, y1), (x2, y2) = windows(stream, 16, generator), windows(mix_with, 16, generator)
            x, y = torch.cat([x1, x2]), torch.cat([y1, y2])
        loss = nn.functional.cross_entropy(model(x).reshape(-1, len(vocab)), y.reshape(-1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()

base = TinyGPT(len(vocab))
print(f"Parameters: {sum(p.numel() for p in base.parameters()):,}")
print(f"Loss after pretraining: {train(base, pretraining_stream, 3000, 3e-3):.3f}")

Parameters: 112,576

Loss after pretraining: 0.401

112,576 parameters, a few thousand more than in the GPT notebook because the vocabulary is larger, 131 tokens. After 3,000 steps the loss is 0.401; as there, it cannot reach zero, because which diagnosis, which advice and which note come next is random.

Step 4: What the pretrained model knows, and what it does not do

---

A pretrained model can only complete text. We give it the beginning of a note, "{diagnosis} is treated with", for each of the 16 diagnoses, and the beginning of a sheet up to "take". Then we ask it a question in the format of step 5 and see how it continues.

def continue_text(model, words, max_new=10):
    words = ['<s>'] + list(words)
    new = []
    model.eval()
    with torch.no_grad():
        while len(new) < max_new:
            token = vocab[model(torch.tensor([[token_id[w] for w in words + new]]))[0, -1].argmax()]
            if token in ('.', '</s>'):
                break
            new.append(token)
    model.train()
    return " ".join(new)

def sheet_drug_accuracy(model, n=200):
    """Share of new discharge sheets in which the model writes the right drug after 'take'."""
    sheet_rng, right = np.random.default_rng(777), 0
    for _ in range(n):
        words, diagnosis = discharge_sheet(sheet_rng)
        right += continue_text(model, words[:words.index(TAKE) + 1]).startswith(cases[diagnosis][0])
    return right / n

note_start = note_templates[0].split('{drug}')[0]
completed = sum(continue_text(base, note_start.format(dx=d).split()) == drug for d, (drug, _) in cases.items())
print(f"Completes '{note_start.format(dx='...')}' with the right drug: {completed} of {len(cases)} diagnoses")
print(f"Writes the right drug in new sheets: {sheet_drug_accuracy(base):.1%}")
question, answer = question_answer('reflux', drug_questions[0])
prompt = question[:question.index(ANSWER[0]) + 2]
print(f"\n{' '.join(prompt)} ...  ->  {continue_text(base, prompt)!r}   (the answer would be {answer!r})")

Completes '... is treated with ' with the right drug: 16 of 16 diagnoses

Writes the right drug in new sheets: 100.0%

question : what should a patient with reflux take ? answer : ...  ->  'heart failure'   (the answer would be 'omeprazole')

The pretrained model knows the facts: it completes the note with the right drug for the 16 diagnoses and writes the right drug in 100.0% of new sheets. But asked a question, it does not answer it: after "answer :" it writes "heart failure", a diagnosis instead of a drug, because it has never seen a question followed by its answer. It is what the book says of a pretrained model: it is not optimized for any task beyond completing text.

Step 5: The fine-tuning data: questions about 10 of the 16 diagnoses

---

Five question templates, three about the drug and two about the warning sign, for the first 10 diagnoses: 50 question-answer pairs. It is a tiny dataset compared with the 8,000 pretraining documents, as the book says. The other 6 diagnoses never appear in a question during fine-tuning: questions about them will tell us whether the model can use what it learned in pretraining.

names = list(cases)
seen, unseen = names[:10], names[10:]
templates = drug_questions + warning_questions
fine_tuning_examples = [question_answer(d, t)[0] for d in seen for t in templates]
fine_tuning_stream = to_stream(fine_tuning_examples)

def question_accuracy(model, diagnoses):
    right = 0
    for d in diagnoses:
        for t in templates:
            question, answer = question_answer(d, t)
            right += continue_text(model, question[:question.index(ANSWER[0]) + 2]) == answer
    return right, len(diagnoses) * len(templates)

print(" ".join(fine_tuning_examples[0]))
print(" ".join(fine_tuning_examples[3]))
print(f"\nFine-tuning examples: {len(fine_tuning_examples)}, about: {', '.join(seen)}")
print(f"Diagnoses kept for the test: {', '.join(unseen)}")

question : what should a patient with pneumonia take ? answer : the antibiotic . </s>
question : when should a patient with pneumonia come back ? answer : shortness of breath . </s>

Fine-tuning examples: 50, about: pneumonia, migraine, hypertension, gastroenteritis, cystitis, asthma, gout, heart failure, hypothyroidism, diabetes
Diagnoses kept for the test: reflux, anemia, shingles, thrombosis, conjunctivitis, otitis

Step 6: Fine-tuning, with and without pretraining text

---

We start from a copy of the pretrained model and train it on the 50 examples for 600 steps, with a learning rate ten times smaller than in pretraining, so as not to move the weights too far. We try it in two ways: with batches made only of questions, and with half of each batch taken from the pretraining text, a common way to keep the model from forgetting what it knew (replay). For each model we count the right answers about the 10 diagnoses of the fine-tuning data and about the 6 new ones, and we check whether it still writes the right drug in discharge sheets.

def report(name, model):
    right_seen, total_seen = question_accuracy(model, seen)
    right_new, total_new = question_accuracy(model, unseen)
    print(f"{name:<44} fine-tuning diagnoses {right_seen:>2}/{total_seen} | new diagnoses {right_new:>2}/{total_new} | "
          f"right drug in sheets {sheet_drug_accuracy(model):6.1%}")

only_questions = copy.deepcopy(base)
train(only_questions, fine_tuning_stream, 600, 3e-4)
with_replay = copy.deepcopy(base)
train(with_replay, fine_tuning_stream, 600, 3e-4, mix_with=pretraining_stream)

report('Pretrained, no fine-tuning', base)
report('Fine-tuned with questions only', only_questions)
report('Fine-tuned with questions + pretraining text', with_replay)
for d in unseen[:3]:
    question, answer = question_answer(d, drug_questions[0])
    prompt = question[:question.index(ANSWER[0]) + 2]
    print(f"   {' '.join(prompt)}  ->  only questions: {continue_text(only_questions, prompt)!r}, "
          f"with replay: {continue_text(with_replay, prompt)!r}  (right: {answer!r})")

Pretrained, no fine-tuning                   fine-tuning diagnoses  0/50 | new diagnoses  0/30 | right drug in sheets 100.0%

Fine-tuned with questions only               fine-tuning diagnoses 50/50 | new diagnoses  9/30 | right drug in sheets   0.0%

Fine-tuned with questions + pretraining text fine-tuning diagnoses 50/50 | new diagnoses 23/30 | right drug in sheets 100.0%
   question : what should a patient with reflux take ? answer :  ->  only questions: 'levothyroxine', with replay: 'omeprazole'  (right: 'omeprazole')
   question : what should a patient with anemia take ? answer :  ->  only questions: 'gout ? answer : colchicine', with replay: 'iron tablets'  (right: 'iron tablets')
   question : what should a patient with shingles take ? answer :  ->  only questions: 'gout ? answer : colchicine', with replay: 'aciclovir twice a day'  (right: 'aciclovir')

Both fine-tuned models answer the 50 questions of their data. The difference is in the 30 questions about the 6 new diagnoses. With questions only, the model gets 9 right and, worse, forgets how to write sheets: the right drug in 0.0% of them. For reflux it answers levothyroxine, the drug of one of its fine-tuning diagnoses. With half of each batch taken from the pretraining text, it gets 23 of the 30 right and keeps writing sheets correctly, 100.0%. Its answers about reflux, anemia and shingles name the right drugs, which it can only have learned in pretraining. The count is strict: "aciclovir twice a day" counts as a mistake.

Step 7: The same questions, trained from scratch

---

Is it the pretraining that makes the difference? We train a model with the same architecture, starting from random weights, on the same 50 questions, 600 steps with the learning rate of pretraining.

torch.manual_seed(0)
from_scratch = TinyGPT(len(vocab))
train(from_scratch, fine_tuning_stream, 600, 1e-3)
report('Trained from scratch with the questions', from_scratch)
question, answer = question_answer(unseen[0], drug_questions[0])
prompt = question[:question.index(ANSWER[0]) + 2]
print(f"   {' '.join(prompt)}  ->  {continue_text(from_scratch, prompt)!r}  (right: {answer!r})")

Trained from scratch with the questions      fine-tuning diagnoses 50/50 | new diagnoses  0/30 | right drug in sheets   0.0%
   question : what should a patient with reflux take ? answer :  ->  'furosemide'  (right: 'omeprazole')

The model trained from scratch learns the 50 questions, 50 of 50, but gets none of the 30 about the new diagnoses right: nothing in its data says what reflux is treated with, so it answers furosemide, one of the drugs it has seen in the questions. And, of course, it cannot write discharge sheets. The difference between 0 and 23 of 30 is what pretraining contributes.

Conclusion

---

• Pretraining on plain text, predicting the next token, puts knowledge into the model, but the pretrained model only completes text: it does not answer questions.
• Fine-tuning with a few examples teaches it a new behaviour, here answering questions, at a small fraction of the cost of pretraining.
• What the model can do after fine-tuning depends on both phases, as the book warns: the fine-tuned model answers about diagnoses that were never in a question because it learned about them in pretraining. A model trained only on the questions cannot know them.
• Fine-tuning can also make the model forget what it learned. Mixing pretraining text into the fine-tuning batches is a simple way to keep it.
• The third type of fine-tuning in the book, alignment with human preferences, is built in the notebook Text generation and learning from preferences of chapter 8.