DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling¶

会议: NeurIPS 2026
arXiv: 2605.23605
代码: https://jmlemercier.github.io/diladiff/
领域: 模型压缩
关键词: 扩散语言模型, 连续隐空间, 掩码扩散, MeanFlow蒸馏, 极速采样

一句话总结¶

针对掩码离散扩散语言模型独立解码假设忽略Token间强相关性导致的多Token并行解码质量崩溃痛点，本文提出了结合语义正则化连续隐空间与MeanFlow轨迹自蒸馏的混合扩散框架DiLaDiff，仅需5步隐空间生成即达成200步教师模型性能，并在保持文本高质量的同时实现7倍推理加速。

研究背景与动机¶

扩散生成模型在连续空间（如图像、视频、音频）已取得巨大成功，具备基于常微分方程（ODE）求解器的严谨采样轨迹以及高效的模型蒸馏能力。受此启发，掩码扩散语言模型（MDLM）被提出用于离散文本生成，其通过逐步去掩码恢复Token序列。然而，离散扩散模型在实际应用中面临极其苛刻的采样质量与吞吐量折中：模型在反向解码时为了并行预测多个Token，不得不采用边际分布独立分解假设，严重忽略了自然语言Token之间高度耦合的上下文依赖关系。特别是在生成初期掩码率极高的阶段，由于可用的干净上下文极为匮乏，真实联合后验分布与独立因子分解后验之间的模型失配被极度放大，导致盲目增加每步并行解码Token数时文本连贯性迅速崩溃。

另一方面，单纯在连续表示空间（如One-hot向量、词嵌入或深度上下文表征）上构建连续扩散文本模型，虽能继承ODE求解器与蒸馏技术，却始终难以媲美离散模型的效果。这根源于连续文本表征的表达能力与可解码性之间的天然矛盾：浅层词嵌入易于通过最近邻反查解码，但无法建模长程语义依赖；而预训练模型深层隐状态虽具备强大语义上下文，但在生成时极难直接无损还原回离散Token序列。

针对离散扩散因独立假设无法高效并行解码、以及连续文本扩散难以兼顾上下文表达与可靠解码的核心矛盾，本文提出了将压缩连续隐空间先验与离散掩码解码器深度耦合的设计路径。核心 idea：利用自编码器构建具备语义平滑性的连续上下文隐空间以显式捕获全局Token相关性，使用基于tanh-logSNR调度的大步长连续扩散建模隐变量先验并以MeanFlow自蒸馏至极少步数，最后仅用极短步数的条件离散掩码扩散完成高质量文本解码。

方法详解¶

整体框架¶

DiLaDiff的整体生成架构分为三个紧密协同的模块：首先，通过微调预训练MDLM并结合BERT上下文特征，构建具备强语义平滑性与鲁棒性的文本自编码器（包含编码器 $\mathcal{E}_\phi$ 与解码器 $\mathbf{x}_\theta$）；其次，引入连续潜变量扩散模型（LaDiff，去噪器为 $\mathbf{z}_\psi$）学习编码器分布的先验，使离散解码器在全局语义潜变量 $\mathbf{z}$ 的条件引导下打破Token独立性限制；最后，使用基于平均速度（MeanFlow）的轨迹自蒸馏算法将连续扩散先验蒸馏为学生模型 $\mathbf{u}_\eta$，仅用5步ODE积分即可生成高质量隐变量。

推理时，系统先从标准高斯噪声出发，由蒸馏后的连续扩散模型快速求解概率流ODE生成潜变量 $\mathbf{z}$；接着将其反归一化并注入离散去噪解码器中，解码器从全掩码序列出发，通过少量步数的反向离散扩散（如 $N_{\text{disc}}=64$）并行恢复完整文本。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["高斯初始隐噪声 z_T ~ N(0, I)"] --> B["MeanFlow隐空间轨迹自蒸馏<br/>学生模型 u_η (N_cont=5 步快速积分)"]
    B --> C["语义连续隐变量 z<br/>坐标级反归一化 (μ_z, σ_z)"]
    D["全掩码离散序列 x_tN ~ [MASK]^L"] --> E["强正则化文本自编码器解码器<br/>零卷积门控交叉注意力注入 z 条件"]
    C --> E
    E --> F["潜在引导的离散掩码扩散采样<br/>q(x_{s}|x_{t}, z) 因子分解并行反向去掩码"]
    F --> G["最终生成文本序列 x_0"]

关键设计¶

1. 强正则化文本自编码器：解决隐空间过度发散与解码器注意力偏置

