跳转至

Continuous Latent Diffusion Language Model

会议: NeurIPS 2026
arXiv: 2605.06548
论文: https://arxiv.org/abs/2605.06548
代码: https://hongcanguo.github.io/Cola-DLM/
领域: LLM 其他
关键词: 连续潜空间扩散语言模型, 分层信息解耦, 块因果 DiT, 流匹配先验, 非自回归生成

一句话总结

提出 Cola DLM,通过 Text VAE 与块因果 DiT 将文本生成严格解耦为连续潜空间全局语义先验输运与局部条件解码两阶段,摆脱严格自回归固化时序偏差并在 2B 规模下展现出超越 AR 与离散扩散的良好扩展潜力。

研究背景与动机

大语言模型在自回归(AR)范式下取得了巨大突破,其通过概率链式法则分解离散文本分布并进行逐词预测,支撑了现代大语言模型的预训练与推理范式。然而,这种从左到右的严格固化顺序带来了难以克服的固有瓶颈:推理过程强绑定串行依赖,导致生成延迟居高不下,同时也将模型的归纳偏置硬性锚定在单一序列顺序上,在非单调生成、局部编辑和全局结构规划等任务中表现出明显的僵化。

离散扩散语言模型(如 LLaDA)尝试通过掩码机制或转移核实现多步去噪非自回归生成,但其迭代状态依然局限在离散词表空间,面临中间掩码状态语义模糊、多步采样计算沉重且难以捕捉全局连续语义几何的问题;而现有的连续扩散方法大多直接在词向量或单纯形空间中进行观测恢复,本质上仍是带噪目标表征的逆向还原,缺乏统一的边缘对数似然视角与明确的分层隐变量概率支撑。

这两种路线与传统自回归的深层矛盾在于:离散表征难以平滑组织高阶语义流形,而词表对齐的扩散又混淆了“全局语义规划”与“局部文本实现”的不同信息层级。本文的切入角度是引入分层隐变量建模,将扩散过程的作用严格限定为连续潜空间先验输运,而非词级别观测恢复。核心 idea:通过分层信息解耦将文本生成拆分为连续潜空间全局语义建模与局部条件解码,利用基于块因果 DiT 的流匹配学习平滑语义先验,实现推理高效、表征可扩展且天然兼容多模态的连续隐扩散语言建模。

方法详解

整体框架

Cola DLM 的核心思想是将离散文本观测与高阶语义表征解耦。系统主要由两个核心模块协同运转:一是严格因果的 Text VAE,负责离散文本序列与紧凑连续潜变量 \(z_0\) 之间的双向双射映射,并约束局部语义平滑;二是块因果扩散 Transformer(Block-Causal DiT),在连续潜空间中通过流匹配(Flow Matching)执行从标准高斯噪声基分布 \(p_1(z_1) = \mathcal{N}(0, I)\) 到潜变量真实先验 \(p_\psi(z_0)\) 的向量场输运。整个生成流程涵盖 Prompt 前缀编码、基于块因果流匹配的先验隐变量逐块生成,以及冻结或协同的条件解码还原。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入离散文本序列 x"] --> B["阶段 1:因果 Text VAE 编码与解码平衡<br/>学习稳定文本-连续隐变量双向映射与语义平滑"]
    B --> C["阶段 2:块因果 DiT 隐空间先验输运<br/>跨块因果依赖 + 块内双向交互的流匹配向量场建模"]
    C --> D["阶段 3:分块生成与干净前缀重绘条件解码<br/>前缀锁定指导先验采样 + 条件因果解码输出文本"]
    D --> E["输出目标文本响应"]

关键设计

