Training it: teaching the model to write Tang poetry

The model is built; now we train it. This lesson's code, train.py, is just over a hundred lines, with the same structure as Module 08's training loop: take a batch of data, compute the loss, backpropagate, update the parameters.

python train.py

Data: shifted by one

A language model's training data needs no human labelling. A piece of text is itself both question and answer:

def get_batch(ids, generator=None):
    """随机截取 BATCH 段长度为 block_size 的文字。目标 y 就是 x 往后错一位：每个位置都要预测下一个字。"""
    starts = torch.randint(len(ids) - cfg.block_size - 1, (BATCH,), generator=generator)
    x = torch.stack([ids[s:s + cfg.block_size] for s in starts])
    y = torch.stack([ids[s + 1:s + cfg.block_size + 1] for s in starts])
    return x, y

All the poems are joined with newlines into one long string, 32 random segments are taken, 128 characters each, as the input x; shift each segment one place later, and that's the answer y. For example, if the input is "白日依山尽，", the answer is "日依山尽，黄": position 1 sees "白" and must guess "日", position 2 sees "白日" and must guess "依". As last lesson said, the causal mask ensures no position can peek at its answer.

The newline has a special meaning here: it marks the end of one poem and the start of the next. The model learns "a full stop is followed by a newline", and during generation it can stop when it produces a newline; the poem is finished.

The last 1,000 poems are held out as a validation set, never seen in training. That way we can use Module 08, Lesson 6's method to see whether the model is learning or reciting.

训练集 34135 首诗，1547255 个词元；验证集 1000 首；词表 6289 个字符
模型：4 层，4 个头，向量维度 128，共 1,614,720 个参数

A few details of the training loop

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.1)
for step in range(1, args.steps + 1):
    for group in optimizer.param_groups:
        group["lr"] = lr_at(step - 1)
    x, y = get_batch(train_ids)
    _, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪：防止偶尔一步梯度太大把训练带飞
    optimizer.step()

Compared with Module 08, there are two additions.

The learning rate changes. For the first 100 steps it rises slowly from 0 to 0.001; this is called warmup. At the start of training the parameters are random and the gradient's direction is unreliable, so steps that are too big easily lead astray. After that, the learning rate falls slowly along a cosine curve to 0.0001; smaller steps late in training can find low-loss regions more precisely.

def lr_at(step, peak=1e-3, warmup=100):
    """学习率：前 100 步从 0 慢慢升上去（预热），然后按余弦曲线降到峰值的十分之一。"""
    if step < warmup:
        return peak * (step + 1) / warmup
    progress = (step - warmup) / max(1, args.steps - warmup)
    return peak * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress)))

Gradient clipping. If the gradient at some step is especially large (from hitting an unusual batch of data), it's scaled down proportionally so its overall size doesn't exceed 1. As Module 08, Lesson 2 showed, one step that's too big can make training diverge outright.

"Warmup + cosine decay + gradient clipping + AdamW" is the most common combination for training Transformers, and LLM training is basically the same, just with different numbers.

Watching it learn to write poetry

Train for 3,000 steps, checking the loss every so often and having it write two poems:

训练前：训练损失 8.779，验证损失 8.778（随便猜的话是 ln(6289) = 8.747）
  训练前随便写的： 耦呌同捐睢寮办猬囊斵樬劫鞑顷饼遁瞩赴冈譀嗾雅踯牧写禧氛迹樯毁尤平螵蔻郓祐菼灺嵊匦

第 300 步（27 秒）：训练损失 5.778，验证损失 5.818
  将南。
  门不独。到李斜，功觅青。谁将风花不可碧，不和山。

第 1000 步（97 秒）：训练损失 4.777，验证损失 4.879
  山里怅望两间道，千里长亭寺断肠。曾梳白兰草，试向玉关鱼。
  此日醉前年，相思高至兹。云如汉陵子，暮水九重宫。

第 2000 步（249 秒）：训练损失 4.410，验证损失 4.583
  秋江野鸟过高楼，野树猿声怨见人。处处秋风满孤照，沧海无端行处闻。
  春山度滟月，多少漫为秋。药罢已开葬，松阴不道开。

第 3000 步（407 秒）：训练损失 4.280，验证损失 4.488
  江山古馆响幽幽，山鸟无人见白头。此时猩猩争得语，又将杯酒醉参差。
  分明人在泪，明月更经过。小谷闲烟树，红潭古石床。归来扶白首，立向卧青山。独有安行处，如何却得还。

训练用了 408 秒，模型存到 .cache/gpt.pt

On my computer (an Apple M4, CPU only), 3,000 steps took under 7 minutes. Each step processes 32×128 = 4,096 characters, so 3,000 steps saw about 12 million characters, roughly 8 passes over the training set.

What it learns appears in stages:

• Before training: a string of random characters, loss 8.78, just guessing.
• Step 300: it has learned to use punctuation and writes common characters, but lines vary in length.
• Step 1,000: most lines are five or seven characters, alternating commas and full stops. But the first poem has two seven-character lines followed by two five-character ones; it doesn't yet know a poem should be consistent.
• Step 2,000: the format is basically right, and lines with coherent imagery start to appear, like "秋江野鸟过高楼，野树猿声怨见人" ("autumn river, wild birds pass the tall tower; wild trees, gibbons' cries resent the sight of men").
• Step 3,000: the first poem is a complete seven-character quatrain, "江山古馆响幽幽，山鸟无人见白头" ("rivers and hills, the old lodge echoes faintly; mountain birds, no one sees the white-haired"); the second is a complete eight-line five-character regulated verse.

