[论文解读] Aligning Text-to-Image Diffusion Models with Reward Backpropagation

Mihir Prabhudesai, Anirudh Goyal+2 more|arXiv (Cornell University)|Oct 5, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

本文提出 AlignProp，一种通过在去噪过程中对可微分奖励梯度进行端到端反向传播，微调文本到图像扩散模型的方法。通过使用低秩适配器（LoRA）和梯度检查点技术，AlignProp 在更少的训练步骤下实现了比强化学习基线更高的奖励性能，从而实现了对图像-文本对齐、美学性、可控性等多种目标的高效优化。

ABSTRACT

Text-to-image diffusion models have recently emerged at the forefront of image generation, powered by very large-scale unsupervised or weakly supervised text-to-image training datasets. Due to their unsupervised training, controlling their behavior in downstream tasks, such as maximizing human-perceived image quality, image-text alignment, or ethical image generation, is difficult. Recent works finetune diffusion models to downstream reward functions using vanilla reinforcement learning, notorious for the high variance of the gradient estimators. In this paper, we propose AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient through the denoising process. While naive implementation of such backpropagation would require prohibitive memory resources for storing the partial derivatives of modern text-to-image models, AlignProp finetunes low-rank adapter weight modules and uses gradient checkpointing, to render its memory usage viable. We test AlignProp in finetuning diffusion models to various objectives, such as image-text semantic alignment, aesthetics, compressibility and controllability of the number of objects present, as well as their combinations. We show AlignProp achieves higher rewards in fewer training steps than alternatives, while being conceptually simpler, making it a straightforward choice for optimizing diffusion models for differentiable reward functions of interest. Code and Visualization results are available at https://align-prop.github.io/.

研究动机与目标

• 为解决预训练文本到图像扩散模型与下游奖励函数（如图像质量、公平性、图文对齐）对齐的挑战。
• 克服扩散模型中长去噪链端到端反向传播带来的极高显存开销问题。
• 开发一种方法，实现通过奖励梯度对扩散模型进行直接、可微分的优化，避免强化学习中高方差梯度的问题。
• 通过 LoRA 适配器的参数平均实现多个奖励函数之间的插值，实现多目标对齐。

提出的方法

• 将去噪过程视为一个可微分的循环策略，将噪声和文本提示映射为图像，从而实现从奖励到模型权重的梯度流动。
• 使用低秩适配器（LoRA）模块，仅微调一小部分参数，与全量微调相比显著降低显存使用。
• 采用梯度检查点技术，按需计算部分导数而非存储它们，避免产生 TB 级别的显存需求。
• 应用截断时间反向传播（BPTT），固定上下文长度 K，以在梯度质量与显存效率之间取得平衡。
• 通过平均在不同奖励函数上微调的模型所得到的 LoRA 权重，实现模型插值，支持多目标生成。
• 使用可微分奖励函数（如基于 CLIP 的图文对齐、美学得分、JPEG 压缩性）引导优化过程。

实验结果

研究问题

• RQ1 能否使扩散采样过程中端到端反向传播的显存开销足够低，以在大规模文本到图像模型中实际应用？
• RQ2 AlignProp 在样本效率和最终奖励性能方面与强化学习基线相比表现如何？
• RQ3 通过组合分别训练的模型所得到的 LoRA 适配器，AlignProp 是否能有效实现多个奖励函数之间的插值？
• RQ4 使用低秩适配器和梯度检查点技术是否能在降低显存与计算成本的同时保持性能？

主要发现

• AlignProp 在 HPSv2 基准测试中平均奖励达到 3.30，显著高于 ReFL 和 DDPO 等基线方法（K=1 时为 2.91）。
• 该方法显著减少了达到高奖励性能所需的训练步数，展现出优于强化学习方法的样本效率。
• 将美学与压缩性奖励微调所得的 LoRA 权重进行插值，可得到一个混合模型，在混合系数 α=0.5 时获得最高总体奖励。
• 消融实验表明，LoRA 微调优于全连接 U-Net 微调；而 EDICT 层虽降低显存使用，但会损害性能并增加推理时间。
• 截断反向传播中 K=10 的性能优于 K=1，表明更长的上下文能提升梯度质量，且不会带来过高的显存开销。
• 该方法能有效实现对复杂多目标奖励的对齐，例如通过参数插值同时提升图文对齐与美学质量。