← 返回 2026-11-05

VA-TDM:速度对抗蒸馏——改进少步扩散模型的统一视角 VA-TDM: A Unified View for Improving Few-Step Diffusion Models by Velocity Adversarial Distillation

Yihong Luo, Tianyu Pang, Tianyang Hu 📅 2026-11-04 👍 10 2026-11-11 18:30
对抗训练 少步蒸馏 扩散模型加速 文生图 文生视频 流匹配

速度对抗蒸馏统一分数蒸馏与矩匹配,4步生成SOTA并扩展至33B视频模型

前置知识

流匹配与速度场

流匹配(如 Rectified Flow)把生成建模为一个连续时间向量场 $u_\tau(x)$ 沿路径把噪声分布传输到数据分布;线性插值路径下速度场为 $u_\tau(x)=\mathbb{E}[\varepsilon-x_0\mid x_\tau=x]$,采样即对该场做 ODE 积分。它与扩散模型在数学上等价,但速度目标比分数目标方差更低、训练更稳定。

本文所有监督信号都构造在速度空间,理解速度场及其与分数的等价关系是读懂 A/B 两种速度形式的前提。

分数蒸馏与分布匹配

以 DMD 为代表的方法用反向 KL 对齐学生与教师分布,梯度形如教师分数与学生分数之差 $\nabla\log p_\theta-\nabla\log q$。由于学生分布是隐式的(样本来自生成器自身),必须额外训练一个辅助分数网络在线跟踪不断变化的学生分布,这一估计的准确性决定了蒸馏上限。

本文的出发点正是该辅助估计在少步中间步上存在分布 mismatch 与高方差问题,读懂动机需要先了解这条技术线。

TDM 轨迹分布匹配

TDM(Luo et al. 2025)不只匹配最终输出,而是沿学生的 $K$ 步采样轨迹,对每一个中间转移做学生与教师边际分布的对齐,让每一步只学一个部分去噪任务,监督因此分摊到整条轨迹,是少步蒸馏的强基线。

VA-TDM 完整保留了 TDM 的逐转移监督框架,仅革新监督信号的构造与在线估计方式,是本文的直接基座与对比对象。

对抗训练与 critic

GAN 通过一个 critic/判别器估计两种分布之间的差异,生成器与判别器进行极小极大博弈;当 critic 达到最优时,其输出给出分布差异的函数表示,可用于指导生成器更新。基于积分概率度量或 f-散度的多种 GAN 变体遵循同一原理。

VAD 把这一思想搬到速度空间:最优 critic 有解析解 $f^\star=u_\tau-\bar{v}^F_{\theta,\tau}$,即速度残差,这是整个框架的理论核心。

去噪分数匹配 DSM

DSM 通过回归条件目标 $\nabla_{x_\tau}\log q(x_\tau\mid x_0)$ 来估计边际分数,是扩散模型训练与在线分数估计的标准做法;回归目标的条件方差决定了辅助网络学习的难易程度,方差越大越难学、监督噪声越大。

论文论证对中间转移输出做 DSM 的条件方差比标准 DSM 更大,因此改用速度目标——这一方差分析是方法动机的关键论据。

研究动机

主流少步蒸馏依赖分数蒸馏/分布匹配(DMD2、TDM 等):更新生成器需要学生分布的分数,因此要训练辅助网络在线跟踪学生分布。问题出在中间步:学生在 $t\to s$ 转移输出的 $x_s$ 保留了继承自上一状态的残差,clean 预测 $\hat{x}_0$ 与该残差相关,因此对 $x_s$ 再加扰动得到的分布 $p^{\text{path}}_{\theta,\tau}$ 并不等价于对 $\hat{x}_0$ 独立加高斯噪声。作者分析:直接对扰动输出做 DSM 的目标条件方差比标准 DSM 更大、辅助网络难学;而 TDM 的 clean-target 去噪加重要性采样实际估计的是另一个分布的分数,与生成器更新真正需要的分数不一致,存在系统性 mismatch。另一被忽视的问题是:少步设置下不同 step index 的转移在同一噪声水平诱导的分布各不相同,但现有方法让所有转移共享同一个辅助分数,理论上就无法精确跟踪学生分布;TDM 用不相交扰动区间部分规避,代价是牺牲监督信号。

