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

为什么

如果你的奖励函数是可微分的，那么直接将奖励模型的梯度反向传播到扩散模型，比使用像 DDPO 这样的策略梯度算法在样本和计算效率上高得多 (25倍)。AlignProp 通过时间进行完整的反向传播，这使得可以通过奖励反向传播来更新去噪过程的早期步骤。

使用 examples/scripts/alignprop.py 入门

alignprop.py 脚本是一个使用 AlignProp 训练器微调 Stable Diffusion 模型的工作示例。这个示例显式配置了与配置对象 (AlignPropConfig) 相关的整体参数的一个小子集。

注意： 推荐使用一个 A100 GPU 来运行。对于内存较低的设置，可以考虑将 truncated_backprop_rand 设置为 False。在默认设置下，这将执行 K=1 的截断反向传播。

几乎每个配置参数都有一个默认值。用户只需提供一个命令行标志参数即可启动并运行。用户应拥有一个Hugging Face 用户访问令牌，该令牌将在微调后用于将模型上传到 Hugging Face Hub。要运行，请输入以下 bash 命令

python alignprop.py --hf_user_access_token <token>

要获取 `stable_diffusion_tuning.py` 的文档，请运行 `python stable_diffusion_tuning.py --help`

在配置训练器时（除了使用示例脚本的用例外），以下是一些需要注意的事项（代码也会为您检查）：

• 可配置的随机截断范围 (--alignprop_config.truncated_rand_backprop_minmax=(0,50))，第一个数字应大于或等于0，而第二个数字应小于或等于扩散时间步数 (sample_num_steps)
• 可配置的截断反向传播绝对步长 (--alignprop_config.truncated_backprop_timestep=49) 该数字应小于扩散时间步数 (sample_num_steps)，仅当 truncated_backprop_rand 设置为 False 时才起作用。

设置图像日志钩子函数

该函数接收一个包含以下键的字典：

['image', 'prompt', 'prompt_metadata', 'rewards']

其中 `image`、`prompt`、`prompt_metadata`、`rewards` 是批处理的。你可以自由地以任何方式记录日志，推荐使用 `wandb` 或 `tensorboard`。

关键术语

• `rewards` : 奖励/分数是与生成图像相关的数值，是引导强化学习过程的关键
• `prompt` : 提示是用于生成图像的文本
• `prompt_metadata` : 提示元数据是与提示相关联的元数据。当奖励模型包含一个 `FLAVA` 设置时，这个字段将不会为空，此时需要提供与生成图像相关的问题和标准答案（请参见这里：https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/rewards.py#L45）
• `image` : 由 Stable Diffusion 模型生成的图像

下面提供了使用 `wandb` 记录采样图像的示例代码。

# for logging these images to wandb

def image_outputs_hook(image_data, global_step, accelerate_logger):
    # For the sake of this example, we only care about the last batch
    # hence we extract the last element of the list
    result = {}
    images, prompts, rewards = [image_data['images'],image_data['prompts'],image_data['rewards']]
    for i, image in enumerate(images):
        pil = Image.fromarray(
            (image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
        )
        pil = pil.resize((256, 256))
        result[f"{prompts[i]:.25} | {rewards[i]:.2f}"] = [pil]
    accelerate_logger.log_images(
        result,
        step=global_step,
    )

使用微调后的模型

假设您已经完成了所有的训练周期，并将模型推送到了 Hub，您可以按如下方式使用微调后的模型

from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipeline.to("cuda")

pipeline.load_lora_weights('mihirpd/alignprop-trl-aesthetics')

prompts = ["squirrel", "crab", "starfish", "whale","sponge", "plankton"]
results = pipeline(prompts)

for prompt, image in zip(prompts,results.images):
    image.save(f"dump/{prompt}.png")

致谢

这项工作深受此处的仓库和相关论文《通过奖励反向传播对齐文本到图像扩散模型》（作者：Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki）的影响。