为使连续隐空间能够被后续生成扩散模型学习，必须打破自编码重建能力与空间平滑性之间的权衡冲突。如果单纯优化重建交叉熵，隐变量会被互相推离形成稀疏流形，导致扩散先验无法拟合。本文基于BERT深层隐特征构建交叉注意力编码器，并施加多重强正则化机制：对BERT特征和潜变量进行基于训练累计统计量的坐标级标准化，并在训练时以各50%概率交替进行特征掩码（掩码率 $p_{\text{mask}}^{\tilde{\mathbf{x}}}=0.7$）或添加高斯噪声（方差 $\sigma_{\text{reg}}^{\tilde{\mathbf{x}}}=0.5$），同时对潜变量施加高比例掩码与噪声扰动。

在解码器侧，为避免从头训练导致的脆弱性，模型直接复用预训练好的MDLM参数。为防止解码器在重构训练中产生注意力偏差而过度依赖潜变量、进而忽略输入掩码序列 $\mathbf{x}_t$ 自身的未掩码上下文，作者在MDLM自注意力层之间插入专门提取隐变量信息的交叉注意力层，并使用零初始化的逐点卷积层（ZeroConv）进行包裹： $$ \mathbf{h} \leftarrow \mathbf{h} + \text{ZeroConv}\left(\text{CrossAttention}\left(\text{ZeroConv}(\mathbf{h}); \mathbf{z}\right)\right) $$ 同时在训练中有75%的概率将隐变量完全替换为高斯随机噪声，强制解码器在保留原生MDLM去掩码能力的同时平滑接入全局隐变量指导。

2. 潜在引导的离散扩散反向分解：严谨消除Token独立解码误差

传统MDLM在推导反向转移后验 $q^\theta(x_{s}|x_{t})$ 时，强制假设去噪器后验在各个Token位置上独立可分，这在掩码率较高的早期阶段导致致命的模型失配。而在DiLaDiff中，引入包含整句全局上下文相关性的隐变量 $\mathbf{z}$ 后，条件反向后验严格满足条件独立性引理： $$ q^\theta_{s|t}(\mathbf{x}{s}|\mathbf{x}}, \mathbf{z}) = \prod_{\ell} q^\theta_{s|t}\left(\mathbf{x{s}^\ell \mid \mathbf{x}}^\ell, \mathbf{x{0}^\ell = \mathbf{x}\theta^\ell(\mathbf{x}_{t}, \mathbf{z})\right) $$ 通过对隐变量先验分布 $p(\mathbf{z})$ 进行边缘化积分，联合反向分布得以显式表达为受 $\mathbf{z}$ 调控的边际乘积。由于全局句意、主题与句法相关性已被隐变量 $\mathbf{z}$ 充分捕获，解码器在每一个扩散步即使同时解除数十甚至上百个Token的掩码，也不会产生词义断裂或局部冗余复读，从根本上释放了离散扩散并行大步长解码的潜力。

针对潜变量先验分布建模，作者采用tanh-logSNR方差保持噪声调度，通过将超参数设为 $d=10$，强制扩散模型将去噪容量集中分配在极高噪声区域，契合了离散文本映射到连续空间时边界陡峭、转移所需扰动尺度偏大的特有几何结构。

3. MeanFlow隐空间轨迹自蒸馏：实现极低步数ODE求解与速度对齐

虽然连续隐空间显著降低了离散解码步数的需求，但若使用标准ODE求解器（如200步连续扩散），连续阶段的反向传播开销在总推理时间中占比将超过110%。为消除连续扩散的计算负担，DiLaDiff利用连续空间的概率流ODE特性，引入MeanFlow自蒸馏技术。不同于传统一致性模型（Consistency Models）仅支持单步映射，MeanFlow旨在训练学生网络 $\mathbf{u}_\eta(\mathbf{z}_t, t, r)$ 预测时间区间 $[r, t]$ 上的平均位移速度： $$ \mathbf{u}(\mathbf{z}t, t, r) := \frac{1}{t-r} \int}^{t} \mathrm{d}\mathbf{z\tau $$ 当 $r \to t$ 时，平均速度收敛为教师模型的瞬时速度 $\mathbf{v}_\psi(\mathbf{z}_t, t)$。在几步采样的粗时间网格上，学生模型通过最小化与切线校正目标之间的均方误差实现轨迹压缩： $$ \mathcal{L}}}(\eta) = \mathbb{E}\left[|\mathbf{u\eta(\mathbf{z}t, t, r) - \text{stopgrad}(\mathbf{u})|}2^2\right], \quad \mathbf{u}}} = \mathbf{v\psi(\mathbf{z}t, t) - (t-r)\left(\mathbf{v}\psi(\mathbf{z}t, t)\partial}}\mathbf{u\eta + \partial_t \mathbf{u}\eta\right) $$ 其中雅可比向量积通过反向模式自动微分高效求解。为了在自蒸馏中稳定继承教师网络的自条件化（Self-Conditioning）能力，作者通过仿射映射 $\Phi(\cdot, t)$ 将瞬时速度反解为干净数据估计，在推理时仅需5步积分（对应10次前向评估NFEs）即可逼近200步教师的生成质量，使隐空间采样的耗时占比降至5%以下。

