跳转至

Continuous Diffusion Scales Competitively with Discrete Diffusion for Language

会议: NeurIPS 2026
arXiv: 2605.18530
代码: https://github.com/NVlabs/RePlaid
领域: LLM 其他
关键词: 连续扩散模型、离散扩散模型、扩展律、变分推断、语言生成

一句话总结

本文重新审视基于变分似然(ELBO)的连续扩散语言模型,通过对齐现代扩散架构与统一步骤评估,提出了 RePlaid,打破了连续扩散在语言建模上扩展性差的固有认知,证明其扩展效率与离散扩散不相上下。

研究背景与动机

扩散模型在语言生成领域近年来展现出巨大的发展潜力,逐步在生成质量与扩展性上逼近自回归(AR)模型。然而,当前语言扩散研究几乎被离散扩散(Discrete DLM,如基于掩码或均匀分布转移的 MDLM、Duo 等)主导。相比之下,连续扩散语言模型(Continuous DLM)虽然在理论上具备极大的吸引力——其平滑连续的隐空间天然支持结构化引导与编辑,且能够直接无缝复用成熟高效的常微分方程(ODE)数值求解器以及单步/少步轨迹蒸馏算法,从而在有限采样步数下提供远优于离散扩散的效率-质量权衡,但在扩展性上一直饱受质疑。

早期的代表性工作(如 Plaid)经验性地报告,连续扩散模型为了达到与自回归模型相同的验证集似然,需要高达 64 倍的训练计算量(FLOPs)。这一巨大计算开销常被学界引证为连续扩散在语言建模扩展性上的根本缺陷。这引发了一个核心矛盾:连续扩散的扩展性劣势究竟是连续建模与离散文本天然脱节所导致的内在数学瓶颈,还是由于早期研究所采用的 Transformer 架构、超参数以及扩展律评估协议与现代离散扩散模型存在不公平的基准偏差所致?

本文重新审视了基于证据下界(ELBO)的连续扩散范式,构建了架构与现代离散扩散完全对齐的 RePlaid,并在标准计算预算协议下系统推导了其 IsoFLOP 扩展律。核心 idea:通过将连续扩散 Transformer 架构全面对齐现代离散扩散标准并保持似然下界(NELBO)端到端优化,连续扩散模型与自回归模型的计算开销差距可大幅收窄至 20 倍(与离散扩散相当),并在过度训练机制下全面反超最强离散基线。

方法详解

整体框架

RePlaid 继承了变分扩散模型(VDM)与 Plaid 的似然优化框架,但在网络架构、参数化以及数值计算路径上进行了深度重构,使其与现代离散扩散模型(如 MDLM、Duo)在架构层面完全对齐。输入离散文本序列首先映射为低维单位球面嵌入,并在连续隐空间中施加保方差高斯扩散过程;去噪骨干网络基于配备双向注意力和 AdaLN-Zero 调制的 Diffusion Transformer,预测离散分类分布并重参数化为预测嵌入,以端到端最小化负证据下界(NELBO)。

整体方法的流程结构如下图所示:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离散文本序列 x<br/>输入 token 序列"] --> B["低维嵌入与前向加噪<br/>映射至 unit sphere 施加 VP 高斯噪声"]
    B --> C["自条件循环增强<br/>25% 批次执行无梯度初筛生成 clean estimate"]
    C --> D["对齐版 DiT 骨干网络<br/>AdaLN-Zero 调制 + 双向自注意力"]
    D --> E["多任务 NELBO 目标优化<br/>先验损失 + 重建损失 + 连续扩散 MSE"]
    E --> F["方差极小化噪声调度<br/>凸显近线性交叉熵与平坦扩散损失"]

关键设计

