驾驭 TIDE：扩散大语言模型的跨架构蒸馏

Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models

arXiv: 2604.26951v1

论文信息

标题: Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models

作者: Gongbo Zhang, Wen Wang, Ye Tian, et al.

发布日期: 2026-04-29

arXiv ID: 2604.26951v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：如何将大规模扩散语言模型（dLLM）的知识压缩到异构的小型模型中，当教师与学生模型在架构、注意力机制甚至分词器上完全不同时，知识传递面临时间可靠性波动、上下文稀缺和词汇对齐三大障碍。
  • 核心方法：提出 Tide 框架，包含三个模块——Tidal（沿训练进度与扩散时间步双轴调节蒸馏强度）、CompDemo（通过互补掩码拆分让教师获得更丰富上下文）、Reverse Calm（以反向贝塔 KL 实现跨分词器稳定的块级概率匹配）。
  • 关键结果：在 8 个基准上，最佳配置将非蒸馏基线平均提升 +1.53 分，HumanEval 达到 48.78，远超同等参数量的自回归模型（32.3），且推理峰值显存仅需 1.4 GB，实现 22 倍的参数压缩（表 1、表 3）。
  • 主要局限：仅验证了 0.6B 参数的块扩散学生模型，未测试更大容量；CompDemo 需要两次教师前传，训练时间增加约 50%；跨分词器对齐方法尚未扩展至连续状态扩散模型。
  • 适合读者：从事大语言模型压缩与部署、扩散生成模型研究的工程师和研究人员，尤其是需要在资源受限环境中部署高质量生成模型的团队。

论文背景和研究动机

扩散语言模型以去噪掩码序列的方式并行生成文本，提供了双向上下文建模能力，近期工作如 LLaDA 2.0 已将参数规模推至 100B，但在 8B-16B 级别才具备竞争力。巨大的参数量和内存开销严重制约了实际部署（图 1）。传统知识蒸馏是压缩模型的有效途径，然而现有 dLLM 蒸馏方法都聚焦于同一架构下的步数压缩，从未尝试跨架构、跨分词器的知识传递。

跨架构蒸馏面临三个原生挑战。第一，扩散模型的训练中教师信号的可靠性随噪声时间步剧烈变化：在低掩码率时教师观察完整序列，预测可信；而高掩码率下教师几乎靠猜测，若不加区分地学习会引入噪声。第二，高掩码率使可用上下文锐减，教师输出的原始 logits 携带的丰富空间信息严重流失。第三，当教师和学生使用不同分词器时，标准 token 级 KL 散度或似然匹配直接失效，需要新的对齐机制。这三个问题分别被归纳为时间、空间和词汇障碍，催生了 Tide 框架的设计。

核心方法和技术细节

Tide 由 Tidal（调度层）、CompDemo（上下文层）和 Reverse Calm（输出层）三个协同模块构成，分别回答何时学、学什么信息、如何投射到学生词汇——形成一条完整的知识蒸馏流水线。

Tidal：时间-迭代双轴调度

传统的蒸馏强度调度仅沿训练进度一个维度变化，在扩散模型中这并不足够。Tidal 将插值系数 λ\lambda 同时建模为训练进度 pp 和扩散时间步 tt 的函数：

λt=λtrain×(1−t)\lambda_t = \lambda_{\text{train}} \times (1-t)

其中 λtrain\lambda_{\text{train}} 按照余弦进度从初始值 λinit=0.1\lambda_{\text{init}}=0.1 逐步升至 λmax⁡=0.9\lambda_{\max}=0.9（式 1、2）。这一设计使得：当 t≈1t \approx 1（高噪声），λt≈0\lambda_t \approx 0，目标主要由学生自身主导，避免学习不可靠的教师信号；当 t≈0t \approx 0（低噪声），λt≈λtrain\lambda_t \approx \lambda_{\text{train}}，学生充分吸收教师分布。双轴调度是 Tide 区别于自回归蒸馏的核心特色——在自回归模型中教师始终获得完整的左侧上下文，不需要此类时间步调制。

蒸馏目标通过对数空间插值得到混合分布 rt\mathbf{r}_t，并计算以温度为 TT 的 KL 散度：

LTIDAL=DKL(rt∥softmax(s/T))×T2\mathcal{L}_{\text{TIDAL}} = D_{\mathrm{KL}}(\mathbf{r}_t \| \mathrm{softmax}(\mathbf{s}/T)) \times T^2

