Прогнозирование порядка следующих токенов улучшает языковое моделирование

Predicting the Order of Upcoming Tokens Improves Language Modeling

August 26, 2025

Авторы: Zayd M. K. Zuhri, Erland Hilman Fuadi, Alham Fikri Aji

cs.AI

Аннотация

Метод Multi-Token Prediction (MTP) был предложен в качестве вспомогательной задачи для улучшения предсказания следующего токена (NTP) при обучении языковых моделей, однако он демонстрирует неоднозначные результаты и уступает в стандартных NLP-бенчмарках. Мы утверждаем, что точное предсказание будущих токенов в MTP является слишком сложной задачей для вспомогательной функции потерь. Вместо этого мы предлагаем Token Order Prediction (TOP), который обучает модели упорядочивать предстоящие токены по их близости, используя функцию потерь, основанную на обучении ранжированию. TOP требует только одного дополнительного слоя для преобразования токенов в сравнении с несколькими трансформерными слоями в MTP. Мы предобучаем модели с 340 млн, 1,8 млрд и 7 млрд параметров, используя задачи NTP, MTP и TOP. Результаты на восьми стандартных NLP-бенчмарках показывают, что TOP в целом превосходит как NTP, так и MTP, даже в масштабе. Наш код доступен по адресу https://github.com/zaydzuhri/token-order-prediction.

English

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 argue that MTP's exact future token prediction is 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, and TOP objectives. Results on eight standard NLP benchmarks show that TOP overall outperforms both NTP and MTP even at scale. Our code is available at https://github.com/zaydzuhri/token-order-prediction