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

1 · 动机:为什么 dLLM 需要跨架构 KD

1.1 历史脉络:dLLM 的"规模诅咒"

Diffusion language model (dLLM) 这条线,从 D3PM (2021) → MDLM (2024) 到最近的 LLaDA / Dream / WeDLM,走的是一条"用 iterative denoising 替代左到右 AR"的路。它的卖点是双向 context + 并行 decoding——对 code、structured generation 这种"全局一致性"任务理论上更友好。

但有个尴尬现实: 想拿到 SOTA 必须烧到 8B 起步。具体看:

• LLaDA-8B / Dream-7B → 至少 8B,LLaDA 2.0 已经冲到 100B
• WeDLM-8B-Instruct → 8B dense,15.5 GB 显存
• LLaDA2.0-mini → 16B MoE,31.3 GB 显存,生成 7.8 tokens/s

对比同尺寸的 AR Qwen3-0.6B (1.2 GB / 51 tokens/s),dLLM 在小尺寸上一直没竞争力。0.6B 的 BD3LM 平均分只有 32.67,被 0.6B AR 的 40.91 甩开 8 分。这是部署侧的硬墙。

2024–2026 年同期,AR 阵营已经把 KD 玩成熟了:DistilBERT (2019)、MiniLM (2020)、MiniLLM (2024)、DistiLLM、GKD (2024)、TAID (2025)。TAID 那条"逐步把 target 从 student 拉到 teacher"的渐进路线特别有启发性,但它只活在 AR 世界。

1.2 为什么"跨架构 KD"是自然解法?— 别的方案为什么不够

关键洞察: dLLM 社区已有的 KD 方法全部假设 teacher = student,只在 step 维度做压缩。这是"多走几步等价于多调几次"的物理观点。但 deployment 真正的痛点是参数量(决定显存 / 移动端可达性),不是 step 数。

1.3 为什么这事不平凡 — dLLM 跨架构 KD 的三个独有难点

把"AR 的 KL-distillation 直接套到 dLLM"会撞上 AR 世界根本不存在的三堵墙。论文把这三堵墙总结为temporal / spatial / vocabulary barriers,我把它翻译成更工程的语言:

难点 ①: Teacher 的可靠性随 timestep 剧烈变化(temporal)

dLLM 训练时 timestep t ∈ [ε, 1) 是随机采样的。t ≈ 0 时几乎不掩码,教师看着完整句子做预测,可靠;t ≈ 1 时几乎全掩码,教师就是瞎猜。如果学生一视同仁地 mimic 教师,在高 t 这一片段实际上是在学噪声。

AR 世界没这个问题——AR 教师永远看到完整 left context,可靠性恒定。

难点 ②: 重度掩码下的 spatial 信号贫乏(spatial)

就算我们只在低 t 区间蒸,中等 t (e.g. 0.5) 时教师看到一半掩码 [M]、一半 token,产出的 logits 仍然抖。教师自己也是在猜——"这个被掩盖位置应该是 ya 还是 sé?"它给的 distribution 信息含量低,蒸不出什么 spatial knowledge。

难点 ③: 不同 tokenizer 让 token-level KL 数学上不存在(vocabulary)

LLaDA2 用 Ling-series 的 V_T,BD3LM/Qwen3 用 Qwen 的 V_S。两个词表不同 → 同一个文本被切成不同 token 序列、不同长度、连分布的支撑集都不同。KL(p_T || p_S) 的两个分布根本不在同一个空间,这是 type error,不是数值精度问题。

2 · 背景速查

2.1 关键术语

2.2 dLLM 训练流程速复习

1. 从 dataset 取 clean sequence x = (x₁, ..., x_L)
2. 采样 timestep t ∈ [ε, 1)
3. 按比例 t 随机掩码 → x_t,被掩位置集合 M
4. 模型 forward x_t,只在 M 位置 log-softmax 预测原 token
5. Loss = CE(model_prediction[M], x[M]) — 即只在 mask 位置算交叉熵

蒸馏额外: 教师 forward 同样的 x_t,产出 logits;学生不仅匹配 ground-truth,还匹配教师 logits。"什么时候、用谁的 logits、怎么算 loss"就是本文要解决的三个问题。