仅对被掩码位置计算，可选地施加中段时间步的高斯权重以强调信息量最大的噪声区间。

CompDemo：互补演示条件去噪

标准 dLLM 蒸馏中，教师和学生对同一个被掩码的输入进行去噪，高掩码率下教师可利用的上下文极度匮乏。CompDemo 借鉴 “演示条件” 思想，将掩码位置 M\mathcal{M} 随机等分为 MA\mathcal{M}_A 和 MB\mathcal{M}_B，交替揭示一部分掩码作为额外上下文进行两次教师前传：一次用 MA\mathcal{M}_A 的真实 token 作为演示，预测 MB\mathcal{M}_B 的 logits；另一次互换角色。最终合并得到每个掩码位置的增强 logits。该方法相当于将教师置于一个有效的更低噪声水平，显著改善预测质量。因为教师冻结（不计算梯度），虽然需要双倍前传次数，但整体训练时间只增加约 50%。

Reverse Calm：跨分词器反向块似然匹配

当教师和学生使用不同的分词器时，Tide 引入 Calm 机制：基于字节对齐将序列分割为最小的 “块”（chunk），每个块包含来自两个分词器的完整 token，从而建立二进制对齐矩阵。将 token 级对数概率投影到块级标量概率 pscp_s^c 和 ptcp_t^c。朴素的前向 Calm 使用二元交叉熵损失 LFwd=−[ptclog⁡psc+(1−ptc)log⁡(1−psc)]\mathcal{L}_{\text{Fwd}} = -[p_t^c \log p_s^c + (1-p_t^c)\log(1-p_s^c)]，但其梯度包含 ptc/pscp_t^c/p_s^c，当学生概率 psc→0p_s^c \to 0 时会产生梯度爆炸。为解决这一问题，提出 Reverse Calm，交换 pscp_s^c 和 ptcp_t^c 的角色：

LRev=−[psclog⁡ptc+(1−psc)log⁡(1−ptc)]\mathcal{L}_{\text{Rev}} = -[p_s^c \log p_t^c + (1-p_s^c)\log(1-p_t^c)]

梯度系数变为 log⁡ptc1−ptc\log\frac{p_t^c}{1-p_t^c}，仅依赖固定的教师概率，天然有界。此外，该目标等效于最小化贝努利 KL 散度 KLBern(psc∥ptc)\mathrm{KL}_{\text{Bern}}(p_s^c \| p_t^c)，具有模式搜寻特性，且通过双重滤波（当 ptc≈0.5p_t^c\approx0.5 或 pscp_s^c 很小时梯度自动衰减）抑制对齐噪声。值得注意的是，由于 Reverse Calm 本身稳定，论文明确指出不应再叠加 Tidal 的进度调度，否则会削弱其自选择机制。

最终，根据分词器兼容性选择不同损失：共享分词器时使用 LCE+wtidalLTIDAL\mathcal{L}_{\text{CE}} + w_{\text{tidal}} \mathcal{L}_{\text{TIDAL}}，并可开启 CompDemo；跨分词器时使用 LCE+wcalmLRev-CalM\mathcal{L}_{\text{CE}} + w_{\text{calm}} \mathcal{L}_{\text{Rev-CalM}}。

创新点和贡献

  • 首次实现 dLLM 跨架构蒸馏：弥补了现有方法仅在同一架构步数压缩内的空白，为将大型 dLLM 压缩到轻量级学生模型提供了可行范式。
  • 双轴调度解决时间可靠性问题：Tidal 同步建模训练进度与扩散时间步，让学生按教师信号质量动态调整信任度，这一思想在扩散蒸馏中具有通用性。
  • 互补掩码提升重度噪声下的上下文质量：CompDemo 利用掩码结构的对称性，以极小额外成本充分挖掘教师潜力，为扩散模型的演示学习提供了新手段。
  • 反向块级匹配克服分词器鸿沟：Reverse Calm 通过交换 BCE 的角色获得稳定梯度，并在理论上建立了与贝努利 KL 的联系，为跨词汇表蒸馏打开了大门。
  • 模块化设计：三个组件职责清晰，可按需组合，实验验证了不同蒸馏管线（跨分词器 vs 共享分词器）确实需要不同的最优策略组合。

实验结果分析