本文的目标是本文的目标是为少步蒸馏构造一套既忠实于学生实际采样过程、又低方差易学的监督信号,并把现有诸多蒸馏目标统一进一个可分析的理论框架。具体而言:其一,摆脱对噪声边际分数的在线估计,改用可以直接从学生采样链 $t\to s\to\tau$ 构造的速度场作为监督载体,从根本上消除 TDM 的分布 mismatch;其二,建立一个速度空间中的对抗博弈框架,使分数蒸馏(反向 KL)、Fisher 散度匹配、矩匹配(MMD)成为同一框架在『选哪种场 × 用哪种梯度更新』下的特例,同时解释 MMD(Salimans et al. 2024)的启发式更新为何有效并给出其目标函数的精确梯度;其三,在 SD3.5-Medium、FLUX.2-Klein 4B 以及 33B 参数的 MiniMax-H3 音视频模型上训练 4 步生成器,在文生图与文生视频上全面超越 DMD2、D-DMD、DPDMD、TDM 与官方 Fast-H3 等基线,显著逼近多步教师的感知质量。

与已有工作不同的是,本文的独特切入是把『蒸馏监督』从分数空间整体搬到速度空间,并显式写成对抗博弈。速度目标的妙处在于 rectified flow 允许路径端点耦合:$t\to s\to\tau$ 复合过程把 clean 预测 $\hat{x}_0$ 与 $x_\tau$ 耦合起来,由此定义的 A-form 速度 $v^A_\tau=(x_\tau-\hat{x}_0)/\tau$ 训练方差低;而 $s\to\tau$ 扰动路径给出 B-form 速度,等价于真实扰动分布的分数,携带全分布信息。作者进一步证明:用一个向量 critic 与速度场做内积的极小极大博弈,最优 critic 恰为速度残差 $f^\star_F=u_\tau-\bar{v}^F_{\theta,\tau}$,残差为零即学生速度处处等于教师速度的纳什均衡;在此框架中改变『用哪种场』与『怎么回传梯度』两个正交选择,就能精确复现 MMD、反向 KL、Fisher 散度等已有方法,并自然导出新目标。这把以往彼此独立、靠启发式拼凑的技巧放进了一个统一且可解析的坐标系。

核心方法

直觉版本:不要去『猜』学生分布的分数,而是在速度空间打一场 GAN——训练 critic 度量教师与学生速度场的差异,把残差反馈给学生。技术路线:(1) 学生在 $x_t$ 处预测 clean 锚点 $\hat{x}_0=x_t-t\,v_\theta(x_t,t)$,经 DDIM 式转移(默认 $\eta=1$)生成 $x_s$,再用与教师前向兼容的高斯核扰动到 $x_\tau$;(2) 从采样链构造两种条件速度:低方差的 A-form $v^A_\tau=(x_\tau-\hat{x}_0)/\tau$ 与携带全分布信息的 B-form(等价于扰动分布的分数);(3) 用共享 backbone、双头的在线场网络 $h_\psi(x_\tau,\tau,s)$ 回归学生平均速度场,critic 参数化为 $(1+\lambda_r)(u_\tau-h_\psi)$;(4) 生成器用 pathwise 与 flow-product 的混合损失更新,critic 输出整体 stop-gradient;(5) 沿 4 步轨迹每个中间转移施加监督,先做 ODE 回归 warm-up。

核心创新有三层。第一,VAD 把少步蒸馏形式化为速度空间的极小极大博弈,并解析证明最优 critic 是速度残差 $f^\star_F=u_\tau-\bar{v}^F_{\theta,\tau}$、残差为零即分布匹配均衡;本质区别是监督由学生自己的采样链定义,而非估计某个替代分布的分数,从根上消除 TDM 的 clean-target mismatch。第二,A/B 双形式与两种梯度约定(pathwise/flow-product)构成 2×2 统一网格(Theorem 5):A-form pathwise 加 detached critic 复现 MMD 的实践更新,A-form flow-product 给出 MMD 期望目标的精确梯度,B-form pathwise 即反向 KL,B-form flow-product 即 Fisher 散度匹配。第三,transition-aware field estimation:让在线场网络额外接收转移目标时刻 $s$,解决不同 step 分布不同却共享辅助网络的盲点,消融显示该设计使 GenEval 从 0.58 提到 0.61。

方法步骤详情

