Reproduce “Multi-Token Prediction Done Right” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Jayden Teoh
Soft-MTP changes multi-token prediction training so MTP heads learn from the model’s own next-token predictions while the shared output head retains cross-entropy training. It recovers this self-distillation signal through gradient routing, without computing a vocabulary-wide KL loss. Across 167M–1.2B models, it improves validation loss and eight-task average accuracy over standard MTP; evidence is limited to tested scales.
Key insights
- The residual component has about four times lower signal-to-noise ratio than the KL component and contributes roughly 80% of the MTP gradient reaching the shared trunk.
- In the reported analysis, steps in the KL component’s direction reduce next-token loss, whereas steps in the residual direction increase it.
- Soft-MTP matches direct KL training on many measured outcomes while preserving standard MTP’s reported memory use and throughput.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.