[论文解读] Directly Fine-Tuning Diffusion Models on Differentiable Rewards

Kevin Clark, Paul Vicol+2 more|arXiv (Cornell University)|Sep 29, 2023

Generative Adversarial Networks and Image Synthesis被引用 7

一句话总结

本文提出直接奖励微调（DRaFT），一种通过扩散采样全过程直接反向传播梯度来微调扩散模型的方法，使用可微奖励函数。通过利用梯度检查点和LoRA，该方法在性能和效率上均表现出色——训练速度比强化学习基线方法快200多倍，且DRaFT-$K$和DRaFT-LV等变体进一步提升了训练稳定性和方差控制。

ABSTRACT

We present Direct Reward Fine-Tuning (DRaFT), a simple and effective method for fine-tuning diffusion models to maximize differentiable reward functions, such as scores from human preference models. We first show that it is possible to backpropagate the reward function gradient through the full sampling procedure, and that doing so achieves strong performance on a variety of rewards, outperforming reinforcement learning-based approaches. We then propose more efficient variants of DRaFT: DRaFT-K, which truncates backpropagation to only the last K steps of sampling, and DRaFT-LV, which obtains lower-variance gradient estimates for the case when K=1. We show that our methods work well for a variety of reward functions and can be used to substantially improve the aesthetic quality of images generated by Stable Diffusion 1.4. Finally, we draw connections between our approach and prior work, providing a unifying perspective on the design space of gradient-based fine-tuning algorithms.

研究动机与目标

• 通过梯度直接优化模型参数，解决预训练扩散模型与期望行为（如美学质量或人类偏好）之间的对齐问题。
• 克服基于强化学习的微调方法在扩散模型对齐中效率低下和样本复杂度高的问题。
• 开发一种高效利用可微奖励函数解析梯度的方法，避免依赖黑箱奖励优化。
• 实现可扩展、高效且可组合的微调，适用于多种奖励函数，包括图像质量、目标检测和对抗样本生成。
• 提供统一的基于梯度的框架，支持多种目标下LoRA权重的插值与组合。

提出的方法

• 提出DRaFT方法，将最终生成图像的奖励梯度反向传播通过整个展开的扩散采样过程，实现扩散模型的端到端微调。
• 使用梯度检查点降低内存开销，采用LoRA仅微调少量参数，保持计算效率。
• 提出DRaFT-$K$，仅将反向传播截断至最后$K$个去噪步骤，以提升训练稳定性并减少梯度爆炸。
• 提出DRaFT-LV，作为DRaFT-1的变体，通过多个噪声样本计算更低方差的梯度估计，提升训练稳定性和速度。
• 将该方法应用于在多种可微奖励（包括LAION美学评分、PickScore和HPSv2）上微调Stable Diffusion 1.4，采用LoRA权重适配。
• 通过混合或缩放在不同奖励目标上训练的LoRA权重，支持模型组合与插值，实现在无需重新训练的情况下进行多目标生成。

实验结果

研究问题

• RQ1 通过完整扩散采样过程的基于梯度的微调，是否能在可微奖励函数上实现优于强化学习的性能？
• RQ2 将反向传播截断至最后$K$步（DRaFT-$K$）与完整反向传播相比，对训练稳定性和收敛性有何影响？
• RQ3 通过多个噪声样本平均（DRaFT-LV）等方差减少技术，是否能显著提升训练效率和收敛速度？
• RQ4 DRaFT在多样化奖励函数（包括图像可压缩性、物体移除和对抗样本生成）上的泛化能力如何？
• RQ5 通过LoRA权重混合是否能实现微调模型的组合或插值，同时在多个目标上保持性能？

主要发现

• 在最大化LAION美学评分时，DRaFT的训练速度比强化学习基线快超过200倍。
• 由于梯度估计方差更低，DRaFT-LV的训练速度约为ReFL（先前的基于梯度的方法）的2倍。
• 当$K=2$时，DRaFT-$K$的性能优于完整反向传播，后者在实践中因梯度爆炸而表现不佳。
• 在PickScore和HPSv2上微调的模型在人类偏好评分上显著提升，且通过LoRA权重缩放有效缓解了过拟合。
• 该方法可有效实现基于不同奖励训练的LoRA权重之间的插值，实现无需重新训练的组合式生成。
• DRaFT在多样化提示下成功提升了图像质量，并支持图像不可压缩性和对抗样本生成等复杂目标。