Nemotron 3.5 Super was released recently and I wanted to understand what changed between 3.0 and 3.5 specifically on Multi Token Prediction (MTP)
A standard language model predicts only the next token. MTP predicts further ahead with a small MTP head that combines the hidden state of the main model at a position with the next token and predicts the token after that.
MTP does two jobs.
a) In training the MTP loss has a weight of 0.3 and pushes the hidden states of the main model to carry information about later tokens.
b) As most folks know at inference the MTP head drafts several tokens and the main model verifies them in one forward pass.
Nemotron 3 Super trains two MTP steps with shared weights and the second MTP step reads the hidden state from the first.
The same MTP head can be chained at inference and with 7 draft tokens Nemotron 3 Super accepts 3.45 tokens per verification on SPEED-Bench.
This however created a gap between training and inference because with 7 draft tokens six drafting steps use hidden states from earlier MTP steps but training exposes the MTP head to this only once.
Nemotron 3.5 Super makes three changes to close this gap.
a) It adds 100B tokens of MTP only training with the main model frozen.
b) It adds hidden state mixing. Details follow in next paragraph.
c) After post-training NVIDIA trains the MTP head for another 25000 steps with the main model frozen. The training uses text the model generated itself and targets the full prediction distribution of the main model. NVIDIA calls this MTP boosting and it raises tokens per verification from 4.287 to 4.718 with greedy decoding and from 3.787 to 4.612 with sampling.
Lets look at hidden state mixing as it is the least obvious change.
1. The MTP head contains an attention layer so when it predicts a token it reads a row with one hidden state for every earlier position. For example if the context is five tokens and the MTP head is drafting its third token then the row holds five hidden states from the main model and two hidden states from earlier MTP steps.
2. Then in this example the row is mixed but in standard training every position in the row comes from the same MTP step.
3. With hidden state mixing each training position receives a hidden state from a randomly selected earlier MTP step. Those hidden states already exist in the training pass so this matches inference better without simulating full drafts.
MTP alignment is how well the MTP head matches the main model and it can get worse during post-training as per the report from Nvidia.
My view is that MTP alignment should be part of every fine tuning pipeline. After fine tuning freeze the updated main model and realign the MTP head.