通过奖励反向传播对齐文本到图像扩散模型

AlignProp

这是我们论文通过奖励反向传播对齐文本到图像扩散模型的官方实现，作者为Mihir Prabhudesai、Anirudh Goyal、Deepak Pathak和Katerina Fragkiadaki。

摘要

文本到图像扩散模型最近在图像生成领域崭露头角，这得益于大规模的无监督或弱监督文本到图像训练数据集。由于弱监督训练，在下游任务中控制它们的行为变得困难，例如最大化人类感知的图像质量、图像-文本对齐或道德图像生成。最近的工作使用普通强化学习来微调扩散模型以适应下游奖励函数，但这种方法以梯度估计的高方差而臭名昭著。在本文中，我们提出了AlignProp，这是一种通过对噪过程的奖励梯度进行端到端反向传播来将扩散模型与下游奖励函数对齐的方法。虽然这种反向传播的简单实现需要大量内存来存储现代文本到图像模型的偏导数，但AlignProp微调低秩适配器权重模块并使用梯度检查点，使其内存使用变得可行。我们测试了AlignProp在微调扩散模型以适应各种目标方面的表现，如图像-文本语义对齐、美学、可压缩性和对象数量的可控性，以及它们的组合。我们展示了AlignProp在更少的训练步骤中实现了更高的奖励，同时概念更简单，使其成为优化扩散模型以适应感兴趣的可微分奖励函数的直接选择。

代码

安装

使用以下命令创建conda环境：

conda create -n alignprop python=3.10
conda activate alignprop
pip install -r requirements.txt

请使用accelerate==0.17.0，其他库依赖可能会有灵活性。

训练代码

Accelerate将自动处理多GPU设置。该代码可以在单个GPU上运行，因为我们会根据CUDA_VISIBLE_DEVICES环境变量中可用的GPU自动处理梯度累积。对于我们的实验，我们使用了4个A100-40GB RAM来运行我们的代码。如果您使用的GPU内存较小，请相应地编辑per_gpu_capacity变量。此外，如果您受到GPU内存的限制，请考虑使用K=1的AlignProp，这将显著减少内存使用。

美学奖励模型

目前我们提前停止代码以防止过拟合，但请根据您的需求随意调整num_epochs变量。

accelerate launch main.py --config config/align_prop.py:aesthetic

如果您受到内存限制，请使用K=1的AlignProp，根据您的内存可用性随意调整trunc_backprop_timestep，使用以下命令。较低的trunc_backprop_timestep值（较高的K值）可以帮助关注更多语义细节：

accelerate launch main.py --config config/align_prop.py:aesthetic_k1

HPSv2奖励模型

accelerate launch main.py --config config/align_prop.py:hps

accelerate launch main.py --config config/align_prop.py:hps_k1

评估和检查点

您可以在这里找到美学奖励函数的检查点，在这里找到Hps-v2奖励函数的检查点。

评估模型检查点，根据配置文件中的resume_from变量。评估包括计算奖励以及将图像存储/上传到本地/wandb。

普通评估

accelerate launch main.py --config config/align_prop.py:evaluate

混合评估

更新resume_from和resume_from_2变量以指定要混合的检查点。将resume_from_2设置为stablediffusion以在resume_from和Stable Diffusion权重之间进行插值。混合系数基于配置文件中可编辑的mixing_coef_1变量。

accelerate launch main.py --config config/align_prop.py:evaluate_soup

致谢

我们的代码库直接建立在DDPO之上。我们要感谢Kevin Black和团队开源他们的代码。

引用

如果您在研究中发现这项工作有用，请引用：

@misc{prabhudesai2023aligning,
      title={Aligning Text-to-Image Diffusion Models with Reward Backpropagation}, 
      author={Mihir Prabhudesai and Anirudh Goyal and Deepak Pathak and Katerina Fragkiadaki},
      year={2023},
      eprint={2310.03739},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}