Advanced Pretraining Objectives

Next-token prediction is a wonderful objective. It is simple, it uses every token as a training signal, and at scale it produces models that can write essays and code. But it has blind spots. A model trained only left to right has never once been asked to write the middle of a function given both ends. It has never been asked to think two tokens ahead. It has never seen a document longer than its training window, or an image. This chapter is about the extra signals that fill those gaps.

None of these replaces next-token prediction. Each one is a cheap modification of the data or a small extra loss that rides alongside it. The theme throughout: the base objective is fixed, so the leverage is in what you feed it and what you ask it to output.

Fill-in-the-middle: teaching a left-to-right model to infill

Open your code editor. Your cursor is in the middle of a function; there is code above and code below. You want the model to write what goes at the cursor. A pure next-token model cannot do this well: it has only ever seen "text so far, then next token", and the code below the cursor is context it has no way to condition on.

You could train a bidirectional model, but then you lose the generation ability that makes decoders useful. Fill-in-the-middle (FIM), from Bavarian et al. (2022) at OpenAI, is the trick that gets both: keep the causal decoder exactly as it is, and rearrange the training documents so that "the middle" comes last.

The PSM transformation

Take a document, pick a random span to be the middle, and split it into prefix, middle, suffix. Then reorder it with three sentinel tokens:

That is it. The model trains with ordinary next-token loss on this rearranged sequence. When it reaches the MID sentinel, everything it has seen is prefix and suffix, and the tokens it must predict are the middle. So a plain causal model learns to infill, with no architectural change. "PSM" names the order: prefix, suffix, middle.

The SPM variant

The alternative order puts the suffix first: suffix, then prefix, then middle. Bavarian et al. use a particular arrangement of the sentinels for it:

Why bother with a second order? Because in SPM the prefix and the middle are adjacent, which matters twice. At the token level, a tokenizer boundary artifact between prefix and middle is less likely to hurt. At the serving level, when a user keeps typing in the editor, the prefix grows and the suffix stays fixed; with SPM the growing part is at the end of the prompt, so a KV cache (see the inference chapter) of the suffix can be reused. Bavarian et al. report training on a 50/50 mix of PSM and SPM.

Why it matters, and the "FIM-for-free" result

The obvious worry: if half your documents are scrambled, does the model get worse at ordinary left-to-right text? Bavarian et al. tested exactly this and found what they call the FIM-for-free property: applying the transformation to 50% of documents leaves left-to-right perplexity and downstream performance essentially unchanged, while adding a strong infilling ability. The intuition is that a rearranged document still contains all the same text, so the model still learns the same statistics; it just also learns the sentinel grammar.

Two details that matter in practice. Document-level FIM (split the document before packing into training sequences) is cleaner than context-level FIM (split after packing) because the sentinels then always describe one coherent document. And the span should be chosen at the character level, then tokenized, rather than at the token level; otherwise the model never sees a middle that starts halfway through a token, which is exactly what a cursor does.

Every serious code model since has adopted FIM: InCoder (Fried et al. 2022) with a related "causal masking" objective, StarCoder, Code Llama, and DeepSeek-Coder all train with it, typically at a 50% rate. If you are training anything that will live in an editor, this is not optional.

Multi-token prediction: think a few tokens ahead

Here is a different blind spot. Next-token loss trains the model to be very good at the immediate next token and rewards it not at all for having a plan for the token after that. Yet a good continuation of "The capital of France is" requires already knowing that "Paris" will be followed by "." rather than ", which". Could we ask the model to predict several future tokens at once, and would that be a better training signal?

The idea and the extra heads

Multi-token prediction (MTP), studied at scale by Gloeckle et al. (2024), attaches $k$ output heads to the shared trunk of the transformer. Head 1 predicts token $t+1$ from position $t$, head 2 predicts $t+2$, and so on to $t+k$. The loss is the sum of the $k$ cross-entropies:

Each head is a small transformer layer plus the shared unembedding; the trunk is the ordinary model. At inference you can throw the extra heads away and decode with head 1 exactly as usual, so MTP is purely a training-time change to what the model is asked to know.

What it buys you

Gloeckle et al. report that at 13B parameters with $k = 4$, MTP improved code benchmarks noticeably (they report about 12% more HumanEval problems solved and 17% more MBPP problems relative to next-token baselines), while small models could be hurt, and gains on natural-language multiple-choice benchmarks were mixed. The benefit grows with model size, which suggests the extra heads act as a regularizer that only helps once the model has capacity to spare. They also found the trained models were better at induction-style and algorithmic reasoning tasks.