1. 低维嵌入与前向加噪:降低词表投影计算开销并约束潜在几何 在连续扩散处理大词表(如 GPT-2 或 LLaMA 词表,大小 \(V \approx 32\text{k} \sim 50\text{k}\))时,若直接在高维 one-hot 空间进行加噪,将引入庞大的矩阵投影开销。RePlaid 采用行向量约束在单位欧氏球面上的低维可学习嵌入矩阵 \(\mathbf{E} \in \mathbb{R}^{V \times d_e}\)(默认 \(d_e = 16\))。给定长度为 \(L\) 的序列 \(\mathbf{x} \in [V]^L\),其连续表示为 \(\mathbf{e} = \mathbf{x}\mathbf{E} \in \mathbb{R}^{L \times d_e}\)。前向过程遵循保方差(Variance-Preserving, VP)高斯扰动: $\(q(\mathbf{z}_t \mid \mathbf{x}) = \mathcal{N}(\alpha_t \mathbf{e}, \sigma_t^2 \mathbf{I}), \quad \alpha_t^2 + \sigma_t^2 = 1, \quad t \in [0, 1]\)$ 其中信噪比定义为 \(\text{SNR}(t) := \alpha_t^2 / \sigma_t^2\)。将 \(V\) 维投影转移至低维 \(d_e\) 维空间,使序列输入到隐藏层维度 \(h\) 的投影 FLOPs 减少了近 50 倍,同时单位球面约束避免了嵌入尺度的任意发散。

2. 自条件循环增强:低成本强化多步生成一致性 为了加速去噪收敛并提升生成质量,RePlaid 深度集成了自条件(self-conditioning)机制。在训练过程中,针对每个批次中 25% 的样本,去噪网络先执行一次无梯度的前向传播,利用当前的带噪隐变量 \(\mathbf{z}_t\) 和时间步先验估计出干净词元概率 \(\hat{\mathbf{x}}^{\text{sc}} \leftarrow \mathbf{x}_\theta(\mathbf{z}_t, \gamma_t, \mathbf{0})\),并将其映射为估计嵌入 \(\hat{\mathbf{e}}_{\text{sc}} = \hat{\mathbf{x}}^{\text{sc}}\mathbf{E}\)。随后,该自条件嵌入与时间步条件、加噪隐变量一同输入网络进行第二次实际预测 \(\hat{\mathbf{x}} \leftarrow \mathbf{x}_\theta(\mathbf{z}_t, \gamma_t, \hat{\mathbf{e}}_{\text{sc}})\)。在采样推理阶段,所有时间步均强制开启自条件反馈,使网络能够显式利用上一步的高置信度重构信息,平抑多步扩散中的误差累积。

3. 对齐版 DiT 骨干网络:消除模型容量与数值计算偏差 传统 Plaid 在网络构建上使用了旧版 Transformer 和较特殊的输出层设计,掩盖了连续扩散的真实潜力。RePlaid 将去噪网络全面迁移至与现代离散扩散模型(如 MDLM)严格相同的 Diffusion Transformer(DiT)架构,包括 RoPE 旋转位置编码、LayerNorm 归一化、带偏置项的 MLP、GELU(tanh) 激活函数以及采用可学习门控机制的 AdaLN-Zero 残差调制。此外,RePlaid 排除了原有在 FP32 下结合输出先验对数密度的数值混淆因素,将最终 Logit 预测头统一置于 bfloat16 计算流中,确保与离散扩散 baseline 处于完全一致的计算与参数量基准线。