1. 因果 Text VAE 编码与解码平衡:建立高保真且具语义平滑性的连续潜空间 传统直接在离散词嵌入上做扩散容易导致隐空间坍缩或几何扭曲,而标准 VAE 在面对强自回归解码器时又极易发生后验坍缩。Cola DLM 设计了严格因果的 Text VAE 架构,不仅保留时序流式生成能力,还引入显式的掩码预测辅助目标 \(\mathcal{L}_{\text{mask}}\),强制编码器抽取的隐变量 \(z_0\) 必须承载足够的语义信息以恢复被掩码的文本,从而彻底阻断了解码器纯靠文本记忆走捷径的坍缩漏洞。同时,模型对潜空间的平滑度进行显式控制,推导出包含重构项、KL 散度与掩码惩罚的变分下界目标: $\(\mathcal{L}_{\text{VAE}} = -\mathbb{E}_{q_\phi(z_0 \mid x)} [\log p_\theta(x \mid z_0)] + \beta \, \mathrm{KL}(q_\phi(z_0 \mid x) \,\|\, p_{\text{base}}(z_0)) + \lambda_{\text{mask}} \mathcal{L}_{\text{mask}}\)$ 更深层的洞察在于,隐空间的信噪比 logSNR 深刻影响着后验几何分布的陡峭程度;通过将 VAE 的 logSNR 设为可学习或维持在适度平滑水平(如 1.5 到 4.5 之间),能使隐空间兼顾微观重构精度与宏观流形平滑性,为扩散模型提供高度鲁棒的去噪空间。

2. 块因果 DiT 隐空间先验输运:跨块因果与块内双向并行的流匹配先验建模 若直接在全序列潜变量上使用完全无因果双向注意力,虽然有利于全局规划,但无法支持自回归式流式推理并带来巨大的显存开销;而完全因果的逐 token 扩散又丧失了双向上下文表征优势。Cola DLM 提出了块因果(Block-Causal)注意力可见集约束:对第 \(b\) 个块 \(z_0^{(b)}\),模型能完全感知历史已知块的无梯度真值 \(\operatorname{sg}(z_0^{(<b)})\),并在当前含噪块 \(z_t^{(b)}\) 内部施加双向全连接注意力: $\(\mathcal{V}_b = \left\{ \operatorname{sg}(z_0^{(<b)}), \, z_t^{(b)} \right\}\)$ 在这一可见集结构下,Cola DLM 摒弃传统的离散去噪转移,使用条件流匹配(Conditional Flow Matching)直接回归最优速度场 \(v_\psi\)。更关键的是,在联合预训练第二阶段,通过引入参考编码器正则项 \(\mathrm{KL}(q_\phi \parallel q_{\phi_{\text{ref}}})\) 限制隐空间漂移,使 VAE 与 DiT 在统一时空下协同演化,达成全局语义压缩与流匹配先验拟合的最佳契合。

3. 分块生成与干净前缀重绘条件解码:克服边界失真的强条件导引推理 在条件文本生成中,首个生成块往往是一个兼具已知 Prompt 隐变量与未知待生成隐变量的“混合区域”。针对该区域如果使用类似图像扩散的局部重绘或简单填充,会导致噪声水平不匹配与注意力混乱。Cola DLM 采用“干净前缀重绘(Clean Condition Repaint)”策略:在整个去噪轨迹中,保持已知 Prompt 区域的潜变量绝对固定为干净的无噪隐表征 \(z^{\text{pre}}\),只对未知的待生成潜变量注入高斯噪声 \(\epsilon^{(b)} \sim \mathcal{N}(0, I)\) 并沿速度场逆向积分输运。每个块生成完毕后,立即固化并作为后续块的因果历史条件。最终,将连续潜空间轨迹直接注入因果解码器: $\(\hat{x}^{\text{res}} \sim p_\theta(x^{\text{res}} \mid z^{\text{pre}}, \hat{z}_0^{(1:B)})\)$ 该机制将序列生成的串行步数缩减了 \(1.6 \times \sim 2.0 \times\)(块大小为 16 时仅需 8~10 步去噪即可饱和),大幅缓解了扩散模型解码沉重的推理瓶颈。

损失函数 / 训练策略

