arXiv · 2503.21801
Efficient Joint Prediction of Multiple Future Tokens
Abstract
In this short report, we introduce joint multi-token prediction (JTP), a lightweight modification of standard next-token prediction designed to enrich hidden state representations by jointly predicting multiple future tokens. Unlike previous multi-token prediction approaches, JTP strategically employs teacher forcing of future-tokens through a carefully designed representation bottleneck, allowing the model to encode rich predictive information with minimal computational overhead during training. We show that the JTP approach achieves a short-horizon belief state representation, while popular alternatives for multi-token prediction fail to do so. We demonstrate the effectiveness of our method on the synthetic star graph navigation task from from Bachmann and Nagarajan [2024], highlighting a significant performance improvement over existing methods. This manuscript presents promising preliminary results intended to stimulate further research.
Explore related subjects
Keep this discovery
Kwangjun Ahn, Alex Lamb, John Langford. 2025-03-24. Efficient Joint Prediction of Multiple Future Tokens. https://arxiv.org/abs/2503.21801
Cite the original work for its findings. Save a collection to share your selection of sources.