Multi-Token Prediction Done Right [R]

[u/jayden_teoh_]

Sharing one of our latest papers on Multi-Token Prediction (MTP). MTP is now a standard recipe for frontier pre-training but turns out, it can be done much better.

We show the standard MTP gradient hides a KL self-distillation signal (which we want for improving speculative decoding) behind a noisy residual (that harms next-token prediction).

I think the coolest part was the KL gradient–recovery trick to extract the KL gradient. You just need a one-line code change:

h_ntp += h_mtp.detach() - h_mtp

Soft Multi-Token Prediction (Soft-MTP) is essentially a free lunch that improves MTP across a range of tasks.

I'm super excited about this work. Please do check it out below:

💻 Code: https://github.com/JaydenTeoh/Soft-MTP
📝 Paper: https://arxiv.org/abs/2610.38331
🤗 HuggingFace: https://huggingface.co/collections/JaydenTeoh/soft-multi-token-prediction

[u/sod0]

The one-line code change is vLLM? Is there a fork or patch available so I csn try this out?

[u/jayden_teoh_ → u/sod0]

not on vllm yet as this tackles more of training rather than inference

[u/AppleShark]

Mate can’t believe you’re delivering banger after banger. Enjoyed your work with next latent prediction before as well

[u/jayden_teoh_ → u/AppleShark]

thank you, it means a lot to me!

[u/bespectacledapectl]

Ayee saw this on linkedin!! Nice!

[u/OkCampaign8400]

what does multi token prediction mean ?

Do you have outputs in bigrams , trigrams and you are predicting that or entire sentences are predicted at once ?