#

ノイズ除去拡散確率モデル (DDPM)

これは、論文「ノイズ除去拡散確率モデル」のPyTorch実装/チュートリアルです。

簡単に言うと、データから画像を取得し、段階的にノイズを追加します。次に、モデルをトレーニングして各ステップでそのノイズを予測し、そのモデルを使用して画像を生成します。

以下の定義と派生は、これがどのように機能するかを示しています。詳しくは論文を参照してください。

転送プロセス

フォワード処理では $x_{0} \sim q (x_{0})$ 、 $T$ タイムステップのノイズがデータに追加されます。

q (x_{t} ∣ x_{t - 1}) = N (x_{t}; 1 - β_{t} x_{t - 1}, β_{t} I) q (x_{1 : T} ∣ x_{0}) = t = 1 \prod T q (x_{t} ∣ x_{t - 1})

$β_{1}, \dots, β_{T}$ 差異スケジュールはどこですか。

以下を使用すると、 $x_{t}$ $t$ 任意のタイムステップでサンプリングできます。

q (x_{t} ∣ x_{0}) = N (x_{t}; \overset{α_{t}}{ˉ} x_{0}, (1 - \overset{α_{t}}{ˉ}) I)

$α_{t} = 1 - β_{t}$ どこと $\overset{α_{t}}{ˉ} = \prod_{s = 1}^{t} α_{s}$

リバースプロセス

逆の処理では $p (x_{T}) = N (x_{T}; 0, I)$ 、4 $T$ つのタイムステップからノイズを除去します。

p_{θ} (x_{t - 1} ∣ x_{t}) p_{θ} (x_{0 : T}) p_{θ} (x_{0}) = N (x_{t - 1}; μ_{θ} (x_{t}, t), Σ_{θ} (x_{t}, t)) = p_{θ} (x_{T}) t = 1 \prod T p_{θ} (x_{t - 1} ∣ x_{t}) = \int p_{θ} (x_{0 : T}) d x_{1 : T}

$θ$ 私たちがトレーニングするパラメータです。

損失

（ジェンソンの不等式から）ELBOを負の対数確率で最適化します。

E [- lo g p_{θ} (x_{0})] \leq E_{q} [- lo g \frac{p _{θ} ( x _{0 : T} )}{q ( x _{1 : T} ∣ x _{0} )}] = L

損失は次のように書き直すことができます。

L = E_{q} [- lo g \frac{p _{θ} ( x _{0 : T} )}{q ( x _{1 : T} ∣ x _{0} )}] = E_{q} [- lo g p (x_{T}) - t = 1 \sum T lo g \frac{p _{θ} ( x _{t - 1} ∣ x _{t} )}{q ( x _{t} ∣ x _{t - 1} )}] = E_{q} [- lo g \frac{p ( x _{T} )}{q ( x _{T} ∣ x _{0} )} - t = 2 \sum T lo g \frac{p _{θ} ( x _{t - 1} ∣ x _{t} )}{q ( x _{t - 1} ∣ x _{t} , x _{0} )} - l o g p_{θ} (x_{0} ∣ x_{1})] = E_{q} [D_{K L} (q (x_{T} ∣ x_{0}) ∥ p (x_{T})) + t = 2 \sum T D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t})) - l o g p_{θ} (x_{0} ∣ x_{1})]

$D_{K L} (q (x_{T} ∣ x_{0}) ∥ p (x_{T}))$ $β_{1}, \dots, β_{T}$ 私たちは一定なので一定です。

コンピューティング $L_{t - 1} = D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t}))$

$x_{0}$ 次の条件で条件付けされたフォワードプロセスの事後処理は

q (x_{t - 1} ∣ x_{t}, x_{0}) \tilde{μ}_{t} (x_{t}, x_{0}) \tilde{β_{t}} = N (x_{t - 1}; \tilde{μ}_{t} (x_{t}, x_{0}), \tilde{β_{t}} I) = \frac{α ˉ _{t - 1} β _{t}}{1 - α _{t} ˉ} x_{0} + \frac{α _{t} ( 1 - α ˉ _{t - 1} )}{1 - α _{t} ˉ} x_{t} = \frac{1 - α ˉ _{t - 1}}{1 - α _{t} ˉ} β_{t}