第 0 步 warm-up:学生 rollout 到 $x_{t_k}$,教师 ODE solver 积分出 $x^{\text{tea}}_{t_{k-1}}$,距离损失回归,严格满足即蕴含轨迹分布匹配。第 1 步采样链:从 $x_{t_K}\sim\mathcal{N}(0,I)$ 逐步执行转移得 $(x_t,x_s,\hat{x}_0)$,经高斯核扰动到 $x_\tau$。第 2 步 critic 学习:在线场 $h^F_\psi$ 以 MSE 回归 $v^F_\tau$,加权重 $\lambda_r$ 的教师速度正则,最优解为师生场加权平均,critic 按 $(1+\lambda_r)$ 校正。第 3 步生成器学习:每形式用 $(1-\lambda)$ pathwise 加 $\lambda$ flow-product 混合损失(对 critic 输出 stop-gradient),再按 $\lambda_m$ 混合双形式,沿轨迹每个中间转移重复。最终配置 $K=4$、$\lambda_m=0.8$、$\lambda=0.7$、$\lambda_r=1$、$\eta=1$。

技术新颖性

技术新颖性体现在四方面。其一,概念统一:此前 score distillation(DMD 系列)、Fisher 散度匹配(Luo et al. 2024; Zhou et al. 2024)与 moment matching(MMD)被视为三种独立方法,VAD 证明它们只是『速度场构造 × 梯度回传方式』在 2×2 网格上的不同位置,且推导给出带精确时间系数的梯度恒等式,而非松散类比。其二,对 MMD 的重新诠释:Salimans et al. 2024 的启发式更新首次被证明是对其期望目标的近似梯度,VAD 的 A-form flow-product 更新给出精确梯度,detached-critic 的 pathwise 更新则从博弈角度解释了原启发式的来源。其三,与 AFM(Lin et al. 2026)同在速度空间做对抗,但本文面向少步蒸馏且避免昂贵的判别器 JVP 计算,训练成本与可行性完全不同。其四,shared-backbone 双头设计与 transition-aware 条件化均为蒸馏文献中首次提出,后者不仅提升 GenEval,还显著改善文字渲染与多样性。

Comparison between detached critic and non-detached critic
Figure 3: Comparison between detached critic and non-detached critic

实验结果

文生图(Table 1):SD3.5-Medium 4 步 GenEval 0.67,与教师持平,高于 TDM 0.61、DPDMD 0.62、DMD2 0.60;OCR 0.59、Aesthetic 5.65、ImageReward 1.02、HPSv2 0.3010、HPSv3 11.89 全面领先基线,LPIPS 0.58。FLUX.2-Klein 4B:GenEval 0.80 超过官方 4 步模型的 0.78,OCR 0.64→0.69,LPIPS 0.47→0.56。文生视频(33B MiniMax-H3,Figure 2):VA-TDM 以 61.33% 偏好率击败 Fast-H3,对 100 NFE 教师偏好率达 45.53%(Fast-H3 仅 24.67%),大幅收窄差距。消融(Table 2):A-form only 0.64、B-form only 0.61、去掉 transition-aware 掉到 0.58;从 $\lambda_m=0.8$ 的 0.63 出发,依次加梯度混合(0.65)、场正则(0.66)、正则校正 critic(0.67)。

Ablation study
Table 2: Ablation study
Comparison of VA-TDM-H3 and competing baselines
Figure 1: Comparison of VA-TDM-H3 and competing baselines
Human preference on video synthesis
Figure 2: Human preference on video synthesis
Qualitative comparison among VA-TDM and competing baselines
Figure 4: Qualitative comparison among VA-TDM and competing baselines
查看结构化数据
任务指标本文基线提升
文生图(SD3.5-Medium)组合提示跟随 GenEval 0.67 TDM 0.61 / DPDMD 0.62 / DMD2 0.60 / D-DMD 0.55;教师 CFG=7 0.67 +0.06 对最强少步基线,持平 100 NFE 教师
文生图(SD3.5-Medium)图像文字渲染 OCR 相似度 0.59 TDM 0.57 / D-DMD 0.56 / DMD2 0.53 +0.02 对最强基线
文生图(SD3.5-Medium)人类偏好 HPSv3 11.89 TDM 11.54 / DMD2 11.18;教师 9.87 +0.35 对最强基线,超过教师
文生图(FLUX.2-Klein 4B)组合提示跟随 GenEval 0.80 官方 4 步模型 0.78;教师 CFG=7 0.80 +0.02 对官方 4 步模型,持平教师
文生图(FLUX.2-Klein 4B)文字渲染与多样性 OCR / LPIPS 0.69 / 0.56 官方 4 步模型 0.64 / 0.47 OCR +0.05,LPIPS +0.09(多样性显著提升)
文生视频(33B MiniMax-H3)人类偏好 偏好率(匿名成对评测) 61.33%(对 Fast-H3);45.53%(对 100 NFE 教师) Fast-H3 38.67%(对 VA-TDM);24.67%(对教师) 对教师偏好率从 24.67% 提升到 45.53%

