Diffusion-DPO（Diffusion-Direct Preference Optimization）の学習方法

カテゴリ：deeplearning

Diffusion-DPO Diffusion Model Alignment Using Direct Preference Optimization は SD3 でも使われた強化学習手法。SD3 ではランク 128 の LoRA として作成している。

ただし、画風や新しい概念の学習のような一般的なタスクは SFT（教師ありファインチューニング）が適している。「○○がうまく描けない」というニーズは○○が言語化できているので SFT を使うべき。SFT はデータセットを用意しやすいし学習負荷も低い。

Diffusion-DPO を使うケース

言語化が困難だが描いてほしくないもの（暴力的・性的表現）がある
ユーザーの選好データを持っている
細部の表現の崩れを改善したい

SFT（教師ありファインチューニング）と Diffusion-DPO との違い

手法	必要なデータ	学習１回に必要な推論回数
SFT	教師画像キャプション	１回
Diffusion-DPO	参照モデルで生成したベスト画像参照モデルで生成したワースト画像キャプション参照モデル	４回
DDPO	強化したい画像のキャプション	画像生成に必要な推論回数＋ VLM 等の推論コスト

参照モデルは強化学習前のベースモデル。

Diffusion-DPO は単純にモデルを２つロードしなければならず、学習１回に必要な推論回数は４回（学習モデルと参照モデルとがそれぞれベスト画像とワースト画像とに対しデノイズを行う）。

教師画像の生成

Diffusion-DPO の学習に使う画像はすべて参照モデルで生成したものを使う。つまり以下のような手順になる。

よく生成するプロンプトのリストを作成する
各プロンプトで複数枚画像を生成し、各プロンプトに対するベストとワーストとをそれぞれ１枚選択して学習画像とする

最終的にプロンプトのリスト x 2 枚の学習画像が必要になる。

SFT 画像を使う場合

ベスト画像をすべて SFT 画像にすると分布ミスマッチによる過学習をおこし、モデルが崩壊する。入れるとしても全データセットの 10～20% にしておく。

ワースト画像の作成

ワースト画像といっても、キャプションからかけ離れている画像や著しく品質の悪い画像は使わない方がいい。

DPO Direct Preference Optimization

Diffusion-DPO の前に DPO の確認をしておく。

報酬モデリング

キャプション $c$ から生成された画像 $x_0$ を人がどのくらい好んでいるかを推定するのは難しい。なぜなら潜在報酬モデル $\mathit{\text{r}}(c, x_0)$ を知ることができないからだ。

手に入れられる情報は画像の選好情報 $x_0^w \succ x_0^l|c$ で好ましい画像を $x_0^w$、好ましくない画像を $ x_0^l$ で表示する。

Bradley-Terry モデルは人の選好を以下のように定義する：

\[ p_{\mathrm{BT}}(x_0^w \succ x_0^l|c) = \sigma (r(c, x_0^w) - r(c,x_0^l)) \]

$\sigma$ はシグモイド関数。$r(c, x_0)$ はニューラルネットワーク $\phi$ でパラメータ化され、二値分類のための最大尤度訓練で推定される：

\[ L_{\mathrm{BT}}(\phi) = - \mathbb{E}_{c, x_0^w, x_0^l} \left [ \mathrm{log} \; \sigma \left (r_\phi (c,x_0^w) - r_\phi(c,x_0^l) \right ) \right ] \]

RLHF Reinforce Learning from Human Feedback

RLHF の目的はキャプション c で条件付けされた条件付き分布 $p_\theta(x_0|c)$ を最適化すること。目的関数は潜在報酬モデル $r(c,x_0)$ で、正則化はリファレンス分布 $p_{\mathrm{ref}}$ を使った KL ダイバージェンスを利用する：

\[ \underset{p_\theta}{ \mathrm{max}} \mathbb{E}_{c \sim \mathcal{D}_c, x_0 \sim p_\theta(x_0|c)}[r(c,x_0)] - \beta \mathbb{D}_{ \mathrm{KL}}[p_\theta(x_0|c)||p_{\mathrm{ref}}(x_0|c)] \]

$\beta$ は正則化を制御するハイパーパラメータ。

DPO の目的関数

唯一の全局最適解 $p_\theta^*$ 以下のようになる：

