跳转至

Token Time Continuous Diffusion for Language Modeling

会议: NeurIPS 2026
arXiv: 2607.14106
OpenReview: https://openreview.net/forum?id=HSVpmk76UW
代码: https://github.com/microsoft/ttcd
领域: LLM 其他
关键词: 连续扩散语言模型, Token时间, 流匹配, 自蒸馏, 快速文本生成

一句话总结

提出 Token Time Continuous Diffusion (TTCD),通过在连续空间流匹配中引入每 token 独立时间与基于初始不确定度的退噪排序,避免了离散扩散的多 token 并行独立采样误差,在极少步数生成下显著超越离散扩散与现有连续扩散基线。

研究背景与动机

当前扩散语言模型(Diffusion LMs)的主流范式是在离散 token 空间中进行迭代式细化,例如通过从掩码 token 解掩码(如 MDLM)或从均匀离散噪声逐步去噪恢复(如 Duo)。尽管这类离散扩散模型在标准多步采样评测中取得了与自回归模型相当的竞争力,但它们在面对高加速(高加速比即 2 至 4 步极少生成步数)场景时存在根本性的结构缺陷:当单个去噪步骤必须同时并行确定多个 token 时,模型只能从各位置的边际分布乘积中独立采样,而无法捕捉真实的底层联合分布,从而引发严重的独立因式分解误差(factorization problem),导致短步数生成文本质量崩溃或严重重复。

为了避开离散并行的因式分解陷阱,连续空间扩散(如将 token 连续嵌入与高斯噪声混合并通过确定性流映射演化)成为重要的探索方向。然而,已有连续扩散语言模型通常为整个序列强行施加单一的全局时间(即全序列共享相同噪声水平)。这引发了两个不可调和的核心矛盾:其一,自然语言序列各位置的预测难度天然极不平衡,部分词汇极易确定而部分词汇极具歧义,但单一全局时间迫使所有 token 经历一致的“从噪声到数据”突变期(transition zone),阻止了易预测 token 提早确立以辅助难预测 token;其二,在提示词条件生成(如问答或补全)场景下,共享全局时间无法区分已知清晰提示词(prompt)与待生成的纯噪声画布(canvas),严重损害了条件建模能力。

本文的切入角度是打破全序列共享同一噪声时钟的传统假设,在连续流匹配框架中显式赋予每个 token 独立的演化时间戳,并通过初始熵感知自动排列各位置的退噪优先级。核心 idea:在连续流匹配中引入基于 Beta 分位数映射的每 token 独立时间(token-time),使确定性高的 token 以更快速率从噪声演变为数据,并在条件生成中将提示词时间常驻为 1,彻底解耦 token 间渐进演化轨迹并消除离散采样因式分解误差。

方法详解

整体框架

TTCD 的核心思想是建立一个在连续嵌入向量空间中演化的确定性扩散生成过程。模型首先从标准高斯噪声分布采样初始连续向量,通过扩散变换器(Diffusion Transformer, DiT)骨干网络逐步沿学习到的速度场退噪,最终在最后一步通过确定性映射(如 argmax 预测或最近邻)映射回离散 token 词表。

在整个生成过程中,序列不再共享单一的时间标量,而是引入了一组局部 token 时间向量 \(\mathbf{t} = [t^1, \dots, t^L]\)。当全局名义时间 \(t_g\) 从 0 匀速推进至 1 时,每个位置的 token 按照其被赋予的专属排序变量(rank)\(u^l \in (0, 1)\),遵循特定的单调分位数函数演化其局部时间 \(t^l = F(t_g, u^l)\)。对于条件生成,输入提示词位置的局部时间恒定置为 1(完全清晰数据),而生成画布位置则根据模型在全噪声状态下的第一步预测熵分配 rank,使高确定性位置在连续空间中更早走向真实数据嵌入,从而为低确定性位置提供更准确的上下文表征。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据:提示词嵌入 + 画布高斯噪声"] --> B["熵感知退噪排序<br/>首遍前向计算画布各位置预测熵并分配 rank"]
    B --> C["Beta 分位数时间调度<br/>各 token 依 rank 映射出异构局部时间 t^l"]
    C --> D["无参数自适应层归一化<br/>局部时间向量独立调制各 token 隐层表征"]
    D --> E["连续流匹配去噪步<br/>按 token 步长沿嵌入凸组合方向确定性演化"]
    E -->|自蒸馏分支| F["快捷自蒸馏学习<br/>多步连续轨迹对齐为 1-2 步极速生成"]
    E --> G["输出:确定性映射至离散词表"]
    F --> G