局限与改进

作者承认的局限:A-form 单独只匹配条件矩,OCR 与 HPSv3 低于 B-form(0.55 对 0.57、11.20 对 11.47),说明矩监督丢失分布信息;B-form 在线场方差高、单独难训;critic 梯度直接回传会崩溃(Fig. 3),实践中必须整体 stop-gradient,偏离理论博弈的原始形式,hybrid 损失更像工程补丁。我的观察:超参偏多($\lambda,\lambda_m,\lambda_r,\eta$ 及 warm-up 长度),最优配置在 SD3.5 上调出,跨 backbone 迁移成本未知;人类评测仅在视频任务进行且未报告样本量与置信区间,45.53% 优势的统计显著性存疑;所有实验固定 4 步与 $\eta=1$,1-2 步极端压缩未探索;方法依赖教师速度场处处准确,教师有伪影或 CFG 失配时的退化未讨论;训练需同时跑学生 rollout、教师 ODE warm-up 与在线场回归,算力开销显著高于纯分数蒸馏。

独立分析的弱点

独立分析的弱点及改进方向:(1) 理论与实践脱节——flow-product 更新用非 detached critic 会崩溃(Fig. 3),最终算法是『stop-gradient + 混合』的经验组合,均衡性质不再严格成立;可分析崩溃机理,设计梯度惩罚或信赖域式稳定方案。(2) 超参敏感——$\lambda_m=0.8$、$\lambda=0.7$、$\lambda_r=1$ 均在 SD3.5 上网格搜索得到,换 backbone 未必最优;可引入按两场残差范数动态调 $\lambda_m$ 的自适应加权。(3) transition-aware 条件化只是把转移身份喂给网络,各步分布差异极大时单一 backbone 容量可能不足;可探索按 step 分组的 LoRA 或混合专家头。(4) B-form 高方差只靠 shared backbone 缓解,本质未解;可借鉴控制变量、antithetic sampling 等方差缩减技术。(5) GenEval/OCR 等可验证指标存在针对评测格式过拟合的风险,建议补充更大规模第三方盲测与失败案例分析。

未来方向

作者提出的方向:把 VAD 确立为设计少步蒸馏方法的原则性基础,后续可在『场的选择』与『更新方式』两个正交轴上继续探索新组合。基于本文成果可延伸的方向:(1) 极端压缩——将 VAD 用于 1-2 步生成,B-form 的全分布监督在步数极少时可能更有价值,并可与一致性模型结合;(2) 更多模态——在音频、3D、分子生成等教师速度场可得的场景复用框架,MiniMax-H3 的成功已证明 33B 级视频模型可行;(3) 与偏好优化衔接——critic 输出的速度残差本身是『学生偏离教师』的逐点度量,可作为奖励信号接入 RLHF 或 preference distillation;(4) 理论深化——研究 hybrid loss 的实际收敛点、非 detached critic 崩溃的条件,以及 $\eta<1$ 随机转移下的博弈均衡;(5) 推广到非高斯前向过程与非常规调度,目前 B-form 的分数等价性依赖高斯核假设;(6) 训练效率——warm-up 阶段教师 ODE 多步积分成本高,可探索用现成轨迹数据或蒸馏式 warm-up 替代。

复现评估

复现评估:代码已在 GitHub 开源(https://github.com/Luo-Yihong/VA-TDM),基座 SD3.5-Medium、FLUX.2-Klein-Base 4B 与 MiniMax-H3 均为公开权重;评测全部基于公开工具——GenEval、OCR 相似度、Aesthetic、ImageReward、PickScore、HPSv2/HPSv3 与 DrawBench prompt 集,多样性用成对 LPIPS,管线可完整重建。算力:SD3.5-Medium 量级蒸馏是主流实验室可达规模,估计需数十张 A100/H100 训练数天;MiniMax-H3 是 33B 音视频模型,复现需大规模集群,普通团队不现实。学习率、batch size、warm-up 步数等细节正文未给出,需查代码与附录。总体难度中高:算法公式清晰、关键超参完整披露($\lambda_m=0.8$、$\lambda=0.7$、$\lambda_r=1$、$\eta=1$、$K=4$),主要门槛在算力与视频模型工程细节;在线场回归还带来额外显存,shared backbone 已有所缓解。

相似论文推荐