3 · 总览:三个组件的角色分工

从 pipeline 角度看,论文做了两条 pipeline 的实验:

• Pipeline A (cross-tokenizer): LLaDA2.0-mini (16B MoE, Ling tokenizer) → BD3LM-0.6B (Qwen tokenizer)。必须用 Reverse CALM。最佳配置 = baseline CALM + Reverse,不开 TIDAL+CompDemo。
• Pipeline B (shared-tokenizer): WeDLM-8B (Qwen tokenizer) → BD3LM-0.6B (Qwen tokenizer)。直接 token-level KL 即可。最佳 = TIDAL + CompDemo,不需要 Reverse(没 cross-tokenizer 不稳定问题)。

4 · TIDAL — Time-Iteration Dual-Axis Lambda 双轴调度

4.1 问题再叙述

TAID (Shing 2025, AR-only) 提出一个干净的想法:学生不要直接 mimic 教师,而是 mimic 一个"插值 target" r = (1−λ)·s + λ·t。λ 从训练开始的 ~0.1 慢慢升到 ~0.9,等价于"先让学生信自己,再慢慢把自信交给教师"——避免训练早期学生还很弱时被教师强行带跑造成 representation collapse。

把 TAID 直接套到 dLLM 的问题: λ 只沿 training-progress 一个轴变,不感知 timestep。dLLM 里有些 batch 是 t=0.05(教师可靠),有些是 t=0.95(教师瞎猜),用同一个 λ 对待是错的。

4.2 TIDAL 的双轴公式

解读:

• 第二式 — λ_train 沿 training progress p ∈ [0,1] 走 cosine schedule,从 0.1 升到 0.9。这是 TAID 的部分。
• 第一式 — 再乘以 (1−t)。t 越大(掩码越重)、教师越不可靠,(1−t) 越小,λ_t 越小,target 越靠学生自己。

4.3 公式细节: Loss

注意几点:

• r_t 走 .detach():保证梯度只通过 student logits 反传(防止"间接更新教师"或者引入虚假梯度路径)
• × T²:Hinton 经典 KD 的温度补偿,使得不同 T 下梯度量级一致
• 只在 mask 位置算 — dLLM 的 standard 做法

4.4 Worked example: 一个具体 token 的 λ_t 计算

训练第 5000 step / 共 50000 step,所以 p = 0.1 → λ_train = 0.1 + (0.9−0.1) × ½(1−cos(0.1π)) ≈ 0.1 + 0.8×0.024 ≈ 0.119。

本 batch 采样到 t = 0.3(掩码 30%)。所以 λ_t = 0.119 × (1−0.3) = 0.083。

假设某 mask 位置(原 token 是"是"):

• student logits s 给出: {"是": 0.4, "为": 0.2, "在": 0.15, ...}
• teacher logits t 给出: {"是": 0.85, "为": 0.05, "在": 0.04, ...}
• 插值 target r = softmax( 0.917 × s + 0.083 × t ),"是" 概率 ≈ 0.4×0.917 + 0.85×0.083 ≈ 0.437

这时学生几乎完全 mimic 自己当前预测,只被教师轻轻"拽"了一点。这正是 p=0.1 早期训练应该的行为——避免还没学好就被强行同化。

到了 p=0.9, t=0.05 的情形: λ_train ≈ 0.876, λ_t ≈ 0.832,target 的 86% 能量来自教师 — 学生此时已经成熟,被教师"猛拽"也学得动。

4.5 反向论证:不要双轴会怎样?

Ablation Table 2 给的硬数据:

移除 timestep 轴 HumanEval 掉 3.05 分——是所有 ablation 里最大的。这反过来说明:dLLM-KD 里 timestep 不是 nice-to-have,是load-bearing。

5 · CompDemo — 互补掩码二次前向

5.1 问题

就算我们用 TIDAL 把高 t 区间下采,中等 t (0.4–0.6) 还是要蒸——这一段教师面对一半 token / 一半 [M],预测有信息但不够锐。能不能想办法让教师在不增加学生输入信息的前提下,看到更多 context?

5.2 解法:互补掩码,跑两次教师