DeepSeek-V3 (2024) used a variant: rather than parallel independent heads, a single sequential MTP module that predicts token $t+2$ conditioned on the trunk output and the embedding of token $t+1$, keeping the causal chain intact. They report it as a training objective that improved their benchmark scores, with the module discarded or reused for speculative decoding at inference.

The speculative-decoding connection

The extra heads have a second life. Speculative decoding (inference chapter) speeds up generation by having a cheap draft model guess several tokens which the real model then verifies in one parallel pass. With MTP, the draft model is free: head 2 through head $k$ produce the guesses. Gloeckle et al. report up to about 3× faster inference with 4 heads on code, and DeepSeek-V3 reports acceptance rates of roughly 85 to 90% for its second-token prediction, translating to about 1.8× tokens per second.

Span corruption and UL2, briefly

Before decoders won, the dominant pretraining objective for text was span corruption, from T5 (Raffel et al. 2019). Mask out random spans of the input, replace each with a sentinel, and train an encoder-decoder to emit the sentinels followed by the missing spans. In tiny form: input "The sat on mat", target " cat the". It is FIM's ancestor, with several middles at once.

UL2 (Tay et al. 2022) argued that no single corruption is best and mixed three "denoisers", each announced by a mode token at the start of the input. The R-denoiser is regular T5 corruption: short spans, about 15% of tokens. The S-denoiser is sequential: corrupt the whole tail, which is just prefix language modeling. The X-denoiser is extreme: very long spans or a high corruption rate, so the model must generate a lot from a little. Training on the mixture with mode switching gave a single model that did well on both understanding and generation tasks. The idea that a decoder can benefit from a mixture of denoising objectives lives on, but for large decoder-only models the community mostly settled on next-token plus FIM, and UL2 is now more a reference point than a recipe.

Long-context training: from 8k to 128k

The next blind spot is length. If the model trained on 8k-token windows, it has never seen position 20,000. Its rotary embeddings (see the positional-encoding chapter) have never rotated that far, and its attention has never had to pick one relevant sentence out of a hundred pages. Feed it a 100k-token document and quality collapses.

Why you cannot just train long from the start

Attention cost grows with the square of sequence length, and activation memory grows linearly with it (see the systems chapter). Training all 15 trillion tokens at 128k would cost many times more than at 8k and would need far more context-parallel communication. Worse, most documents are short, so most of that long context would be padding or unrelated concatenated documents. So the standard recipe trains short and extends late.

The staged approach

Grattafiori et al. (2024) describe Llama 3's long-context stage in Section 3.4.2. After the main pretraining at 8k, they increased the context in six stages, from 8k up to 128k, training on roughly 800 billion tokens in total for this phase. At each stage they waited for two signals before moving on: short-context benchmark performance had fully recovered, and the model solved needle-in-a-haystack retrieval perfectly at the new length. They also report choosing a large RoPE base frequency of 500,000 for the architecture, which spreads the rotary frequencies so that positions far apart remain distinguishable at long range.

Other levers used across the field: position interpolation (Chen et al. 2023), which rescales positions so that a 32k input looks like a stretched 4k one to the existing rotary frequencies; YaRN (Peng et al. 2023), a more careful frequency-dependent rescaling; and simply raising the RoPE base, which Code Llama did (to 1,000,000) before its long-context fine-tuning. All of them are cheap ways to make the model's positional machinery tolerate lengths it never saw, so that a small amount of long-context training finishes the job.

The data you need

Long-context training needs genuinely long documents: books, long code repositories, long-form articles, multi-turn transcripts. Concatenating short documents to fill the window teaches nothing about long-range dependence, and worse, it teaches the model to attend across document boundaries that carry no information. Llama 3 reports using an attention mask that prevents tokens from attending across documents within a packed sequence, and notes it mattered mostly for the long-context stage. Upsampling long documents in the mix is standard; some recipes also synthesize long-range tasks (multi-document question answering, long summarization) for the extension phase.

Needle in a haystack

How do you know the extension worked? The simplest probe is the needle-in-a-haystack test (Kamradt 2023): hide one distinctive sentence ("the best thing to do in San Francisco is…") at a chosen depth in a long pile of unrelated text, then ask about it. Sweep context length and needle depth, and plot retrieval accuracy as a heatmap. A model with a real 128k context is green everywhere; a model with a nominal 128k context often shows a red band in the middle depths (the "lost in the middle" effect, Liu et al. 2023) and a cliff beyond its true effective length.

