面向语言模型的并行词元预测

Parallel Token Prediction for Language Models

arXiv: 2512.21323v1

论文信息

标题: Parallel Token Prediction for Language Models

作者: Felix Draxler, Justus Will, Farrin Marouf Sofian, et al.

发布日期: 2025-12-24

arXiv ID: 2512.21323v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：这篇论文要解决大语言模型自回归解码的串行瓶颈，即每生成一个 token 都需要一次完整的前向传播，导致长文本生成延迟极高。
  • 核心方法：提出并行 Token 预测框架（PTP），将采样过程中使用的辅助随机变量直接作为模型输入，使模型能够在单次 transformer 调用中联合预测多个相互依赖的 token。
  • 关键结果：在 Spec-Bench 基准测试上，O-PTP 在 Vicuna-7B 上平均每步接受超过 4 个 token（表 3），在 CodeContests 上使用 TinyLlama-1.1B 达到 7.0 个可接受 token（表 2），实现了并行解码的 SOTA 性能。
  • 主要局限：有限模型容量限制了单次前向传播能生成的一致性文本长度，需要结合错误校正机制（论文第 2.3 节）；更大规模的从零训练 PTP 模型尚未探索。
  • 适合读者：对 LLM 推理加速、投机解码、并行序列生成感兴趣的研究者和工程师，尤其是有 transformer 基础、对 token 生成机制有理解的读者。

论文背景和研究动机

自回归 transformer 是当代大语言模型的基础架构。其文本生成过程存在根本的串行瓶颈：每个 token 的生成依赖于前面所有已生成的 token，因此每生成一个新 token 都需要一次完整的模型前向传播。当输出较长序列时，推理延迟呈线性增长，相比单次 transformer 调用所能实现的理论速度，实际效率远未达到极限。

近年来，学术界提出了两大类加速方案。第一类是投机解码，采用系统架构思路，用轻量级模型生成候选 token 序列，再由大模型并行验证。第二类是一次性预测多个 token，包括多 token 预测和离散扩散模型，但这些方法假设未来 token 之间彼此独立，导致模型表征能力受限——采样的 token 是从边缘分布中提取的，无法捕捉条件依赖关系（见论文附录 G 的形式化推导）。

本文的核心突破是认识到：自回归采样虽然形式上表现为迭代过程，但一旦辅助随机变量 ui∼U[0,1]u_i \sim \mathcal{U}[0,1] 被抽取完毕，从累积分布函数中选择 token 的过程就变成了一个确定性的函数 ti=Pick⁡(ui,Pi)t_i = \operatorname{Pick}(u_i, P_i)。这意味着，如果模型能够提前获得这些辅助变量作为输入，它就能"预知"未来将采样哪些 token。

这一思想源于标准化流文献中的两个概念：Inverse Autoregressive Flows 实现了连续多维度并行采样，Free-form Flows 则从教师模型中蒸馏出快速生成网络。本文将这些连续空间的并行生成技巧移植到离散序列采样领域。

核心方法和技术细节

并行采样机制

经典自回归模型中，给定历史 token t<it_{<i}，下一个 token 的分布为 Pi:=P(ti∣t<i)P_i := P(t_i | t_{<i})。采样时抽取 ui∼U[0,1]u_i \sim \mathcal{U}[0,1]，通过累积分布函数查找对应 token：

ti=Pick⁡(ui,Pi)≡min⁡j∈{1,…,V}{j:Fij>ui},Fij=∑l=1jPilt_i = \operatorname{Pick}(u_i, P_i) \equiv \min_{j \in \{1,\dots,V\}} \{j : F_{ij} > u_i\}, \quad F_{ij} = \sum_{l=1}^{j} P_{il}

要害在于，uiu_i 一旦确定，tit_i 就是确定性结果。将此逻辑递归展开，就得到了核心定理（定理 1，见论文附录 A.1 的完整证明）：未来任意 token tkt_k 都可以表达为历史 token t<it_{<i} 和所有中间辅助变量 ui,…,uku_i, \dots, u_k 的确定性函数。

O-PTP 与 C-PTP 两种模式

基于上述思路，论文提出两种并行 Token 预测模式：

One-Hot Parallel Token Prediction (O-PTP) 将所有未来 token 对应的辅助变量直接输入模型，模型输出退化为单点分布，通过 argmax⁡\operatorname{argmax} 直接提取预测 token（公式 5）。这种方式不暴露原始条件概率，无法调节温度或进行不确定性量化。

Categorical Parallel Token Prediction (C-PTP) 则采用更精细的设计：预测 token tkt_k 时，仅以 ui,…,uk−1u_i, \dots, u_{k-1} 为条件，明确排除 uku_k 本身。这样做的好处是，过去的辅助变量已经确定性地编码了已采样历史 t<kt_{<k}，而保留 uku_k 的不确定性就能完整恢复原始自回归条件分布 P(tk∣t<k)P(t_k | t_{<k})（定理 2，见论文附录 A.2 的归纳证明）。这使得 C-PTP 可以进行温度调参等后处理操作。

训练范式与错误校正

