Multi-Token Prediction
Traditional LLMs are trained using Next Token Prediction. At any given context position \(t\), the model is optimised via a cross-entropy to predict only the single immediate next token \(x_{t+1}\).
Multi-Token Prediction (MTP) restructures the training objective. At every single position \(t\) in the training corpus, the model is instructed to predict \(n\) future tokens simultaneously: \(x_{t+1}, x_{t+2}, \ldots, x_{t+n}\).
\[L_1 = -\sum_t \log P_\theta(x_{t+1} | x_{t:1}) \qquad (1)\]
where \(P_\theta\) is the LLM being trained and \(x_{t:1} = x_1, \ldots, x_t\).
MTP adapts the learning task to predict \(n\) future tokens at once. This translates to:
\[L_n = -\sum_t \log P_\theta(x_{t+1:t+n} | x_{t:1}) \qquad (2)\] While written as a joint probability, the architecture factorises this objective assuming conditional independence across the heads. Because the \(n\) heads operate in parallel without knowing each other’s outputs, the true joint distribution is approximated as the product of independent probabilities, which translates to the sum of individual losses.
The LLM \(P_\theta\) employs a shared trunk to produce a latent representation \(z_{t:1}\) of the observed context \(x_{t:1}\), then fed into \(n\) independent heads to predict in parallel each of the \(n\) future tokens.
The architecture is decomposed as:
Shared Transformer Trunk, \(f_s\): This is the body of the transformer, it produces a latent hidden representation vector \(z_t\).
Independent Output Heads, \(f_{h_i}\): The model constructs \(n\) independent prediction heads. Each head \(i \in \{1, \ldots, n\}\) consists of a dedicated final transformer layer. They are arranged in parallel. Head 1 operates on \(z_t\) to forecast \(x_{t+1}\). Head 2 operates on \(z_t\) to forecast \(x_{t+2}\), without knowing what Head 1 predicted.
Shared unembedding matrix; \(f_u\): Maps the output for each head back to vocabulary space, all heads share a single, identical unembedding weight matrix, \(f_u\).
The conditional probability for the r-th future token is formalised as
\[P_\theta(x_{t+r} \mid x_{1:t}) = \text{softmax}(f_u(f_{h_r}(f_s(x_{1:t})))) \qquad (3)\]
Computing loss for MTP
During training, we have the entire sequence. At token position \(t\), the model has the exact future tokens \((x_{t+1}, x_{t+2}, \ldots)\). Consequently, cross entropy loss can be computed for each head independently. For a model trained with a lookahead horizon of \(n\) tokens, the total multi token prediction loss at a specific sequence position \(t\) is the sum of cross entropy losses across all heads
\[L_{\text{total}} = -\sum_{i=1}^{n} L_i(P_{h_{i}}(x_{t+i}|x_{t:1}), x_{t+i})\]
where \(L_i\) is the standard cross entropy loss
\(x_{t+i}\) is the ground truth token label
\(P_{h_{i}}\) is the model probability distribution output by head \(i\)
Optimising Memory
A naive parallel setup would require computing & storing logit tensors of size \(b \times s \times |V|\) (where \(b\) is batch size, \(s\) is the sequence length, \(|V|\) is the vocabulary size) for all \(n\) heads concurrently.
To mitigate this, the MTP approach implements a memory saving schedule during forward & backward passes:
The shared transformer trunk performs its forward pass to compute \(z_t\)
Head 1 executes a forward pass, compute its specific next token cross entropy loss and perform its backward pass immediately. This computes & accumulates the gradients directly into the shared trunk - and lets us discard Head 1’s logits & activations from memory.
We do this for all the other \(n-1\) heads.
Once Head \(n\) finishes, the shared trunk contains the accumulated gradient sum from all \(n\) independent tasks. Then the optimiser takes a single step to update the trunk of the transformer.
By chaining execution, the model keeps the peak GPU memory overhead identical to standard single token training.
For illustration, suppose the model trunk produced a shared output \(z_t = \begin{bmatrix} 1 \\ 2 \end{bmatrix}\) and let there be 2 heads. Let their losses be:
head 1 loss: \(L_1 = 0.319\)
head 2 loss: \(L_2 = 0.319\)
Next, compute gradient vector for each head. For illustration, let the gradient for each head be
\[\frac{\partial L_1}{\partial z_t} = \begin{bmatrix} -0.5 \\ 0.2 \end{bmatrix}, \quad \frac{\partial L_2}{\partial z_t} = \begin{bmatrix} 0.1 \\ -0.6 \end{bmatrix}\]
Each gradient is generated sequentially. When \(\frac{\partial L_1}{\partial z_t}\) is computed, it is written to a temporary gradient buffer assigned to the shared trunk. After this, head 1’s temporary activations are deleted.
Next, head 2 runs the backward pass to compute \(\frac{\partial L_2}{\partial z_t}\). Because head 2 is looking for a different token trajectory, its gradient looks different.
The hardware adds head 2’s gradient directly into the existing trunk buffer, as opposed to storing both vectors simultaneously.
\[\frac{\partial L_{total}}{\partial z_t} = \frac{\partial L_1}{\partial z_t} + \frac{\partial L_2}{\partial z_t} = \begin{bmatrix} -0.5 \\ 0.2 \end{bmatrix} + \begin{bmatrix} 0.1 \\ -0.6 \end{bmatrix} = \begin{bmatrix} -0.4 \\ -0.4 \end{bmatrix}\]
Note: Each independent head executes its forward pass and immediate backward pass _only down to the latent representation layer \(z_t\).
The final accumulated vector \([-0.4, -0.4]^T\) is what backpropagates through the shared transformer trunk.
Inference Strategy
MTP allows 2 options for inference time token generation:
Standard autoregressive decoding: Involves discarding heads \(2, \ldots, n\) only use head 1. The model behaves like a standard LLM.
Self-Speculative Decoding: Involves predicting \(n\) tokens at a time. Keep all the heads active, predict multiple tokens at a time. Since the heads are parallel, we don’t know what each token guesses - hence a draft & verify routine.
Suppose we have a \(n=4\) token prediction horizon. Our input is \(x_1, \ldots, x_t\) and the predicted tokens are \(\hat{x}_{t+1}, \hat{x}_{t+2}, \hat{x}_{t+3}, \hat{x}_{t+4}\).
Step 1: Draft Phase
The model looks at the context \(x_1, \ldots, x_t\). In a single forward pass, the shared trunk evaluates, and all 4 heads generate their predictions simultaneously:
head 1: \(\hat{x}_{t+1}\)
head 2: \(\hat{x}_{t+2}\)
head 3: \(\hat{x}_{t+3}\)
head 4: \(\hat{x}_{t+4}\)
We now have a 4 token draft proposal.
Step 2: Verification
All drafted tokens are appended to the prompt sequence \(x_1, \dots, x_t, \hat{x}_{t+1}, \hat{x}_{t+2}, \hat{x}_{t+3}, \hat{x}_{t+4}\).
The sequence is passed to the transformer, this time using only the first head. The output tensor from head 1 contains softmax probabilities for the next token at every step in the sequence. We only care about drafting new candidates, look exclusively rows corresponding to the end of the sequence. The output matrix is \(L \in \mathbb{R}^{s \times |V|}\), where \(s=t+4\).
\[L = \begin{bmatrix} \cdots & \cdots & \cdots & \cdots \\ l_{t,1} & l_{t,2} & \cdots & l_{t,|V|} \\ l_{(t+1),1} & l_{(t+1),2} & \cdots & l_{(t+1),|V|} \\ l_{(t+2),1} & l_{(t+2),2} & \cdots & l_{(t+2),|V|} \\ l_{(t+3),1} & l_{(t+3),2} & \cdots & l_{(t+3),|V|} \\ l_{(t+4),1} & l_{(t+4),2} & \cdots & l_{(t+4),|V|} \end{bmatrix}\] We can map the rows of \(L\) to the tokens they are supposed to give probabilities for. For instance,
\[\hat{x}_{t+5} : [l_{(t+4),1} \quad l_{(t+4),2} \quad \cdots \quad l_{(t+4),|V|}]\] \[\hat{x}_{t+4} : [l_{(t+3),1} \quad l_{(t+3),2} \quad \cdots \quad l_{(t+3),|V|}]\]
\[\cdots\]
\[\hat{x}_{t+1} : [l_{t,1} \quad l_{t,2} \quad \cdots \quad l_{t,|V|}]\]
Each row contains token probabilities, let \(\bar{x}_i\) denote the token which is the token that should be selected based on the softmax probabilities.
Next, we compare the drafted tokens \(\hat{x}_{t+i}\) and the selected tokens \(\bar{x}_{t+i}\), we accept tokens that match & discard when they don’t.
Compare \(\hat{x}_{t+1}\) and \(\bar{x}_{t+1}\), if they match, retain token \(\hat{x}_{t+1}\) and move to \(\hat{x}_{t+2}\) and compare with \(\bar{x}_{t+2}\) and so on. When a mismatch occurs, \(\hat{x}_{t+k} \neq \bar{x}_{t+k}\), then:
Accept the valid draft prefix tokens that matched prior to the mismatch: \(\hat{x}_{t+j}\) are accepted for all \(j < k\)
Accept the bonus correction token. Although we don’t accept \(\hat{x}_{t+k}\), we do have \(\bar{x}_{t+k}\) so we accept \(\bar{x}_{t+k}\)
Discard all tokens \(\hat{x}_{t+j}\) for \(j > k\)
If all draft tokens match, we accept all draft tokens plus the one extra token.
How does generated sequence end?
If using self speculative decoding, during the verification phase all draft tokens are considered at once. The moment an end of sequence token, \(\langle\text{EOS}\rangle\), is detected at any point in the sequence, the process halts and all tokens past \(\langle\text{EOS}\rangle\) are discarded.