StepX-OPD:重新思考扩散模型中的多教师在线策略蒸馏 StepX-OPD: Rethinking Multi-Teacher On-Policy Distillation in Diffusion Models

完整流程见 Alg. 1。(1) 抽取训练源 $d\in\{$GenEval, OCR, PickScore$\}$ 并取一个同源 mini-batch;前 200 步 warmup 中目标 $v_{tgt}$ 为冻结的初始学生(自蒸馏),之后由静态查找表 $r$ 把 $d$ 路由到对应专家 $v_{te}^{r(d)}$。(2) 每个训练步重抽随机 rollout 网格 $1=\sigma_0>\dots>\sigma_4=0$(段数固定、两端钉死、内部节点随机,所有并行进程共享同一抽样以保证梯度一致),学生用 Euler 更新 $z_{\sigma_{i+1}}=z_{\sigma_i}+(\sigma_{i+1}-\sigma_i)v_\theta$ 走完 4 步。(3) 均匀选一段 $[\sigma_\ell,\sigma_u]$,在其内均匀抽 $\tau_{base}$,对两端状态插值 $x_{base}=\gamma z_{start}+(1-\gamma)\,\mathrm{sg}[z_{end}]$,其中 $\gamma=\frac{\tau_{base}-\sigma_\ell}{\sigma_u-\sigma_\ell}$;$\mathrm{sg}$ 为 stop-gradient,Lemma A.1 证明 $x_{base}$ 恰是从 $z_{start}$ 出发到 $\tau_{base}$ 的精确 Euler 子步。(4) 在 $[\tau_{base},\sigma_u]$ 内抽监督时刻 $\tau$,用矩匹配的前向 SDE 转移精确再加噪:$z_\tau=\kappa x_{base}+\beta\epsilon$,其中 $\kappa=\alpha_\tau/\alpha_{\tau_{base}}$、$\beta=\sqrt{\max(\sigma_\tau^2-(\kappa\sigma_{\tau_{base}})^2,0)}$,使 $z_\tau$ 精确携带 $x_{base}$ 的条件分布。(5) 按默认 TDM 更新 fake 速度 $v_\psi$;路由专家在该点做条件/无条件两次前向并 CFG 合成目标 $v^{cfg}_{tgt}=v_{tgt}(z_\tau,\tau,\varnothing)+w[v_{tgt}(z_\tau,\tau,c)-v_{tgt}(z_\tau,\tau,\varnothing)]$。(6) TDM 项用伪 Huber 损失:$\mathcal{L}_{TDM}=\rho(z_{end}-\mathrm{sg}[z_{end}-\tau\Delta_\tau])$,其中 $\Delta_\tau=v^{cfg}_{tgt}-v_\psi(z_\tau,\tau,c)$。(7) 低噪速度回归:当训练步 $s\ge s_v=500$ 且 batch 平均时刻 $\bar{\tau}\le\tau_v=0.5$ 时,$\mathcal{L}_v=\frac{1}{N_{elt}}\sum_j\|v^{stu}_\theta(z_{\tau_j},\tau_j,c)-\mathrm{sg}[v^{cfg}_{tgt,j}]\|^2$,权重 $w_v=0.5$。(8) 总目标 $\mathcal{L}_{total}=\mathcal{L}_{TDM}(v^{cfg}_{tgt})+w_v g_v\mathcal{L}_v$。训练用 8 块 NVIDIA H800(FSDP)、全局 batch 256、rank-32 LoRA 学生与 fake 适配器,学习率分别为 $2\times 10^{-6}$ 与 $1\times 10^{-5}$,200 步优化器 warmup。

