[控场AI]
· 5 分钟阅读· 2,721 字

SMD:简化在线蒸馏,4步生成赶超多步扩散模型

SMD:简化在线蒸馏,4步生成赶超多步扩散模型

SMD 仅用在线速度回归加 Pivot 残差代理,实现 4 步少步扩散生成。

Simplified Matching Distillation(SMD)针对扩散与流匹配模型采样成本高的问题,提出一种极简的少步蒸馏方案:摒弃现有方法普遍依赖的辅助 score 模型、对抗训练和特殊网络结构,仅靠在线速度回归即可支撑 4-NFE 生成。其核心设计包含两个部分:一是 Pivot 残差代理,在学生预测粗时间区间均值的同时建模区间内部的残差波动,且只在训练阶段存在、不改变部署时的生成器结构;二是排斥弱教师预测的轻量机制,避免学生继承教师不可靠区域的输出。实验覆盖图像、3D、视频、音视频生成,规模最高 33B 参数,4-NFE 下超过强少步基线,部分任务甚至超越多步教师,展现了跨模态、跨架构的通用性。

扩散模型和流匹配模型在图像、视频、3D 内容生成上已经证明了自己的能力,但采样成本始终是落地时绕不开的痛点。少步蒸馏(few-step distillation)试图把几十步的采样压缩到 4 步以内,代价通常是引入辅助 score 模型、对抗训练或特殊网络结构——这些附加组件让训练流程变重,也限制了方法的通用性。

最近一篇工作提出了 Simplified Matching Distillation(SMD),核心主张相当克制:仅靠简单的在线速度回归(on-policy velocity regression),就足以支撑 4-NFE 的生成质量。

rss source: Simplifying and Improving On-Policy Distillation for Few-Step Diffusion

摒弃辅助模块的蒸馏思路

现有少步蒸馏方法大致分为两类。一类依赖额外的 score 模型来估计教师分布,另一类引入对抗损失来拉近学生与教师的输出分布。两者都能提升效果,但代价是训练管线复杂、超参敏感,并且往往绑定特定架构。

SMD 的出发点是回到最朴素的回归目标:让学生模型在粗粒度时间步上预测教师的速度场均值。这个做法本身并不新,但作者发现,只要处理得当,简单的在线速度回归已经是一个很强的基线。

关键问题在于,教师的真实速度在每个粗时间区间内并非恒定,直接用区间均值做监督会丢失局部变化信息。

少步蒸馏与 NFE

NFE(Number of Function Evaluations)指生成一张图像(或一段视频、一个 3D 资产)时,模型需要调用网络前向计算的次数。扩散模型的标准采样需要 20 到 50 次前向,因为每一步只沿着学到的速度场前进一小段,步长必须足够小才能保证轨迹稳定;流匹配模型同理,虽然轨迹更接近直线,但通常也需要十几到几十次积分。少步蒸馏(few-step distillation)的目标就是把 NFE 从几十压缩到 4 步甚至 1 步,让推理成本下降一个数量级。这里的难点在于,学生网络必须在极粗的时间网格上直接预测从某一时刻跳到下一时刻所需的位移,而教师在这段区间内的速度并不是恒定的——这正是 SMD 需要引入 Pivot 残差代理的根本原因。

Pivot:训练时的残差代理

为了在粗粒度转移中保留教师速度的变化,作者引入了一个名为 Pivot 的残差代理。它的作用是把局部修正从学生的区间均值预测中分离出来:学生负责预测整个区间的平均速度,Pivot 则负责建模区间内部的残差波动。

这个代理只在训练阶段存在,部署时完全不改变生成器的结构。换句话说,推理时的计算图和普通少步模型没有任何区别,没有额外分支、没有额外参数量。这也是 SMD 与许多蒸馏方法的重要差异——它把复杂度留在了训练端,把简洁留给了推理端。

速度场、在线回归与 "On-Policy"

速度场 v(x,t) 是流匹配/扩散模型的核心量:在状态 x、时刻 t 处,它指出数据应该往哪个方向、以多快的速度移动。训练学生的常见做法是离线回归——先用教师预先采一批轨道,再让学生去拟合这些固定样本,但学生会访问到教师从未走过的状态,导致误差累积。在线(on-policy)回归则不同:每步让学生自己从噪声出发采样出一条轨迹,一路上记录自己到达的状态,再让教师在这些学生真实访问到的状态上给出速度作为监督信号。这样监督分布和学生自身的推理分布是对齐的,避免了训练与推理之间的"分布漂移"。SMD 采用的正是这种在线方式,而 "Pivot 残差代理" 就是在这一框架下额外引入的、用于吸收区间内部速度波动的辅助预测头,它只在训练时被调用,学生主干在部署时不会多出任何计算。

排斥弱教师预测的轻量机制

除了拟合教师速度,SMD 还加入了一个反向约束:让学生远离教师较弱的预测。

这个设计的直觉是,教师模型在不同时间步上的预测质量并不均匀,某些区域的输出本身就不够可靠。如果学生只是一味拟合,可能会继承这些不理想的模式。通过施加排斥信号,学生被推向教师预测中更可靠的那一侧,从而避免落入不良模式。

作者将这一机制描述为轻量级的,因为它不需要对抗训练中常见的判别器,也不需要在线 score 估计,只是对训练目标做了一项额外约束。

为什么需要"排斥"而不是"对齐"

传统的知识蒸馏假设教师给出的目标总是可靠的,学生只需尽力逼近即可。但在少步场景下,教师在某些时间段(尤其是 t 接近 0 或接近 1 的端点)的预测本身噪声很大,如果学生无条件拟合,反而会把这种不确定性固化下来。SMD 的做法是在损失中加一项排斥项:当教师预测的置信度较低、或方差较大时,让学生远离这一预测,从而把学习信号集中在教师更可靠的时间区域。这一思路与对比学习中的负样本、以及一些能量模型里的"把概率质量推离低质量区域"有相通之处。相比对抗训练需要一个额外的判别器去在线估计真假,这种基于教师自身置信度的排斥不需要任何辅助网络,训练开销几乎为零,因而被称为轻量机制。

跨模态的通用性验证

SMD 的实验覆盖范围相当广:图像生成与编辑、3D 生成、视频合成、音视频联合生成,模型规模最高扩展到 33B 参数。

在 4-NFE 的设置下,SMD 在多个基准上超过了现有的强少步基线,甚至在部分任务上超越了多步教师模型。这意味着学生模型用 4 步采样,就能达到教师用更多步数才能取得的生成质量。

方法还支持跨模型蒸馏、多教师蒸馏以及因果视频蒸馏,说明它并不绑定某一种教师架构或某一种模态。

为什么这件事值得关注

少步蒸馏的工程价值在于推理成本。当生成模型进入产品化阶段,4 步和 20 步之间的差距直接体现在延迟和算力账单上。SMD 的意义不只是又一个蒸馏方法,而是它展示了一条更简洁的路径:不需要对抗训练,不需要辅助 score 模型,纯粹基于回归就能达到有竞争力的少步生成效果。

对于需要部署扩散模型生成能力的团队来说,训练端可以复杂,推理端必须简单——SMD 在这一点上的取舍,比单纯的指标提升更值得参考。

论文标题为《Simplifying and Improving On-Policy Distillation for Few-Step Diffusion》,核心贡献可以概括为 Pivot 残差代理与排斥机制的联合设计,以及由此带来的跨模态、跨规模的可扩展性。

#少步蒸馏#扩散模型#在线蒸馏#4-NFE生成#SMD#速度回归#模型加速#few-step diffusion
分享:

相关推荐