What a model is trained to predict, how frontier labs turn the raw web into trillions of clean tokens, and the number formats, optimizers, schedules, and stability tricks that let a model learn from them without falling over.

Chapter 6 ended with a small GPT that writes passable sonnets after a few minutes on a laptop. Between that model and a frontier one sit three questions, and this chapter takes them in order. What should the model be trained to predict? What text should it read? And how do you keep an optimizer stable for months, across thousands of GPUs, in number formats with only a handful of bits? This section answers the first question. Sections 8.2 to 8.4 build the data, from a raw web crawl to a deduplicated, balanced mixture. Sections 8.5 to 8.8 cover the number formats, optimizers, schedules, and stability tricks that let a large run train at all, and §8.9 takes three real runs apart.

The first question sounds settled. A language model predicts the next token, done. But here is a puzzle. In 2019 the most careful comparison of pretraining objectives anyone had run, the T5 study (Raffel et al. 2020), trained the same-sized models on the same data with different objectives and architectures. Next-token prediction with a decoder-only model came out worst of all, 73.8 on GLUE against 83.3 for an encoder-decoder trained to fill in blanks. Four years later, essentially every frontier model was a decoder-only model trained on next-token prediction. Both facts are true. By the end of this section you will be able to explain why, and to implement every objective that took part in the contest.

What an objective decides

Every pretraining objective we'll meet is a recipe with three parts. A corruption turns a clean text into what the model sees. An attention pattern decides which positions may read which. A loss mask decides which predictions are graded. Next-token prediction is the case with no corruption, a causal pattern, and every position graded. Masked language modeling corrupts 15 percent of the tokens, lets every position see every other, and grades only the corrupted ones. The rest of the section varies these three knobs.

The widget below shows six objectives side by side on one sentence. Switch between them and watch three things, what goes in, which positions are graded, and the attention pattern.

Causal language modeling is maximum likelihood

Write a text as tokens $\mathbf{x} = (x_1, \dots, x_T)$ from a vocabulary $\mathcal{V}$ of size $V$, with $x_{\lt t} = (x_1, \dots, x_{t-1})$ the tokens before position $t$. A model with parameters $\theta$ assigns the text a probability by the chain rule (§6.1),

Training maximizes the log-likelihood of the corpus, which is the same as minimizing the average negative log-likelihood per token,

In plain words, read the text left to right and pay $-\log p$ for every token, where $p$ is the probability the model gave that token just before seeing it. This is the cross-entropy loss of §6.3, averaged over every position.

Why is this the right thing to minimize? Let $p_{\text{data}}$ be the true distribution of texts. The expected loss of a model $q = p_\theta$ splits into two pieces,

The step is one line. Add and subtract $\sum p_{\text{data}}\log p_{\text{data}}$, then group. The entropy $H(p_{\text{data}})$ doesn't depend on the model, and the KL divergence is never negative and is zero only when $q = p_{\text{data}}$ (§0.5). So minimizing the next-token loss is exactly fitting the model to the distribution of text, and the loss can never go below the entropy of text itself. The chain rule turns the left side into a sum of per-token terms, which is why one loss per position is all you need.

Let's check the identity on a world small enough to enumerate, three-token texts over a three-token vocabulary, so 27 possible texts.

The three numbers agree to the precision shown, and the sampled average lands within about 0.01 of the exact expectation. Training on samples estimates exactly this quantity.

Three properties of causal language modeling matter for everything that follows.

1. Every position is graded. A window of $T$ tokens yields $T$ predictions in one forward pass, because the causal mask lets each position see only its past (§4.4).
2. It is a complete generative model. You can sample a text token by token, and you can score any text exactly. Nothing else in this section gives you both for free.
3. Any task can be written as a continuation. A translation, a classification, or an answer is just more text after a prompt, which is what lets one pretrained model do in-context learning (§6.8).

Its weakness is the one BERT was invented to fix. Position $t$ never sees anything to its right, even when the whole text is already written.

Masked language modeling, and why it is not a language model