PTP 支持两条训练路径。一是从已有的自回归教师模型蒸馏：对训练序列的每个 token 反向推导兼容的辅助变量边界 uk∈[Fk,tk−1,Fk,tk)u_k \in [F_{k,t_k-1}, F_{k,t_k})，然后分别训练学生模型（使用交叉熵或 KL 散度损失）。二是反向自回归训练，无需教师模型，通过迭代地让当前模型预测下一 token 的 logit 并用该 logit 推导辅助变量来实现自监督。

由于有限模型容量导致长序列预测可能出现错误，论文引入错误校正机制：PTP 一次性生成候选 token 序列后，用教师模型（或无教师时的自验证）重新采样并逐 token 比对，只接受完全匹配的前缀。这样保证输出与高质量自回归模型完全一致。

创新点和贡献

本文的核心创新在于打破了"token 必须逐个生成"的固有认知。通过将辅助变量的采样环节嵌入模型设计，实现了：

  1. 理论上的普适性：O-PTP 和 C-PTP 在模型容量足够的条件下，可以表征任意自回归序列分布，不存在建模能力的根本损失（定理 1、2）。

  2. 方法上的统一性：PTP 框架既覆盖有教师模型蒸馏的迁移学习场景，也支持无教师的从零训练场景，具有良好的灵活性。

  3. 实践上的高性能：在投机解码框架中将草稿模型升级为 PTP，可直接利用现有系统架构实现显著加速（图 4 显示 PTP 草稿模型在不同规模下均优于自回归草稿模型）。

  4. 对独立预测局限性的有力证据：通过对比实验清楚地表明（图 5、表 2），不提供辅助变量的独立 token 预测会产生大量语义不一致的 token 组合，如 def numpy 这类无效 Python 代码片段。

实验结果分析

从零训练的有效性验证

在 NYC 出租车数据集（将城市划分为 25 个区域预测接客位置）上，C-PTP 经过反向自回归训练后，困惑度（19.88）几乎与相同架构的自回归模型（19.81）持平（表 1），证明了无教师训练路径的可行性。

代码生成性能突破

在 CodeContests 编程挑战数据集上，以 TinyLlama-1.1B-Chat 为教师模型蒸馏 O-PTP 时，PTP 草稿模型每步可接受 7.0 个 token，而独立预测基线仅 6.2 个（表 2）。图 4 展示了不同草稿模型规模下的扩展性分析，PTP 在所有参数规模下均保持明显优势。

通用自然语言生成的 SOTA 表现

在 SpecBench 基准测试（涵盖多轮对话、翻译、摘要、问答、数学推理、检索增强生成共六类任务）上，O-PTP 平均每步接受 4.18 个 token，超过 SAMD（3.90）、Eagle-2（3.86）等竞争方法（表 3）。该项实验以 Vicuna-7B 为教师，通过 gated LoRA 适配器进行参数高效微调，体现了 PTP 在大模型上的实用潜力。

辅助变量嵌入设计的工程选择

论文通过一系列消融实验阐明了几个关键工程决策（见附录 C）。辅助变量嵌入采用将 IEEE-754 浮点数比特位直接映射为 32 维二值向量再线性变换的方案（称为"arithmetic coding"嵌入），相比全连接网络或简单阈值嵌入，在 O-PTP 上达到 46.1 个正确 token，显著优于其他方案的 35.9-40.9 个（表 5）。

实践建议

对于希望将 PTP 集成到现有 LLM 推理系统的工程师，以下几点值得关注：

  1. 从投机解码切入：PTP 与投机解码框架天然兼容，最简单有效的落地方式是将现有的自回归草稿模型替换为 O-PTP 草稿模型。论文已在 TinyLlama 和 Vicuna 上提供了端到端验证，性能提升在 7B 模型上明确可见（表 3、图 4）。

  2. 辅助变量嵌入的选择：使用 bits(u) 将 float32 的 IEEE-754 位表示映射为二值向量的方法（公式 14）在不同任务和模型尺度上均表现稳定，论文建议直接采用此方案。

  3. 教师数据生成策略：蒸馏 PTP 时，从教师模型采样训练序列（而非使用原始训练数据集）可获得最佳效果（见论文附录 C.2 表 4），以教师分布为提议分布可以降低训练方差。

  4. 条件分布保真度的权衡：如果需要保留温度调节、top-k/top-p 采样等后处理能力，应选择 C-PTP（公式 6-7）；若只追求最大程度的并行速度提升，O-PTP 的 one-hot 输出更易训练且预测准确率更高。

  5. 硬件资源利用：论文附录 E 提供了利用并行计算资源进一步降低延迟的思路——同时运行多个独立的 PTP（抽取不同辅助变量），择优选择预测更稳定的序列，可在不引入偏差的前提下提升正确 token 数量（表 7 展示了 M=1024M=1024 时平均正确 token 从 45.36 提升到 59.79）。

值得提醒的是，目前 PTP 主要在中等规模模型（<=7B 参数）上验证，论文明确指出，更大规模模型的从零训练以及推理时的计算最优方案仍待后续工作探索。实际部署时应评估模型容量是否足以在多 token 并行预测中维持输出质量。