最近,集中模型(Diffusion Model)在图象生成领域取得了显著的进展,为图象生成和视频生成任务带来了前所未有的发展机遇。尽管取得了令人印象深刻的结果,集中模型在推理流程中天然存在的多步数迭代去噪特性导致了较高的计算成本。近期出现了一系列集中模型蒸馏算法来减速集中模型的推理流程。这些方式大致可以分为两类:i) 轨迹保持蒸馏;ii) 轨迹重构蒸馏。然而,这两类方式会分别受到动机天花板有限或者或输出域变化这两个问题的限制。
为了解决这些问题,字节跳动技术团队提出了一种名为 Hyper-SD 的轨迹分段一致性模型。Hyper-SD 的开源也得到了Huggingface首席执行官 Clem Delangue的肯定。
该模型是一种新颖的集中模型蒸馏框架,结合了轨迹保持蒸馏和轨迹重构蒸馏两种策略的优点,在压缩去噪步数的同时保持接近无损的功能。与现有的集中模型减速算法相比,该方式取得了卓越的减速动机。经过大量实验和用户评测的验证,Hyper-SD 在 SDXL 和 SD1.5 两种架构上都能在 1 到 8 步生成中兑现 SOTA 级别的图象生成功能。
项目主页:https://hyper-sd.github.io/
论文链接:https://arxiv.org/abs/2404.13686
Huggingface 链接:https://huggingface.co/ByteDance/Hyper-SD
单步生成 Demo 链接:https://huggingface.co/spaces/ByteDance/Hyper-SDXL-1Step-T2I
实时画板 Demo 链接:https://huggingface.co/spaces/ByteDance/Hyper-SD15-Scribble
引言
现有用于集中模型减速的蒸馏方式大致可以分为两大类:轨迹保持蒸馏和轨迹重构蒸馏。轨迹保持蒸馏技术旨在维持集中对应的常微分方程(ODE)的原始轨迹。其原理是通过迫使蒸馏模型和原始模型产生相似的输出来减少推理步骤。然而需要注意的是,尽管能够兑现减速,由于模型容量有限以及训练拟合流程中不可避免的误差,这类方式可能导致生成质量下降。相比之下,轨迹重构方式则直接利用轨迹上的端点或真实图象作为监督的主要来源,忽略了轨迹的中间步骤,能够通过重建更有效的轨迹来减少推理步骤的数量,并在有限的步骤内探索模型的潜力,将其从原始轨迹的约束中解放出来。然而,这通常会导致减速模型与原始模型的输出域不一致,从而得到不理想的结果。
本论文提出了一种结合轨迹保持和重构策略优点的轨迹分段一致性模型(简称 Hyper-SD)。具体而言,该算法首先引入轨迹分段一致性蒸馏,在每个段内强制保持一致性,并逐渐减少段的数量以兑现全时一致性。这一策略解决了由于模型拟合能力不足和推理误差累积导致的一致性模型功能次优的问题。随后,该算法利用人类反应学习(RLHF)来提升模型的生成动机,以弥补减速流程中模型生成动机的丧失,使其更好地适应低步数推理。最后,该算法使用分数蒸馏来增强一步生成功能,并通过统一的 LORA 兑现理想化的全时间步数一致集中模型,在生成动机上取得了卓越的成果。
方式
1. 轨迹分段一致性蒸馏
一致性蒸馏(CD)[24] 和一致性轨迹模型(CTM)[4] 都旨在通过一次性蒸馏将集中模型转换为整个时间步范围 [0,T] 的一致性模型。然而,由于模型拟合能力的限制,这些蒸馏模型往往达不到最优性。受到 CTM 中引入的软一致性目标的启发,我们通过将整个时间步范围 [0, T] 划分为 k 段并逐步执行分段一致模型蒸馏来细化训练流程。
在第一阶段,我们设置 k=8 并使用原始集中模型来初始化 和。起始时间步是从中均匀随机采样的。然后,我们对结束时间步进行采样,其中计算以下:
训练丧失计算以下:
其中通过公式 3 进行计算,表示学生模型的指数滑动平均(EMA)。
随后,我们恢复上一阶段的模型权重并继续训练,逐渐将 k 减少到 [4,2,1]。值得注意的是,k=1 对应于标准 CTM 训练方案。对于距离度量 d,我们采用了对抗性丧失和均方误差 (MSE) 丧失的混合。在实验中,我们观察到,当预测值和目标值接近时(例如,对于 k=8, 4),MSE 丧失更为有效,而对抗性丧失则随着预测和目标值之间的差异增加而变得更加精确(例如,对于 k=2, 1)。因此,我们在整个训练阶段动态增加对抗性丧失的权重并减少 MSE 丧失的权重。此外,我们还集成了噪声扰动机制来增强训练稳定性。以两阶段轨迹分段一致性蒸馏(TSCD)流程为例。以下图所示,我们第一阶段在 和 时间段内执行独立一致性蒸馏 ,然后基于之前的两段一致性蒸馏结果,进行全局一致性轨迹蒸馏。
完整的算法流程以下:
2. 人类反应学习
除了蒸馏之外,我们进一步结合反应学习以提高减速集中模型的功能。具体来说我们通过利用人类审美偏好和现有视觉感知模型的反应来提高减速模型的生成质量。对于审美反应,我们利用 LAION 审美预测器和 ImageReward 中提供的审美偏好奖励模型来引导模型生成更具美感的图象,以下所示:
其中是审美奖励模型,包括 LAION 数据集和 ImageReward 模型的审美预测器,c 是文本提示,与ReLU函数一起作为铰链丧失 。除了来自审美偏好的反应之外,我们注意到嵌入有关图象的丰富先验知识的现有视觉感知模型也可以作为良好的反应提供者。根据经验,我们发现实例宰割模型可以指导模型生成结构合理的物体。具体来说,我们首先将潜在空间中图象上的噪声集中到,之后,类似于 ImageReward,我们执行迭代去噪,直到 特定时间步并直接预测。随后,我们利用感知实例宰割模型通过检查真实图象实例宰割标注与去噪图象的实例宰割预测结果之间的差异来评估结构生成的功能,以下所示:
其中是实例宰割模型(例如 SOLO)。实例宰割模型可以更准确地捕获生成图象的结构缺陷并提供更有针对性的反应信号。值得注意的是,除了实例宰割模型之外,其他感知模型也适用。这些感知模型可以作为主观审美的补充反应,更多地关注客观生成质量。因此,我们用反应信号优化集中模型可以定义为:
3. 一步生成强化
由于一致性丧失的固有限制,一致性模型框架内的一步生成并不理想。正如 CM 中分析的那样,一致性蒸馏模型在引导位置处的轨迹端点方面表现出卓越的准确性。因此,分数蒸馏是一种合适且有效的方式来进一步提升我们的 TSCD 模型的一步生成动机。具体来说,我们通过优化的分布匹配蒸馏(DMD)技术来推进一步生成。DMD 通过利用两个不同的评分函数来增强模型的输出:来自教师模型分布和来自假模型的。我们将均方误差 (MSE) 丧失与基于分数的蒸馏结合起来,以提高训练稳定性。在这个流程中,前面提到的人类反应学习技术也被集成进来,用来微调我们的模型以有效地生成具有保真度的图象。
通过集成这些策略,我们的方式不仅能够兑现在 SD1.5 和 SDXL 上都兑现卓越的低步数推理动机(并且无需 Classifier-Guidance),同时能够兑现理想的全局一致性模型,无需针对每个特定的步数训练 UNet 或者 LoRA 兑现统一的低步数推理模型。
实验
在 SD1.5 和 SDXL 上和目前现有的各种减速算法的定量比较,可以看到 Hyper-SD 显著优于当前最先进的方式
此外,Hyper-SD 能够用一个模型来兑现各种不同低步数的推理,上面的定量指标也显示了我们方式在使用统一模型推理时的动机。
在 SD1.5 和 SDXL 上的减速动机可视化直观地展示了 Hyper-SD 在集中模型推理减速上的优越性。
大量的 User-Study 也表明 Hyper-SD 相较于现有的各种减速算法的优越性。
Hyper-SD 训练得到的减速 LoRA 能够很好地兼容不同的风格的文生图底模。
同时,Hyper-SD 的 LoRA 也能适配现有的 ControlNet,兑现低步数下高质量的可控图象生成。
总结
论文提出了 Hyper-SD,一个统一的集中模型减速框架,可以显著提升集中模型的在低步数情况下的生成能力,兑现基于 SDXL 和 SD15 的新 SOTA 功能。该方式通过采用轨迹分段一致性蒸馏,增强了蒸馏流程中的轨迹保存能力,兑现接近原始模型的生成动机。然后,通过进一步利用人类反应学习和变分分数蒸馏提升模型在极端低步数下的潜力,从而产生了更优化、更高效的模型生成动机。论文还开源了用于 SDXL 和 SD15 从 1 到 8 步推理的 Lora 插件,以及专用的一步 SDXL 模型,旨在进一步推动生成式 AI 社区的发展。