Pre-training objectives

Why this matters

Architecture is the skeleton; the training objective is the soul. Why does ChatGPT generate text but BERT does not? Why does FIM (Fill-in-the-middle) make code models so much better? Why are encoder-decoder models still the right choice for translation? Each of these questions reduces to what loss the model was trained against.

This lesson surveys the objectives that built modern NLP — and tells you when each one is the right choice.

Learning objectives

1. Distinguish causal LM, masked LM, prefix LM, span-corruption, and FIM objectives.
2. Implement causal LM and FIM losses in PyTorch.
3. Reason about the scaling trade-offs of each objective.
4. Pick the right objective for a given downstream task.

1. Causal language modeling (CLM) — GPT family

The objective: at each position, predict the next token.

loss = -1/T sum_t log p(x_t | x_<t)

Used by every modern chat / general LLM (GPT, Llama, Mistral, Claude, Gemini). The model is autoregressive — generation is the same operation as training.

# x: (B, T)  input token ids
# during training, predict x[t] from x[<t]
inputs = x[:, :-1]
targets = x[:,  1:]
logits = model(inputs)
loss = F.cross_entropy(logits.reshape(-1, V), targets.reshape(-1))

Pros: matches generation exactly; simple; scales beautifully. Cons: each token only sees the past — not the strongest possible representation for retrieval or classification.

2. Masked language modeling (MLM) — BERT family

The objective: mask 15% of tokens, predict them given bidirectional context.

[I, ate, [MASK], pizza, yesterday]   →   model predicts "good"

Used by BERT, RoBERTa, ModernBERT (2024), DeBERTa, mBERT. Strong for representation learning; bad for generation (the model has never been asked to produce text autoregressively).

Today MLM lives on in:

• Embedding models for retrieval (BGE, GTE, E5, ModernBERT-Embed).
• Reranking models (Cohere Rerank, BGE-Reranker).
• Classifiers / NER / token tagging.

You will rarely train one yourself, but every RAG system depends on a model trained with this objective.

3. Prefix LM and span corruption — T5 / UL2

Prefix LM (UniLM, GLM)

Bidirectional attention on the prefix, causal on the target. Combines understanding (encoder-like) with generation (decoder-like) in one architecture.

Span corruption (T5)

Replace contiguous spans with <X>, <Y>, <Z> sentinels. Train the model to generate the original spans separated by the same sentinels.

Input:  "Thank you <X> me to your party <Y> week."
Target: "<X> for inviting <Y> last <Z>"

Strong for translation, summarisation, and any input → output task. T5 / mT5 / FLAN-T5 are still SOTA on some translation benchmarks.

UL2 (Tay 2022)

A mixture of denoising objectives (causal, prefix, span). Trained to recognise which mode it is in via prompts ([NLU], [NLG], [S2S]). Underrated; some labs still train UL2-style.

4. Fill-in-the-Middle (FIM) — code models

A simple but transformative trick (Bavarian et al., OpenAI 2022). Reorder a portion of training examples so the model learns to fill gaps:

ORIGINAL:  prefix + middle + suffix
FIM ARRANGED: <PRE> prefix <SUF> suffix <MID> middle

The model still does next-token prediction — but now it can be prompted with <PRE> .. <SUF> .. <MID> at inference to fill in arbitrary middles. Crucial for code completion: IDEs need to fill code between the existing context and the file's end.

Used by:

• StarCoder / StarCoder2
• DeepSeek-Coder
• Qwen-Coder
• Code Llama
• Most modern code-specialised LLMs.

# Random FIM mix during pretraining (50% probability per doc)
def fim_transform(ids, fim_rate=0.5):
    if random.random() > fim_rate:
        return ids
    L = len(ids)
    a, b = sorted(random.sample(range(L), 2))
    pre, mid, suf = ids[:a], ids[a:b], ids[b:]
    return [PRE_ID] + pre + [SUF_ID] + suf + [MID_ID] + mid + [EOS_ID]

Tokenizers reserve <PRE>, <SUF>, <MID>, <EOS> special tokens.

5. Multi-token prediction (MTP) — DeepSeek-V3

DeepSeek-V3 (2024) introduced training the model to predict the next two tokens jointly via auxiliary heads:

loss = CE(t1) + λ * CE(t2 | t1_truth)

Reported to improve sample efficiency and downstream quality. Used at inference for speculative-decoding-style speedups (covered in Phase 5).

6. Choosing an objective

For 95% of you, causal LM with optional FIM is the right answer.

7. Loss curves and what they tell you

When pretraining, watch:

• Train loss: should descend smoothly through warmup → cosine. Sudden spikes = lr too high or instability (gradient clipping helps).
• Validation loss: should track train (no overfitting in single-epoch pretraining).
• Per-domain loss: split val into code, English, math, multilingual; ensure none plateau early.
• Gradient norm: spikes 10× the median often precede divergence; many labs auto-skip such steps.
• Token throughput (tokens/s/GPU): the only number that matters for timing.

You will use WandB or similar to monitor — Phase 5 covers infra.

Hands-on lab (3 hours)

objectives_lab.ipynb:

1. Take your tiny_shakespeare GPT from Lesson 2.4. Train it for 3000 steps without label shifting (predict same token). Show that loss does not decrease meaningfully (it should plateau at log(V)).
2. Re-train with proper targets = x[:,1:]. Should converge.
3. Add a 30% FIM transform to the data loader. Train for 3000 steps. After training, prompt with <PRE> the king is <SUF> . and sample. Check it produces a plausible middle.
4. Implement a tiny MLM loss on the same model: mask 15% of tokens with a <MASK> token id, compute loss only on masked positions. Train. Note this works only if you allow bidirectional attention (set is_causal=False). Verify by setting causal back and seeing the loss collapse to near zero (the model just copies x from the past).
5. Bonus: implement T5-style span corruption with one sentinel token. Train a small encoder-decoder.

Common pitfalls

1. Computing CE loss over the full sequence in MLM (including unmasked positions) → trivially low loss. Use ignore_index for unmasked tokens.
2. Forgetting to shuffle documents during pretraining — the model overfits to the order.
3. Training a model with FIM tokens that were never added to the tokenizer. Always extend the vocab and retrain the head.
4. Mixing causal and bidirectional masks in the same training run without conditioning — model gets confused.

Self-check

1. Difference between CLM and MLM in one sentence each.
2. Why does FIM matter for code models?
3. What objective do retrieval embedding models start from?
4. What does MTP add to vanilla CLM?
5. When is encoder-decoder still preferable to decoder-only?

References

• Radford et al. (2018), "Improving Language Understanding by Generative Pre-Training" (GPT-1, CLM).
• Devlin et al. (2018), "BERT: Pre-training of Deep Bidirectional Transformers" (MLM).
• Raffel et al. (2019), "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer" (T5).
• Tay et al. (2022), "UL2: Unifying Language Learning Paradigms."
• Bavarian et al. (2022), "Efficient Training of Language Models to Fill in the Middle" (FIM).
• DeepSeek-AI (2024), "DeepSeek-V3 Technical Report" (multi-token prediction).