Cola DLM 的训练采用由浅入深的两阶段策略: - 第一阶段(Text VAE 预训练):固定基分布优化目标式 (3.16),通过 \(\mathcal{L}_{\text{mask}}\) 与 KL 散度约束稳定编码器 \(q_\phi\) 与解码器 \(p_\theta\) 的重建映射,不急于拟合最终先验。 - 第二阶段(VAE 与 DiT 联合预训练):联合优化目标 \(\mathcal{L}_{\text{stage2}}\),结合 VAE 重建正则、块因果流匹配损失 \(\mathcal{L}_{\text{FM}}\) 以及锚定初始 VAE 的参考约束: $\(\mathcal{L}_{\text{stage2}} = \lambda_{\text{VAE}} \left( -\mathbb{E}_{q_\phi} [\log p_\theta] + \beta \, \mathbb{E}_{q_\phi} [\log q_\phi] + \lambda_{\text{mask}} \mathcal{L}_{\text{mask}} \right) + \lambda_{\text{fm}} \mathcal{L}_{\text{FM}} + \lambda_{\text{ref}} \mathbb{E}_{p_{\text{data}}} \left[ \mathrm{KL}(q_\phi(z_0 \mid x) \,\|\, q_{\phi_{\text{ref}}}(z_0 \mid x)) \right]\)$ 优化器采用 AdamW,学习率线性预热至 \(1.5 \times 10^{-4}\) 并按余弦退火衰减,训练噪声时间调度采用 Logit-Normal 分布(\(\text{loc}=1.0\))对齐潜变量语义信息密度。

实验关键数据

主实验

在严格对齐模型参数量(非 Embedding 参数约 1.8B,总参数量约 2.0B)、训练 Token 数及优化超参的条件下,Cola DLM 与基准自回归 LLaMA-2B 及离散扩散代表 LLaDA-2B 在 8 个基准上进行了少样本生成式评测。

模型类别 / 架构 LAMBADA (Cont.) MMLU (MC Gen.) SIQA (MC Gen.) SQuAD (Gen.) 任务平均 (Task Avg)
LLaMA-2B (严格对齐 AR) 36.8 10.4 22.1 18.2 19.5
LLaDA-2B (严格对齐离散扩散) 35.2 9.8 21.4 22.5 20.1
Cola DLM (本文连续隐扩散) 37.1 11.9 24.8 23.4 22.6

(注:数据采自论文严格控制 2000 EFLOPs 预算下的评测表现。多选题采用严苛的 Few-shot 生成式字符串匹配协议,因而全系绝对分数偏低,但横向相对可比性高度一致)

消融实验

1. 隐空间演化与优化机制消融(EFLOPs = 117 算力预算) 分析 Text VAE 在第二阶段是否更新、如何更新对下游生成任务平均性能的影响:

潜空间更新配置 LAMBADA MMLU SIQA 平均准确率 (Avg.)
Joint DiT \(\times 1\) (预训练初始化+联合演化) 34.6 10.1 21.6 22.10
Joint DiT \(\times 0.01\) (微小步长演化) 31.8 8.2 18.4 19.47
Fix VAE (冻结第一阶段 VAE) 30.4 7.7 18.4 18.83
Interval (交替冻结/微调) 28.5 7.1 16.9 17.50
All Scratch \(\times 1\) (无预训练直接冷启动) 20.9 5.4 7.6 11.30

2. 首块生成引导与填充策略消融 分析在首个包含 Prompt 提示与待生成内容的混合块中,不同引导机制的生成准确率:

首块引导策略 LAMBADA MMLU SIQA 平均得分 (Avg.)
Clean Condition Repaint (全程干净前缀锁定) 37.1 11.9 24.8 24.6
Right Padding (右侧填充对齐) 24.7 11.5 13.8 16.7
Left Padding (左侧填充对齐) 24.6 8.4 14.9 16.0
Partial Repaint (\(t=3, m=1.0\)) 7.0 7.6 13.3 9.3
Partial Repaint (\(t=1, m=1.0\)) 8.5 7.9 8.8 8.4

