论文

通过奖励倾斜分布匹配强化少步生成器

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

摘要

几步扩散 蒸馏 的最新进展已经实现了高效的图像生成,但使这些模型与人类偏好保持一致仍然具有挑战性。我们提出奖励倾斜分布匹配 蒸馏 (RTDMD),这是一个两阶段框架,它将分布匹配 蒸馏 与奖励引导的强化学习统一起来,用于少步流生成器。我们证明,将 KL 散度最小化为奖励倾斜的教师分布可以自然地分解为分布匹配项和奖励最大化项。在第一阶段,我们引入环境一致性分布匹配蒸馏(AC-DMD),它执行子区间分布匹配,并使用一致性正则化器增强假分数目标,以帮助假分数模型在有限更新下跟踪移位生成器分布。在第二阶段,我们联合优化这两项:对于奖励最大化项,我们推导出一个混合策略梯度,它将用于随机中间转换的 GRPO 式估计器与通过确定性最终步骤的直接奖励反向传播相结合,并进一步引入步骤子集 GRPO (SubGRPO) 以减少方差。 SD3、SD3.5 和 FLUX.2 上的实验表明,RTDMD 仅用 4 个推理步骤就在偏好、审美和构图指标方面建立了新的最先进结果,优于之前的几个步骤的文本到图像生成方法。代码和模型可在 https://github.com/Harahan/RTDMD 获取。