损失函数 / 训练策略¶

训练分为三阶段串行流水线： 1. 自编码器微调：从预训练MDLM权重初始化，冻结嵌入层，使用BERT作为特征抽取器。在16张H100上以Batch Size 512微调200k步，学习率5e-5，目标为带掩码输入的加权交叉熵损失（将传统ELBO权重设为常数 $-1$ 以降低方差）。 2. LaDiff连续扩散先验训练：冻结自编码器，使用DiT结构学习潜在扩散先验，在8张H100上以Batch Size 512训练500k步，学习率2e-4，采用 $d=10$ 的tanh-logSNR调度及50%概率的自条件化。 3. DiLaDiff自蒸馏：冻结LaDiff教师，在16张H100上以Batch Size 2048蒸馏25k步，学习率5e-5，加入25%概率纯流匹配损失（$r=t$），并采用损失归一化因子5与切线项10k步线性预热。

实验关键数据¶

主实验¶

在OpenWebText基准上（序列长度 $L=1024$，词表大小30522），以GPT2-Large评估生成困惑度（GenPPL）与MAUVE评分，并与官方基线MDLM与DUO进行多步与少步生成对比。

模型	连续步数 $N_{\text{cont}}$	离散步数 $N_{\text{disc}}$	GenPPL (↓)	Token熵 (↑)	MAUVE (↑)	相对MDLM加速比
MDLM (官方基线)	–	1024	104.5	5.63	–	1.0×
DUO (官方基线)	–	1024	71.7	5.22	–	1.0×
MDLM (作者复现)	–	1024	90.5	5.50	0.74	1.0×
LaDiff (本文全量)	200	1024	62.9	5.40	0.82	0.93×
LaDiff ($\zeta=0.7$)	200	1024	47.8	5.26	–	0.93×
MDLM (快速基线)	–	64	120+	–	0.62	16.0×
LaDiff (快速采样)	200	64	68.3	–	0.81	7.0×
DiLaDiff (蒸馏极速)	5 (10 NFE)	64	108.7 (76.3@$\zeta=0.7$)	5.39	0.79	15.2×

在单卡NVIDIA RTX 6000 Pro Blackwell上对连续扩散计算开销进行墙钟时间分析（Batch Size=32）： - 在 $N_{\text{disc}}=1024$ 时，LaDiff（$N_{\text{cont}}=200$）的连续扩散开销仅占离散解码时间的 7.0%。 - 在 $N_{\text{disc}}=64$ 时，未蒸馏的LaDiff（$N_{\text{cont}}=200$）连续扩散开销骤增至 111.6%（连续部分成为新瓶颈）。 - 经过MeanFlow蒸馏后的DiLaDiff（$N_{\text{cont}}=5$），连续扩散开销剧降至离散解码时间的 5.0%，真正实现了近乎零额外负担的连续扩散引导。

消融实验¶

自编码器设计与正则化强度对表示学习及最终生成质量的影响（$N_{\text{cont}}=200, N_{\text{disc}}=1024$，压缩比2×）：

实验配置	掩码与噪声正则参数	重建PPL (↓)	词恢复率 % (↑)	生成GenPPL (↓)	生成MAUVE (↑)	说明
LaDiff (默认强正则)	$\sigma=0.5, p_{\text{mask}}=0.7, p_{\text{drop}}=75\%$	$\le 3.25$	69.4%	62.9	0.82	隐空间最平滑，生成质量最优
MildAug	$\sigma=0.4, p_{\text{mask}}=0.5, p_{\text{drop}}=75\%$	$\le 2.25$	75.2%	100.1	0.71	正则偏弱，隐空间变硬
SoftAug	$\sigma=0.3, p_{\text{mask}}=0.3, p_{\text{drop}}=75\%$	$\le 1.92$	81.1%	98.5	0.24	扩散先验无法有效拟合，MAUVE崩溃
Dropout-50%	$\sigma=0.5, p_{\text{mask}}=0.7, p_{\text{drop}}=50\%$	$\le 1.21$	92.2%	101.5	0.21	解码器过度依赖潜变量，生成多样性崩塌
BigEnc (6层编码器)	深度扩展 (2层 $\to$ 6层)	$\le 2.88$	76.7%	65.1	0.78	参数量增加但生成提升不显著
BigDec (3层交叉注意力)	交叉层扩展 (1层 $\to$ 3层)	$\le 3.01$	75.6%	64.5	0.80	浅层注入即可充分传递全局语义

