💬 NLP

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

本文介绍了潜在扩散语言模型（LDLM），该模型联合训练潜在编码器、扩散模型和解码器以构建高质量的连续潜在空间，相较于现有的离散和连续扩散模型，实现了更优越的文本生成性能和显著更快的推理速度。

原作者： Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin, Ilya Koziev, Alexander Korotin, Dmitry Vetrov

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin, Ilya Koziev, Alexander Korotin, Dmitry Vetrov

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个机器人写故事。旧的方法（称为“自回归”）就像一个学生从左到右一次写一个单词。如果他们在早期犯了错误，就无法回头修正，除非重写整个内容。而且，他们一次只能写一个单词，这非常缓慢。

本文提出的新方法称为LDLM，就像一位雕塑家在一块大理石上工作。雕塑家不是一次雕刻一个单词，而是从一块粗糙、充满噪声的石块（故事的“噪声”版本）开始，逐渐凿去噪声，显露出最终的形状。他们可以同时审视整块石块，并行地进行 refinement，瞬间修正故事中任何位置的错误。

然而，这里有一个陷阱：要有效地进行雕塑，你需要合适种类的大理石。如果石头太硬或太软，雕塑家就无法完成工作。在人工智能术语中，这种“大理石”被称为潜在空间——这是模型在将其转换回单词之前所处理的一种隐藏的、数学化的文本表示。

问题：“冻结”的蓝图

此前尝试这种“雕塑”方法时，使用的是预先制作好的大理石蓝图。他们采用了一个智能的、预训练的语言模型（就像一个已经知道单词如何组合的词典），并将其冻结在原地。他们试图利用这个固定的蓝图来雕塑文本。

本文的作者意识到，这就像试图用为另一种石材设计的蓝图来雕塑雕像。预训练模型并未针对“雕塑”过程（扩散）进行优化，因此结果显得笨拙、缓慢或质量低下。

解决方案：联合训练（“团队合作”）

作者构建了一个名为LDLM（潜在扩散语言模型）的新系统。他们不使用冻结的蓝图，而是创建了一个三人团队，从零开始共同学习：

编码器：一个将单词转换为“大理石”（潜在空间）的翻译器。
雕塑家（扩散模型）：负责去除噪声以 refine 故事的部分。
解码器：一个将 refine 后的大理石转换回可读单词的翻译器。

神奇之处在于这三者共同学习。随着雕塑家在去除噪声方面变得更好，它会告诉编码器：“嘿，我需要大理石被塑造成这种形状才能更好地工作。”编码器随后调整其形状以帮助雕塑家。这是一个反馈循环，整个团队不断进化，以完美地相互适配。

成功的“配方”

作者发现，仅仅将这三者放在一起并不能立即奏效；团队会互相争执，无法学习。他们必须遵循特定的“配方”才能使它们协同合作：

热身：在最初阶段，他们让编码器和解码器在没有雕塑家干扰的情况下练习单词翻译。这给了编码器在雕塑家开始尝试塑造它之前建立稳定基础的机会。这就像让乐队开始演奏之前，先让音乐家调试他们的乐器。
“噪声”技巧：在训练过程中，他们故意在解码器的输入中添加少量静态噪声。这迫使编码器变得稳健并高效地存储信息，而不是依赖完美但脆弱的信号。这就像让跑步者背着沉重的背包训练，这样当他们不背背包跑步时，会感到轻盈而快速。
智能时机：他们不会以恒定的速度训练模型。他们根据任务在那一刻的难易程度调整何时进行练习，确保模型在困难部分的学习效果与轻松部分一样好。

结果：更快、更智能

当他们在大型文本数据集（如新闻文章和书籍）上测试这个新团队时：

质量：生成的故事比以前的方法更连贯、更多样化。
速度：因为该模型是在“大理石”（隐藏的数学表示）上工作，而不是在每一步都试图从巨大的词典中挑选单个单词，所以它比竞争对手快 2 到 13 倍。
效率：它在生成高质量文本和保持文本多样性（不重复）之间取得了更好的平衡。

总结

简而言之，本文介绍了一种教 AI 写作的新方法，即让“翻译器”、“雕塑家”和“精修者”作为一个整体共同学习，而不是使用预先制作好的、冻结的工具。通过遵循特定的训练配方，他们创造了一个系统，能够比以前快得多地写出更好的文本。

技术摘要：如何联合潜在空间训练潜在扩散语言模型

问题陈述

自回归模型目前主导着文本生成，但受限于其从左到右的序列性质，无法并行生成文本或修正之前的错误。扩散语言模型通过并行迭代细化来生成文本，提供了一种替代方案。然而，现有方法面临重大障碍：

离散扩散：直接在词元（token）空间操作，但受限于因子化的联合分布，在每一步独立预测词元。这限制了同时生成多个词元的能力，使得高质量的多步生成变得困难。
连续扩散：在潜在空间操作，允许对所有位置进行渐进式细化。虽然优于离散方法，但当前的连续方法依赖冻结的预训练编码器（例如基于 Transformer 的编码器）来构建潜在空间。这些编码器是在代理任务（如下一词元预测）上训练的，并未针对扩散去噪过程的具体需求进行优化。反之，即使针对扩散进行了调整，浅层词元嵌入（one-hot 或固定向量）产生的潜在空间也是次优的。