What it hasn't learned is just as clear: the meaning often doesn't connect ("此时猩猩争得语", "at this moment the orangutans compete to speak"), and it pays no attention to tonal patterns or rhyme. For 1.61 million parameters and 7 minutes of training, this is already pretty good.

Throughout training, training loss and validation loss both fall and stay close (4.28 and 4.49 at the end). By Module 08, Lesson 6's account, this means there's no obvious overfitting: it does almost as well on poems it hasn't seen.

Is it writing poems or reciting them?

Earlier modules of this course have said it again and again: good-looking output doesn't mean the model really learned; it may just have memorised the training data. So after training, the script runs a check: generate 100 poems and see how many lines are identical to some line in the training set.

train_sentences = {s for p in train_poems for s in p.replace("。", "，").split("，") if s}
torch.manual_seed(1)
generated = sample(100)
sentences = [s for p in generated for s in p.replace("。", "，").split("，") if s]
copied = sum(s in train_sentences for s in sentences)

生成 100 首诗，共 576 句，其中 0 句（0%）和训练集里的某一句一模一样
整首和训练集里某一首一模一样的：0 首

576 lines, not one copied. What it learned is how Tang poems are "written": the format, the common characters and words, which characters tend to appear together; then it combines them into new lines.

Give it only 300 poems

Module 08, Lesson 6 said a model recites when data is too scarce. Now try the same model with only 300 poems:

python train.py --poems 300 --steps 1500 --out small.pt

训练集 300 首诗，13612 个词元；验证集 1000 首；词表 6289 个字符

第 150 步（22 秒）：训练损失 5.079，验证损失 6.584
第 500 步（70 秒）：训练损失 0.421，验证损失 8.864
第 1000 步（145 秒）：训练损失 0.103，验证损失 9.817
  雕鹗途程在碧天，彩衣东去复何言。二千宾客旧知己，十二山河新故园。吟看桂生溪月上，醉听鲲化海涛翻。好期圣代重相见，莫学袁生老竹轩。
第 1500 步（230 秒）：训练损失 0.071，验证损失 9.984

生成 100 首诗，共 593 句，其中 454 句（77%）和训练集里的某一句一模一样
整首和训练集里某一首一模一样的：40 首

Training loss fell to 0.07, far lower than the 4.28 with all the data. But validation loss rose all the way to 9.98, higher even than the 8.78 of random guessing before training.

The seven-character regulated verse written at step 1,000 is beautifully neat, because it was recited word for word from the training set. The final check confirms it: 77% of lines were copied, and 40 of the 100 poems were copied whole.

With 1.61 million parameters and only 13,000 characters to learn, the model is entirely capable of memorising them all. Memorising is the easiest way to lower the training loss. On poems it hasn't seen, it actually does worse than having learned nothing at all: it's too "confident" about what it memorised, and given another poem, the probability it gives the right answer is pitifully low.

The same model, the same code, with data cut from 34,000 poems to 300, and the result changes from "learned to write poetry" to "recites poetry". This is Module 08, Lesson 6's conclusion replayed on a language model, and it's why LLMs need massive data.

What a loss of 4.49 means

The final validation loss is 4.49. Cross-entropy is -log(probability of the correct character), so on average the model gives the correct next character a probability of about e to the power -4.49, roughly 1.1%.

That sounds low, but remember: when writing poetry the next character genuinely has many reasonable choices. After "白日依山尽，黄河入海" ("the white sun sinks behind the mountains; the Yellow River into the sea"), the model gives some probability to "无", "间", "多" and "深" (you'll see the actual numbers next lesson). The original poem has "流" (flows), but writing another character isn't necessarily wrong. The loss can't fall to 0, for the same reason as Module 08, Lesson 1's "the data has random fluctuation, so the loss can never reach 0".

One yardstick is to compare with random guessing: picking evenly from 6,289 characters gives a loss of 8.75. 4.49 is equivalent to narrowing the choice from 6,289 characters down to about e to the power 4.49, roughly 89 characters.

Exercises

1. Change --steps to 6,000. Does the loss keep falling? Do the generated poems get better?
2. Set the number of training poems to 1,000, 3,000 and 10,000 (--poems), train each for 1,500 steps, and record the validation loss and the share of "copied lines" in a table.
3. Remove the learning rate warmup and decay (have lr_at simply return 0.001), train for 3,000 steps as before, and compare the final validation loss.

Self-check

1. When training a language model, what are the input and the answer?

A segment of text is taken as the input, and the same segment shifted one place later is the answer. Each position has to predict the next character from itself and the characters before it. A 128-character segment provides 128 prediction questions at once.

2. What problems do learning rate warmup and gradient clipping each solve?

Warmup: at the start of training the parameters are random and the gradient's direction is unreliable, so a large learning rate from the start easily leads astray; the learning rate starts small and rises gradually. Gradient clipping: occasionally a step's gradient is especially large, and updating with it directly would send the parameters too far and might make training diverge, so overlarge gradients are scaled down proportionally.

3. When training on only 300 poems, why is the training loss very low while the validation loss is higher than random guessing?

The model has enough parameters to memorise all 300 poems, so the training loss becomes very low. But what it learned is those 300 poems, not the patterns of writing poetry. It's very confident about what it memorised, so on poems it hasn't seen it gives the right answer a very low probability, and the cross-entropy ends up higher than random guessing. That 77% of the generated lines were copied confirms it.