关键发现¶

重构性能与生成性能的尖锐反转：纯重构指标最好（恢复率92.2%、重构PPL 1.21）的模型在实际生成中MAUVE得分仅有0.21；只有施加高强度掩码、高斯噪声注入以及高概率隐通道丢弃（强迫自编码器学习具有高容错率的紧凑平滑流形），连续扩散模型才能学好先验分布。
降温采样稳定性：LaDiff在将解码器Softmax温度降低到 $\zeta=0.7$ 时，不仅GenPPL从62.9大幅优化至47.8，且Token熵仅从5.40微降至5.26，证明句子级语义多样性已被连续潜变量牢牢锁定；而原生MDLM降低温度时熵急剧跌落并迅速陷入单一模式复制。
时间调度偏置特性：超参数搜索显示 $d=10$ 的tanh-logSNR噪声调度在所有配置下均最优，证实离散文本在嵌入连续空间后各Token簇之间具备明显的离散断裂带，扩散学习必须将绝大部分步数分配在高噪声状态以跨越离散分界。

亮点与洞察¶

零卷积门控融合：在微调MDLM作为解码器时，引入ZeroConv包裹的交叉注意力层，使得训练初期新增模块完全透传，彻底规避了引入连续隐变量扰动对原有离散掩码解码能力的灾难性破坏。
连续空间ODE蒸馏破解离散扩散并行难题：离散扩散模型本身难以直接应用微分方程蒸馏技术，而DiLaDiff另辟蹊径，把蒸馏对象精准限定在具有严谨ODE流的连续潜变量空间，用连续蒸馏降维打击离散生成瓶颈。
自条件化速度逆映射：在MeanFlow多步蒸馏中，通过推导瞬时速度与干净数据的显式代数映射 $\Phi(\mathbf{v}, t)$，将复杂的自条件化预测转化为零额外开销的确定性闭式计算。

局限与展望¶

离散解码端尚未蒸馏：本文的蒸馏完全聚焦于连续潜变量阶段（$N_{\text{cont}}$ 从200降至5），但离散解码器仍采用传统的逐步去掩码，导致极端低步数（如 $N_{\text{disc}} \le 16$）场景下的表现仍落后于专门做离散蒸馏的方案（如DUO+DCD）。
学生与教师仍存在小幅性能差距：DiLaDiff在5步下MAUVE达0.79，虽逼近教师的0.81，但在GenPPL上仍有一定差距，表明非欧几何文本隐空间的少步蒸馏仍有优化空间。
正则化超参数高度经验化：当前使隐空间具备可扩散性的数据增强策略依赖于网格搜索经验启发，尚未建立从词表曲率到先验扩散收敛界的严格理论解析。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [首次成功在上下文文本隐空间实现ODE轨迹蒸馏，架构设计严谨且巧妙]
实验充分度: ⭐⭐⭐⭐⭐ [涵盖重构、生成、多样性熵分析、墙钟耗时剖析及多维度消融]
写作质量: ⭐⭐⭐⭐⭐ [数学推导详尽完备，核心动机阐述清晰，图表极具说服力]
价值: ⭐⭐⭐⭐⭐ [为非自回归扩散语言模型摆脱低吞吐困境提供了兼具理论与实效的全新范式]

实验配置	掩码与噪声正则参数	重建PPL (↓)	词恢复率 % (↑)	生成GenPPL (↓)	生成MAUVE (↑)	说明
LaDiff (默认强正则)	\(\sigma=0.5, p_{\text{mask}}=0.7, p_{\text{drop}}=75\%\)	\(\le 3.25\)	69.4%	62.9	0.82	隐空间最平滑，生成质量最优
MildAug	\(\sigma=0.4, p_{\text{mask}}=0.5, p_{\text{drop}}=75\%\)	\(\le 2.25\)	75.2%	100.1	0.71	正则偏弱，隐空间变硬
SoftAug	\(\sigma=0.3, p_{\text{mask}}=0.3, p_{\text{drop}}=75\%\)	\(\le 1.92\)	81.1%	98.5	0.24	扩散先验无法有效拟合，MAUVE崩溃
Dropout-50%	\(\sigma=0.5, p_{\text{mask}}=0.7, p_{\text{drop}}=50\%\)	\(\le 1.21\)	92.2%	101.5	0.21	解码器过度依赖潜变量，生成多样性崩塌
BigEnc (6层编码器)	深度扩展 (2层 \(\to\) 6层)	\(\le 2.88\)	76.7%	65.1	0.78	参数量增加但生成提升不显著
BigDec (3层交叉注意力)	交叉层扩展 (1层 \(\to\) 3层)	\(\le 3.01\)	75.6%	64.5	0.80	浅层注入即可充分传递全局语义