关键设计

1. 局部 Token 时间与 Beta 分位数时间调度:打破全局时钟同质化约束 传统流匹配假定所有 token 共享同一时间 \(t \in [0, 1]\),导致无法模拟 token 间难度异构与提示词前缀固定特征。TTCD 针对这一痛点,引入局部时间函数 \(t^l = F(t_g, u^l)\)。为了保证物理意义自洽,函数 \(F\) 必须严格满足四大原则:边界一致性(\(F(0, u)=0, F(1, u)=1\))、对全局时间单调递增、对 rank 变量 \(u\) 单调递增(高 rank 保证全程局部时间领先),以及全期望守恒(\(\mathbb{E}_u[F(t_g, u)] = t_g\),即局部时间的期望严格等于全局进度)。

论文证明了通过 Beta 分布的分位数函数可完美闭式解析地满足上述性质: $\(t^l = F(t_g, u^l) := Q_{t_g}(u^l)\)$ 其中 \(Q_{t_g}\) 为参数为 \(\alpha(t_g) = \frac{1}{1 - t_g}\) 与 \(\beta(t_g) = \frac{1}{t_g}\) 的 \(\mathrm{Beta}(\alpha, \beta)\) 分布的逆累积分布函数(分位数函数)。这一设计使得高 rank token 的时间曲线呈现凹函数特征,在全局时间很小时即快速接近 1,从而实现确定性高的 token 提早退噪。

2. 初始熵感知退噪排序:自适应确定先验生成次序 为了让模型决定哪些 token 应当更早去噪,TTCD 在推理之初利用一次前向传播动态确定排序。在初始状态(画布位置全为高斯噪声且局部时间为 0,提示词位置为清晰嵌入且局部时间为 1)下,模型输出各画布位置的 token 预测概率分布 \(p_\theta^l(\cdot \mid \mathbf{z}_{\mathrm{inp}}, \mathbf{t})\)。算法计算各位置预测分布的香农熵 \(\eta^l = H[p_\theta^l]\),并将从低到高排列的序号与 rank 空间单调映射对齐——预测熵越低(即确定性越高)的 token 被赋予越高的 rank \(u^l\)。一旦排序建立,在后续所有去噪步中保持固定,无需重复重排,使得高置信度 token 作为锚点率先具象化。

3. 无额外参数的 per-token 自适应层归一化:精确时间条件注入 在扩散变换器(DiT)中,传统的自适应层归一化(adaLN)仅接受一个标量全局时间向量嵌入并调制整句所有 token。TTCD 修改了 adaLN 的计算逻辑,将输入维度调整为逐 token 独立的时间向量:每个 token 位置利用其对应的局部时间 \(t^l\)(通过正弦位置编码投影)独立进行 scale、shift 和 gate 的仿射调制。该改造复用了原有 adaLN 变换参数,不引入任何新增可学习权重,却赋予了 DiT 在层内部精细区分“当前 token 是噪声还是清晰数据”的强条件感知能力。

4. 快捷自蒸馏学习:拉平高加速步数下的生成质量退化 针对少步数极速推理场景,TTCD 进一步结合了基于流模型的快捷模型(Shortcut Models)自蒸馏技术。在蒸馏阶段,模型被要求直接预测从起始全局时间 \(t_{\mathrm{start}}\) 到目标全局时间 \(t_{\mathrm{end}}\) 之间的平均流速。为保持训练稳定并防止大步长下各 token rank 信息模糊,TTCD 创新地在输入时间条件中级联拼接了起点、中点与终点的局部时间向量 \((\mathbf{t}_{\mathrm{start}}, \mathbf{t}_{\mathrm{mid}}, \mathbf{t}_{\mathrm{end}})\),以二步展开输出概率为监督目标,最小化一步跳跃输出与二步未截断输出之间的 KL 散度: $\(\mathcal{L}(\phi) = \mathrm{KL}\left[\mathrm{stopgrad}(\mathrm{target}) \mid\mid g_\phi(\cdot \mid \mathbf{z}, (\mathbf{t}_{\mathrm{start}}, \mathbf{t}_{\mathrm{mid}}, \mathbf{t}_{\mathrm{end}}))\right]\)$ 辅以在极小步长 \(\epsilon\) 下与真实 token 的交叉熵正则,实现了在 1 至 4 步去噪下保持极低困惑度且文本熵不崩溃。