关键发现

  • 隐空间绝不能从零冷启动,但亦不可彻底冻结:All Scratch 冷启动导致隐变量严重坍缩为简单的单向离散漂移;而 Fix VAE 虽然早期稳定但很快遭遇性能饱和。以高质量预训练为基座并在第二阶段以相同学习率进行 Joint DiT 协同演化,能够随算力扩展激发出最强扩展潜力。
  • 似然估计 PPL 与生成质量严重脱节:直接通过连续正规化流估计的边缘负对数似然(PPL)常因为局部高维点密度锐度问题呈现天文数字(如 \(10^6\) 量级),但其生成的文本语义不仅完全正确甚至优于平滑隐空间;过度追求微观 PPL 会迫使模型退化为生成通用平庸的高频词(如 "in", "the", "went")。
  • 块尺寸与推理去噪步数存在极佳甜点:块尺寸设为 16 展现出比逐 token(块大小 1)和超大块(64、128)更强的性能;推理去噪步数在 10~16 步即达到性能饱和,以单块 16 个 token 计算,实现了接近 2 倍于自回归的有效加速。

亮点与洞察

  • 扩散本质是隐先验输运而非观测还原:Cola DLM 厘清了连续扩散在 NLP 中的理论定位,证明扩散路径应当用来在连续紧凑流形上搬运语义概率密度,把离散符号表层生成的繁重语法负担剥离给自回归解码器。
  • 揭示最优时间偏移与潜变量维度的系统性漂移规律:实验证实潜空间并非独立因子的集合,最优调度参数 \(\text{loc}\) 随潜变量维度 \(d\) 的增大单调增大(\(d=16 \to 1.0\), \(d=128 \to 2.3\)),为连续隐式语言模型的扩散信噪比校准提供了坚实的理论依据。
  • 天然打通多模态交互界面:通过将离散文本转变为连续潜向量,消除了文本与视觉等天然连续模态在表征形态上的根本鸿沟,使用共享的块因果 MMDiT 即可自然统一文本续写、以图生文与文生图任务。

局限与展望

  • 序列长度压缩面临非整除边界敏感问题:在探索 Patch Size 为 2 的序列级隐变量压缩时,奇数长度 Prompt(Mod1)在拼接边界由于语义位移导致准确率暴跌至接近零,尚未形成稳健的变长序列连续压缩机制。
  • 高维隐空间推断存在变分差距:由于流匹配先验评估高度依赖 Hutchinson 迹估计器,条件边缘似然的计算仍极为昂贵且方差较大,限制了传统基于困惑度的轻量化在线评测。
  • 多模态与更强潜变量架构的统一潜力待释放:目前基于 VAE 的文本隐空间仍存在局部信息损失,未来有望替换为表征自编码器(RAE)或漂移模型(Drifting Model),进一步推高大参数量下的扩展上限。

相关工作与启发

  • vs LLaDA / 离散掩码扩散:LLaDA 直接在离散 Token 空间进行掩码擦除与逐步还原,状态空间缺乏连续梯度流形;Cola DLM 将去噪搬到连续潜空间,通过高斯流匹配实现了更紧凑平滑的全局语义规划。
  • vs Plaid / 词向量对齐连续扩散:Plaid 等模型在无损词向量或单纯形上做连续扩散,本质是连续空间里的观测还原;Cola DLM 引入了分层隐变量概率模型,明确解耦了全局语义组织与局部词级解码。
  • vs 经典自回归 (Llama / Qwen):自回归单向绑定从左到右的时序马尔可夫链;Cola DLM 在块内引入非自回归双向并行,在生成灵活性与全局一致性上展现出独特的扩展斜率。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从分层隐变量解耦与流匹配先验输运的理论高度构建并跑通了 2B 规模的连续潜扩散语言模型。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大科学设问、8 个基准测试、多达 2000 EFLOPs 的扩展曲线消融及多模态拓展实验。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导完备自洽,从马尔可夫路径与率失真理论切入,逻辑链条严丝合缝。
  • 价值: ⭐⭐⭐⭐⭐ 为破除大模型自回归依赖提供了最具可行性与理论美感的连续扩散语言建模路线。