Finetuning Stable Diffusion with DDPO via TRL
Hugging Face has integrated Denoising Diffusion Policy Optimization (DDPO) into the trl library via the DDPOTrainer, allowing Stable Diffusion models to be fine-tuned using reinforcement learning (RL) to better align with human preferences and aesthetic standards.
DDPO vs. Reward-Weighted Regression
DDPO improves upon previous RL methods for diffusion models, such as Reward-weighted regression (RWR), by reducing computational overhead and eliminating approximation errors. While RWR reuses the denoising loss function and weights it based on the reward of final samples—effectively ignoring intermediate steps—DDPO treats the entire denoising process as a multistep Markov Decision Process (MDP).
By framing the process as an MDP and using a fixed sampler, DDPO allows the agent policy to be an isotropic Gaussian. This enables the calculation of the exact likelihood of each denoising step rather than relying on the approximate likelihood of the final sample, leading to better performance and the ability to handle more complex objectives.
The DDPO Algorithm and RLHF Workflow
DDPO utilizes a policy gradient method, specifically Proximal Policy Optimization (PPO), to tackle the optimization problem. The primary customization in the DDPO implementation is the trajectory collection portion of the PPO algorithm.
When integrated into a Reinforcement Learning from Human Feedback (RLHF) workflow, the process for aligning diffusion models is streamlined into three steps:
- Pretrained Model: Start with a pretrained Diffusion Model (e.g., Stable Diffusion).
- Reward Model: Gather preference data and train a reward model to serve as a signal.
- DDPO Fine-tuning: Fine-tune the model using DDPO with the reward model as the signal.
In the case of improving image aesthetics, Hugging Face utilized a frozen CLIP model with a trainable regression head trained on the Aesthetic Visual Analysis (AVA) dataset to act as the reward signaller.
Implementation and Training Guidelines
Training Stable Diffusion with DDPO requires significant hardware resources; an NVIDIA A100 GPU is the minimum requirement to avoid out-of-memory (OOM) errors. The implementation is handled via the DDPOTrainer and DDPOConfig classes in the trl library.
Recommended Hyperparameters
For single GPU training, the following hyperparameters are recommended:
| Parameter | Recommended Value |
|---|---|
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 |
Key Findings and Lessons Learned
Experimental results from Hugging Face indicate that DDPO-tuned models generalize well across a wide variety of prompts, even when the training prompt set is minimally sized.
LoRA vs. Full Fine-tuning
- LoRA: Recommended and proven stable across multiple tests.
- Full Fine-tuning (non-LoRA): Can produce more intricate images than LoRA but is significantly more challenging to stabilize. For non-LoRA runs, it is recommended to set the learning rate lower (approximately
1e-5) and setmixed_precisiontoNone.
Current Limitations
The current implementation of DDPOTrainer in the trl library is limited to fine-tuning vanilla Stable Diffusion models. While LoRA is the primary focus and works well, full training is possible but requires more precise hyperparameter tuning.