arXiv:2410.20898cs.CVcs.AI2024-10ICML被引 9

小模型通过强化学习优化,用极短时间超越大模型生成效果。

David and Goliath: Small One-step Model Beats Large Diffusion with Score Post-training

  • 用基于得分的正则化替代传统KL散度,提升训练稳定性和性能。
  • 2.6亿参数单步模型在1024×1024分辨率下,仅用1.88%推理时间胜过50步120亿参数大模型。
  • 适合追求高效、高质图像生成的应用场景,尤其关注推理效率的开发者。

我们提出Diff-Instruct*(DI*),一种数据高效的后训练方法,用于改进单步文本到图像生成模型的人类偏好,且无需图像数据。该方法将对齐问题建模为在线人类反馈强化学习(RLHF),在最大化人类奖励函数的同时,通过与参考扩散过程保持接近进行正则化。不同于传统方法依赖KL散度,我们引入新型通用得分基散度正则化,显著提升性能与训练稳定性。尽管该目标函数理论上不可直接优化,我们推导出一个严格等价的可计算损失函数,能高效求得梯度。我们构建了 extit{DI*-SDXL-1step},一个2.6B参数、1024×1024分辨率的单步模型,基于DMD2对齐SDXL。在Parti提示基准测试中,该模型在ImageReward、PickScore和CLIP分数上均优于50步120亿参数的FLUX-dev模型,且仅使用1.88%的推理时间。结果表明,经过恰当后训练,小单步模型可击败大型多步扩散模型。代码已开源:https://github.com/pkulwj1994/diff_instruct_star。期望本工作推动以人类为中心的机器学习技术发展。

原文摘要 · Abstract (English)

We propose Diff-Instruct* (DI*), a data-efficient post-training approach for one-step text-to-image generative models to improve its human preferences without requiring image data. Our method frames alignment as online reinforcement learning from human feedback (RLHF), which optimizes the one-step model to maximize human reward functions while being regularized to be kept close to a reference diffusion process. Unlike traditional RLHF approaches, which rely on the Kullback-Leibler divergence as the regularization, we introduce a novel general score-based divergence regularization that substantially improves performance as well as post-training stability. Although the general score-based RLHF objective is intractable to optimize, we derive a strictly equivalent tractable loss function in theory that can efficiently compute its \emph{gradient} for optimizations. We introduce \emph{DI*-SDXL-1step}, which is a 2.6B one-step text-to-image model at a resolution of $1024\times 1024$, post-trained from DMD2 w.r.t SDXL. \textbf{Our 2.6B \emph{DI*-SDXL-1step} model outperforms the 50-step 12B FLUX-dev model} in ImageReward, PickScore, and CLIP score on the Parti prompts benchmark while using only 1.88\% of the inference time. This result clearly shows that with proper post-training, the small one-step model is capable of beating huge multi-step diffusion models. Our model is open-sourced at this link: https://github.com/pkulwj1994/diff_instruct_star. We hope our findings can contribute to human-centric machine learning techniques.

图像生成强化学习小模型高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。