跳转至

LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling

会议: NeurIPS 2026
arXiv: 2604.11748
OpenReview: https://openreview.net/forum?id=fUzRdnKBK4
代码: https://github.com/nealchen2003/LangFlow
领域: LLM 其他
关键词: 连续扩散模型、Flow Matching、语言建模、噪声调度、自条件化

一句话总结

LangFlow 通过将词嵌入空间连续扩散与 Bregman 散度流匹配建立理论关联,并结合基于对数信噪比的 Gumbel 噪声调度与自条件化机制,使连续扩散语言模型在困惑度与生成质量上首次全面比肩甚至超越离散扩散模型。

研究背景与动机

连续扩散模型在图像、视频以及物理控制等连续模态领域取得了巨大成功,具备灵活的潜在轨迹、确定性 ODE 采样以及成熟的模型蒸馏加速工具箱。然而在自然语言等离散范式中,连续扩散的发展却远落后于离散扩散模型。主流的离散扩散虽然通过连续时间马尔可夫链实现了与自回归模型相竞争的表现,但其舍弃了表达力丰富的连续潜在空间,难以进行灵活的轨迹编辑与极简少步生成加速;而单纯在概率单纯形(simplex)上构建连续扩散又受到词表高维稀疏空间的严重制约。

这促使研究者重新审视基于连续词嵌入空间的扩散语言模型。然而,早期嵌入空间扩散模型长期受困于缺乏严密的理论基础:其训练目标大多依赖启发式设计(如直接用均方误差 MSE 回归嵌入向量,导致词向量空间严重坍缩)或实现复杂的动态分块损失;在评测层面,该领域始终缺失严格基于确定性概率流常微分方程(ODE)的负对数似然(NLL)/ 困惑度(PPL)评估上界,仅能退化使用随机 SDE 变分下界;评估基准与严谨目标的缺失反过来又使连续扩散在语言模态下的噪声调度与训练技巧缺乏清晰的优化方向。

针对上述瓶颈,本文的核心切入点是将嵌入空间扩散与 Flow Matching 框架深度对齐,证明在词表离散空间预测下交叉熵损失在数学上严格等价于基于 Bregman 散度的连续后验匹配。核心 idea:将词嵌入扩散重塑为 Bregman 散度流匹配,推导首个严谨的 ODE 评估困惑度上界,并通过基于信息均匀原则的 Gumbel 噪声调度与自条件化重构训练流程,使连续扩散在语言建模上首次匹敌离散扩散。

方法详解

整体框架

LangFlow 构建于词嵌入空间的概率流常微分方程(PF-ODE)之上。在训练阶段,输入的离散文本序列首先映射为连续的词嵌入向量,随后根据对数信噪比加噪,网络以加噪嵌入为输入预测 clean token 的后验类别概率,并以交叉熵损失驱动监督学习;在采样阶段,模型从高斯先验噪声出发,沿时间反转 ODE 连续数值积分,每一步通过网络预测的离散分布加权生成连续去噪方向,并在积分终止时通过 argmax 解码还原离散文本。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离散文本序列 x<br/>长度 L, 词表大小 V"] --> B["Bregman 散度流匹配<br/>词嵌入映射 z 与 gamma 路径加噪"]
    B --> C["信息均匀 Gumbel 噪声调度<br/>根据熵衰减率拟合最佳噪声分布"]
    C --> D["非对称自条件化机制<br/>注入前一步预测概率嵌入 z_SC"]
    D --> E["概率流 ODE 确定性求解<br/>数值积分至目标噪声并 argmax 解码"]
    E --> F["输出高质量生成文本"]

关键设计

1. Bregman 散度流匹配与交叉熵理论对齐:消除嵌入坍缩并实现端到端训练