Needle retrieval is necessary but far from sufficient. A model can find one sentence and still fail to reason over a long document. Harder probes such as RULER (Hsieh et al. 2024) add multiple needles, variable tracking, and aggregation tasks, and show that many models' "effective" context is well short of their advertised one.

Continued pretraining and domain adaptation

Now suppose you already have a good general base model and you want it to be excellent at code, or mathematics, or Japanese. You could train a new model from scratch on a domain-heavy mix, but that throws away trillions of tokens of learning. Continued pretraining (CPT) instead resumes next-token training from the existing checkpoint on new, domain-focused data.

Examples that worked

Code Llama (Rozière et al. 2023) is the canonical case: start from Llama 2 and continue on about 500 billion tokens of mostly code, then specialize further for Python and for instruction following. Llemma (Azerbayev et al. 2023) continued Code Llama on a mathematics-heavy corpus. Multilingual extensions routinely continue a mostly-English model on a target-language mix. In each case the adapted model far outperforms the base on the target domain at a small fraction of the from-scratch cost.

Catastrophic forgetting, and replay

Here is the catch. Train only on code for long enough and the model gets worse at English. This is catastrophic forgetting: gradient steps on the new distribution overwrite what was learned on the old one. Two things make it worse: a high learning rate (the model moves far from its starting point) and a large distribution shift (the new data is very different from the old).

The standard mitigation is replay: mix a fraction of the original pretraining distribution back into the CPT data. Ibrahim et al. (2024) studied this systematically and report that a small replay fraction, a few percent for a mild shift and more for a strong one, combined with re-warming the learning rate and then decaying it again, recovers most of the from-scratch quality on the union of both distributions. The learning-rate part matters: resuming at the tiny end-of-schedule learning rate learns the new domain too slowly, while jumping straight to the peak rate causes a loss spike and heavy forgetting, so warm up again from a low value (Gupta et al. 2023).

Curriculum and annealing, revisited

Everything above changes what the model sees. There is also leverage in when. Curriculum, in the loose sense used in LLM pretraining, means ordering the data: broad web text for most of training, then a shift toward higher-quality and more targeted data at the end, while the learning rate anneals to zero.

The last few percent of training at a decaying learning rate is disproportionately influential, because the weights settle into their final basin during it. Llama 3 reports (Section 3.1.3 and 3.4.3) that annealing the learning rate on a small amount of high-quality code and mathematical data at the very end improved the 8B model markedly on GSM8K and MATH, while the 405B model saw negligible gain, suggesting that the largest models already have that knowledge and the trick mostly helps smaller ones. The warmup-stable-decay schedule from MiniCPM (Hu et al. 2024) makes this explicit: hold the learning rate constant for most of training, then decay sharply on the best data. The data engineering chapter covers the mixing and quality-filtering side.

Model merging, warm-starting, and depth up-scaling

Three related ways to avoid starting from scratch, each in a sentence or two. Model merging averages the weights of several models trained from a shared starting point; "model soups" (Wortsman et al. 2022) showed that averaging fine-tuned checkpoints can beat any single one, and Llama 3 reports averaging models from different post-training runs. Warm-starting initializes a new model from an old one, either directly (continued pretraining is warm-starting on new data) or by growing the architecture. Depth up-scaling is a specific growth recipe: SOLAR 10.7B (Kim et al. 2023) took a 32-layer Mistral 7B, duplicated it, dropped the last 8 layers of one copy and the first 8 of the other, stacked them into a 48-layer model, and continued pretraining. The duplicated model starts much closer to a good solution than random initialization does, and the authors report strong results for the compute spent. The general lesson: weights are expensive and reusable.

Multimodal pretraining

A text model has never seen an image. Making it see is, surprisingly, mostly a data-and-adapter problem rather than a new-architecture problem, and the recipe now looks similar across labs.

Vision encoders and adapters

Take a pretrained vision encoder, usually a Vision Transformer trained contrastively like CLIP (Radford et al. 2021), which turns an image into a grid of patch embeddings. Add an adapter that maps those embeddings into the language model's token space, then feed them into the LLM as if they were text tokens. LLaVA (Liu et al. 2023) used a single linear projection (later an MLP) and trained it in two stages: first only the projection on image-caption pairs, then the projection plus the LLM on instruction-style visual conversations. That is enough to get a model that can describe and answer questions about images.

