Pretraining on Unlabeled Data

A single loss value is the entire training signal. Point an optimizer at it for ten epochs and watch what happens.

That model can generate — but always the same way. To control how it picks the next token, we reach past argmax.

None of this matters if the weights are bad. Our 20 KB model only memorized. The fix isn't more clever decoding — it's real pretraining, which we can borrow.

The lab

Drive the real thing. These are OpenAI GPT-2's actual next-token logits — warm up the temperature, tighten top-k, and sample. The full continuations underneath were generated by the real model at build time.

Scrub the training run epoch by epoch. Watch the training loss dive while the validation loss flattens — the signature of a model memorizing its tiny corpus.

A bonus from Appendix D: the training-loop refinements chapter 5 leaves out. Warm the learning rate up before decaying it, and clip gradients that grow too large — both pure functions, running live.

The real code

The demos run a TypeScript port; here is the Python they mirror. The loss is one call — cross-entropy over the flattened logits and the shifted targets:

def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch, target_batch = input_batch.to(device), target_batch.to(device)
    logits = model(input_batch)
    loss = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1), target_batch.flatten())
    return loss
 
perplexity = torch.exp(loss)   # the "effective vocabulary size" of the model's doubt

Decoding gets two new dials on top of chapter 4's greedy loop — temperature scaling and top-k filtering:

def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]
        with torch.no_grad():
            logits = model(idx_cond)
        logits = logits[:, -1, :]
 
        if top_k is not None:                       # keep only the k highest logits
            top_logits, _ = torch.topk(logits, top_k)
            logits = torch.where(
                logits < top_logits[:, -1], torch.tensor(float("-inf")), logits)
 
        if temperature > 0.0:                       # sample from a tempered softmax
            probs = torch.softmax(logits / temperature, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
        else:                                        # or take the argmax (greedy)
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)
 
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

def train_model_simple(model, train_loader, val_loader, optimizer, device,
                       num_epochs, eval_freq, eval_iter, start_context, tokenizer):
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1
 
    for epoch in range(num_epochs):
        model.train()
        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()          # compute gradients
            optimizer.step()         # update weights
            tokens_seen += input_batch.numel()
            global_step += 1
 
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader, device, eval_iter)
                train_losses.append(train_loss)
                val_losses.append(val_loss)
        generate_and_print_sample(model, tokenizer, device, start_context)
    return train_losses, val_losses, track_tokens_seen

# Save the trained weights — and the optimizer state, so training can resume
# exactly where it left off (AdamW carries per-parameter momentum):
torch.save({
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
}, "model_and_optimizer.pth")
 
checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

# Same GPTModel class, real OpenAI weights (context 1024, qkv_bias=True):
settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")
gpt = GPTModel(NEW_CONFIG)
load_weights_into_gpt(gpt, params)   # copy every tensor into place
 
token_ids = generate(
    model=gpt, idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=25, context_size=1024, top_k=50, temperature=1.5)

OpenAI released GPT-2 in four sizes; §5.5 loads the smallest, but the same GPTModel class scales to all of them by changing three config numbers (§5.5):

Four things to remember

1. Loss is cross-entropy on next-token prediction. It's the mean negative log-probability the model gives the correct token; perplexity is its exponential — the effective number of tokens the model is choosing between.
2. A tiny model on tiny data memorizes. Ten passes over one short story drove the training loss under 0.4 while validation stalled near 6.5 — the model learned the text, not the language.
3. Temperature and top-k trade safety for diversity. Temperature reshapes the distribution before sampling; top-k clips the nonsense tail. Together they let a model be creative without being incoherent.
4. Pretraining is expensive; loading it is not. The same architecture, filled with OpenAI's 124M weights, turns gibberish into fluent English — no training required.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original notebook (Apache 2.0). The loss, perplexity, and sampling math on this page run a TypeScript port verified against its executed Python with npm run verify:sampling; the training trajectory and pretrained continuations were precomputed from real seed-123 runs.