Directly Fine-Tuning Diffusion Models on Differentiable Rewards

ReFL

📌 核心贡献

提出 DRaFT（Direct Reward Fine-Tuning），通过将可微奖励函数的梯度直接反向传播穿过完整扩散采样过程来微调扩散模型，并提出截断反向传播变体 DRaFT-K 和低方差变体 DRaFT-LV，统一了 ReFL 等方法作为 DRaFT-1 的特殊情况，比 RL 方法快 200 倍以上。

📖 摘要

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.

🔍 关键信息

字段	内容
机构	Google DeepMind
发表	2023-09-29
分类	cs.CV, cs.LG
链接	arXiv

📝 我的笔记

动机与问题

基于 RL 的扩散模型微调方法（如 DDPO）将采样过程视为黑盒，忽略了奖励模型可微的特性。如果奖励函数可微，理论上可以直接通过反向传播来优化，效率远高于 RL 方法的方差估计。然而直接反传面临 T 步 UNet 反向传播的显存消耗挑战。

核心方法

DRaFT 目标函数

直接对采样过程做梯度上升：

$J(\theta )={\mathbb{E}}_{c\sim p(c),x_T\sim \mathcal{N}(0,I)}[r(\text{sample}(\theta ,c,x_T),c)]$

通过 LoRA 参数化（冻结预训练权重，注入低秩矩阵 $h=W_{0}x+BAx$）减少可训练参数量，结合 gradient checkpointing 降低显存。

DRaFT-K：截断反向传播

只对最后 K 步反向传播梯度，前面步骤使用 stop_gradient：

$J_K(\theta )=\mathbb{E}[r({\text{sample}}_K(\theta ,c,\text{sg}[x_{T-K}]),c)]$

关键发现：K=1 时性能最好，因为梯度范数随 K 指数增长导致优化不稳定。当 K=1 时等价于 ReFL（ImageReward 论文中的算法）。

DRaFT-LV：低方差 K=1 变体

对 DRaFT-1 生成的图像 $x_0$ 进行 $n$ 次前向扩散加噪：$x_1={\alpha }_1\cdot \text{sg}[x_0]+{\sigma }_1\cdot \epsilon$，然后平均多次梯度估计以降低方差，约 2 倍快于 DRaFT-1。

统一框架

该论文的核心理论贡献是将 DRaFT、DRaFT-K、DRaFT-LV 和 ReFL 统一为同一算法的不同 stop_gradient 配置，揭示了基于梯度的扩散模型微调方法的设计空间。

关键实验结果

训练效率对比：

方法	相对速度	说明
DRaFT (full)	基准	全采样反向传播
DRaFT-K (K=1)	~200x vs RL	比 DDPO 等 RL 方法快 200 倍以上
DRaFT-LV	~2x vs DRaFT-1	低方差版本

截断步数 K 的 Ablation：

K 值	1	5	10	30	50
性能	最优	良好	良好	下降	最差

梯度范数随 K 指数增长，解释了较大 K 值性能下降的原因。

多样化奖励函数应用：

• 美学评分（LAION Aesthetic）
• 人类偏好（HPSv2, PickScore）
• JPEG 可压缩/不可压缩性
• 目标检测（OWL-ViT）
• 对抗样本生成（ResNet-50）

LoRA 权重组合：

独立训练的 LoRA 权重可线性组合实现多奖励混合，无需重新训练。缩放参数 $\alpha <1$ 可在预训练和微调模型之间平滑插值。

总结

DRaFT 的核心洞察是”可微奖励就应该用梯度来优化”，绕过了 RL 的高方差问题。论文最重要的理论贡献是统一框架，揭示了 ReFL 是 DRaFT-1 的特殊情况。Published at ICLR 2024。

🔗 相关论文

基于/改进自： ImageReward

同方向： DDPO, AlignProp