\[ \begin{split} p_\theta^*(x_0|c) &= \dfrac{p_{\mathrm{ref}}(x_0|c)\mathrm{exp}(\dfrac{r(c,x_0)}{\beta})}{Z(c)} \\ Z(c) &= \sum_{x_0} p_{\mathrm{ref}}(x_0|c)\mathrm{exp}(\dfrac{r(c,x_0)}{\beta}) \end{split} \]

$Z(c)$ は分配関数。上式の報酬関数は以下のように書き換えられる：

\[ r(c,x_0) = \beta \; \mathrm{log} \dfrac{p_\theta^*(x_0|c)}{p_{\mathrm{ref}(x_0|c)}} + \beta \; \mathrm{log} Z(c) \]

$L_{\mathrm{BT}}(\phi)$ の式に代入すると、$\beta \; \mathrm{log} Z(c)$ は引き算で消え、DPO の目的関数は以下のようになる：

\[ L_{\mathrm{DPO}}(\theta) = -\mathbb{E}_{c, x_0^w, x_0^l}\left [ \mathrm{log} \; \sigma \left ( \beta \; \mathrm{log} \dfrac{p_\theta(x_0^w|c)}{p_{\mathrm{ref}}(x_0^w|c)} - \beta \; \mathrm{log} \dfrac{p_\theta(x_0^l|c)}{p_{\mathrm{ref}}(x_0^l|c)} \right ) \right ] \]

DPO を拡散モデルに適用する

学習に使うデータセット $\mathcal{D} = {(c, x_0^w, x_0^l)}$ は、キャプション c とモデル $p_{\mathrm{ref}}$ を使用して生成されたものを使う。

まず解決しなければならないのは、$p_\theta (x_0|c)$ は計算不可能（not tractable）という問題だ。$p_\theta$ を計算するには x_0 に至るすべての可能な拡散経路（x_1,..., x_T）を周辺化（marginalize）する必要がある。

そこで ELBO（evidence lower bound）を使う。latents $x_{1:T}$ を導入し、全経路の報酬として $R(c, x_{0:T})$ を定義することで $r(c,x_0)$ を以下のように定義できる：

\[ r(c,x_0) = \mathbb{E}_{p_\theta(x_{1:T}|x_0,c)} \left [ R(c,x_{0:T}) \right ] \]

KL 正則化項については結合 KL ダイバージェンス $\mathbb{D}_{\mathrm{KL}}[p_\theta (x_{0:T}|c)||p_{\mathrm{ref}}(x_{0:T}|c)]$ の上界を最小化する。RLHF の式に代入すると以下の式になる：

\[ \underset{p_\theta}{\mathrm{max}} \mathbb{E}_{c\sim \mathcal{D}_c, x_{0:T}\sim p_\theta(x_{0:T}|c)} [r(c,x_0)] - \beta \; \mathbb{D}_{\mathrm{KL}}[p_\theta (x_{0:T}|c)||p_{\mathrm{ref}}(x_{0:T}|c)] \]

この目的関数は RLHF の式の並列定式化だが、経路 $x_{0:T}$ 上で定義されている。この目的はオリジナルの参照逆拡散過程の分布にマッチングさせつつ、逆拡散過程の $p_\theta(x_{0:T})$ の報酬を最適化することだ。

最適化すると以下の式（キャプションの c が省略されていることに注意）になる：

\[ \begin{eqnarray} L_{\mathrm{DPO-Diffusion}}(\theta) = -\mathbb{E}_{(x_0^w, x_0^t)\sim \mathcal{D}} \mathrm{log} \; \sigma \left ( \beta \; \mathbb{E}_{x_{1:T}^w\sim p_\theta (x_{1:T}^w|x_0^w) \\ x_{1:T}^l\sim p_\theta (x_{1:T}^l|x_0^l)} \left [ \mathrm{log} \dfrac{p_\theta(x_{0:T}^w)}{p_{\mathrm{ref}(x_{0:T}^w)}} - \mathrm{log} \dfrac{p_\theta(x_{0:T}^l)}{p_{\mathrm{ref}(x_{0:T}^l)}} \right ] \right ) \end{eqnarray} \]