作者在两条异构蒸馏管线上进行了全面评估：管线 A 用 16B MoE 的 LLaDA2.0-mini 蒸馏至 0.6B 的 BD3LM（跨分词器），管线 B 用 8B 稠密的 WeDLM 蒸馏至同一学生（共享分词器），所有参数都相对于非蒸馏 BD3LM 基线进行比较。

主结果（表 1）：跨架构蒸馏在所有配置下均优于未蒸馏基线（平均 32.67）。使用跨分词器的原生策略（Reverse Calm）达到平均最高分 34.20，较基线提升 +1.53；共享分词器下借助 Tidal + CompDemo 达到 33.55。值得注意的是，即便最简单的蒸馏目标（KL 或前向 Calm）也已超过基线，说明跨架构知识迁移本身是可行的。更突出的原因是，蒸馏后模型在代码生成任务上表现卓越：HumanEval 最高达 49.39（Tide-Shared），远优于同等 AR 模型的 32.30，且 MBPP 也有 2 分的优势。这说明扩散式的并行生成特别适合维持代码的全局语法与语义一致性。

消融实验（表 2）：对共享分词器管线逐步移除 Tidal 的时间步轴、训练进度轴（仅保留时间步调度作为基线）以及 CompDemo。完整方法取得最高平均分 33.14。移除时间步轴导致平均下降 0.26，且 HumanEval 骤降 3.05，验证了时间步调制的关键性；去掉 CompDemo 也导致 0.17 的平均损失，尤其在代码和知识任务上退化明显。同时，仅使用时间步调度的基线较完整框架低，说明双轴调度和 CompDemo 的补充不可或缺。

推理效率（表 3、表 4）：蒸馏学生在 H100 上的峰值显存仅 1.4 GB，相较 16B 教师减少 22 倍，生成 256 token 延迟 6.25 秒，吞吐 41.0 token/s。相比非蒸馏 BD3LM，蒸馏带来的吞吐下降仅约 2.6%，且在不同基准评测中延迟增加均匀可控。在保持与 AR 基线大致 80% 吞吐的情况下，蒸馏模型获得了显著的质量提升，展现出良好的部署性价比。

实践建议

基于 Tide 的设计与实验结果，对实际部署跨架构扩散模型蒸馏可形成几条指导原则：

  1. 根据分词器兼容性选择蒸馏策略：若教师和学生使用相同或同系列分词器，优先采用 Tidal（带 CompDemo）的 token 级蒸馏，其双轴调度在高质量信号上充分利用教师分布，在代码和推理任务上效果更突出（见 HumanEval 提升至 48.78）。若分词器不同，应选用 Reverse Calm 进行块级匹配，避免梯度爆炸并利用其噪声过滤能力，此时不建议叠加 Tidal 的进度调度。

  2. 合理设计训练调度：对于任何 dLLM 蒸馏，都应让蒸馏强度随扩散时间步衰减。最简单的方法就是以 1−t1-t 为系数进行调制，并配合余弦进度课程从较低的初始 λ\lambda（如 0.1）缓慢爬升，这能有效防止早期表示坍塌，并在后期最大化知识吸收。

  3. 在高掩码率下增强教师上下文：若计算资源允许额外 50% 的训练时间，强烈推荐启用 CompDemo。它几乎不需要调整超参数（掩码分割比 ρ=0.5\rho=0.5 即可），且对结构化生成任务（代码、数学推理）收益显著，适合对质量要求高的场景。

  4. 部署时的硬件评估：蒸馏后的 0.6B BD3LM 可在 1.4 GB 显存下运行，适合消费级 GPU。采用块扩散和 staircase attention 的推断延迟约为 6 秒/256 token，吞吐约为 AR 基线的 80%。在需要并行生成或双向上下文的应用（如代码补全、数学求解）中，这种以少许吞吐换取生成质量的权衡是值得的。工程实践中可进一步结合采样步数优化（如使用少于 256 步），在可接受精度下降范围内提升速度。

  5. 警惕跨分词器对齐中的特殊符号：跨词典蒸馏时务必去除聊天模板等特殊 token，仅对齐内容部分，否则会引入虚假的对齐块，破坏 Reverse Calm 的稳定性。在实验中，对齐矩阵构建仅针对纯文本，这一实践应当延续。未来扩展至长序列时，需关注 chunk 数量增长可能带来的梯度噪声累积。