把 mask 集合 M 随机切成两半 M_A ⊔ M_B,|M_A| ≈ |M_B| ≈ |M|/2。然后:

• Pass 1: 揭示 M_A 的真值给教师当 context,只让教师预测 M_B 的 logits → t⁽¹⁾[M_B]
• Pass 2: 揭示 M_B 的真值给教师,让教师预测 M_A 的 logits → t⁽²⁾[M_A]
• 合并: t_final[M_B] ← t⁽¹⁾[M_B], t_final[M_A] ← t⁽²⁾[M_A]

关键:每个被蒸的位置,教师都比学生少看了 ~一半的 [M](因为另一半被揭开了)。等价于教师在更低的有效 timestep 下做预测。学生还是面对原始 x_t,没作弊。

5.3 反向论证

移除 CompDemo 在 ablation 里掉 0.17 平均分,HumanEval 掉 2.44 分。结构化输出(代码)对"教师是否给出锐利预测"特别敏感——code 的 syntax 经常需要严格匹配,教师在重度掩码下乱猜会污染学生的 code 认知。

6 · Reverse CALM — 反向 chunk-level BCE

6.1 问题再细化

LLaDA2 (Ling tokenizer) 和 BD3LM (Qwen tokenizer) 编码同一段文本:

"I don't wanna"
Ling tokenizer:  ["I", "don", "'t", "wanna"]      → 4 tokens, V_T 词表
Qwen tokenizer:  ["I", "do", "n't", "wan", "na"]  → 5 tokens, V_S 词表

KL(p_T(·|x) || p_S(·|x)) 数学上不存在——两个 token 不在同一空间。怎么办?

6.2 CALM 的核心 idea: chunk-level scalar 概率

用 tokenkit 在字节层面对齐两条 token 序列,找出最小的"对齐单位":两边都能被整数个 token 切完的最小 byte span,称为 chunk。

对每个 chunk,定义其概率为 chunk 内所有 token 的联合 log-probability 之和(再 exp/temperature 化):

关键:p_s^c 和 p_t^c 都是 [0,1] 之间的标量(不是分布),代表"模型生成这个 chunk 的概率"。两边都是同样的物理量,可以直接比较。

6.3 Forward CALM 的 bug

自然 baseline = 在 p^c 上做 BCE,把 teacher 的 chunk 概率当 soft label:

梯度有个问题:对 p_s^c 求导,系数 ∝ p_t^c / p_s^c − (1−p_t^c)/(1−p_s^c)。Cross-tokenizer 对齐本身有噪声(byte 切分不完美 / 教师对某 chunk 概率高但学生还没学到)→ p_t^c → 1, p_s^c → 0 的情况频繁,gradient 直接爆炸。

6.4 Reverse CALM:把 BCE 翻过来

把 p_s 和 p_t 的角色互换。对 p_s^c 求导:

这个梯度系数只依赖于 fixed teacher 的 p_t,与 p_s 无关 → 天然有界。

6.5 Worked example 数字

设 chunk c = "don't",T=1。某 step 时:

• 教师概率 p_t^c = 0.92
• 学生概率 p_s^c = 0.05(初期还没学好)

Forward CALM:

• L = −(0.92·log(0.05) + 0.08·log(0.95)) ≈ 2.76
• ∂L/∂p_s = −0.92/0.05 + 0.08/0.95 ≈ −18.32 ← 巨大,容易炸

Reverse CALM:

• L = −(0.05·log(0.92) + 0.95·log(0.08)) ≈ 2.40
• ∂L/∂p_s = −(log(0.92) − log(0.08)) ≈ −2.44 ← 有界

学生在不同 p_s 下,Reverse 的梯度始终 ≈ 2.44(只看 teacher),而 Forward 的梯度可以从 −∞ 变到 +很大,数值不稳定。

6.6 第二个 bonus: 噪声过滤

如果某 chunk 在 cross-tokenizer 对齐里是"伪 chunk"(byte 切分歧义),教师对它的 p_t^c 应该接近 0.5(教师也不确定)。Reverse CALM 此时 logit(p_t) ≈ 0 → 梯度 ≈ 0,该 chunk 不参与训练。这就是免费的 noise filter。

