近年来,语言模型的显著进展主要得益于大规模文本数据的可获得性以及自回归训练方法的有效性。
在这种自回归的训练框架中,每一个 token 都作为其前文上下文的预测目标。
这一方法无需显式标注,使得自回归模型在训练阶段具有明显优势,因而成为主流范式。
然而,在推理阶段,自回归生成本质上是顺序执行的,计算开销大,因为每一步解码都需要完整地运行一次模型。相比之下,人类在表达时常常会先在句子层面组织思路,再逐词输出。
那么在大模型推理阶段,能否让大模型摆脱单个 token 顺序预测的瓶颈,也和人类的逻辑一样,从更大层面考虑输出内容呢?
假如能够打破自回归模型的顺序执行本质,就产生了如网友所说的「具有时间跳跃能力的 LLM」。
受到启发,苹果的研究人员为此展开思考,借助类似策略,开发了一个框架,使预训练的自回归大型语言模型能够执行多 token 预测,在保持生成质量的同时,为代码和数学任务提供高达 5.35 倍的推理加速,以及为一般任务提供约 2.5 倍的推理加速。
AI 工程师 Jackson Atkins 认为这是一项具有开创性的工作,「最令人兴奋的是这一切都可以通过对现有模型进行 LoRA 微调实现」。
想象一下,将 AI 运行成本削减数倍,或是在轻量设备上流畅运行强大的实时助手,这对推理优化而言是一次颠覆性的变革。
如果你手中有一个运行速度提升 5 倍的大语言模型,第一件事会做什么?
论文标题:Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential
论文链接:https://www.alphaxiv.org/abs/2507.11851
首先探讨一个基本问题:语言模型能否在单个推理步骤中生成多个 token?
令人鼓舞的是,答案是肯定的。
现有的推测解码研究已经探索了这个方向以加快生成速度。推测解码方法利用一个草稿模型生成多个 token,然后通过一个验证器检查它们与标准自回归输出的一致性。虽然这种方法提供了加速,但它仍然根本上依赖于自回归生成。
在这项工作中,研究者们提出了一个更深层次的问题:我们能否训练真正非自回归的语言模型?
研究者们通过设计完全非自回归的训练算法来探讨这一问题,例如基于扩散机制的语言模型。然而,此类方法通常需要构建全新的建模与训练流程。
于是研究者们进一步提出一个问题:是否可以在尽可能少改动的前提下,适配现有的自回归训练与推理框架?目标是实现在保留自回归模型核心优势的同时,实现高效的多 token 生成。
为进一步论证这一点,研究者们首先观察到,自回归模型虽然并未明确针对未来 token 进行训练,但实际上已经在一定程度上编码了关于未来 token 的信息。
例如,给定提示词 “what is two plus two?”,一个预训练模型在标准自回归解码过程中通常会生成 “two plus two equals four”。为了检验模型是否具备对未来 token 的感知能力,研究者们在提示后添加占位 token(图中以 ↔ 表示),并分析输出的 logits,如图 1 (左)所示。
令人惊讶的是,未来 token 的正确序列出现在前 200 个 logits 中,说明模型在某种程度上已隐含地掌握了即将生成的 token 信息。
图 1:自回归模型能够隐式预判未来的 token。
基于上述观察,研究者们进一步探究能否引导模型发挥其潜在能力,从而更好地组织对未来 token 的预测。
为此,研究者在提示词末尾引入若干 mask token,并对模型进行训练,使其能够直接预测这些 token。如图 1 (中)所示,经过微调后的模型能够将正确的 token 提升至前 10 个 logits 中。
最后,为生成连贯的多 token 输出,研究者们引入了一个轻量级的采样模块:一个两层感知机(two-layer perceptron),该模块在预测每个 token 时,结合先前已采样出的 token 作为条件,具体结构如图 1 (右)所示。
与之前的一些方法不同,本文训练模型填充 mask token,以预测未来 token。该模型在推理这些 token 时,充分利用了自身的全部深度与表示能力,并结合整个序列的上下文信息,从而在效果上显著优于现有的多 token 预测方法。
此外,得益于一种简单而有效的技术 —— 门控 LoRA 适配(gated LoRA adaptation),本方法在生成质量上无任何下降。
在较为轻量的监督微调下,研究者们相较于传统回溯式生成模型实现了加速效果。如图 2 所示,在 Tulu3-8B 模型基础上微调以预测 8 个额外 token,即可获得显著性能提升。
图 2:在通过监督微调训练门控 LoRA 和采样头后获得的加速效果。
方法
为在尽量减少重新训练的前提下实现多 token 生成,引入了一类特殊的 token,称为 mask。设原始序列为 X = [x1, …, xn],基本思路是在该序列末尾附加 k 个唯一的 mask token,从而构造出扩展序列 Xm = [x1, …, xn, m1, …, mk]。其中,mask token [m1, …, mk] 的表示向量被初始化为随机向量,并添加至模型的嵌入表中。
在本文中,将模型对标准下一个 token 的预测称为 NTP(Next Token Prediction),而对 mask token 的预测称为 MTP(Mask Token Prediction)。
本文所提出的 MTP 模型的整体架构如图 3 所示,其中展示了在微调时使用 k=2 个 mask 的模型在推理阶段的工作流程。在图中 box-1(左上角),扩展后的序列 Xm 被输入至解码器,模型生成的潜在表示中,[z1, …, zn] 对应于 NTP token 的表示,而 [zn+1, …, zn+k] 则对应于 MTP token 的表示。
图 3:MTP 模型的组成部分。框 1(左上)显示了带有门控 LoRA 参数的自回归模型。框 2(左下)说明了采样头。框 3(右)展示了门控 LoRA 模块的框图。
图中 框 2(左下)展示了采样器头(sampler head)的结构。第一个(NTP)token 通过标准的反嵌入层(unembedding layer)以自回归方式生成,即由 z_n 预测出 y_n+1。而剩余的(MTP)token 则由采样器模块依次生成。在每一步中,采样器根据 z_n+k+1 与前一个已生成的 token y_n+k 来预测 y_n+1+k,确保每个生成的 token 同时融合模型的潜在表示与先前已采样的 token 信息。
为了在微调时保留预训练模型的行为,研究者们在解码器层中引入了门控 LoRA 模块(gated LoRA)。微调过程中,仅更新 LoRA 参数与采样器头参数,原始解码器权重保持冻结状态。门控 LoRA 模块通过对 NTP 与 MTP token 分别采取不同的计算路径,确保微调过程不会影响 NTP token 的生成行为,如图中 框 3(右侧)所示。该差异通过引入二进制 mask 实现,二进制 mask 被作为额外输入传入解码器层。
详细的模型训练过程请参阅原论文。
实验
研究者们在 Tulu3-8B SFT 模型上开展了实验。该模型属于 LLaMA-3 系列,并基于 Tulu3 数据集进行了有监督微调。
生成质量评估
研究者们首先在微调过程中跟踪了模型的准确率变化。由于所用模型为 SFT 模型,其对额外训练非常敏感。图 6 (a) 展示了模型在 ARC-Challenge 基准上的 zero-shot 准确率,该评估通过 Harness 库完成。
如图 6 (a) 中的虚线所示,gated LoRA 能够保持准确率稳定,原因在于其保证了 NTP token 的输出不会受到微调影响。
如图 6 (b) 中所示,在标准 LoRA 与 gated LoRA 两种训练方式都实现了有效的收敛。
接下来,研究者们在图 6 (c) 中分析了 NTP token 上的交叉熵损失,该指标与生成质量直接相关。如图所示,使用标准 LoRA 训练的模型,其 NTP 损失呈上升趋势,表明生成质量逐渐下降。相比之下,采用 gated LoRA 的模型在整个训练过程中保持了几乎恒定的 NTP 损失,这是因为在 gated LoRA 机制下,梯度不会在 NTP token 上传播,从而有效避免了对原始生成能力的干扰。
图 6:使用标准 LoRA 和门控 LoRA 适配器训练的模型收敛性分析。
加速能力分析
为了评估所提方法在生成速度上的提升,研究者们采用了自推测解码算法(self-speculative decoding)。
实验最多运行 100 个生成步骤,若模型在第 100 步之前生成了句子结束符(end-of-sentence token),则提前终止生成过程。设总生成步骤数为 T,总生成 token 数为 G,则计算接受率(acceptance rate)为 G/T。该指标反映了模型平均每一步生成多少个被接受的 token,进而衡量所提方法带来的加速效果。
接受率的理论最小值为 1,因为每一步至少会生成一个 token(即标准的下一个 token 预测)。最大理论值为 k+1=9,其中 k=8 是训练时使用的 mask token 数量。
表 1 报告了模型在五个不同任务领域下的接受率,包括:知识问答、数学、编程、对话和安全。
如表所示,所提出的多 token 生成算法在不同任务上实现了约 1.5 到 5.2 倍不等的加速,具体取决于任务类型及插入的 mask 数量。值得注意的是,在编程和数学任务中,加速效果更为显著,可能是因为这些领域中未来 token 的可预测性较高。
表 1:Tulu-3 在不同领域中通过多标记预测获得的
加速。
消融实验
表 1 中的结果对应于本方法的最优配置,包含以下三个关键组件:
(1)使用采样器 MLP 头(sampler MLP head);
(2)训练过程中引入 LCM 损失(LCM loss);
(3)生成阶段采用二次解码算法(quadratic decoding algorithm)。
研究者们对各个组件的贡献进行了消融实验。
图 7: MTP 模型的简单配置和最先进配置所达到的平均加速效果。基础版本使用线性解码,不包含 LCM 损失或采样头,而高级版本则包含 LCM 损失、采样头和二次解码。
研究者们训练了多个不同 LoRA 秩的模型,以进一步探索其影响。图 8 总结了相关结果:左侧与中间子图分别展示了在未使用与使用采样器头的条件下的接受率;右侧子图则展示了采样器头与 LoRA 参数所带来的显存开销。
图 8:LoRA 秩的影响:无采样头加速(左),带采样头加速(中),以及 LoRA 和采样头的内存开销(右)。
总结
研究者们评估了自回归模型在语言模型有监督微调阶段对多 token 预测任务的适应能力。未来值得探索的一个方向,是在预训练阶段或下游任务自适应阶段引入该方法,以进一步检验其适用性与效果。另一个具有前景的研究方向是将基于扩散的生成方法应用于多 token 预测任务。研究者们认为,多 token 预测位于完全自回归生成与完全扩散生成之间,能够在两者之间取得优势的平衡,兼具效率与质量的潜力。
更多信息,请参阅原论文。