$Σ_{θ} (x_{t}, t) = σ_{t}^{2} I$ $σ_{t}^{2}$ $β_{t}$ 用紙セットは定数またはに設定されています。 $\tilde{β_{t}}$

次に、 $p_{θ} (x_{t - 1} ∣ x_{t}) = N (x_{t - 1}; μ_{θ} (x_{t}, t), σ_{t}^{2} I)$

$ϵ \sim N (0, I)$ 与えられたノイズに対して $q (x_{t} ∣ x_{0})$

x_{t} (x_{0}, ϵ) x_{0} = \overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ = \frac{1}{α _{t} ˉ} (x_{t} (x_{0}, ϵ) - 1 - \overset{α_{t}}{ˉ} ϵ)

これにより、

L_{t - 1} = D_{K L} (q (x_{t - 1} ∣ x_{t}, x_{0}) ∥ p_{θ} (x_{t - 1} ∣ x_{t})) = E_{q} [\frac{1}{2 σ _{t}^{2}} ∥ ∥ \tilde{μ} (x_{t}, x_{0}) - μ_{θ} (x_{t}, t) ∥ ∥^{2}] = E_{x_{0}, ϵ} [\frac{1}{2 σ _{t}^{2}} ∥ ∥ \frac{1}{α _{t}} (x_{t} (x_{0}, ϵ) - \frac{β _{t}}{1 - α _{t} ˉ} ϵ) - μ_{θ} (x_{t} (x_{0}, ϵ), t) ∥ ∥^{2}]

ノイズを予測するためのモデルによる再パラメータ化

μ_{θ} (x_{t}, t) = \tilde{μ} (x_{t}, \frac{1}{α _{t} ˉ} (x_{t} - 1 - \overset{α_{t}}{ˉ} ϵ_{θ} (x_{t}, t))) = \frac{1}{α _{t}} (x_{t} - \frac{β _{t}}{1 - α _{t} ˉ} ϵ_{θ} (x_{t}, t))

$ϵ_{θ}$ $ϵ$ $(x_{t}, t)$ 与えられたものを予測する学習済み関数はどこか

これにより、

L_{t - 1} = E_{x_{0}, ϵ} [\frac{β _{t} ^{2}}{2 σ _{t}^{2} α _{t} ( 1 - α _{t} ˉ )} ∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]

つまり、ノイズを予測するためのトレーニングを行っています。

簡易損失

$L_{simple} (θ) = E_{t, x_{0}, ϵ} [∥ ∥ ϵ - ϵ_{θ} (\overset{α_{t}}{ˉ} x_{0} + 1 - \overset{α_{t}}{ˉ} ϵ, t) ∥ ∥^{2}]$

これにより、 $- l o g p_{θ} (x_{0} ∣ x_{1})$ $t = 1$ $L_{t - 1}$ $t > 1$ ウェイトを廃棄するタイミングと廃棄するタイミングを最小限に抑えることができます。 $L_{t - 1}$ $\frac{β _{t} ^{2}}{2 σ _{t}^{2} α _{t} ( 1 - α _{t} ˉ )}$ 重みを捨てると、 $t$ （ノイズレベルが高い）高いほうに与えられる重みが増え、サンプルの品質が向上します

。

このファイルには、トレーニング中に画像を生成するために使用する損失計算と基本的なサンプリング方法が実装されています。

$ϵ_{θ} (x_{t}, t)$ これがコードを提供してトレーニングするUNetモデルです。このファイルでは、トレーニング済みのモデルからサンプルと補間を生成できます

。

162from typing import Tuple, Optional
163
164import torch
165import torch.nn.functional as F
166import torch.utils.data
167from torch import nn
168
169from labml_nn.diffusion.ddpm.utils import gather

#