Predicting the Order of Upcoming Tokens Improves Language Modeling

TL;DR

Token order prediction (TOP) is proposed, which trains models to order upcoming tokens by their proximity using a learning-to-rank loss, which requires only a single additional unembedding layer compared to MTP's multiple transformer layers.

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 code is available at https://github.com/zaydzuhri/token-order-prediction

Figures (11)

• Figure 1: An overview of token order prediction (TOP). Given an input token sequence, a vocabulary, a sequence length of 4 and window size of 4, a TOP target sequence is constructed via Algorithm \ref{['alg:seq_to_top']}. The output hidden representation of the final layer goes to two separate unembedding heads for NTP and TOP. The final loss to optimize is a sum of the NTP and TOP loss.
• Figure 2: Left: Training loss of a MTP transformer with 16 MTP heads predicting tokens at $t+1,...,t+16$ offsets. Right: Training loss of the MTP model averaged over all 16 heads, compared to the training loss of a same-sized TOP model with window size 16.
• Figure 3: Illustration of a star graph training sample with $d=3$ and $l=3$ due to bachmann2024pitfalls.
• Figure 4: Test set accuracy during training of NTP, TOP, MTP, DS-MTP models on the star graph pathfinding task. Star graph setups: $G(3,3)$ (top left), $G(3, 5)$ (top right), $G(5, 3)$ (bottom left), $G(5, 5)$ (bottom right).
• Figure 5: Pretraining loss of 340M parameter base models.