Large-scale Reinforcement Learning for Diffusion Models

This is the official implementation of our paper Large-scale Reinforcement Learning for Diffusion Models by Yinan Zhang, Eric Tzeng, Yilun Du, and Dmitry Kislyuk.

Text-to-image diffusion models are a class of deep generative models that have demonstrated an impressive capacity for high-quality image generation. However, these models are susceptible to implicit biases that arise from web-scale text-image training pairs and may inaccurately model aspects of images we care about. This can result in suboptimal samples, model bias, and images that do not align with human ethics and preferences. In this paper, we present an effective scalable algorithm to improve diffusion models using Reinforcement Learning (RL) across a diverse set of reward functions, such as human preference, compositionality, and fairness over millions of images. We illustrate how our approach substantially outperforms existing methods for aligning diffusion models with human preferences. We further illustrate how this substantially improves pretrained Stable Diffusion (SD) models, generating samples that are preferred by humans 80.3% of the time over those from the base SD model while simultaneously improving both the composition and diversity of generated samples.

• Model Weights
• Other Reward functions

The hyperparameters and the training setting for human preference fine-tuning can be found in configs/imagereward_train_configs.yaml. Run the following command to train SD2 model with ImageReward:

accelerate launch --multi_gpu ./train.py --config_name image_reward

We use the opensourced ImageReward model for optimizing human preferences. The teaser video on our website is built using the animation script provided by the authors of DDPO.

We would like to thank Kevin Black, DDPO and ImageReward teams for opensourcing their code.

If you find this work useful in your research, please cite:

@misc{zhang2024largescale,
      title={Large-scale Reinforcement Learning for Diffusion Models}, 
      author={Yinan Zhang and Eric Tzeng and Yilun Du and Dmitry Kislyuk},
      year={2024},
      eprint={2401.12244},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}