另一头:学生对该 chunk 还没信心(p_s^c 小),∂p_s^c/∂θ 小,grad 小——再过滤一次。这是"dual-end" 的含义。

6.7 数学小注:Reverse CALM = Bernoulli KL

论文 Appendix C 推导:

所以 Reverse CALM 是mode-seeking 的(student 的 chunk 概率应该 collapse 到 teacher 的某一个)。这与 Forward(mode-covering)对应,呼应 MiniLLM 那条"AR-KD 应该用 reverse KL"的研究线。但这里换到了 scalar Bernoulli 空间,而不是 token-level distribution 空间。

7 · 实验关键结果与读图

7.1 主表精读 (8 benchmark, two pipelines)

关键观察

1. Reasoning task 上 BD3LM 学生远不如同尺寸 AR(GSM8K 45.56 vs 59.60,MATH 13.08 vs 32.40)。蒸馏只能挽回一部分(GSM8K 到 52.24,MATH 几乎不动)。dLLM 在小尺寸上的 reasoning 短板 KD 救不全。
2. Code task 上 BD3LM 学生反而压过 AR(HumanEval 46.34 vs 32.30 已经赢了,蒸到 48.78 进一步扩大优势)。这是 dLLM 真正发光的领域。
3. 每个 pipeline 偏爱自己的 native 策略。Cross-tokenizer 用 Reverse CALM(TIDE-Cross 列)拿 34.20;Shared-tokenizer 用 TIDAL+CompDemo(TIDE-Shared 列)拿 33.55。交叉用 (TIDE-Cross 在 shared,TIDE-Shared 在 cross)各掉 ~2.7 / ~0.4 分,符合方法论预期。

7.2 推理效率 (Table 3)

0.6B 学生 vs 16B 教师: 22× 显存压缩 + 5× 加速。但相比同尺寸 AR (51.3 tokens/s) 还慢 ~20%——这是 BD3LM 多步 denoising 的开销,KD 改不了。

8 · 为什么 HumanEval 涨幅最大?

这是论文最有意思的结果,但作者的解释比较短(就一段:"parallel generation 维护 global coherence,适合 structured output")。我想细化一下:

1. Code 是高度结构化的 — bidirectional context 是真有用。函数体的开头和结尾要配对(def f(): 和 return ...),括号要平衡。AR 模型从左到右生成时,生成 return 时还看不到后续可能的 except:。BD3LM 在一个 block 内同时知道前后位置,优势天生。
2. Code 的 token-level distribution 比自然语言更尖锐。"for i in" 之后大概率是 "range"——教师概率分布有明显 mode。Reverse CALM 的 mode-seeking 性质正中这个特性。
3. HumanEval 的 prompt+test 长度短 (256 tokens 内能搞定),恰好是 BD3LM 512 序列长度的甜点。MATH 和 GSM8K 的推理过程很长,会走出 BD3LM 训练的 block 范围,multi-block 的 AR 跳跃式生成削弱 bidirectional 优势。
4. CompDemo 在 code 上收益最大(ablation 移除 CompDemo,HumanEval 掉 2.44 — 八个任务里第二大)。原因猜测:code 的 syntax 严格,教师在重度掩码下乱猜带来很大损失;给教师补一半 context,教师才能给出锐利的 code-specific 概率分布。
5. 蒸馏后的 0.6B BD3LM 在 HumanEval 上 (48.78) 甚至超过了 8B 教师 WeDLM?论文没明示 WeDLM 的 HumanEval,但从 Table 1 推测教师并非压倒性强 — 教师本身 + 学生的 BD3LM 结构 + dataset 重组合,最终甚至能提纯 code 能力。

9 · 与同类工作对比

9.1 与 EAGLE-3 多层 KD 的具体差异

有读者可能会问: EAGLE-3 也是 KD,也是大教师 → 小学生(draft head 比 base 小得多),为什么不能直接套?