与最接近工作的差异:(1) 相对 DiffusionOPD、Flow-OPD、DanceOPD,监督从全步状态分布搬到学生自己的 4 步 rollout,并把"选哪个教师"与"训练数据源"绑定——Flow-OPD 的任务路由发生在教师合并阶段,DanceOPD 虽按样本路由速度场但在全步 MSE 下进行,而本文在少步轨迹上按 mini-batch 源路由,静态查找表不需要任何进程间通信,数据并行完全友好。(2) 相对 CDM,随机网格内的监管点不是从单点速度外推,而是用两个已访问节点插值,论文证明二者计算的是同一个量(精确 Euler 子步),但插值路线不需要 rollout 之外的额外速度评估;再加噪系数由矩匹配精确导出,监管点携带的正是 rollout 的真实条件分布。(3) 相对 DMD/TDM,教师分数不再是单一固定教师,而是按源路由的 CFG 合成专家速度,并新增了由反向 KL 分解直接推导出的低噪门控速度回归。(4) 副作用量化本身也是贡献:提示词复现计数、四个提示分层、slogan/描述语域测试构成了一套可复用的测量协议。

(1) 教师生态封闭且静态:路由表只覆盖 GenEval/OCR/PickScore 三个源,新增能力(如中文文字渲染、风格控制)需要训练新专家适配器并重新设计初始化 soup,改进方向是学习式软路由(由学生状态预测专家混合权重)或对专家增量做持续合并。(2) 提示词复现的治理是补丁式的:精确教师对短 rollout 上的引导速度取平均,只把 slogan 复现率从 0.820 压到 0.610,且根因(OCR 奖励与能力整合的相互作用)未被解释;改进方向是把复现率作为正则项直接进损失(训练时用 OCR 检测器惩罚检测结果与提示词的重叠),或构造显式的"无文字"负例训练源。(3) 优化稳定性对调度敏感:速度回归从第 0 步与第 500 步开启结果接近(0.8895 对 0.8901),但第 1000 步才开启掉到 0.8652;源采样比例在 fresh-LoRA 与 merged-init 两种学生下最优值相反(1:2:1 对 1:1:1),说明超参之间存在复杂耦合,改进方向是自适应损失加权与按验证分数的源重加权。(4) 只验证了 4 步与 512×512:更激进的 1-2 步学生、更高分辨率、视频扩散上的有效性未知,可在 DMD2 式单步设置与多分辨率数据上复测。(5) Average 的轴归一化使跨方法比较依赖指标选择,缺少人类评估与更细粒度的失败模式分析。

作者方向:把精确教师目标以更低成本集成(如缓存平均速度或将其蒸馏为单次前向);进一步压低提示词复现,弄清 OCR 奖励与复现行为之间的因果链;在更多骨干模型上验证阶段融合的普适性。基于其成果可延伸的:把静态源路由升级为对状态的软路由/门控混合专家,统一 DanceOPD 的逐样本路由与本文的逐 batch 路由;把该框架与强化学习后训练(Flow-GRPO 教师,Table 9 已初步验证)组合,实现"少步 + 偏好对齐"一步到位;将反向 KL 分解推广到含对抗损失(DMD2)或矩匹配(MM)的目标族,为更多蒸馏损失提供统一的权重结构解释;研究复现现象在其他能力(非 OCR 奖励)整合下是否同样出现;并扩展到视频生成的时间步蒸馏。

论文给出 GitHub、HuggingFace 与项目页链接(正文为占位,实际开放程度需确认);学生与 fake 网络均为 rank-32 LoRA,冻结专家直接复用 DiffusionOPD 等前作公开发布的适配器,训练提示集沿用 Flow-GRPO/DiffusionOPD 的划分,评测协议完全公开(GenEval 官方检测器、PaddleOCR 归一化匹配、PickScore 除以 26、Average 公式,以及各基准图像数量:GenEval 8848、OCR 4072、PickScore 8192、DeQA 等 750 张)。算力门槛中等:8 块 H800、全局 batch 256、约 72.7 GPU 小时,单节点一天内可完成,但对没有 H800 级集群的团队仍不轻;理论部分(附录 A)只依赖高斯核运算与广义 de Bruijn 恒等式,易于核验。总体复现难度中等偏上:关键是拿到三个冻结专家适配器、复刻随机网格的多进程一致性以及伪 Huber 的 TDM 项实现;若只想验证消融趋势,LoRA 加单节点多卡即可。