以“排序”替代“预测”：Token Order Prediction（TOP）如何改进语言模型训练

在当下以 Next-Token Prediction（NTP）为主导范式的 LLM 训练中，研究者一直在探索如何通过辅助目标提升内部表征，从而带来更好的下游泛化。本文聚焦一篇来自 MBZUAI 的早期预印本，提出以“排序未来词元的接近度”替代“精确预测未来词元”的新思路，将 Multi-Token Prediction（MTP）的“难度过高”转化为 Learning-to-Rank 的“近似排序”问题。核心结论显示：在 340M、1.8B 与 7B 三个模型规模上，作为 NTP 的辅助目标，Token Order Prediction（TOP）整体优于 NTP 与 MTP，且在标准 NLP 基准上更稳健。

该研究的关键创新在于将未来词元的“顺序信息”转为可学习的排名分布，用 Listwise 排序损失对单个额外的 unembedding 头进行监督，避免为每个未来偏移位置增设独立 Transformer 层的额外参数与计算开销。与 MTP 相比，TOP 训练更稳定、扩展性更好，并在多项通用任务上呈现出一致的提升趋势。

1. 基本信息

• 论文《Predicting the Order of Upcoming Tokens Improves Language Modeling》
• 作者为 Zayd M. K. Zuhri、Erland Hilman Fuadi 与 Alham Fikri Aji，研究单位为 MBZUAI。
• 该工作为早期预印本，版本 v1 发表于 2025-08-26，arXiv 链接为：https://arxiv.org/abs/2508.19228。
• 作者开源了实现与训练代码，仓库地址为：https://github.com/zaydzuhri/token-order-prediction

2. 研究背景

以 Next-token Prediction, NTP 为中心的自回归训练范式在语言理解与生成上取得了卓越成绩，但其局限与改进空间同样被持续讨论。近年来，一类代表性尝试是以Multi-token prediction, MTP 为辅助目标，让模型在共享 Backbone 的末端分出多个分支，分别预测   的精确词元，从而促使主干表征“向前看”。这类方法在代码生成、摘要等需要较强“前瞻性”的生成任务上常有收益，同时还可被用于一定形式的自我推测解码，以提高推理速度。

然而，MTP 的泛化改善在通用 NLP 基准上并不稳定，尤其在小模型上常见无效甚至退化。经验上，当未来偏移步数   增加时，训练难度显著提升，且最优   难以一刀切地确定。

这从侧面表明，“精确预测多步未来词元”的学习目标可能过难，难度的不匹配反而削弱了其作为辅助目标的普适价值。该研究据此提出一个关键判断：与其强迫模型“精确命中远期词元”，不如让模型“学会排序哪些词元更快出现”，即把“多步精确预测”放宽成“近似顺序学习”的 Learning-to-Rank 问题，从难题转化为“更可学”的目标。

因此，研究的动机在于：在不牺牲“向前看”这一有益归纳偏置的前提下，构建更温和、参数与计算开销更友好、且在通用任务上更稳定的辅助训练目标。Token Order Prediction, TOP 即在此语境下提出，用单一额外的 unembedding 头和一个 listwise 排序损失，诱导 Backbone 习得对“未来词元接近度”的结构化感知。

3. 方法

该方法的设计出发点是充分保留“未来结构”的监督信号，同时回避 MTP 在多步精确命中上的固有难度与扩展瓶颈。MTP 的每一个未来步都需配备一层单独的 Transformer 头，既增加参数又加重计算；更关键的是，越远的未来步难度越高、梯度信号越噪，容易拖累整体优化。Token Order Prediction, TOP 将“精确分类到某个未来偏移位置”的任务，替换为“对全词表进行‘下一个出现时间’的接近度打分”，把硬分类变成软排序，使训练信号更平滑、覆盖面更广、且随着窗口大小变化不会线性地膨胀参数。

整体框架上，模型主干仍采用标准自回归 Transformer，仅在输出层并联两个线性头：用于常规 NTP 的   与用于排序监督的  。训练时联合最小化两者的损失之和，推理时移除  ，仅保留  ，因此不会改变推理时的架构与接口。

具体实现分为目标构造与损失定义两部分。

• 首先，给定输入序列  、词表大小   与窗口大小  ，对每个时间步   构造长度为   的“接近度向量”  。对任意词元  ，若其在区间   的首次出现位置与   的距离为  ，则令  ；若在窗口内不可达，则  。这意味着离当前更近的“即将出现词元”将获得更高分值，从而形成对“未来顺序”的隐式排序监督。该目标可以用如下伪代码（与论文一致）刻画其自右向左的一次扫描构造过程：

Input: token sequence x (length T+W), vocab size V, window size W
Output: target tensor y of shape (T, V)

Initialize y[:] = -∞
Initialize next[v] = T + W for all v in [0, V-1]