此前连续语言扩散常沿用图像连续扩散的 MSE 损失来直接回归连续去噪嵌入 \(\hat{\bm{z}}\)。然而,MSE 损失的梯度会同时将所有其他 token 的嵌入向目标 token 牵引,导致词嵌入表在几何上聚集坍缩(最近邻距离趋近于 0),严重损害模型的泛化与表达能力。LangFlow 从凸函数 \(f\) 定义的 Bregman 散度出发: $\(\mathcal{D}_{f}(\bm{p}, \bm{q}) = f(\bm{p}) - f(\bm{q}) - \nabla f(\bm{q}) \cdot (\bm{p} - \bm{q})\)$ 通过设定 \(f(\bm{p}) = \bm{p} \cdot \log \bm{p}\),当目标分布 \(\bm{p} = \bm{1}_{x}\) 为 one-hot 标签时,Bregman 散度严格退化为单 token 的负对数似然。因此,直接监督预测分布 \(\hat{\bm{x}}_{\bm{\theta}}(\bm{z}_{\gamma}, \gamma)\) 的交叉熵损失在数学上直接等价于在连续隐变量 \(\bm{z}_{\gamma}\) 下最小化与真实后验的 Bregman 散度,从而既保留了分类监督对嵌入几何结构的保护,又自然给出了可微的连续速度场: $\(\hat{\bm{z}}_{\bm{\theta}}^{(i)}(\bm{z}_{\gamma}, \gamma) = \sum_{k=1}^{V} \hat{x}_{\bm{\theta}}^{(i, k)}(\bm{z}_{\gamma}, \gamma) \bm{e}_{k} = \mathbf{E}^{\top} \hat{\bm{x}}_{\bm{\theta}}^{(i)}(\bm{z}_{\gamma}, \gamma)\)$

2. 严谨的 ODE 困惑度(PPL)评估上界:建立连续流模型理论基准

传统连续语言模型因依赖随机 SDE 难以计算准确且确定性的负对数似然,且随机性天然阻碍了一致性模型(Consistency Models)等少步流蒸馏技术的应用。LangFlow 沿保方差概率流 ODE(PF-ODE)推进,通过瞬时变量替换定理(Instantaneous Change of Variables)显式追踪连续状态密度变化,推导出离散序列 \(\bm{x}\) 的严谨对数似然变分下界: $\(\log p(\bm{x}) \geq \mathbb{E}\left[ \frac{LD}{2} - \frac{\|\bm{z}_{b}\|^{2}}{2\sigma_{b}^{2}} + \sum_{i=1}^{L} \log \hat{\bm{x}}_{\bm{\theta}}^{(i, x^{(i)})}(\bm{z}_{a}, a) - \int_{a}^{b} \frac{\alpha_{\gamma}}{2} \nabla \cdot \hat{\bm{z}}_{\bm{\theta}}(\bm{z}_{\gamma}, \gamma) \mathrm{d}\gamma \right]\)$ 该界由纯高斯先验对数密度项、终止步的离散解码交叉熵项以及沿 ODE 反向轨迹的速度场散度积分项共同构成,散度项可通过 Hutchinson 迹估计高效计算。这使得连续扩散模型首次拥有了能与离散扩散的 ELBO 及自回归困惑度直接对标的评估手段。

3. 信息均匀原则与可学习 Gumbel 噪声调度:破除图像均匀调度的低效瓶颈

