「DRLX」による「Stable Diffusion」の「RLHF」を試したので、まとめました。

1. DRLX

「DRLX」は、「Cyper」が開発したStable Diffusion用のRLHFライブラリです。

「Carper」では、LLMをRLHFで学習するために「TRLX」を開発しました。そして今度は、diffusersモデルをRLHFで学習するために「DRLX」を開発しました。初期リリースでは、「DDPOトレーナー」とさまざまな「報酬モデル」が提供されています。

現在の制限事項は、次のとおりです。

2. 報酬モデル

「報酬モデル」は、画像生成モデルの強化学習で報酬信号を生成するために使用されます。 通常、画像を取得し、何らかの報酬を返します。報酬を生成する際にプロンプトを使用するものもありますが、必須ではありません。

2-1. Toy Rewards

テスト用の報酬モデルです。

2-2. Aesthetics

美的スコアが高い画像に報酬を与えるモデルです。

2-3. Pickscore (WIP)

「PickAPic」の「PickScore」モデルを使用した報酬モデルです。

3. 学習の実行

Colabでの学習手順は、次のとおりです。

(1) パッケージのインストール。

# パッケージのインストール
!pip install ftfy regex tqdm
!pip install git+https://github.com/openai/CLIP.git
!pip install git+https://github.com/CarperAI/DRLX.git
!pip install datasets

(2) 「ddpo_sd.yml」のダウンロード。
「DRLX」の設定ファイルのサンプルです。

# ddpo_sd.ymlのダウンロード
!wget https://raw.githubusercontent.com/CarperAI/DRLX/main/configs/ddpo_sd.yml

(3) 「ddpo_sd.yml」の編集。
学習時間がかかるため、total_samplesを100分の1 (5.0e+2) に変更します。

train:
  total_samples: 5.0e+4 # 50k

↓

train:
  total_samples: 5.0e+2

(4) 学習の実行。
A100で1時間半かかりました (Accelerateでもっと高速に学習する方法ありそう)。
最後にモデルを保存してます。

from drlx.reward_modelling.aesthetics import Aesthetics
from drlx.pipeline.pickapic_prompts import PickAPicPrompts
from drlx.trainer.ddpo_trainer import DDPOTrainer
from drlx.configs import DRLXConfig

# 学習の実行
pipe = PickAPicPrompts()
config = DRLXConfig.load_yaml("ddpo_sd.yml")
trainer = DDPOTrainer(config)
trainer.train(pipe, Aesthetics())
trainer.save_pretrained("./drlx_model")

「DRLX」で保存されたモデルは、元のパイプラインと互換性があり、他のオブジェクトと同様にロードできます。

・drlx_model

4. 推論の実行

Colabでの推論の実行手順は、次のとおりです。

(1) 推論の実行。

from diffusers import StableDiffusionPipeline

# 画像生成の実行
pipe = StableDiffusionPipeline.from_pretrained("./drlx_model")
prompt = "A mad panda scientist"
image = pipe(prompt).images[0]
image.save("test.jpeg")

美的スコアは効果がわかりにくいので、まずはAverageBlueRewardでテストすればよかった気がする。