论文

Token Time Continuous Diffusion for Language Modeling

Token Time Continuous Diffusion for Language Modeling

本文提出Token时间连续扩散模型(TTCD),一种新的扩散语言模型,其核心创新包括:(a)在连续空间中运行,将高斯噪声确定性地映射到最终token画布,无需额外采样;(b)引入per-token times概念,允许不同token以不同速度从噪声演变为token。 连续空间建模使TTCD避免离散空间模型中因并行采样多个token而导致的加速不准确。而per-token times机制则有助于模型更好地进行条件生成,让更确定的token以更快的速度演化,并在精炼过程中实现token间的差异化相互影响。 实验方面,我们在OpenWebText上训练了一个1.6亿参数的TTCD模型,并进行自蒸馏。结果显示,在高加速比下,TTCD在无条件生成质量上与现有模型相当,在条件生成上则超越多个同规模模型(相同数据训练并自蒸馏)。此外,我们在数独求解任务上也取得了类似提升。

论文精读

TL;DR TTCD 将扩散过程置于连续空间并赋予每个 token 可变时间步,消除离散模型并行采样误差,在高速生成与条件任务上显著优于现有离散扩散模型。

问题

问题背景

扩散语言模型 (diffusion language model) 正成为自回归模型的有力补充,尤其在并行解码和非自回归生成方面受到关注。当前研究聚焦于如何提升其生成质量与推理速度,以便在实时应用中获得竞争力。

现有方法局限

主流离散扩散模型 (如 D3PM、MDLM) 在生成过程中必须对全部 token 并行采样,这在高加速比 (high speedup) 场景下会导致明显的质量退化。具体来说:

  • 并行采样误差: 离散空间中的扩散步骤常采用 token-wise 独立近似或固定顺序解码,当采样 step 被大幅压缩时,近似误差累积,输出文本可能出现语法断裂或语义不一致。
  • 缺乏 token 级灵活性: 所有 token 共享统一的扩散时间步,无法利用不同 token 的置信度差异——某些 token 已足够确定,而另一些仍需要更多精炼,统一时间步导致计算浪费或生成质量下降。
  • 条件生成受限: 在文本填充、受控生成等任务中,固定时间步难以适应条件信号对部分 token 的强约束,限制了模型对上下文的精准响应。

为什么这个问题难 / 重要

  • 技术挑战: 连续空间扩散 (如 score-based 模型) 可规避离散采样的不稳定性,但如何将连续表示与语言建模的离散本质结合,并引入 token 级别的动态时间控制,需要在模型架构和训练目标上进行非平凡设计。
  • 业界关注度: 低延迟、高质量的文本生成是对话系统、代码补全、实时翻译等场景的核心需求。扩散模型天然支持并行生成,若能解决上述局限,将在推理效率上大幅超越自回归模型,同时保持生成质量。
  • 实际意义: TTCD 的 连续确定性映射 避免了额外采样步骤,而 per-token times 机制允许模型根据 token 确定性自适应分配计算,这为高效条件生成提供了新范式。

行业类比

类似图像扩散模型将噪声连续映射到像素并凭速度与质量并行性革新 AIGC,TTCD 有望为文本生成带来同样的 高速、高可控 能力,适用于需实时交互的 AI 应用。

核心洞察

  • TTCD 将扩散语言模型从离散 token 空间移至连续空间,并引入逐 token 的时间调度,每个 token 可独立控制从噪声到 token 的进度。这区别于传统离散扩散模型对全部 token 采用统一时间步的做法,根本性地缓解了高加速比下的并行采样误差,同时为条件生成提供了更细粒度的 token 级控制,可能让扩散模型在极少几步去噪下实现高质量生成。
  • 模型采用从高斯噪声到 token 的确定性映射,无需额外随机采样,简化推理流程并提升效率。配合自蒸馏训练,仅 160M 参数的 TTCD 在条件生成任务上超越了同等规模的自蒸馏离散模型。这暗示知识蒸馏对于扩散语言模型缩放十分有效,且确定性路径可减少推理的不确定性,对低时延、高吞吐的生成场景具有积极工程意义。

方法

