新扩散语言模型引入Token时间连续机制

背景

近年来，扩散模型在图像、音频等连续域取得了巨大成功，但在自然语言处理中，语言建模通常依赖离散 token 空间。已有的扩散语言模型大多在离散空间迭代，通过并行采样多个 token 来加速生成。然而，这种并行采样在高加速比下会引入显著的不准确性——因为 token 之间相互依赖，同时采样多个 token 会破坏这种依赖关系，导致生成质量下降。此外，传统扩散语言模型对所有 token 采用统一的加噪-去噪时间表，无法区分不同 token 的不确定性程度，也难以建模条件生成中的差异性。为了解决这些问题，本文提出了一种新型扩散语言模型——token time continuous diffusion (TTCD)，通过引入连续空间建模和每个 token 独立的时间概念，在保持高质量生成的同时实现了高加速比。

核心内容

本文提出 token time continuous diffusion (TTCD) 模型，其核心创新包含两点：

1. 连续空间建模：TTCD 在连续空间中操作，从高斯噪声出发，通过确定性映射（无需额外采样）直接生成最终的 token 画布（token canvas）。这与传统离散空间扩散模型不同，后者需要在每一步并行采样多个 token。连续空间建模避免了并行采样带来的不准确性，因为生成过程是确定性的，且每个 token 的演化路径不依赖同一时刻的其他 token 的随机采样结果。

2. 逐 token 时间概念：TTCD 引入了每个 token 独立的时间参数（per-token times），即不同 token 从噪声到最终 token 的演化速度可以不同。有些 token 可以以更快的速率完成从噪声到 token 的转变，而其他 token 则较慢。这种设计带来三方面好处：

   • 更好地建模条件生成：在条件生成中，某些 token 可以更快确定（例如已知前缀），而其他 token 需要更多步骤。
   • 更确定（更确信）的 token 可以以更快的速率推进，减少不必要的迭代。
   • 在 refinement（精炼）过程中，不同 token 之间可以实现差异化的相互影响：较慢的 token 可以从较快确定的 token 处获得更准确的上下文信息。

在实验中，作者在 OpenWebText 数据集上训练了一个 160M 参数的 TTCD 模型，然后对其自蒸馏（self-distill）。在高加速比（即生成步数很少）的设置下，TTCD 在无条件下生成质量与多个同类大小、相同数据训练并自蒸馏的离散模型相当；在有条件生成（条件生成）任务上，TTCD 显著优于这些模型。此外，在数独求解任务上，TTCD 也取得了类似的性能提升。

关键要点

• TTCD 在连续空间中操作，从高斯噪声到最终 token 画布采用确定性映射，无需任何额外采样步骤，从根本上避免了离散空间并行采样带来的不准确性。
• 首次在扩散语言模型中引入逐 token 时间概念，允许不同 token 以不同速度从噪声演化到最终 token，突破了传统统一时间表的限制。
• 逐 token 时间带来的优势：更好地建模条件生成、允许更确定的 token 快速通过、在精炼阶段实现 token 间差异化影响。
• 模型参数量 160M，在 OpenWebText 上训练后自蒸馏，高加速比下无条件生成质量与同类离散模型相当，有条件生成表现更优。
• 在数独求解任务上同样观察到一致的性能提升，说明 TTCD 的泛化能力不仅限于自然语言。

意义与影响

TTCD 为扩散语言模型提供了一种新的设计范式：结合连续空间和逐 token 自适应时间表，在提升生成速度的同时保持了高质量，尤其在高加速比场景下优势明显。这一工作有望推动扩散模型在自然语言生成（如文本完成、条件故事生成）以及结构化生成（如数独、代码等）中的应用。逐 token 时间概念可以启发未来研究：模型可以动态决定哪些 token 需要更多 refinement，从而更高效地分配计算资源。此外，确定性映射避免了采样噪声，使生成过程更可控，便于与其他模块（如强化学习、规划）结合。尽管目前只在 160M 参数规模验证，但该方向为更大规模、更高效率的语言扩散模型开辟了道路。