for t from T+W-1 down to 0:
    if x[t] in vocab:
        next[x[t]] = t
    if t < T:
        for each v in [0, V-1]:
            d = next[v] - t
            if 0 < d ≤ W:
                y[t, v] = W - d

损失函数借鉴 listwise Learning-to-Rank 思想。设主干最后一层隐藏状态为  ，NTP 头与 TOP 头分别为线性映射   与  。标准 NTP 损失为

TOP 的 listwise 排序损失将   视作“软目标分数”，在归一化后与预测打分的归一化分布求交叉熵：

最终优化目标是两者之和：

从表征学习角度看，  以“接近度排名”为监督，迫使   捕获“短期将出现哪些词元及其大致先后”的结构性信息；这种信息与 NTP 的“下一个词元概率”目标高度一致，二者在“靠近下一个词元”的方向上形成合力，从而以较小的额外参数与显著低于 MTP 的计算开销，增强主干的建模能力。实现上，作者使用了融合的 Triton kernel 将 unembedding 与损失计算在 block 级别一次完成，几乎不引入额外吞吐损失；由于 TOP 仅新增一个与   同形的线性层，参数与显存开销远低于随   线性增长的 MTP 多层头。

值得强调的是，若仅以 TOP 训练而移除 NTP，则推理时只能进行贪心生成，缺乏概率采样的灵活性；因此该工作定位 TOP 为“辅助目标”，而非替代 NTP 的主目标。

4. 实验与发现

实验在三个模型规模上系统比较了 NTP、MTP 与 TOP：约 340M、1.8B 与 7B 参数。训练数据来自 FineWeb-Edu 的“sample-100BT”子集，340M 训练 52B tokens，1.8B 与 7B 训练 104B tokens。统一采用序列长度 4096、RoPE  、词表大小 32k、未绑权重（untied embeddings），优化器为 AdamW，余弦学习率调度与适当 warmup，并在 MTP 设置中使用 4 个未来步。为保证可重复性，论文详细给出了每个规模的层数/隐藏维度/头数、学习率与 batch 配置；实现基于 Flame 与 flash-linear-attention。

评测覆盖 8 个标准 NLP 基准：LAMBADA（准确率与困惑度）、HellaSwag、ARC Challenge、PIQA、SciQ、Social IQa、NaturalQuestions Open、TriviaQA（Exact Match）。核心发现可以概括为三个层面。第一，在几乎所有规模与多数任务上，TOP 相较基线 NTP 与 MTP 都带来一致提升。例如，LAMBADA 上 TOP 在 340M/1.8B/7B 的准确率分别优于 NTP（36.35%→37.07%、49.58%→50.34%、55.89%→57.03%），困惑度也同步下降（340M：30.34→28.76；1.8B：11.38→11.19；7B：7.97→7.64）。HellaSwag 的归一化准确率在三个规模上亦全面提升（340M：42.53%→43.57%；1.8B：60.05%→60.45%；7B：67.44%→68.73%）。第二，MTP 在通用理解类任务上并不稳定，尤其 7B 规模时常出现退化；TOP 则随规模增大持续受益，在 TriviaQA（EM）上差距尤为明显（1.8B：11.85→18.93；7B：24.28→30.90）。第三，尽管 TOP 的训练阶段 NTP 头上记录到的训练损失略高于纯 NTP（提示正则化效应与更少过拟合），但在评测困惑度与准确率上表现更佳，指向更强的泛化。

从统计与实际意义讨论来看，这些改进不仅体现在平均数上，也具备一致的跨任务、跨规模可迁移性。尤其是当任务偏向理解与检索式问答（如 NQ Open、TriviaQA）时，TOP 在较大规模上展现出明显优势，说明“顺序接近度”的结构性监督与语言模型的长程一致性建模、知识寻址与答案定位存在内在耦合。相对地，个别社会常识类任务（如 Social IQa）在 7B 上略有波动，提示未来可以在窗口机制、权重共享或损失权重上做更细粒度的调节与消融。

值得注意的是，MTP 的小规模模型在该复现实验中并非处处落后，这一结果与部分先前报告形成互补；但随着规模扩大，MTP 在通用任务上的弱势更加明显。TOP 则展现出“越大越好”的单调趋势，符合“排序监督更易优化、对主干更友好”的设计初衷。

5. 结论与展望

该研究以 Token Order Prediction（TOP）为核心贡献，给出了一个轻量、可扩展且与 NTP 强一致的辅助训练目标。通过将“未来词元精确预测”的目标，替换为“未来词元接近度排序”的 listwise 学习，TOP 在相近或更小的额外开销下，更稳定地提升了 LLM 在通用 NLP 基准上的表现。实验显示，随着参数规模增大，TOP 的收益进一步扩大，这为大模型预训练中的辅助目标设计提供了新的方向。

说明：本文基于作者公开的早期预印本撰写，实验细节与扩展结果以后续版本为准。