图像生成中普遍采用的时间 \(t \in [0, 1]\) 均匀噪声调度在语言模态下存在严重失效。实验分析表明,由于离散 token 嵌入在几何上具有孤立性,当 \(t > 0.2\)(即中低噪声区间)时交叉熵损失几乎迅速逼近 0,造成 80% 的采样和训练步数被浪费在几乎没有信息增益的区域。LangFlow 引入基于对数信噪比(logNSR)的 \(\gamma = \log(\sigma_{t}^{2}/\alpha_{t}^{2})\) 参数化,将高噪区域线性放大;进一步提出“信息均匀原则”,即让采样的每一步在后验熵 \(H_{\gamma}\) 的衰减率 \(H_{\gamma}' = \mathrm{d}H_{\gamma}/\mathrm{d}\gamma\) 上获得相等的熵减量。经实证,该信息增益曲线呈现右偏单峰分布,天然契合 Gumbel 分布形式: $\(H_{\gamma} = H_{+\infty} \cdot \exp\left( -\exp\left( -\frac{\gamma - P_{\mu}}{P_{\beta}} \right) \right)\)$ LangFlow 将位置参数 \(P_{\mu}\) 与尺度参数 \(P_{\beta}\) 设为可学习参数并由调度损失 \(\mathcal{L}_{\text{Scheduler}} = (\mathrm{stopgrad}(\ell_{\text{CE}}) - H_{\gamma})^{2}\) 动态拟合,将离散生成的生成困惑度(Gen. PPL)从约 1000 骤降至 154.2。

4. 连续流中的非对称自条件化机制:同时驱动困惑度与采样质量

自条件化(Self-Conditioning, SC)将前一步预测的去噪嵌入反馈至当前网络输入中。在离散掩蔽扩散(如 MDLM)中,启用自条件化虽略微改善生成多样性,却会损害模型自身的对数似然(PPL 上升),导致离散文献中普遍弃用自条件化评测。而在连续扩散的向量空间中,自条件化起到了显著的语义纠错与漂移阻断作用。分析表明,随着噪声增加,模型的后验预测倾向于从正确语义退化至高频虚词;自条件化能够锚定长程语义一致性,在训练时以 \(p=0.25\) 概率前向二次预测,使得验证集困惑度从 49.0 剧降至 30.0,同时生成困惑度从 154.2 优化至 81.5,展现出与离散扩散截然相反的协同正向增益。

损失函数 / 训练策略

模型整体训练通过联合优化交叉熵损失与调度器损失进行: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{CE}}(\bm{\theta}) + \mathcal{L}_{\text{Scheduler}}(P_{\mu}, P_{\beta}, H_{+\infty})\)$ 其中主网络采用与离散扩散对齐的标准 12 层 Diffusion Transformer (DiT,参数量 130M),配备 RoPE 旋转位置编码与自适应层归一化(AdaLN)。在 LM1B(上下文 128)和 OpenWebText(上下文 1024)上以 batch size 512 训练 1M 步,优化器使用 AdamW,最大学习率 \(3\times 10^{-4}\),在前 5000 步线性引入嵌入重构偏置项 \(r\)。

实验关键数据

主实验

在标准语言建模基准 LM1B 和 OpenWebText (OWT) 上,将 LangFlow 与自回归模型、经典吸收态离散扩散、均匀态离散扩散以及以往连续扩散代表方法进行全面评测。

模型类别 模型名称 LM1B Gen. PPL ↓ LM1B PPL ↓ OWT Gen. PPL ↓ OWT PPL ↓
自回归基准 AR Transformer 66.7 22.8 35.9 17.5
离散吸收态扩散 D3PM Absorb – 76.9 – –
离散吸收态扩散 SEDD Absorb 115.9 32.0 – 24.1
离散吸收态扩散 MDLM 103.9 31.0 104.9 23.2
离散均匀态扩散 SEDD Uniform – 40.3 103.6 29.7
离散均匀态扩散 UDLM 99.8 33.8 – 27.4
离散混合态扩散 Duo 97.6 33.6 77.6 25.2
连续扩散 Diffusion-LM – 118.6 – –
连续扩散 Plaid 77.3 32.4 – –
连续扩散 FLM (1024步) 96.9 – 62.2 –
连续扩散 (本文) LangFlow (Ours) 92.2 30.0 36.5 24.6

在跨领域零样本迁移泛化(Zero-Shot Transfer)实验中,在 OWT 上预训练的 LangFlow 表现优异:

模型类别 模型名称 PTB ↓ Wikitext ↓ LM1B ↓ Lambada ↓ AG News ↓ Pubmed ↓ Arxiv ↓
自回归模型 Transformer 82.05 25.75 51.25 51.28 52.09 49.01 41.73
离散扩散 SEDD Absorb 100.09 34.28 68.20 49.86 62.09 44.53 38.48
离散扩散 MDLM 95.26 32.83 67.01 47.52 61.15 41.89 37.37
离散扩散 Duo 89.35 33.57 73.86 49.78 67.81 44.48 40.39
连续扩散 (本文) LangFlow (Ours) 81.20 32.28 68.21 46.93 69.41 46.74 38.47

消融实验

自条件化(Self-Conditioning, SC)在离散扩散 MDLM 与连续扩散 LangFlow 上的消融对比(LM1B 数据集):