• EAGLE-3 的 draft head 共享教师的 hidden states(dim 4096+),它本质是"借用教师的 representation 做 next-token 预测"。Draft 不是独立模型,是教师的"附属预测头"。
• TIDE 的学生是独立的 0.6B 模型,有自己的 4096 → 1024 维 hidden,不依赖教师 forward。两个模型可以分别部署。
• EAGLE-3 应用场景是同架构 inference 加速(教师永远在线);TIDE 应用场景是替换教师(教师永远离线)。

9.2 与同期 dLLM-KD 工作的真正新意

2024–2025 的 dLLM-KD(CDLM/DDD/LSD/SDTT)全部假设 student = teacher 同架构(只压 step)。本文是第一个承认 cross-architecture 才是更值钱的方向 + 把它跑通。这是 framework 级别的开拓,不是 incremental 改进。

10 · 局限 / 个人 take / 待验证问题

论文承认的局限

• 学生只跑到 0.6B。Future work 提到 "scaling student capacity" — 1.5B/3B 是否还能维持比例增益未验证。
• 只在 discrete diffusion(masked dLLM)上验证。Continuous-state diffusion LM(类似 SEDD)需要不同设计。
• 序列长度只到 512 + block size 32。长 context (4k+) 下 BD3LM 的 attention pattern 是否仍能享受蒸馏增益未知。
• Pipeline B 用的 WeDLM-8B 本身是从 Qwen3-8B 转过来的,与学生 BD3LM(也从 Qwen3-0.6B 转过来)的 representation space 已经很对齐。这个 "shared-arch within shared-tokenizer" 设定其实有"作弊"嫌疑。真正"完全异构 + 完全 cross-tokenizer"只有 LLaDA2 那条 pipeline。

我的 take

1. 论文最 underrated 的贡献其实是把 ALM 的 forward 翻成 reverse 这个微小改动。Forward CALM 是 ALM 原版,本文做的核心新事情就是分析 forward 的梯度爆炸 + 用 reverse 修。这个 trick 可以直接搬到任何 cross-tokenizer KD(不只是 dLLM)。
2. TIDAL 的两个轴本质上是 TAID(progress 轴)+ "教师可靠性 mask"(timestep 轴)。timestep 这个轴是 dLLM specific 的,但思想是普适的:任何"教师可靠性随时间变化"的 KD 都可以加这一轴。例如 RL-distillation 里教师在某些 state 上不可靠,就该降 λ。
3. CompDemo 是非常 dLLM-native 的 trick,搬不走。但成本 +50%(虽然论文说"忽略不计")在 100B+ 教师上会很贵。
4. 整篇论文最反直觉的发现:cross-tokenizer pipeline 用 native 策略只赢 0.37 分,shared-tokenizer 用 native 策略赢 2.76 分。这意味着 Reverse CALM 鲁棒性强(任何 pipeline 都能接近最佳),而 TIDAL+CompDemo 的好处更"挑剔"——只有当 token-level KL 本身可以算时才能完全发挥。

想验证的问题

1. 把 Reverse CALM 单独搬到 AR 跨 tokenizer KD(e.g. Llama → Qwen),是否依然有显著增益?预期 yes,这是 free transfer。
2. TIDAL 的 timestep 轴是否对所有 dLLM 教师都通用?对 reasoning-heavy 教师(教师在 high-t 上不那么差),timestep 轴的好处会缩小吗?
3. CompDemo 的 ρ=0.5 是默认,但 ρ=0.25 / 0.75 没扫。极端 ρ(几乎全揭/几乎不揭)的退化模式是什么?
4. 能否在 RL post-training 之后再做 TIDE?论文只做 SFT 蒸馏。如果教师是 RLHF 后的(更像最终目标),学生能否进一步提升?
5. 当前 pipeline B 是 same-arch within same-tokenizer-family(Qwen3-8B WeDLM → Qwen3-0.6B BD3LM)。如果换成真正异构(LLaDA2 → BD3LM 同 tokenizer 假想版本),数字会怎样?
6. HumanEval 48.78 vs MBPP 38.60 的差距:为什么 MBPP 涨幅小很多?MBPP 的 prompt 风格更口语化,可能 chunk 对齐噪声更大?

记忆点

精读笔记 v1 · 配套论文 PDF: /data/szhang967/papers/paper-notes/models/TIDE_2604.26951.pdf