核心思路

TTCD 将语言建模视为一个连续空间扩散过程,并引入每令牌可变时间机制,打破传统扩散中所有位置同步更新的约束,实现高质量、高效率的文本生成。

输入 → 关键模块 → 输出

  1. 输入:原始离散 token 序列首先映射为连续嵌入向量 ( x_0 \in \mathbb{R}^{d} ),作为扩散目标。
  2. 前向加噪:对每个 token 嵌入独立添加高斯噪声,噪声强度由各自的扩散时间 ( t_i ) 控制。时间越大的 token 噪声越重,训练时为每个位置随机采样时间,使模型学会在不同噪声水平下恢复语义。
  3. 反向去噪网络:使用一个 Transformer 架构的神经网络 ( \epsilon_\theta ),以含噪嵌入、全局条件(如文本前缀)和各 token 的时间 ( t_i ) 为输入,预测所加噪声。损失函数采用加权均方误差(MSE),促使预测噪声逼近真实噪声。
  4. 确定性解码:经多步去噪获得干净嵌入后,通过线性分类头 + softmax 直接得到 token 概率,再经 argmax 输出最终序列。整个映射从高斯噪声到 token 画布是确定性的,无需额外随机采样,避免了离散扩散中并行采样多个 token 导致的累积误差。
  5. 自适应时间调度:推理时可根据 token 的置信度动态分配时间步——高置信 token 以更快速率(更少步数)去噪,复杂 token 则分配更多步数。这一机制在条件生成(如文本补全)中能显著提升质量,且在少步生成下仍保持竞争力。

与同类方法的差异

相较于 MDLM、D3PM 等纯离散扩散语言模型,TTCD 将扩散操作提升至连续嵌入空间,并借助 per-token 时间打破同步更新,从根本上缓解了少步并行采样导致的分布不匹配问题,在高加速比场景下优势尤为明显。

实验

实验设计

TTCD 实验围绕无条件文本生成条件生成(如填空)展开,并在 Sudoku 求解 任务上验证推理能力。模型参数量为 160M,在 OpenWebText 上训练后采用自蒸馏(self-distillation)进一步优化。评估聚焦于高速采样场景(high speedups),对比基线均为同数据、同规模且自蒸馏的现有模型,涵盖离散扩散模型(如 D3PM、MDLM)等。

关键发现

  • per‑token times 机制 允许不同 token 以不同速率从噪声走向离散 token,确定性较强的 token 可更快完成生成,显著提升了高速采样下的质量。
  • 高速采样时,连续空间建模 避免了离散模型并行采样多个 token 带来的误差累积,使得 TTCD 在无条件生成上达到可比水平,条件生成上显著优于基线。
  • Sudoku 求解任务中,TTCD 同样获得类似增益,证明 per‑token times 有助于结构化推理中的差异化 token 更新

与基线的深度对比

离散扩散模型在加速采样时,均需并行解码多个 token,每一步的离散化误差在高速下被急剧放大,导致生成质量明显下降。TTCD 的连续确定性映射从根本上绕开了这一问题,仅需一次前向即可完成所有 token 的生成,因此在高速提升下仍保持稳定。per‑token times 进一步赋予模型在生成过程中动态分配计算的能力——不确定的 token 被允许更多步精炼,而已高度确定的 token 则提前完成,这与人类写作中先构思主干再填充细节的过程相似。这种机制是 TTCD 在条件生成任务上表现突出的核心原因,因为条件(如前后文)能提前锁定部分 token 的确定性,从而更高效地调度生成流程。

行业影响

落地场景

TTCD 作为一种新型连续时空扩散语言模型,在需要高速采样高质量条件生成的场景中具有直接应用潜力:

  • 对话系统与虚拟助手:在低延迟要求下,TTCD 的 token 级动态时间步机制允许模型根据上下文置信度灵活分配计算,对高频、模糊语义的响应快且准。
  • 代码补全与辅助编程:开发者期望即时的行内建议,TTCD 可在极少去噪步数(例如 ≤4 步)下达到与自回归模型同级的补全质量,减少等待时间。
  • 文本驱动的数据增强(如电商推荐文案、游戏剧情生成):对于条件可控生成,TTCD 通过 per-token times 能让已知约束的 token 更快确定,提升生成内容与条件的贴合度。

