[论文] 通过自蒸馏的多标记预测

Source: arXiv - 2602.06019v1

Overview

一支研究团队提出了一种出人意料的简便方法，能够将任何预训练的自回归语言模型（例如 GPT‑2、LLaMA）转变为 multi‑token predictor，在推理时的速度提升可达三倍。研究者并未添加额外的“投机者”模型或构建复杂的解码流水线，而是通过 online self‑distillation 目标，训练原始模型在一次前向传播中预测多个后续 token。实验表明，在数学推理基准（GSM8K）上，准确率下降不足 5 %，同时实现了显著的加速，且无需更改模型架构或部署代码。

关键贡献

• 自蒸馏用于多标记预测：引入一种在线蒸馏损失，使预训练模型输出一小段标记序列，而不是单个下一个标记。
• 零改动部署：最终模型使用与原单标记模型完全相同的检查点和推理代码——无需额外的验证器、预测器或自定义运行时。
• 经验性的速度‑精度权衡：在 GSM8K 上实现 > 3 倍的解码速度提升，且相对精度损失 < 5 %，缩小了投机解码与朴素单标记生成之间的差距。
• 广泛适用性：适用于任何自回归语言模型，无论规模或预训练数据如何，都是现有服务的即插即用升级。

方法论

1. Baseline model – 从一个冻结的预训练自回归语言模型开始（例如，仅解码器的 Transformer）。
2. Online teacher – 在训练期间，同一模型以常规的单 token 模式运行，为接下来的 k 个 token 生成“教师”预测（k 为 4 或 8 等小整数）。
3. Student head – 在最终隐藏状态上附加一个轻量级的额外头部，训练其一次性直接输出 k 个 token。
4. Distillation loss – 对学生的 logits 与教师的 logits（交叉熵）在所有 k 位置上的偏差进行惩罚，同时保留首个 token 的原始语言模型损失。
5. Curriculum – 随着训练进行，逐步增大 k 的取值，使模型能够在不失稳的情况下适应更长的预测视野。
6. Inference – 推理时，模型直接调用新头部一次性生成 k 个 token，然后向前移动上下文并重复该过程，无需额外的验证器来检查多 token 预测是否正确。

由于教师和学生是同一个网络，整个过程是 self‑distilling，可以在线完成（无需单独的教师模型或数据集生成步骤）。

结果与发现

• 速度：解码速度大致随 k 线性提升，直至某一点；当 k = 8 以上时，准确率下降更为明显。
• 质量：适度的准确率损失主要源于偶发的“漂移”，即早期令牌错误在多令牌块中传播，但大多数错误可在后续块中恢复。
• 兼容性：相同方法已在 GPT‑2‑medium 和 1.3 B LLaMA 检查点上测试，显示出一致的加速且无需任何架构更改。

实际影响

• 更快的 API：云服务提供商只需通过自蒸馏损失对模型进行微调，即可提升现有 LLM 接口的吞吐量——无需额外服务器或自定义推理内核。
• 成本节约：每生成一个 token 所需的前向传播次数减少，直接降低 GPU/TPU 使用率，为高并发应用（聊天机器人、代码补全等）削减推理费用。
• 边缘部署：计算资源受限的设备（移动端、物联网）可以通过一次推理生成多个 token，从而更流畅地运行更大的模型，扩大了设备端可用的模型规模。
• 简化流水线：与投机解码不同，无需维护单独的“验证器”模型或编排投机‑回退逻辑，降低了工程工作量和潜在的 bug。

限制与未来工作

• 错误传播：当模型在块的早期 token 上预测错误时，错误会影响该块的其余部分，导致偶尔出现低质量输出的突发。
• 固定块大小：当前方法使用静态 k；基于置信度的自适应块长度可能会提升准确性‑速度的权衡。
• 基准范围：实验仅聚焦于 GSM8K（数学推理）；需要更广泛的评估（对话、代码生成、长文本）来验证通用性。
• 训练开销：自蒸馏微调相比纯推理部署会增加额外计算，尽管相对于预训练成本仍属适度。

未来的研究方向包括在块内部集成基于置信度的提前停止，将自蒸馏与量化或剪枝相结合以实现更紧的延迟预算，以及探索多模态扩展（例如视觉‑语言模型）。

作者

• John Kirchenbauer
• Abhimanyu Hans
• Brian Bartoldson
• Micah Goldblum
• Ashwinee Panda
• Tom Goldstein

论文信息

• arXiv ID: 2602.06019v1
• 类别: cs.CL, cs.LG
• 出版日期: 2026年2月5日
• PDF: 下载 PDF