Understanding Multi-Token Prediction in DeepSeek‑R1 Architecture

Introduction

The fourth part of the DeepSeek‑R1 series focuses on the Multi‑Token Prediction (MTP) technique, an innovation introduced from the V3 version onward. Unlike conventional next‑token prediction, MTP asks the model to predict several future tokens simultaneously, improving learning efficiency and generation speed.

Next‑Token Prediction

Large language models such as GPT and Llama are trained with a next‑token prediction loss, minimizing cross‑entropy for the probability of the next token x_{t+1} given the history x_{1:t}:

Meta’s Multi‑Token Prediction

Meta’s 2024 paper “Better & Faster Large Language Models via Multi‑token Prediction” extends the next‑token task by adding n parallel output heads that each predict a future token, still using a cross‑entropy loss. Experiments in the paper identify an optimal n and show that MTP improves training performance and accuracy.

The Meta MTP network consists of a shared transformer backbone with four parallel heads. For an input token t_i, the heads predict t_{i+1}, t_{i+2}, t_{i+3}, and t_{i+4} respectively. Each head is an independent transformer layer (MHA + 2 FFN) followed by a shared vocabulary projection.

DeepSeek’s Adaptation of MTP

DeepSeek adopts a similar multi‑head structure but with a more complex module chain that preserves causal connections. The architecture includes a shared embedding layer, a shared output head, a transformer block TRMk(·), and a linear projection M_k ∈ ℝ^{d×2d} for each of the D sequence modules.

The backbone is a decoder‑only multi‑layer transformer that encodes the input token sequence x_{1:t} into hidden states z_{1:t}.

Four parallel heads are attached to z_{1:t}. Head 1 predicts the next token, Head 2 predicts the token after that, and so on.

Each head contains its own transformer layer (MHA + FFN) and feeds into the original model’s vocabulary projection (shared matrix + softmax).

During training, each head computes a cross‑entropy loss; the losses are weighted by a coefficient λ and averaged to form the final loss:

Training vs. Inference

“Our MTP strategy is primarily for improving the main model’s performance; during inference we can drop the MTP modules and let the main model run independently.” – DeepSeek‑V3

Thus MTP is used only in the training phase; the inference pipeline remains unchanged.

Conclusion

Multi‑Token Prediction represents a significant advance in LLM training, allowing simultaneous prediction of multiple future tokens, which enhances context modeling, speeds up training, and improves downstream generation quality, especially for large‑scale models.