通过 TRL 使用 DDPO 微调 Stable Diffusion
Hugging Face 已通过 DDPOTrainer 将去噪扩散策略优化 (Denoising Diffusion Policy Optimization, DDPO) 集成到了 trl 库中，允许使用强化学习 (RL) 对 Stable Diffusion 模型进行微调，从而更好地符合人类偏好和审美标准。
DDPO vs. Reward-Weighted Regression
DDPO 通过减少计算开销和消除近似误差，改进了以往的扩散模型强化学习方法，例如奖励加权回归 (Reward-weighted regression, RWR)。RWR 复用去噪损失函数并根据最终样本的奖励进行加权——实际上忽略了中间步骤——而 DDPO 将整个去噪过程视为一个多步马尔可夫决策过程 (MDP)。
通过将该过程构建为 MDP 并使用固定采样器，DDPO 允许智能体策略为各向同性高斯分布。这使得计算每个去噪步骤的精确似然成为可能，而不是依赖于最终样本的近似似然，从而带来更好的性能并能够处理更复杂的目标。
DDPO 算法与 RLHF 工作流
DDPO 利用策略梯度方法，特别是近端策略优化 (Proximal Policy Optimization, PPO) 来解决优化问题。DDPO 实现中的主要定制部分是 PPO 算法的轨迹收集部分。
当集成到人类反馈强化学习 (RLHF) 工作流中时，对齐扩散模型的流程被简化为三个步骤：
- 预训练模型：从预训练的扩散模型（例如 Stable Diffusion）开始。
- 奖励模型：收集偏好数据并训练一个奖励模型作为信号。
- DDPO 微调：使用 DDPO 并以奖励模型作为信号来微调模型。
在提高图像美感的情况下，Hugging Face 使用了一个冻结的 CLIP 模型，并配有一个在 Aesthetic Visual Analysis (AVA) 数据集上训练的可训练回归头，作为奖励信号器。
实现与训练指南
使用 DDPO 训练 Stable Diffusion 需要大量的硬件资源；为了避免显存溢出 (OOM) 错误，最低配置要求是 NVIDIA A100 GPU。实现是通过 trl 库中的 DDPOTrainer 和 DDPOConfig 类来完成的。
推荐超参数
对于单 GPU 训练，建议使用以下超参数：
| 参数 | 推荐值 |
|---|---|
num_epochs |
200 |
train_batch_size |
3 |
sample_batch_size |
6 |
gradient_accumulation_steps |
1 |
sample_num_steps |
50 |
sample_num_batches_per_epoch |
4 |
per_prompt_stat_tracking |
True |
per_prompt_stat_tracking_buffer_size |
32 |
mixed_precision |
True |
train_learning_rate |
3e-4 |
主要发现与经验教训
来自 Hugging Face 的实验结果表明，即使训练提示词集规模很小，经过 DDPO 微调的模型也能在广泛的提示词上表现出良好的泛化能力。
LoRA vs. 全量微调
- LoRA：推荐使用，并在多次测试中被证明是稳定的。
- 全量微调 (non-LoRA)：可以产生比 LoRA 更复杂的图像，但稳定性的挑战要大得多。对于非 LoRA 运行，建议设置较低的学习率（约
1e-5）并将mixed_precision设置为None。
当前局限性
目前 trl 库中 DDPOTrainer 的实现仅限于微调原生的 Stable Diffusion 模型。虽然 LoRA 是主要关注点且效果良好，但全量训练是可能的，只是需要更精确的超参数调优。