Predicting the Order of Upcoming Tokens Improves Language Modeling

Multi-token prediction (MTP) has been proposed as an auxiliary objective to improve next-token prediction (NTP) in language model training but shows inconsistent improvements, underperforming in standard NLP benchmarks. We found MTP's exact future token prediction to be too difficult as an auxiliary loss. Instead, we propose token order prediction (TOP), which trains models to order upcoming tokens by their proximity using a learning-to-rank loss. TOP requires only a single additional unembedding layer compared to MTP's multiple transformer layers. We pretrain models of 340M, 1.8B, and 7B parameters using NTP, MTP, DeepSeek MTP (DS-MTP) and TOP objectives. The results of nine standard NLP benchmarks show that TOP overall outperforms NTP, MTP, and DS-MTP even at scale. TOP models with continued training on math and code also perform better on 4 relevant benchmarks. On the synthetic star graph task, TOP enables pathfinding on graphs where NTP, MTP, and DS-MTP fail. Our implementation and training code is available at https://github.com/zaydzuhri/token-order-prediction.

Large language models learn to generate text by predicting the next word in a sentence, a training method known as next-token prediction. More recently, researchers proposed multi-token prediction, where the model learns to predict several future words at once. While this can help on tasks like coding and mathematics, it does not consistently improve general language understanding. We believe this is because predicting several exact future words is too difficult as a learning task. Instead, we propose a simpler and more flexible approach called token order prediction. Rather than guessing the exact future words, the model only learns which words are likely to appear in what approximate order in the upcoming text. This objective encourages the model to “look ahead” without requiring perfect predictions far into the future. Across a wide range of experiments, including language understanding, coding, and mathematics, our method consistently improves performance over existing approaches while also requiring fewer additional model components during training.