ここで $x_{1:T}\sim p_\theta(x_{1:T}|x_0)$ をサンプリングする必要があるが非効率的だ。イェンセンの不等式を使って E を外側に出し、T = 1,000 回のサンプリングを T を掛け算することで近似して以下の式になる。詳細は論文の S.2 Details of the Primary Derivation を参照。

\[ \begin{eqnarray} L_{\mathrm{DPO-Diffusion}}(\theta) = -\mathbb{E}_{(x_0^w, x_0^t)\sim \mathcal{D}, t\sim \mathcal{U}(0, T), \\ x_{t-1,t}^w \sim p_\theta (x_{t-1,t}^w|x_0^w) \\ x_{t-1,t}^l \sim p_\theta (x_{t-1,t}^l|x_0^l)} \mathrm{log} \; \sigma \left ( \beta T \; \mathrm{log} \dfrac{p_\theta(x_{t-1}^w|x_t^w)}{p_{\mathrm{ref}}(x_{t-1}^w|x_t^w)} - \beta T \; \mathrm{log} \dfrac{p_\theta(x_{t-1}^l|x_t^l)}{p_{\mathrm{ref}}(x_{t-1}^l|x_t^l)} \right ) \end{eqnarray} \]

$p_\theta (x_{t-1}, x_t|x_0,c)$ からサンプリングするのは不可能なので、$p_\theta (x_{1:T}|x_0)$ を $q(x_{1:T}|x_0)$ で近似して以下の式を得る：

\[ \begin{eqnarray} L(\theta) = -\mathbb{E}&_{(x_0^w,x_0^l)\sim \mathcal{D}, t \sim \mathcal{U}(0,T), x_t^w \sim q(x_t^w|x_0^w), x_t^l \sim q(x_t^l|x_0^l)} \mathrm{log} \; \sigma (-\beta T \\ &+ \mathbb{D}_{\mathrm{KL}}(q(x_{t-1}^w|x_{0,t}^w)||p_\theta (x_{t-1}^w|x_t^w)) \\ &- \mathbb{D}_{\mathrm{KL}}(q(x_{t-1}^w|x_{0,t}^w)||p_{\mathrm{ref}} (x_{t-1}^w|x_t^w)) \\ &- \mathbb{D}_{\mathrm{KL}}(q(x_{t-1}^l|x_{0,t}^l)||p_\theta (x_{t-1}^l|x_t^l)) \\ &+ \mathbb{D}_{\mathrm{KL}}(q(x_{t-1}^l|x_{0,t}^l)||p_{\mathrm{ref}} (x_{t-1}^l|x_t^l))) \end{eqnarray} \]

最終的に以下の式になる：

\[ \begin{split} L(\theta) &= -\mathbb{E}_{\substack{(\mathbf{x}_0^w, \mathbf{x}_0^l) \sim \mathcal{D}, \\ t \sim \mathcal{U}(0, T), \\ \mathbf{x}_t^w \sim q(\mathbf{x}_t^w \mid \mathbf{x}_0^w), \\ \mathbf{x}_t^l \sim q(\mathbf{x}_t^l \mid \mathbf{x}_0^l)}} \left[ \log \sigma \left(-\beta \omega(\lambda_t) \left( \\ \left\| \boldsymbol{\epsilon}^w - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t^w, t) \right\|_2^2 - \left\| \boldsymbol{\epsilon}^w - \boldsymbol{\epsilon}_{\text{ref}}(\mathbf{x}_t^w, t) \right\|_2^2 \\ - \left( \left\| \boldsymbol{\epsilon}^l - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t^l, t) \right\|_2^2 - \left\| \boldsymbol{\epsilon}^l - \boldsymbol{\epsilon}_{\text{ref}}(\mathbf{x}_t^l, t) \right\|_2^2 \right) \right ) \right) \right] \end{split} \]

$\lambda_t = \alpha_t^2/\sigma_t^2$ は SNR。$\omega(\lambda_t)$ はウェイト関数で、T は β に入れている。

この損失関数は $\epsilon_\theta$ が $x_t^l$ よりも $x_t^w$ の改善に反応するようになっている。

学習について

クラウドで同じ GPU を４台使うのが一番簡単。

疑似コード

実装は SalesforceAIResearch/DiffusionDPO を参照。

