AI Engineer Exhaustive
EasyPlusDeepSeek-V3

Multi-Token Prediction

DeepSeek-V3

Easy

Predict multiple future tokens simultaneously

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Next-token prediction gives one loss term per position. Multi-Token Prediction (MTP) asks the same trunk to also forecast tokens further ahead, densifying the gradient and, the authors argue, forcing hidden states to plan beyond t+1. Gloeckle et al. (2024) attach independent heads that predict the extra tokens in parallel. DeepSeek-V3 instead stacks sequential modules so each extra token is predicted with a full causal chain: depth k sees the previous depth's representation plus the embedding of the token that actually sits at that offset. Released V3 uses $$ — one extra future token besides the main next-token head. MTP is a training objective; the extra module can be dropped at inference or reused as a speculative-decoding draft head.

How it works

Depth k owns a Transformer block \operatorname{TRM}_k and a mix-in matrix M_k\in\mathbb{R}^{d\times 2d}. It shares the main model's embedding \operatorname{Emb} and output head \operatorname{OutHead}. For position i at depth k, concatenate the RMS-normalized previous-depth state with the RMS-normalized embedding of token t_{i+k}:

\mathbf{h}_i^{\prime k}=M_k\bigl[\operatorname{RMSNorm}(\mathbf{h}_i^{k-1});\operatorname{RMSNorm}(\operatorname{Emb}(t_{i+k}))\bigr].

When k=1, \mathbf{h}_i^{0} is the main-model hidden state at i. The block then runs causally on the shortened sequence (length T-k):

\mathbf{h}_{1:T-k}^{k}=\operatorname{TRM}_k(\mathbf{h}_{1:T-k}^{\prime k}), \qquad P_{i+k+1}^{k}=\operatorname{OutHead}(\mathbf{h}_i^{k}).

So depth 1, at position i, is trained to put mass on t_{i+2} — one token beyond the main head's t_{i+1}. Each depth k has its own cross-entropy; the extra training term is the average, scaled by \lambda:

\mathcal{L}_{\mathrm{MTP}}^{k}=-\frac{1}{T}\sum_{i=2+k}^{T+1}\log P_i^{k}[t_i], \qquad \mathcal{L}_{\mathrm{MTP}}=\frac{\lambda}{D}\sum_{k=1}^{D}\mathcal{L}_{\mathrm{MTP}}^{k}.

The paper sets \lambda=0.3 for the first 10T tokens and \lambda=0.1 for the remaining 4.8T. After training, discarding the MTP stack leaves a normal 61-layer causal LM. The same stack can draft tokens for speculative decoding; the authors report that use as an optional inference bonus, not as the training goal.

The released weight layout documents one MTP module as model.layers.61 (after num_hidden_layers=61 main layers), plus enorm, hnorm, and eh_proj for the two RMSNorms and M_k. embed_tokens and shared_head alias the main embedding and lm_head. Unique MTP parameters are about 11.5B; activations including the shared embed/head are about 2.4B.

Official code

The public inference demo inference/model.py builds only the 61-layer main Transformer. It does not instantiate MTP blocks. Weight meaning is specified in README_WEIGHTS.md: num_nextn_predict_layers (1 for the open V3 dump), layer index 61, fields enorm, hnorm, eh_proj, and a shared output head. Do not invent a starter signature from that README — it describes tensors, not a training API.

Watch-outs

Sources