Flamingo (Alayrac et al. 2022) took the other route: keep the LLM frozen and insert gated cross-attention layers that let text tokens attend to the visual features. Llama 3 (Section 7) follows the cross-attention pattern, reporting that they trained image adapters on top of the finished text model so that adding vision could not degrade text quality. The choice between "inject as tokens" and "cross-attend" trades simplicity against protecting the text model.

Interleaved data

Caption pairs teach "this image is a dog". Interleaved documents, web pages with images placed among the paragraphs that discuss them, teach the harder skill of using an image in context and of handling several images in one document. Flamingo made interleaved data central; most modern recipes mix captions, interleaved documents, OCR-heavy documents, charts, and synthetic visual question answering, staged from easy to hard.

Audio in a paragraph

Speech follows the same template with a different encoder. A pretrained speech encoder (Whisper-style, Radford et al. 2022, or a conformer trained on unlabeled audio) converts a waveform into frame embeddings; an adapter compresses and projects them into the LLM's token space; the LLM then reads them as tokens. Llama 3 (Section 8) reports a speech adapter trained this way with the text model kept fixed, so that a single text model can accept typed or spoken input. Generating audio is a separate problem usually handled by a decoder on the output side.

Reasoning-oriented pretraining data

A last lever, and one where claims run ahead of evidence, so let us be careful. Models trained on more code and mathematics tend to be better at multi-step reasoning, even on tasks that are neither code nor math. The Llama 3 authors report a mix of roughly a quarter mathematical and reasoning data and 17% code, up from earlier generations, and attribute part of their reasoning gains to it. DeepSeekMath (Shao et al. 2024) showed that a large, carefully filtered corpus of web mathematics (they report about 120 billion tokens) continued into a code model produced a strong math model at 7B.

Synthetic data goes further. The phi series (Gunasekar et al. 2023, "Textbooks Are All You Need") trained small models on LLM-generated textbook-style explanations and exercises and reported strong coding results for the size. Synthetic chain-of-thought, where a stronger model writes step-by-step solutions that go into the pretraining mix, is now common, and it clearly helps on benchmarks whose format it matches.

The balanced view. In favour: reasoning data is dense in "structure per token" and models do transfer some of it. Against: synthetic data narrows the distribution, can leak benchmark-like content and inflate scores (see the evaluation chapter on contamination), and training repeatedly on model outputs can degrade diversity (the "model collapse" concern of Shumailov et al. 2024). The practical recipe most labs converge on is a natural-data majority with a minority of high-quality synthetic reasoning data, concentrated in the annealing phase where it has the most effect.

Practice

Further reading

• Bavarian et al. (2022). Efficient Training of Language Models to Fill in the Middle. PSM/SPM, FIM-for-free, the practical details.
• Fried et al. (2022). InCoder: A Generative Model for Code Infilling and Synthesis. The causal-masking cousin of FIM.
• Gloeckle et al. (2024). Better & Faster Large Language Models via Multi-token Prediction.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report. Sequential MTP module and its use for speculative decoding.
• Raffel et al. (2019). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (T5). Span corruption.
• Tay et al. (2022). UL2: Unifying Language Learning Paradigms. The mixture of denoisers.
• Grattafiori et al. (2024). The Llama 3 Herd of Models. Section 3.4 for staged long-context training and annealing; Sections 7 and 8 for vision and speech.
• Chen et al. (2023). Extending Context Window of Large Language Models via Positional Interpolation; Peng et al. (2023). YaRN.
• Liu et al. (2023). Lost in the Middle: How Language Models Use Long Contexts; Hsieh et al. (2024). RULER.
• Rozière et al. (2023). Code Llama; Azerbayev et al. (2023). Llemma. Continued pretraining in practice.
• Ibrahim et al. (2024). Simple and Scalable Strategies to Continually Pre-train Large Language Models; Gupta et al. (2023). Continual Pre-Training of Large Language Models: How to (re)warm your model?
• Hu et al. (2024). MiniCPM. The warmup-stable-decay schedule and annealing on high-quality data.
• Kim et al. (2023). SOLAR 10.7B: Depth Up-Scaling; Wortsman et al. (2022). Model Soups.
• Liu et al. (2023). Visual Instruction Tuning (LLaVA); Alayrac et al. (2022). Flamingo; Radford et al. (2022). Whisper.
• Gunasekar et al. (2023). Textbooks Are All You Need; Shao et al. (2024). DeepSeekMath; Shumailov et al. (2024). The Curse of Recursion (model collapse).