[控场AI]
· 4 分钟阅读· 2,414 字

DMD蒸馏新突破:批评者模型需与教师模型同步缩放

DMD蒸馏新突破:批评者模型需与教师模型同步缩放

DMD蒸馏中批评者需与教师同步扩容并以生成步为条件,方能稳定提升少步生成质量。

本文系统研究了分布匹配蒸馏(DMD)中批评者的容量与信息需求。控制实验发现,当教师模型增大而批评者规模固定时,学生质量反而下降甚至训练崩溃,说明单纯放大教师并不够。为此作者提出联合缩放策略,让批评者与教师同步扩大,从而防止崩溃并使学生真正受益于更大教师。在信息层面,传统批评者建模与步数无关的混合分布,信息损失严重;作者提出步感知批评者STAC,以学生生成步为条件估计步特定分数,提供更精准的训练信号。两者结合在类别条件ImageNet、SDXL与FLUX.2文生图、Wan2.1文生视频上均取得一致改进,表明批评者的容量与信息设计是DMD蒸馏的通用瓶颈。

分布匹配蒸馏(Distribution Matching Distillation, DMD)一直是少步生成模型的关键技术,它通过教师模型的评分与在线批评者(critic)对学生分布评分的差异来训练学生模型。但批评者到底需要什么样的容量和信息才能提供可靠的训练信号?一篇最新研究给出了系统性答案:批评者也需要“与时俱进”。

DMD蒸馏研究示意图

批评者容量不足会导致训练崩溃

研究者在ImageNet上使用七种不同规模的SiT模型进行了控制实验,发现了一个反直觉的现象:当教师模型增大而批评者保持小规模固定时,学生模型的质量反而下降,甚至出现训练崩溃。这意味着单纯增大教师模型并不能自动带来更好的蒸馏效果——批评者如果“跟不上”教师的表达能力,整个训练过程就会失衡。

为此,团队提出了**联合缩放(joint scaling)**策略:让批评者与教师同步扩大规模。实验表明,这种策略能够防止崩溃,并使学生模型真正从更大的教师中获益。这一发现挑战了此前“教师越大越好”的简单假设,强调了蒸馏系统中各组件容量匹配的重要性。

SiT 模型简介

SiT(Scalable Interpolant Transformers)是一种基于 Transformer 骨干的可扩展扩散/流匹配生成模型,它把扩散模型的训练统一到插值(interpolant)框架下,将前向加噪过程视为数据与噪声之间的插值路径,并允许灵活选择插值方式、采样器和时间步加权。由于 SiT 提供了一系列不同参数规模的模型(从较小到较大),且架构与训练配方统一,因此非常适合做“控制变量”式的缩放实验。本文正是借助 SiT 的多种规模,在保持其他条件不变的前提下单独改变教师与批评者的容量,从而清晰地分离出“教师变大”与“批评者是否同步变大”这两个因素对蒸馏效果的影响。

步感知批评者:提供更精准的分数估计

除了容量问题,批评者需要什么样的信息同样关键。传统DMD训练中,批评者被要求建模一个与步数无关的学生输出混合分布,这相当于让批评者“盲猜”不同生成步的得分。研究者认为这种做法信息损失严重,于是提出了步感知批评者(STAC, Step-Aware Critic):让批评者以学生的生成步为条件,估计步特定的分数。

这种设计让批评者能够区分不同步数下的分布特征,从而提供更精确的训练信号。实验显示,STAC与联合缩放结合后,在多个任务上均带来了一致的改进,包括类别条件ImageNet生成、基于SDXL和FLUX.2的文本到图像生成,以及基于Wan2.1的文本到视频生成。

分数匹配与评分函数(score function)背景

在扩散模型与DMD这类蒸馏方法中,“评分”(score)指的是对数概率密度对输入的梯度,即 ∇ₓ log p(x)。扩散模型的反向生成过程本质上是在每一步沿着当前噪声水平对应的评分方向去噪,因此学生模型要学会教师在各噪声水平上的评分行为。DMD 不直接回归教师的输出,而是用一个可学习的批评者去估计“学生分布”的评分,再用教师评分与该估计之差作为梯度信号更新学生,这属于分数蒸馏(score distillation)的思路。批评者的作用相当于一个随训练在线更新的判别器:它越能准确刻画学生在各噪声水平下的分布,梯度信号就越接近真实分布匹配的方向。这也解释了为什么批评者的容量与条件信息会直接决定蒸馏成败——它估计的是分布层面的量,而非单个样本的回归目标。

多任务验证:从图像到视频的普适性

该研究的验证范围覆盖了三大类生成任务,体现了方法的通用性。在类别条件ImageNet生成中,联合缩放和STAC的组合显著提升了少步蒸馏的效果;在文本到图像任务中,SDXL和FLUX.2两个主流模型均受益;在文本到视频任务中,Wan2.1的实验同样证实了方法的有效性。这种跨模态的一致性表明,批评者的容量与信息设计是DMD蒸馏的通用瓶颈,而非特定任务的偶然现象。

SDXL、FLUX.2 与 Wan2.1 简介

SDXL 是 Stable Diffusion 系列的扩大版文生图模型,采用更大的 UNet 骨干与双文本编码器,是开源文生图的代表性基线;FLUX.2 属于基于流匹配(flow matching)与 Transformer 架构的新一代文生图模型,代表了当前高质量开源生成模型的技术路线;Wan2.1 则是面向文本到视频生成的开源模型系列,需要在时间维度上保持一致性,蒸馏难度显著高于图像。将 DMD 蒸馏分别应用于这三类模型,意味着方法要同时面对 UNet 与 Transformer 两种骨干、扩散与流匹配两种训练范式、以及图像与视频两种数据模态。跨这些差异取得一致改进,是本文主张“批评者容量与信息设计是通用瓶颈”的关键证据。

对蒸馏实践的启示

这项研究给从事少步生成模型研发的团队提了个醒:优化蒸馏流程时,不能只盯着教师模型或学生架构,批评者的设计同样需要精心考量。联合缩放提供了一种简单可行的容量匹配方案,而STAC则从信息层面改进了训练信号的 quality。两者叠加,有望成为未来DMD蒸馏的标准配置。

对于资源有限的团队,联合缩放意味着在扩大教师模型时,必须同步为批评者分配相应的计算预算,否则可能适得其反。而STAC的实现复杂度不高,却能在不增加太多开销的情况下带来稳定收益,值得优先尝试。

总结

DMD蒸馏的性能不仅取决于教师模型的大小,更取决于批评者能否“配得上”教师。联合缩放与步感知批评者STAC从容量和信息两个维度解决了这一匹配问题,在图像和视频生成任务上均验证了有效性。这项工作为少步生成模型的训练提供了新的设计准则。

#DMD#分布匹配蒸馏#少步生成模型#批评者模型#联合缩放#STAC#扩散模型蒸馏#生成模型
分享:

相关推荐