损失函数 / 训练策略

在预训练阶段,给定真实文本序列 \(\mathbf{x} \sim \mathcal{D}\),初始高斯噪声 \(\mathbf{z}_0 \sim \mathcal{N}(0, \mathbf{I})\),随机采样的全局时间 \(t_g \sim \mathrm{Unif}(0, 1)\) 及各 token 独立 rank \(u^l \sim \mathrm{Unif}(0, 1)\)。连续空间插值向量为 \(z^l = t^l z_1^l + (1 - t^l) z_0^l\)(其中 \(z_1^l = e_\theta(x^l)\) 为经过 \(L_2\) 范数截断为 1 的真实词嵌入)。模型目标为最小化各位置对真实 token 的交叉熵损失: $\(\mathcal{L}(\theta) = -\sum_{l=1}^L \log p_\theta^l(x^l \mid \mathbf{z}, \mathbf{t})\)$ 在推理去噪时,单步速度场更新由下式计算并执行一阶欧拉演化: $\(z^l \leftarrow z^l + \frac{\Delta^l}{1 - t^l}\left(\sum_{x \in \mathcal{V}} e_\theta(x) \cdot p_\theta^l(x \mid \mathbf{z}, \mathbf{t}) - z^l\right)\)$ 其中 \(\Delta^l = F(t_g + \Delta, u^l) - F(t_g, u^l)\) 为该 token 在当前全局步长下的位移增量。

实验关键数据

主实验

论文在 OpenWebText 数据集上训练了 160M 参数规模的 DiT 模型,评估了无条件生成与前缀条件生成(Prefix-Conditioned Generation,给定 992 词前缀生成最后 32 词),主要指标包括 GPT2-Large 评估的生成困惑度(Gen PPL,越低越好)以及文本熵(Entropy,越高说明多样性好,避免退化重复)。下表展示了极少步数(4 步与 8 步)生成设置下的核心对比(均取最佳温度与采样参数):

方法类别 / 模型 4 步无条件 PPL (Ent) 8 步无条件 PPL (Ent) 4 步前缀条件 (32词) PPL (Ent) 8 步前缀条件 (32词) PPL (Ent)
离散掩码: MDLM 155.42 (4.38) 76.49 (4.57) 186.86 (3.14) 77.90 (3.17)
离散均匀: Duo 78.23 (4.06) 47.42 (4.52) 110.96 (3.09) 60.44 (3.16)
离散蒸馏: Duo w/ DCD 84.43 (4.46) 48.04 (4.92) 93.16 (3.05) 57.82 (3.11)
连续流模型: FLM 362.22 (5.30) 218.25 (5.46) 329.42 (2.17) 331.00 (2.22)
连续流蒸馏: FMLM 96.52 (5.29) 70.39 (5.24) 529.98 (2.51) 348.51 (2.41)
TTCD (本文, 基础版) 193.18 (5.07) 105.26 (5.17) 221.28 (3.09) 112.24 (3.16)
TTCD w/ Shortcut (本文, 蒸馏版) 138.38 (5.41) 108.52 (5.40) 274.45 (3.25) 193.91 (3.24)

注:离散模型 Duo 在少步数下通过压低文本熵(低于 4.5)获得较低 PPL,实质存在短语重复问题;而连续基线 FLM/FMLM 在前缀条件生成中发生严重的熵崩溃(Ent < 2.5),TTCD 在维持高多样性(熵维持在 3.25+ 与 5.4+)的前提下展现了出色的条件生成稳定性。

消融实验

为了严格验证 per-token 时间机制与熵排序设计的必要性,论文在算法推理密集的 9x9 数独(Sudoku)求解任务(已知约 25 格,补全其余 56 格)上进行了系统消融(评估不同步数下的全盘解出准确率 %):