核心挑战在于构建一个同时具备可解码性（能映射回词元）和扩散友好性（易于去噪）的潜在空间，同时避免将编码器训练与扩散训练分离所导致的不匹配问题。

方法论：LDLM

作者提出了潜在扩散语言模型（LDLM），这是一个编码器、扩散模型和解码器联合训练的框架。这使得扩散目标能够直接塑造潜在空间。

架构

两阶段编码：
- 冻结词元编码器：输入词元 $w$ 首先由冻结的预训练词元编码器（GPT-2）处理，以提取上下文隐藏状态 $h$ 。
- 可训练潜在编码器：一个可训练的编码器 $E_\theta^z$ 将这些隐藏状态 $h$ 映射到扩散潜在空间 $z_0$ 。这重塑了表示形式，使其适合扩散，同时保持可解码性。
两阶段解码：
- 潜在解码器：一个可训练的解码器 $D_\theta^h$ 将潜在变量 $z_0$ （在训练期间被噪声扰动）映射回隐藏状态 $\hat{h}$ 。
- 词元解码器：一个词元解码器 $D_\theta^w$ 将 $\hat{h}$ 转换为词元 $w$ 。
训练目标：
模型最小化组合损失： $L(\theta) = L_{diff} + L_h + L_w$ $L (θ) = L_{d i f f} + L_{h} + L_{w}$ 。
- $L_{diff}$ ：标准的连续扩散损失（MSE），用于从噪声观测中恢复干净的潜在变量。
- $L_h$ （隐藏状态 MSE）：从潜在空间重建冻结编码器的隐藏状态 $h$ 。关键在于，这里使用**均方误差（MSE）**而非交叉熵（CE）。
- $L_w$ （词元 CE）：从隐藏状态重建词元。梯度在重建的隐藏状态处停止，以防止影响潜在编码器。

关键训练配方（“四部分”解决方案）

作者指出，简单的联合训练会失败（产生低多样性的生成），并提出了一套具体的配方使其生效：

MSE 解码器损失：与以往使用交叉熵的工作不同，LDLM 使用 MSE 来重建隐藏状态。CE 迫使潜在变量充分分离以区分词元，这与扩散所需的平滑几何结构相冲突。MSE 要求较宽松，允许如果底层隐藏状态相似，潜在变量可以彼此接近，从而在保持平滑流形的同时防止潜在坍塌。
扩散到编码器的预热：在训练开始时，重建损失（鼓励大的潜在幅度）和扩散损失（鼓励小的幅度以简化去噪）将编码器拉向相反的方向。作者实施了一种预热机制，其中扩散梯度由系数 $\gamma(s)$ 缩放，该系数从接近 0 开始增加到 1。这使得编码器能够先学习一个可解码的潜在空间，然后再由扩散目标对其进行塑造。
自适应时间步采样：由于潜在空间在训练过程中不断演变，固定的噪声调度是次优的。作者动态调整时间步采样分布，确保去噪损失随时间步线性增长，使所有时间步对不确定性减少的贡献均等。
解码器输入噪声：在训练期间，向潜在解码器输入注入高斯噪声（ $\sigma_{dec}\epsilon$ $σ_{d ec} ϵ$ ）。这有三个目的：
- 鼓励编码器高效地存储信息（防止未使用的维度破坏信号）。
- 充当数据增强，使解码器对扩散误差具有鲁棒性。
- 稳定潜在坐标的方差，改善扩散参数化。

结果

作者在OpenWebText (OWT) 和 LM1B 上评估了 LDLM，将其与最先进的离散（MDLM, Duo）和连续（CANDI, FLM）扩散基线进行比较。

质量 - 多样性权衡：LDLM 在生成困惑度（Gen. PPL）和熵之间实现了最佳的帕累托前沿。在可比熵水平下，其 Gen. PPL 显著优于基线，接近真实文本的质量。
潜在空间比较：消融研究表明，LDLM 联合学习的潜在变量优于：
- 浅层嵌入（从头训练或冻结）。
- 来自冻结编码器的上下文嵌入（TEncDM 风格）。
- 来自带有显式平滑目标的自编码器训练的潜在变量（COSMOS 风格）。
  这证实了潜在空间必须针对扩散目标进行专门调整。
计算效率：LDLM 在采样期间比竞争方法快2–13 倍。与在每个步骤都投影到词表的词元级扩散模型不同，LDLM 在整个生成过程中在低维潜在空间中操作，仅在最后一步投影到词元。

意义与主张

该论文主张，联合学习潜在空间是使潜在扩散在文本生成中具有竞争力的关键且此前未被解决的步骤。

克服简单的联合训练：这项工作表明，仅仅将编码器和扩散模型一起训练是不够的；需要特定的技术（MSE 损失、预热、自适应采样、解码器噪声）来稳定训练并实现高质量。
上下文潜在变量的优越性：结果表明，高层序列信息（由冻结的 GPT-2 编码器捕获）至关重要，但必须由一个专门针对扩散的可训练编码器进行重塑。
效率：通过避免每步的词表投影，LDLM 提供了显著的速度提升，解决了基于扩散的文本生成中的一个主要瓶颈。

作者得出结论，直接用去噪目标塑造潜在空间对于构建有效的潜在文本扩散模型至关重要，这既实现了最先进的生成质量，又保持了优越的推理速度。