以下のコードは数式と符号が一致しないように見えるが、Diffusion-DPOでは「損失（loss）」ではなく「対数尤度 log-likelihood」ベースで式が定義されており、それを loss に変換する場合は符号処理を注意する必要がある。

損失は小さい方がいいが、対数尤度は大きい方がいいので符号は逆になる。

# 初期化：モデル、参照モデル (frozen)、オプティマイザ、データローダなど
model = …          # 学習対象モデル θ  
ref_model = …      # 参照モデル θ_ref（凍結）  
dataset = PreferenceDataset(...)  # 勝ち／負けペア (x⁺, x⁻), 各プロンプト付き  
dataloader = DataLoader(dataset, batch_size=…)  
optimizer = …  
beta = ...  # 論文では SD 1.5 = 2000, SDXL = 5000。T と beta とを合算しているのでこのような大きな値になる

for epoch in range(num_epochs):
    for batch in dataloader:
        # batch contains e.g. prompts, winning_images, losing_images
        prompts, x_win, x_loss = batch

        # ノイズ付加／拡散モデルの入力準備
        # https://github.com/SalesforceAIResearch/DiffusionDPO の実装では論文とは異なりノイズは共通
        noise_win = sample_noise(...)
        noise_loss = sample_noise(...)
        noisy_win = add_noise(x_win, noise_win, t)
        noisy_loss = add_noise(x_loss, noise_loss, t)

        # モデルおよび参照モデルからの予測（例えば UNet の εθ 予測）
        pred_win = model(noisy_win, t, prompts)
        pred_loss = model(noisy_loss, t, prompts)
        ref_pred_win = ref_model(noisy_win, t, prompts)
        ref_pred_loss = ref_model(noisy_loss, t, prompts)

        # 損失（例えば MSE や ELBO に相当）を計算
        err_win = mse(pred_win, noise_win)          # モデルの勝ちサンプル誤差
        err_loss = mse(pred_loss, noise_loss)        # モデルの負けサンプル誤差
        ref_err_win = mse(ref_pred_win, noise_win)  # 参照モデルの勝ち誤差
        ref_err_loss = mse(ref_pred_loss, noise_loss) # 参照モデルの負け誤差

        # DPO-風の差分項を計算  
        # 例："inside_term" = β * ((err_loss − err_win) − (ref_err_loss − ref_err_win))
        diff_model = err_loss − err_win
        diff_ref = ref_err_loss − ref_err_win
        inside = beta * (diff_model − diff_ref)

        # ロスを −logσ(inside) によって定義
        loss = - log_sigmoid(inside).mean()

        # 逆伝播 & 最適化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

ノイズを共通化すべきか

論文の数式は勝ちサンプルと負けサンプルとで独立したノイズを付与している。しかし論文の疑似コードと GitHub のコードはノイズを共通化している。共通ノイズは勾配推定の分散を下げる（共通乱数を用いることで誤差差分の揺らぎが相殺される）が、一般には母平均との一致性が損なわれる可能性がある。

Calibrated Multi-Preference Optimization for Aligning Diffusion Models Supplementary Materials#A.2 Diffusion preference optimization では以下のように述べている。

Diffusion-DPO の論文では共通ノイズを使うように提案しているが、CaPO では独立したノイズを使用した。こちらの方が理論的に正当（grounded）、経験上性能もわずかに良いようである。

TailorPO

A Tailored Framework for Aligning Diffusion Models with Human Preference

手法が SPO: Aesthetic Post-Training Diffusion Models from Generic Preferences with Step-by-step Preference Optimization に類似しているので、論文自体はリジェクトされているが２つの重要な指摘をしている。

Diffusion-DPO の学習画像は最終結果のみで評価されるが、それは途中のデノイズ過程が優れていることを意味しない
キャプションは共通しているが、優れた画像と良くない画像とでベースとなる画像が異なるので、その軌跡を直接引き算するのは問題がある

Inversion-DPO

Inversion-DPO: Precise and Efficient Post-Training for Diffusion Models

Inversion DDIM で SFT 画像から初期ノイズを推定し、その初期ノイズをデノイズした画像を使う方法。ただしデノイズ後の画像の品質が常に優れているとは限らない。

Inversion-DPO は計算負荷は変わらないが、Diffusion-DPO 比で学習が２倍以上早いと主張している。論文では 11,140 ペアの画像で性能の改善を確認している。