4. 多任务 NELBO 目标优化与分类重参数化:严格保持对数似然界 与近期引入启发式交叉熵(CE)或流匹配损失的连续扩散方法不同,RePlaid 坚持采用严格的负证据下界(NELBO)。通过分类重参数化 \(\mathbf{e}_\theta(\mathbf{z}_t, t) := \mathbf{x}_\theta(\mathbf{z}_t, t)\mathbf{E}\),NELBO 严密分解为先验匹配项、边界重建项和沿时间积分的连续扩散项: $\(\mathcal{L}_{\text{NELBO}}(\mathbf{x}) = \text{KL}(q(\mathbf{z}_1 \mid \mathbf{x}) \parallel p(\mathbf{z}_1)) + \mathbb{E}_{\mathbf{z}_0 \sim q(\mathbf{z}_0 \mid \mathbf{x})}\left[\sum_{l=1}^L -\log \langle \mathbf{x}_\theta^l(\mathbf{z}_0, 0), \mathbf{x}^l \rangle \right] - \frac{1}{2} \mathbb{E}_{t, \mathbf{z}_t}\left[\text{SNR}'(t) \|\mathbf{e}_\theta(\mathbf{z}_t, t) - \mathbf{e}\|^2\right]\)$ 在训练实现中,动态自适应划分批次以追踪重建损失与扩散损失的方差比例,保证各个项协同收敛,并规避了额外判别性损失对隐空间低秩几何结构的破坏。

5. 方差极小化噪声调度:自发诱导近线性信息衰减 RePlaid 将噪声调度函数参数化为单调递增的多层感知机网络,端点 \(\gamma_0, \gamma_1\) 直接优化扩散损失,而内部曲线 \(\tilde{\gamma}(t)\) 则通过特殊的梯度钩子专门优化扩散损失的蒙特卡洛方差 \(\text{Var}_t[\ell_{\theta, \gamma}(t)]\)。理论证明表明,当方差最优化时,瞬时扩散损失在所有时间步恒为常数 \(\kappa\);在贝叶斯最优去噪器假设下,互信息 \(I(\mathbf{e}; \mathbf{z}_t)\) 随时间精确线性衰减,使得每一步的解码交叉熵 \(\text{CE}(t)\) 呈现近线性变化,自然平衡了不同时步去噪的内在难度,无需任何启发式的人工时间重参数化。

损失函数 / 训练策略

总训练目标直接采用无偏且方差缩减的 NELBO。在训练期间,输入批次大小为 256(序列长度 2048),采用 AdamW 优化器,学习率针对骨干网络从 \(2 \times 10^{-4}\) 余弦衰减至 \(2 \times 10^{-5}\)(权重衰减 0.1),而针对低维嵌入矩阵和单调调度网络则采用专用的 \(1 \times 10^{-2}\) 初始学习率以加速几何结构的自组织学习。调度网络在内部计算中维持 float64 精度,确保单调性与梯度的数值稳定性。

实验关键数据

主实验

在语言建模评测中,遵循离散扩散标准基准,在 OpenWebText(OWT,序列长度 1024)和 One Billion Word(LM1B,序列长度 128)上评测困惑度(PPL 依据 NELBO 上界计算)。对于 RePlaid,评估了是否启用自条件(s.c.)的性能。

数据集 模型类别 模型名称 自条件 (s.c.) PPL (250K步) PPL (1M步)
OWT 自回归 (AR) Transformer - 17.8 17.5
OWT 离散扩散 (Mask) MDLM (low var.) - 25.0 23.1
OWT 离散扩散 (Uniform) Duo - 27.1 25.2
OWT 连续扩散 (Continuous) LangFlow ✗ - 38.4
OWT 连续扩散 (Continuous) LangFlow ✓ - 32.2
OWT 连续扩散 (Continuous) Plaid (原版) ✗ 27.5 25.7
OWT 连续扩散 (Continuous) Plaid (原版) ✓ 25.4 24.4
OWT 连续扩散 (Continuous) RePlaid (本文) ✗ 26.8 23.6
OWT 连续扩散 (Continuous) RePlaid (本文) ✓ 24.9 22.1
LM1B 离散扩散 (Mask) MDLM (low var.) - - 30.8
LM1B 离散扩散 (Uniform) Duo - - 33.7
LM1B 连续扩散 (Continuous) RePlaid (本文) ✓ - 31.6

消融实验

在 OpenWebText 上对 0.1B 规模的 RePlaid(s.c.,训练 1M 步)进行核心模块消融,验证各组件对最终 NELBO 困惑度上界的贡献:

配置 OWT 测试集 PPL 相对恶化量 (\(\Delta\)) 核心结论与机制分析
RePlaid (s.c.) 完整模型 22.1 - 整合所有设计,刷新连续扩散 SOTA 并逼近 AR
去掉输出先验 (w/o output prior) 22.5 +0.4 缺少闭式高斯对数密度正则,末层 logit 约束减弱
去掉自条件 (w/o self-conditioning) 23.6 +1.5 缺少步间 clean estimate 循环前馈,采样依赖单步直接预测
去掉可学习噪声调度 (w/o learnable schedule) 24.4 +2.3 无法自适应平坦化瞬时扩散损失,去噪时步难度失衡
冻结为随机词嵌入 (w/o learnable embeddings) 39.4 +17.3 隐空间失去低秩词性结构聚类,离散-连续映射严重退化

关键发现

  • 低维嵌入学习是决定性支柱:消融表明将嵌入冻结为随机分布会导致 PPL 从 22.1 灾难性暴跌至 39.4,这证明了连续隐空间中自发涌现的低秩语义流形对扩散至关重要,也解释了为何固定 one-hot 嵌入的流匹配方法在似然上难以突破。
  • 扩展律与过度训练优势:在 IsoFLOP 曲线拟合中,RePlaid (s.c.) 达到 AR 相当性能所需的计算开销倍数从历史认知的 64 倍降低至 20 倍,显著优于离散扩散的 Duo(22 倍)。更为重要的是,在模型参数受限且需要大量训练以压缩推理成本的过度训练(over-trained)机制下,RePlaid 全面反超离散扩散 MDLM(例如在 66M 参数下,RePlaid 仅需 3.1 倍最优计算即可达到 MDLM 需 6.9 倍最优计算才能实现的损失水平)。
  • ODE 求解器展现高阶优势:在无温度条件采样时,二阶 DPM-Solver++(2M) 在极少采样步数(\(T < 32\))下显著优于一阶 DDIM 与随机 DDPM 采样器;当采样步数达到 64 步以上时,DDPM 凭借随机性平抑确定性误差,在 MAUVE 评测指标上达到平台期峰值。

亮点与洞察

  • 严格遵循似然下界的优越性:论文揭示了纯交叉熵(CE)损失对隐空间嵌入具有过度的“判别性排斥力”,破坏了嵌入在 PCA 和 t-SNE 上呈现的低秩且符合词性(POS)的几何凝聚态;而基于 ELBO 的均方误差形式天然兼顾了词表重构与紧凑表示。
  • 方差极小化与线性信息衰减的内在等价:利用 I-MMSE 恒等式在理论上证明了,最小化 ELBO 估计方差得到的理想连续噪声调度,在贝叶斯最优去噪器下天然保证了互信息线性衰减与解码交叉熵线性增长,为扩散时间重参数化提供了严谨坚实的理论依据。
  • 打破离散 vs 连续的壁垒:通过将现代 DiT 架构引入连续扩散,澄清了过往学界对于“连续扩散在语言任务中难以扩展”的认知偏差,为后续结合 ODE 快速采样和流形蒸馏开辟了通用可扩展的道路。

局限与展望

  • 与自回归模型的固有计算差距:尽管 RePlaid 将计算差距从 64 倍缩小至 20 倍,但在达到同等验证集似然时,仍然需要比自回归模型高出量级的训练 FLOPs,在大规模数十亿(Billion)参数下的可扩展性仍需进一步实证。
  • 自条件在密集采样时的熵坍缩:自条件机制在极高采样步数(\(T \ge 512\))时,因误差累积和过高的确定性反馈,导致生成文本的 unigram 熵过低,多样性有所受损,亟需更鲁棒的正则化或退出机制。
  • 上下文感知嵌入扩展:当前模型基于静态查找表嵌入矩阵,未来将扩展到序列级上下文编码器(如结合轻量预训练 Encoder),有望进一步大幅缩小与 AR 模型的建模差距。

相关工作与启发

  • vs Plaid: 原版 Plaid 采用传统 Transformer 架构,未对齐现代离散扩散的优化规范,得出了 64 倍计算差距的悲观结论;RePlaid 全面对齐 DiT 模块、AdaLN-Zero 与现代超参数,并在 OWT 上将 PPL 从 24.4 刷新至 22.1。
  • vs MDLM / Duo: MDLM 和 Duo 是现代离散掩码/均匀扩散的标杆模型,MDLM 在基准训练量下保持较小算力差距(14倍),但 RePlaid 在过度训练场景和参数紧凑性上表现更优,以更小参数量(少约 1.8 倍)取得更优或相当的下界。
  • vs FLM / LangFlow: 近期连续扩散工作放弃了真正的 ELBO,转而使用启发式交叉熵或单步流匹配,容易诱发嵌入发散或依赖复杂的 ODE 似然近似;RePlaid 证明了坚持原生 NELBO 配合方差最小化噪声调度,不仅数学理论自洽,且在实际文本质量与 PPL 上均显著胜出。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 重新审视经典似然扩散框架,澄清了语言连续扩散的扩展性迷思,并给出了方差极小化与线性信息衰减的优美理论闭环。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖了严格的 IsoFLOP 扩展律推导、跨数据集 PPL 评估、多采样器轨迹对比及完整的消融实验。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层层递进,动机纯粹,理论分析与实证分析高度互为印证。
  • 价值: ⭐⭐⭐⭐⭐ 为非自回归语言模型中的连续扩散流派重塑了信心,证明了连续扩散具备完全不输离散扩散的规模扩展前景。