透過獎勵反向傳播對齊文字到影像擴散模型

為什麼

如果你的獎勵函式是可微分的，那麼直接將獎勵模型的梯度反向傳播到擴散模型，比使用像 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）的影響。