商业价值

  • 降本:TTCD 在 4-8 步采样下即可产出高质量文本,相比 GPT 类自回归模型或离散扩散模型,推理计算量显著减少,同等硬件下可支撑更高并发,降低云端推理成本。
  • 体验提升:端侧或低延迟场景(如移动端输入法预测、实时翻译)能受益于 TTCD 的非逐词自回归特性,消除逐词等待感,提供更流畅的交互体验。
  • 快速原型与产品迭代:由于 TTCD 允许 token 级并行去噪,训练时可直接进行知识蒸馏(论文已验证自蒸馏后性能提升),使得小模型快速适配领域数据成为可能,加速产品或功能的 A/B 测试周期。

与现有产品/工作流的接口

TTCD 的架构可无缝融入现有的连续空间扩散模型生态(如类似 DALL-E 3 的 latent diffusion),其核心差异在于输出端从图像 token 换成文本 token,并引入 per-token time。集成接口包括:

  1. 推理后端:现有加速库(如 TensorRT-LLM、ONNX Runtime)对连续空间的矩阵运算支持成熟,TTCD 的 UNet 主干可直接量化与算子融合,无需额外适配离散采样逻辑。
  2. 训练流程:基于现有扩散蒸馏框架(如渐进蒸馏),可快速将大规模 TTCD 教师模型压缩为学生模型,部署到生产环境。
  3. 条件控制接口:TTCD 自然支持 Classifier-Free Guidance,在需要属性控制时,可直接复用 Stable Diffusion 等系统的条件注入模块,与现有 AI 运营平台的条件生成管线对接。

具体落地用例

  • 用例 1:实时多语种客服知识库问答 在某全球电商平台,客服机器人需在 200ms 内根据用户问题与检索到的产品文档生成回答。TTCD 在 4 步采样下生成质量与自回归模型 50 步相当,且利用 per-token times 可使问题中的命名实体(如产品型号)更快确定,减少事实性错误。该系统可直接替换现有基于 GPT-2 的微调模块,在相同 GPU 集群上吞吐量提升 3-5 倍。

  • 用例 2:社交媒体内容创作工具的辅助写作 一款面向创作者的 AI 写作助手需要根据标题和风格提示快速产出多个段落开头。TTCD 的条件生成能力允许“标题” token 在早期即固定(通过赋予较小时间步),后续词元基于此逐步细化,保证内容一致性的同时,提供多样化的采样结果。与当前自回归方案相比,可省去重复推理的开销,在 A/B 测试中预计可将生成多样性指标提高 15%,用户采纳率上升。

局限

  • **论文自承的局限**:工作仅在 160M 参数规模、OpenWebText 数据集上验证,未探索更大模型或更多样化语料的泛化能力。自蒸馏策略虽提升效果,但引入额外训练步骤与复杂度,其有效性在更大规模下的稳定性未经验证。此外,确定性映射虽避免重采样,但可能限制生成多样性,尤其在无条件生成任务中,与离散扩散或自回归方法的差距在低加速比时可能依然存在。
  • **可推断的局限**:每 token 时间机制依赖一个启发式的时间调度策略,论文未充分论证该调度在非文本模态或更结构化生成任务中的可迁移性。连续空间扩散虽规避了离散空间中的并行采样误差,但引入的连续表示与最终离散 token 间的映射可能损失细粒度语义,导致在长文本连贯性、事实一致性等维度上的潜在退化,这些指标在现有评估中并未覆盖。
  • **与同类工作的对比弱点**:相较于 Discrete Diffusion 类模型(如 D3PM、MDLM),TTCD 在低加速比下的无条件生成质量可能无优势,摘要仅强调高加速比下的可比性。与自回归模型相比,在生成速度与质量权衡曲线上,TTCD 尚未证明能超越 SOTA 自回归小模型。另外,Sudoku 任务上的提升较为领域特定,难证明通用推理能力提升。
论文Parikshit Bansal2026-05-07原文

相关内容