用奖励倾斜分布匹配提升4步图像生成质量,更符合人类审美。
Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

- 将奖励学习与分布匹配结合,分两阶段优化生成器。
- 仅用4步推理即在多个指标上超越现有方法。
- 适合追求高效高质图像生成的研究者和开发者。
近期少步扩散蒸馏进展实现了高效的图像生成,但与人类偏好对齐仍具挑战。本文提出奖励倾斜分布匹配蒸馏(RTDMD),一种两阶段框架,统一分布匹配蒸馏与奖励引导强化学习,用于少步流生成器。我们证明最小化到奖励倾斜教师分布的KL散度可自然分解为分布匹配项与奖励最大化项。第一阶段引入环境一致分布匹配蒸馏(AC-DMD),进行子区间分布匹配,并通过一致性正则化增强假分数模型对有限更新下生成器分布变化的追踪能力。第二阶段联合优化两项:针对奖励最大化项,推导出混合策略梯度,结合GRPO风格估计随机中间步骤与直接通过确定性最后一步反向传播奖励;进一步提出步子集GRPO(SubGRPO)以降低方差。在SD3、SD3.5和FLUX.2上的实验表明,RTDMD仅用4步推理即在偏好、美学与构图指标上达到新最佳性能,优于以往少步文生图方法。代码与模型见https://github.com/Harahan/RTDMD。
原文摘要 · Abstract (English)
Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。