PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models

📄 arXiv: 2402.08714v2 📥 PDF

作者: Fei Deng, Qifei Wang, Wei Wei, Matthias Grundmann, Tingbo Hou

分类: cs.LG, cs.AI

发布日期: 2024-02-13 (更新: 2024-03-27)

备注: CVPR 2024. Project page: https://fdeng18.github.io/prdp

💡 一句话要点

提出PRDP以解决视觉领域大规模奖励微调问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 奖励微调 扩散模型 强化学习 图像生成 监督学习 大规模训练 计算机视觉

📋 核心要点

  1. 现有基于强化学习的奖励微调方法在视觉领域大规模训练中不稳定，无法处理复杂的未见提示。
  2. 提出近端奖励差异预测（PRDP），通过奖励差异预测（RDP）目标实现对扩散模型的稳定微调。
  3. 实验结果显示，PRDP在小规模训练中与强化学习方法相当，在大规模训练中生成质量显著提升。

📝 摘要（中文）

奖励微调已成为将基础模型与下游目标对齐的有效方法。在语言领域，通过强化学习最大化反映人类偏好的奖励取得了显著成功。然而，在视觉领域，现有的基于强化学习的奖励微调方法在大规模训练中不稳定，无法有效泛化到复杂的未见提示。本文提出了近端奖励差异预测（PRDP），首次在超过10万提示的大规模数据集上实现了对扩散模型的稳定黑箱奖励微调。我们的关键创新是奖励差异预测（RDP）目标，它与强化学习目标具有相同的最优解，同时享有更好的训练稳定性。我们通过理论证明，完美的奖励差异预测模型恰好是强化学习目标的最大化者。实验表明，PRDP在小规模训练中能够匹配成熟的基于强化学习的方法，并在大规模训练中实现了对复杂未见提示的优越生成质量。

🔬 方法详解

问题定义：本文解决的是在视觉领域进行大规模奖励微调时，现有强化学习方法的不稳定性和泛化能力不足的问题。现有方法在复杂的未见提示上表现不佳，限制了其应用。

核心思路：论文提出的PRDP方法通过引入奖励差异预测（RDP）目标，利用监督回归的方式，使扩散模型能够预测生成图像对的奖励差异，从而实现稳定的奖励微调。这样的设计使得模型在训练过程中更具稳定性，避免了传统强化学习方法的波动。

技术框架：PRDP的整体架构包括数据预处理、模型训练和在线优化三个主要阶段。在数据预处理阶段，使用大规模提示数据集；在模型训练阶段，采用RDP目标进行监督学习；在线优化阶段则通过近端更新算法进行稳定优化。

关键创新：PRDP的核心创新在于RDP目标的引入，它与传统的强化学习目标具有相同的最优解，但在训练过程中表现出更好的稳定性。这一创新使得扩散模型能够在大规模数据集上有效训练。

关键设计：在损失函数设计上，RDP目标采用监督回归的方式，模型需要预测生成图像对的奖励差异。此外，近端更新算法的设计确保了模型在训练过程中的稳定性，避免了大规模训练中的常见问题。

📊 实验亮点

实验结果表明，PRDP在小规模训练中能够与成熟的强化学习方法相媲美，而在大规模训练中，PRDP在生成质量上显著优于强化学习方法，尤其是在处理复杂未见提示时，生成质量提升明显，展示了其强大的应用潜力。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉中的图像生成、图像修复和图像编辑等任务。通过实现稳定的奖励微调，PRDP能够提升生成模型在实际应用中的表现，满足用户对高质量图像生成的需求，未来可能在艺术创作、广告设计等领域产生深远影响。

📄 摘要（原文）

Reward finetuning has emerged as a promising approach to aligning foundation models with downstream objectives. Remarkable success has been achieved in the language domain by using reinforcement learning (RL) to maximize rewards that reflect human preference. However, in the vision domain, existing RL-based reward finetuning methods are limited by their instability in large-scale training, rendering them incapable of generalizing to complex, unseen prompts. In this paper, we propose Proximal Reward Difference Prediction (PRDP), enabling stable black-box reward finetuning for diffusion models for the first time on large-scale prompt datasets with over 100K prompts. Our key innovation is the Reward Difference Prediction (RDP) objective that has the same optimal solution as the RL objective while enjoying better training stability. Specifically, the RDP objective is a supervised regression objective that tasks the diffusion model with predicting the reward difference of generated image pairs from their denoising trajectories. We theoretically prove that the diffusion model that obtains perfect reward difference prediction is exactly the maximizer of the RL objective. We further develop an online algorithm with proximal updates to stably optimize the RDP objective. In experiments, we demonstrate that PRDP can match the reward maximization ability of well-established RL-based methods in small-scale training. Furthermore, through large-scale training on text prompts from the Human Preference Dataset v2 and the Pick-a-Pic v1 dataset, PRDP achieves superior generation quality on a diverse set of complex, unseen prompts whereas RL-based methods completely fail.