这篇论文提出DualHeads方法,旨在解决少步蒸馏扩散模型难以对齐任务需求和人类偏好的问题。现有面向多步扩散模型的对齐方法依赖可微奖励或分数场与随机轨迹,而蒸馏生成器往往不具备这些条件。作者提出仅用外部奖励评估和参考样本即可适配少步生成器,具体做法是在冻结的扩散骨干上添加两个轻量二分类头:奖励头学习同一提示下哪些样本得分更高,参考头区分参考图像与生成图像,生成器则最大化二者logits的加权和。当头最优时,该目标等价于对学习到的奖励做KL正则化最大化,头权重之比决定正则强度。实验表明,在规则型和模型型奖励下,少步DualHeads生成器的任务奖励接近40步强化学习方法,同时保持图像质量,且不增加额外推理成本。
| DualHeads | 本文提出的方法,在冻结扩散模型上添加奖励头和参考头两个轻量分类器,用于少步生成器的奖励对齐。 |
| 少步扩散模型 | 经过蒸馏、能在少量采样步骤内生成高质量图像的扩散模型。 |
| 奖励倾斜生成 | 通过奖励信号调整生成模型,使其输出更符合任务要求或人类偏好的生成范式。 |
| KL正则化最大化 | 在最大化奖励的同时,用KL散度约束生成分布不过度偏离参考分布的目标函数。 |
| 冻结扩散骨干网络 | 参数固定不参与训练的预训练扩散模型主体,仅在其上训练轻量级头。 |
📱 每天一份 AI 前沿日报
关注公众号,每天 09:00 推送 · 不错过任何重磅