配置 / 变体 2 步求解率 (%) 4 步求解率 (%) 8 步求解率 (%) 16 步求解率 (%) 说明
离散掩码 (随机解掩码) 0.62 2.73 8.79 20.09 缺乏启发式排序
离散掩码 (熵解掩码, MDLM) 11.65 68.29 92.90 97.30 2步下严重受限于因式分解独立采样
连续模型 + 序列级全局时间 (Sequence-time) 0.00 0.00 3.03 20.96 提示词与空位混同,连续退噪完全失效
连续模型 + 全局时间 + 误差时间规整 (w/ Warping) 0.00 0.16 4.69 15.04 单一时间规整无法挽救条件混同
连续模型 + Token时间 (随机 rank) 1.01 22.27 36.75 37.89 引入异构时间即带来显著增益
TTCD 完整模型 (Token时间 + 熵排 rank) 31.51 61.33 65.85 68.46 2步极速下达 31.5%,远超离散基线 11.6%

关键发现

  • 极速生成摆脱因式分解坍塌:在 2 步数独这类强约束逻辑任务中,离散掩码扩散在每步并行猜测大量格子时由于乘积边际独立性破坏了联合分布约束,成功率骤降至 11.65%;而 TTCD 借助连续流空间的平滑演化与局部时间差异,在 2 步下斩获 31.51% 的成功率,提升近 3 倍。
  • 天然契合前缀条件生成:传统的全局时间连续扩散(FLM/FMLM)无法固定已知输入前缀,导致条件生成出现灾难性的熵坍塌(文本熵低于 2.5,大量无意义字符重复);TTCD 通过将前缀时间硬编码为 1,完全杜绝了前缀退化。
  • 采样超参鲁棒性极佳:改变初始高斯噪声方差与 softmax 温度时,以往连续模型(FLM)的 PPL 和熵曲线剧烈震荡,而 TTCD 表现出极其平滑的 Pareto 最优前沿。

亮点与洞察

  • 将全局标量时间拓展为张量局部时间:巧妙利用 Beta 分布分位数函数严格证明并实现了单调性、边界与期望守恒,为连续扩散模型处理序列内部的非均匀信息密度提供了优美的数学工具。
  • 零新参数的条件架构适配:仅仅通过改变 DiT 的 adaLN 输入广播逻辑,将一维全局时间向量改为逐位置时间向量输入,完全无需引入新增参数即可实现细粒度条件调制。
  • 破解离散生成在少步数下的硬伤:揭示了离散扩散在极限加速(1-4 步)下必然遭遇的多变量独立采样误差,证明连续空间流匹配配合异构时间退噪是实现真正极速文本扩散生成的极具前景的路径。

局限与展望

  • 模型参数规模有待扩大:论文当前实验验证主要集中在 160M 级轻量模型,尚未在 1B 至 7B 级以上的大规模语言模型上与先进自回归多 token 预测器(MTP)展开直接比拼。
  • 长步数上限与离散启发式模型仍有微小差距:在 16 步以上充分采样的数独任务中,离散掩码模型凭借成熟的离散确定规则可达到 97% 准确率,而连续流匹配在映射回离散边界时累积的微小浮点误差使得最终上限稳定在 68% 左右。
  • 离散扩散的 token-time 泛化空间:作者指出将 token-time 的异构推进理念迁移至均匀离散加噪模型也是未来值得深挖的研究方向。

相关工作与启发

  • vs MDLM / Duo (离散扩散模型):离散模型直接在 token 空间执行 mask 或 uniform 替换,每步必须确定性地 commit 离散选择,少步数并行时因独立分解引发严重错误;TTCD 全程在连续嵌入空间平滑演化,直到最后一步才映射为离散 token,从物理机制上规避了该误差。
  • vs FLM / FMLM (全局时间连续流语言模型):FLM 采用全序列共享的单一时间与复杂的经验性时间规整(time warping);TTCD 指出时间规整无法解决不同词汇难度不均的问题,通过基于 Beta 分布的 token-time 和初始熵排序替代全局规整,在前缀条件生成与超参稳定性上大幅占优。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统提出连续扩散中的 per-token 异构时间概念与 Beta 分位数映射,数学性质自洽优雅。
  • 实验充分度: ⭐⭐⭐⭐ 涵盖数独逻辑推理、OpenWebText 无条件/条件语言建模以及分子生成任务,消融与超参分析详实。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,动机叙述清晰直击痛点,理论证明完备。
  • 价值: ⭐⭐⭐⭐ 为攻克扩散语言模型极速推理(少步数生成)提供了非常关键的方法学启示。