You met BERT's objective in §7.1, so here is the one-paragraph recap. Select a random set $\mathcal{M}$ of about 15 percent of the positions. Replace 80 percent of them with [MASK], 10 percent with a random token, and leave 10 percent unchanged, giving the corrupted input $\tilde{\mathbf{x}}$. Run a Transformer with no causal mask, and grade it only on the selected positions,

The 80-10-10 split exists because [MASK] never appears at fine-tuning time. The unchanged and random cases force the model to keep a useful representation at every visible position, since any of them might be graded and any might be fake. Exercise 7.1b implements it.

Two consequences of this formula shape the rest of the story.

Only about 15 percent of positions teach anything. The forward and backward pass cost the same as a causal model's, but a 512-token sequence gives about 77 graded predictions instead of 512. Each prediction is better informed, since it sees both sides, so the comparison is not a plain 6.7 to 1, but the signal is sparse.

The masked tokens are predicted independently. Mask both words of "[MASK] [MASK] is a city in California" and the model scores each blank separately given the same corrupted input. It may well rate "San" and "Angeles" as the best choices for the two blanks, each plausible alone and wrong together. The model learns conditionals of the form $p(x_t \mid \tilde{\mathbf{x}})$, not a joint distribution over texts. So there is no cheap way to sample a paragraph from it or to compute the probability of a sentence. (You can score a sentence by masking each token in turn, a pseudo-log-likelihood, but that costs $T$ forward passes and still isn't a true probability.)

That is the first half of the answer to the puzzle. MLM makes excellent representations for understanding tasks, which is what GLUE measures, and it gives up generation entirely.

Variants of masking

Each of the next three ideas fixes a specific weakness of token-level masking.

Whole-word masking. BERT masks WordPiece tokens, so it can hide ##ing while leaving play visible, and the blank is trivial (§2.4). Whole-word masking selects words and then masks every piece of each selected word. Google released whole-word-masking BERT checkpoints in 2019. The idea matters even more for Chinese, where BERT splits text into single characters and most words span several tokens.

SpanBERT (Joshi et al. 2020) goes further and masks contiguous spans of whole words. It samples each span length $\ell$ (in words) from a geometric distribution $\ell \sim \operatorname{Geo}(0.2)$ clipped at 10, picks a random start, and repeats until 15 percent of the tokens are covered, applying 80-10-10 to whole spans. It also adds a span boundary objective, predicting every token inside a span from the two tokens just outside it plus a position embedding, which pushes the boundary representations to summarize the span. The paper quotes a mean span length of 3.8, and we can check it.

The mean is 3.80. A geometric law makes short spans common (22 percent are a single word) while still showing the model spans of up to ten words. SpanBERT's gains were largest on span-selection tasks such as extractive question answering and coreference.

import numpy as np

def whole_word_mask(tokens, word_selected):
    """tokens: list of str; word_selected: (n_words,) bool. Returns (n_tokens,) bool."""
    return np.zeros(len(tokens), dtype=bool)

ELECTRA (Clark et al. 2020) attacks the 15 percent problem head on. A small masked language model, the generator $p_G$, fills each masked position with a sample from its prediction. The main model, the discriminator, sees the result and decides for every position whether its token is original or replaced. With $\hat y_t = \sigma(z_t)$ the discriminator's probability, from its logit $z_t$, that position $t$ was replaced,

In plain words, the generator plays BERT, the discriminator spots its forgeries, and every one of the $T$ positions now contributes to the main model's loss. The label is about the token, not the position. When the generator happens to sample the original token, that position counts as original, and the paper found this convention moderately better. The generator is trained by maximum likelihood, not to fool the discriminator, since you cannot backpropagate through the sampling step, and the discriminator's gradient never reaches it. ELECTRA used $\lambda_{\text{disc}} = 50$, a generator a quarter to half the discriminator's size, and shared token embeddings between them. After pretraining, the generator is thrown away.

The results were striking for small budgets. ELECTRA-Small, trained on one GPU for four days, beat GPT on GLUE although GPT used about 30 times more compute, and ELECTRA-Large matched RoBERTa and XLNet with less than a quarter of their compute. The ablation explains where the gain comes from. On GLUE, a discriminator graded only on the 15 percent of masked positions ("ELECTRA 15%") scored 82.4, barely above BERT's 82.2, while full ELECTRA scored 85.0. So most of the improvement comes from grading every token, not from the forgery game itself.

import numpy as np

def electra_corrupt(ids, masked, samples):
    corrupted = np.where(masked, samples, ids)
    labels = masked.astype(float)
    return corrupted, labels

def disc_loss(logits, labels, valid):

Span corruption, T5's objective

T5 (Raffel et al. 2020) asked the question "what if the corrupted text goes into an encoder, and a decoder writes out what was removed?". Its span corruption objective drops about 15 percent of the tokens in contiguous spans, replaces each dropped span with a sentinel token that is unique within the example, and makes the target the dropped spans, each introduced by its sentinel, followed by one final sentinel. The paper's example, with its sentinels written <X>, <Y>, <Z> (in the released vocabulary they are <extra_id_0>, <extra_id_1>, and so on, 100 of them),

original  Thank you for inviting me to your party last week .
input     Thank you <X> me to your party <Y> week .
target    <X> for inviting <Y> last <Z>

Each design choice has a reason.

Unique sentinels tell the decoder which hole each piece fills. With one shared mask symbol, a target of "for inviting last" would be ambiguous about where the break falls.

Only the dropped tokens are targets. BART (Lewis et al. 2020), a contemporary, makes the decoder rewrite the entire original text, which costs a decoder pass over every token. T5's target is about a fifth as long, and in T5's own ablation, predicting only the corrupted spans did as well as or slightly better than reconstructing everything.

Spans of mean length 3. T5 compared corrupting tokens independently with spans of average length 2, 3, 5, and 10, all at 15 percent. Length 3 was slightly but significantly better than independent tokens on most tasks, and 10 slightly worse. It also shortens the targets, because each span costs only one sentinel. Corruption rates of 10, 15, and 25 percent performed about the same, and 50 percent was worse.

The span lengths are not sampled independently. T5 fixes the number of corrupted tokens, $\operatorname{round}(rL)$ for corruption rate $r$ and length $L$, and the number of spans, $n_s = \operatorname{round}(rL/\mu)$ for mean span length $\mu$, then splits both the corrupted and the kept tokens into $n_s$ pieces at random and interleaves them, starting with a kept piece. Here is that procedure, and the arithmetic that sets T5's sequence lengths.

To fill a 512-token encoder input exactly, T5 starts from 568 raw tokens. Of those, 85 are corrupted in 28 spans, so the encoder reads $568 - 85 + 28 + 1 = 512$ tokens and the decoder writes $85 + 28 + 1 = 114$. Per raw token, the model is graded on about 0.2 targets, which is denser than BERT's 0.15, though still a fifth of what a decoder gets.

import numpy as np

def span_corrupt(ids, noise, sentinels):
    """returns (inputs, targets) as 1-D int arrays"""
    ids, noise = np.asarray(ids), np.asarray(noise, bool)
    inputs = np.where(noise, sentinels[0], ids)        # wrong: one sentinel per token, all the same
    targets = ids[noise]
    return inputs, targets

Prefix language modeling

There is a middle ground between a causal decoder and an encoder-decoder. Split each text at some position, call the first $P$ tokens the prefix, and let the prefix attend bidirectionally within itself while the rest stays causal. Grade only the predictions of tokens after the prefix. This is a prefix language model, also called a non-causal decoder. It is one stack of weights, like GPT, but the prompt is read with full context, like an encoder. UniLM (Dong et al. 2019) trained one Transformer with bidirectional, causal, and prefix masks at once, and T5 compared a prefix LM directly with its encoder-decoder.

Here are the three attention patterns for a 6-token window with a 3-token prefix. A 1 means "may attend".

The loss mask has a classic off-by-one, and it's worth slowing down. With inputs $x_0, \dots, x_{T-1}$ and targets shifted by one, $y_t = x_{t+1}$, position $t$ predicts the next token. In a prefix LM, position 0 can see $x_1$ and $x_2$, because they are in the prefix, so its prediction of $y_0 = x_1$ is a copy, not a prediction. Grading it would teach the model to copy and leak information. Position $P - 1 = 2$ predicts $x_3$, the first suffix token, which no prefix position can see. So the graded positions are exactly $t \ge P - 1$.

import numpy as np

def prefix_lm_masks(T, P):
    P = np.asarray(P)
    i, j = np.arange(T)[:, None], np.arange(T)[None, :]
    allowed = np.broadcast_to(j <= i, (len(P), T, T))
    graded = np.ones((len(P), T), bool)
    return allowed, graded

UL2's mixture of denoisers

By 2022 there were three families of objective, each good at something. Span corruption with short spans teaches knowledge and understanding. Prefix LM teaches fluent continuation. And the T5 study had found that corrupting 50 percent of the text works poorly on its own. UL2 (Tay et al. 2023) asked whether one model should train on all of them, and framed them as one family of denoisers, parameterized by the mean span length $\mu$, the corruption rate $r$, and the number of spans.

That's seven denoisers, mixed in roughly equal shares. A mode token, [R], [S], or [X], is prepended to the input so the model knows which game it's playing, and at fine-tuning or prompting time you add the mode token that matches the task. Their ablations preferred a modest share of S-denoising, about 20 percent, over raising it to half.

The idea that matters most for this section is that the knobs interpolate. With $\mu = 1$ and $r = 0.15$, you have something close to BERT-style masking. With one span covering the end, you have a prefix LM. With $r = 0.5$ and long spans, the model writes half the text from the other half. The widget lets you move along that line and watch the encoder and decoder lengths, which set the cost.

The paper's architecture conclusions are worth reading carefully, because they are more nuanced than the eventual industry choice. In UL2's experiments at about 167M parameters and 32B tokens, the GPT-style setup, a causal decoder with next-token prediction, was beaten by every alternative they tried, and the authors recommended always training with at least a prefix LM. They also argued that the objective matters more than the backbone, since a UL2-trained decoder beat the T5 encoder-decoder. UL2 20B, an encoder-decoder with 32 layers on each side, width 4,096, and the T5 vocabulary, was then trained on 1 trillion tokens of C4.

The most influential follow-up was not a new model. U-PaLM (Tay et al. 2022) took the finished causal-decoder PaLM and continued training it on the UL2 mixture for about 0.1 percent of its original compute. At 540B parameters, the result matched final PaLM's performance at roughly half its compute, a saving the authors put at about 4.4 million TPUv4 hours. So mixed denoising objectives were not abandoned. They were folded into a causal decoder. Google's PaLM 2 report later described training on "a tuned mixture of different pre-training objectives".

Fill-in-the-middle for code

Code completion has a shape that left-to-right models handle badly. You are editing line 40 of a file. The 39 lines above are the prefix, and the 200 lines below, which call the function you are writing and import things you might need, are the suffix. A causal model can condition only on the prefix. An encoder-decoder could use both, but its training spans are a few tokens long, far shorter than the multi-line completions people want.

Bavarian et al. (2022) found a fix that needs no new architecture at all, only a data transformation. Take a document, cut it at two random positions into a prefix, a middle, and a suffix, and move the middle to the end, with special tokens marking the pieces,

PSM   <PRE> prefix <SUF> suffix <MID> middle <EOT>

A causal model trained on such documents learns to write the middle after reading both the prefix and the suffix. At inference time, you feed <PRE> prefix <SUF> suffix <MID> and sample until the model emits <EOT>, which is how it says that the middle now joins up with the suffix. Several details of the recipe matter.

• Cut in characters, not tokens. The two cut points are uniform over the characters of the document, before tokenization, so each piece has an expected length of a third of the document. Users stop typing mid-token, and a model that has only seen cuts at token boundaries handles that badly. The paper recommends always including character-level spans.
• Keep the loss on everything. The prefix, the suffix, and the middle are all graded, and so is <EOT>. Grading only the middle would throw away two thirds of the ordinary training signal.
• Apply it to a fraction of documents. The FIM rate $p_{\text{fim}}$ is the probability that a document is transformed. Bavarian et al. used 50 percent for their main models.
• Transform after packing. Document-level FIM cuts a document before it is chunked into training contexts, so a long document's prefix and suffix can end up in different contexts. Context-level FIM splits each packed context at its <EOT> tokens, transforms the documents inside, and trims or pads the result back to the context length. The paper found it slightly better and used it in all its main runs.

There is a second format, SPM, which puts the suffix first,

SPM   <PRE> <SUF> suffix <MID> prefix middle <EOT>

Two things make it attractive. With SPM, the prefix and the middle are contiguous text, so the model is literally continuing the prefix, and the place where they join is tokenized naturally (Code Llama encodes prefix and middle together in SPM for this reason). And SPM helps caching in an editor. When you type, the prefix grows while the suffix stays fixed, and since the suffix comes first, its keys and values stay valid in the KV cache (§6.5). The odd-looking leading <PRE> <SUF> is deliberate. With it, an SPM example is exactly what a PSM example looks like when the random prefix happens to be empty, so the two formats share structure and transfer to each other. The paper trains on both, half and half, and found SPM slightly stronger at evaluation time.

Here are both formats for a tiny function, cut at character positions.

Both cuts land inside a token, one in the middle of return and one in the middle of the number. That is exactly what users create when they stop typing mid-word, and why character-level cuts matter.

Why FIM costs nothing on left-to-right

The headline result of the paper is what the authors called the FIM-for-free property. They trained eight model sizes from 50M to 6.9B parameters on 100B tokens of code and of natural language, with and without FIM at a 50 percent rate. The left-to-right test loss of the FIM models matched the plain models at every size, and so did HumanEval and the language benchmarks, even though the FIM models saw ordinary documents only half the time. In a separate ablation, FIM rates up to 90 percent still left the left-to-right loss unchanged, while 100 percent hurt it. And learning FIM later by fine-tuning was expensive, needing far more compute to reach the quality that pretraining gave for free.

This is an empirical finding, and the paper's intuition for it is simple. A FIM document is three ordinary documents glued together in a new order. Inside each piece, every token is still predicted from the tokens just before it, which are its true left context. Only the predictions near the two seams have an unusual context. Let's count how many predictions keep their true local context.

With a document of 2,000 tokens, 99.7 percent of predictions see the same previous token as in the original, 98 percent the same previous 16 tokens, and 88 percent the same previous 128. Locally, FIM data is ordinary data. What changes is the global pattern. When predicting the middle, the model can attend to text that in the original comes after it, which is precisely the new skill.

"Free" deserves one caveat. Code Llama retested the claim at 7B and 13B parameters on 500B tokens of code, with a 90 percent FIM rate. The left-to-right test loss again held, but the infilling-trained models lost 0.6 (7B) and 1.1 (13B) percentage points on average across HumanEval and MBPP. So FIM is nearly free, and the rate is a dial between infilling skill and a small downstream cost. Production settings span the range. Code Llama used 90 percent, half PSM and half SPM. DeepSeek-V3, a general model, used a rate of 0.1 in PSM format, applied at the document level, with the tokens <|fim_begin|>, <|fim_hole|>, and <|fim_end|> in the roles of <PRE>, <SUF>, and <MID>.

import numpy as np

def fim_context(ctx, eot, pre, suf, mid, rate, spm_rate, rng):
    return np.asarray(ctx).copy()

Encoder, decoder, or both

Now we can take the contest apart. Three architectures carry these objectives.

The fifth row trips people up. In an encoder-decoder, each token passes through either the encoder or the decoder, not both. With $N$ parameters per stack, a forward pass over $a$ input tokens and $b$ target tokens costs about $2Na + 2Nb$ FLOPs, the same as a decoder-only model with $N$ parameters reading the concatenation of $a + b$ tokens (§5.6 derives the $2N$ FLOPs per token). So at matched compute, an encoder-decoder has about twice the parameters, which UL2's authors point out is a kind of sparsity. (Cross-attention adds a little, since each decoder block carries about $16d^2$ weights instead of $12d^2$.) Comparisons between the two can be compute-matched or parameter-matched, and they disagree.

Let's put numbers on the signal, for one training window of 512 tokens.

Now the evidence, in the order it arrived.

T5, 2019. Every model was fine-tuned on each task. On GLUE, an encoder-decoder with denoising scored 83.28, a prefix LM with denoising 81.82, and a decoder-only model with next-token prediction 73.78. Two more rows of the same table separate the objective from the architecture. An encoder-decoder whose two stacks share weights, so that it has the decoder's parameter count, still scored 82.81 with denoising, and the prefix LM trained with a plain language-modeling objective scored 79.68, well above the causal decoder's 73.78. Both the denoising objective and the bidirectional view of the input helped. The paper concluded that "using a denoising objective always results in better downstream task performance compared to a language modeling objective". That conclusion was correct for what it measured, fine-tuned understanding tasks at about 220M parameters.

Wang et al., 2022. This study asked a different question. Which architecture and objective give the best zero-shot generalization, with no task-specific fine-tuning? They trained causal decoders, non-causal (prefix) decoders, and encoder-decoders at matched compute, 4.8B parameters for the decoders and 11B for the encoder-decoders, on 168B tokens. Immediately after pretraining, the causal decoder with next-token prediction was best (44.2 average accuracy on their EAI-Eval suite, against 43.5 for the non-causal decoder and 39.9 for the encoder-decoder). After multitask prompted fine-tuning, the ranking flipped, and the encoder-decoder trained with span corruption was best. They also showed that conversion is cheap in either direction. Adapting a causal model into a non-causal masked model took a third of the tokens of training one from scratch, and adapting a non-causal MLM into a causal LM took 105B extra tokens to match 168B tokens of training from scratch.

UL2 and U-PaLM, 2022. The objective matters more than the architecture, pure next-token prediction is not optimal, and a small dose of mixed denoising late in training measurably improves a causal decoder.

So why did decoder-only causal LM win at the frontier? The evidence points to a combination of reasons.

1. The goal changed. Once the product became a model you prompt, rather than a model you fine-tune per task, zero-shot and few-shot ability became the metric, and that is where Wang et al. found the causal decoder best. Every task, including chat, is a continuation.
2. Dense signal and simplicity. Every token is graded, and one stack with one attention rule is easier to scale, shard, and optimize. GPT-3 showed in 2020 that this recipe kept improving with scale, and the infrastructure followed.
3. Inference is natural. A conversation has no fixed split into input and output. A decoder caches the keys and values of everything said so far and appends. An encoder-decoder has to decide what counts as the input and re-encode it as the conversation grows.
4. The rivals' strengths can be absorbed. FIM gives a causal model infilling. UL2R gives it a dose of denoising. Instruction tuning grades only the response, which is a prefix-LM loss mask on a causal model (§14.3). And when you need bidirectional representations, a small encoder trained separately is often the better tool (§7.1).

The honest summary is that encoder-decoders did not lose on quality at matched compute for the tasks they were measured on. They lost on generality, simplicity, and fit with how models came to be used. Hyung Won Chung, who worked on Flan-T5 and PaLM at Google, offers one reading of this history in the talk linked with this section. In his view, the encoder-decoder's extra structure, separate weights for input and output and a bidirectional input, is an inductive bias whose value shrinks as scale grows. Many practitioners share that view, but it is an interpretation of the evidence, not a theorem.

Multi-token prediction

Next-token prediction grades each position on one token. Gloeckle et al. (2024) asked what happens if each position also predicts the tokens two, three, and four steps ahead. Their architecture has a shared trunk producing the hidden state $\mathbf{h}_t$ at each position, $K$ independent output heads, each one Transformer layer, and a shared unembedding. Head $k$ predicts $x_{t+k}$ from the trunk's state at position $t$. With $K$ heads, the loss averages over heads,

In plain words, every position must commit to a guess about the next $K$ tokens using only what it has read so far, so the trunk is pushed to represent where the text is going, not only what comes next. (Gloeckle et al. sum the heads instead of averaging, which only rescales the loss.) Head 1 is the ordinary next-token head, and at inference you can use it alone.

Their findings, all empirical.

• The benefit grows with model size. Small models got worse with multi-token prediction, and large ones better. At 13B parameters on code, the multi-token model solved 12 percent more HumanEval problems and 17 percent more MBPP problems than a next-token baseline (relative improvements).
• $K = 4$ was best for 7B models on 200B tokens of code. On byte-level models, predicting 8 bytes was best, and far better than predicting one byte at a time.
• The extra heads make the model up to three times faster at inference, through self-speculative decoding. The heads propose the next few tokens, the next-token head verifies them in one pass, and accepted tokens are free (the general method is speculative decoding).
• On natural language, the gains were smaller and showed up mostly on generative tasks, with little change on multiple-choice benchmarks.

The obvious cost is memory. Each head produces logits of size $V$ at every position, and with $V$ in the hundreds of thousands, $K$ copies of those logits and their gradients would dominate memory. The paper's trick is to run the heads one at a time. Compute the trunk once, then for each head do its forward and backward pass, add its gradient into the trunk's output gradient, and free its logits before starting the next head. Peak memory drops from $O(KV + d)$ to $O(V + d)$ per position, with no extra compute. The torch block below implements this and checks that it gives the same gradient.

DeepSeek-V3 (DeepSeek-AI 2024) changed the design in one important way. Its MTP modules are sequential. Module $k$ takes the representation from depth $k - 1$ at position $t$, combines it with the embedding of the true token $x_{t+k}$, and runs one more Transformer block,

where $[\cdot\,;\cdot]$ stacks the two $d$-vectors into one of length $2d$ and $\mathbf{E}_{x_{t+k}}$ is the embedding of token $x_{t+k}$. The output of the block goes through the shared output head to predict $x_{t+k+1}$. In plain words, each extra prediction is made knowing the true intermediate tokens, so the chain of predictions stays causal, the way a speculative-decoding draft model works. DeepSeek-V3 used one extra depth (it predicts two tokens per position), shared the embedding and output head with the main model, and weighted the MTP loss by $\lambda_{\text{mtp}} = 0.3$ for the first 10T tokens and 0.1 for the remaining 4.8T. At inference, the module can be dropped, or used to draft a second token. DeepSeek reports that the second token is accepted 85 to 90 percent of the time, giving about 1.8 times the generation speed. Over 14.8T training tokens, a single extra depth adds roughly 14.8T more graded predictions, at the price of one extra Transformer block and a second pass through the output head.

import numpy as np

def mtp_loss(H, W, ids):
    raise NotImplementedError

In PyTorch

Here is a tiny decoder that supports a causal or prefix-LM mask and $K$ multi-token heads on a shared trunk. It computes the prefix-LM loss with the ignore_index convention, then the multi-token loss twice, once naively and once with the paper's head-by-head backward, and checks that the trunk gradients agree.

With real models, the same pieces are a few flags away. Hugging Face's collator does BERT masking, T5's preprocessing does span corruption, and FIM is usually a dataset transform written in a few dozen lines like Exercise 8.1e. These need downloads, so they are shown and not run.

Choosing an objective

Read more

• Raffel et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. Section 3.2 compares architectures and Section 3.3 compares objectives, corruption rates, and span lengths, the most thorough study of its kind.
• Clark et al. (2020). ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators. Replaced token detection, with the ablation (Table 5) that shows where the gains come from.
• Joshi et al. (2020). SpanBERT: Improving Pre-training by Representing and Predicting Spans. Geometric span masking and the span boundary objective.
• Tay et al. (2023). UL2: Unifying Language Learning Paradigms. The mixture of denoisers, mode switching, and a careful discussion of decoder-only versus encoder-decoder comparisons.
• Bavarian et al. (2022). Efficient Training of Language Models to Fill in the Middle. PSM and SPM, the FIM-for-free property, and recommended hyperparameters (Section 8.1 and Appendix D).
• Wang et al. (2022). What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization? The compute-matched comparison of causal, non-causal, and encoder-decoder models, and adaptation between them.
• Gloeckle et al. (2024). Better & Faster Large Language Models via Multi-token Prediction. Parallel heads, the memory-efficient backward, and self-speculative decoding.
• DeepSeek-AI (2024). DeepSeek-V3 Technical Report. Section 2.2 for sequential multi-token prediction, and Section 4.1 for its FIM setting.

Watch next

▶

▶