Predicting the Order of Upcoming Tokens Improves Language Modeling

Zayd Muhammad Kawakibi Zuhri, Erland Hilman Fuadi, Alham Fikri Aji
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:169012-169029, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-zuhri26a, title = {Predicting the Order of Upcoming Tokens Improves Language Modeling}, author = {Zuhri, Zayd Muhammad Kawakibi and Fuadi, Erland Hilman and Aji, Alham Fikri}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {169012--169029}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/zuhri26a/zuhri26a.pdf}, url = {https://proceedings.mlr.press/v306/zuhri26a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Predicting the Order of Upcoming Tokens Improves Language Modeling %A Zayd Muhammad Kawakibi Zuhri %A Erland Hilman Fuadi %A Alham Fikri Aji %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-zuhri26a %I PMLR %P 169012--169029 %U https://proceedings.mlr.press/v306/zuhri26a.html %V 306 %X 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.
APA
Zuhri, Z.M.K., Fuadi, E.H. & Aji, A.F.. (2026). Predicting the Order of Upcoming Tokens Improves Language Modeling. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:169012-169029 Available from https://proceedings.mlr.press/v306/zuhri26a.html.

Related Material