模型配置 Gen. PPL ↓ Δ (Gen. PPL) PPL ↓ Δ (PPL) 说明
MDLM (无 SC) 103.9 – 31.0 – 离散基线默认配置
MDLM (+ 自条件化) 94.9 -9.0 32.7 +1.7 采样改善但似然恶化
LangFlow (无 SC) 154.2 – 49.0 – 连续扩散初始状态
LangFlow (+ 自条件化) 81.5 -72.7 30.0 -19.0 采样与似然实现双重飞跃

关键发现

  • 连续与离散的范式反转:在自条件化机制下,离散掩码扩散(MDLM)会产生似然损失(PPL 恶化 1.7),而连续扩散 LangFlow 获得了大幅正向收益(PPL 从 49.0 下降 19.0 点至 30.0,Gen. PPL 减半),一举抹平了连续与离散扩散的性能鸿沟。
  • 高质量生成能力:在大型语料库 OpenWebText 上,LangFlow 取得了 36.5 的生成困惑度(Gen. PPL),大幅领先现有所有离散扩散模型(如 Duo 的 77.6 与 MDLM 的 104.9),逼近自回归模型的 35.9。
  • 跨域零样本迁移突破:LangFlow 在 PTB (81.20)、Wikitext (32.28) 和 Lambada (46.93) 上均创下了扩散模型的最优纪录,在 7 个基准中有 4 个超越了自回归 Transformer,证明了连续潜在空间表征在泛化迁移中的优越性。

亮点与洞察

  • 理论层面的桥梁搭建:巧妙利用 Bregman 散度与负对数似然的一致性,将离散交叉熵与连续 Flow Matching 严密统一,消除了困扰连续扩散已久的嵌入塌缩问题。
  • 模态差异驱动的噪声重构:指出图像与语言在连续扩散上的本质差异在于几何离散性带来的熵衰减集中性,提出信息均匀原则与 Gumbel 分布自适应调度,摆脱了照搬视觉经验的误区。
  • 开辟少步确定性蒸馏前景:基于确定性 PF-ODE 建立完整的生成与似然边界体系,规避了 SDE 连续随机注入所破坏的确定性双射,为未来将一致性模型蒸馏扩展至高效少步语言生成奠定了理论基石。

局限与展望

  • 采样熵偏低现象:LangFlow 生成文本的词分布熵略低于部分离散扩散基准(例如 OWT 采样熵 5.25 对比数据的 5.44),深入分析表明这源于模型对高频主题词全局复现率略高,而非局部重复退化。
  • 更大规模验证空缺:当前验证基于 130M 参数级别的模型规模,在数十亿参数(如 7B/14B)上的可扩展性及在预训练长文本推理任务上的能力仍待进一步验证。
  • 少步加速潜能未完全释放:虽然建立了 ODE 理论,但在更极端的少步推理(如 4~8 步生成)下,未经蒸馏的直接数值求解性能仍有下滑空间,需跟进流蒸馏算法。

相关工作与启发

  • vs Plaid: Plaid 采用 SDE 与 MSE 目标训练,导致嵌入空间坍缩且需要复杂的动态批次采样;LangFlow 采用 Bregman 散度交叉熵目标与确定性 ODE 评估,在保证几何不坍缩的前提下将 OWT 困惑度大幅推进到 SOTA 水平。
  • vs MDLM / Duo: MDLM 与 Duo 基于离散吸收/掩蔽状态建模,虽然在困惑度上有优势,但受限于离散状态转移跳跃而无法拥有平滑可微的轨迹编辑空间;LangFlow 证明了纯连续嵌入流同样能够匹敌甚至超越其困惑度与生成质量。
  • vs FLM: FLM 作为同期的连续流模型探索,并未推导 ODE 下的似然(PPL)评估上界;LangFlow 给出了严格的 NLL 变分上界公式,填补了连续流语言建模理论评测的空白。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将 Bregman 散度与 Flow Matching 结合统一连续语言建模,提出了创新的信息均匀 Gumbel 调度与 ODE 似然边界。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 LM1B、OWT 及 7 个下游零样本迁移数据集,消融全面且与主流离散/连续扩散基线进行了严格的同架构重训对比。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从动机分析到模态几何差异的洞察极其深刻,图表严谨。
  • 价值: ⭐⭐⭐⭐⭐ 为连续扩散语言模型树立了全新的坚实基线,有力打破了离散扩散在非自回归语言建模中的单一垄断局面。