Diffusion Models

拡散モデルは、段階的なノイズ付加プロセスを逆転させる方法を学習してデータを生成する生成AIモデルです。学習時には、画像などのデータに多数の小さなステップでガウスノイズを加え、各ステップのノイズを予測するようニューラルネットワークを学習させます。生成時には、モデルは完全なランダムノイズから開始し、新しいサンプルが現れるまで段階的にノイズを取り除きます。Jascha Sohl-Dicksteinらは2015年にDeep Unsupervised Learning using Nonequilibrium Thermodynamicsでこの手法を提案し、Jonathan Ho、Ajay Jain、Pieter Abbeelは2020年にノイズ除去拡散確率モデル（DDPM）で実用的な性能を実現し、CIFAR-10で最先端のFIDスコア3.17を達成しました。2021年、OpenAIの研究者は拡散モデルがImageNet画像生成でGANを上回ることを示しました。拡散モデルとその関連手法であるフローマッチングは現在、テキストから画像を生成するシステムのStable Diffusionをはじめ、動画生成器、画像編集ツール、科学向けモデルを支えています。

拡散モデルの仕組み#

順方向プロセスは学習されるのではなく、固定されています。T個のステップそれぞれで少量のガウスノイズを加えます。DDPMではT = 1,000とし、1ステップあたりのノイズ分散を0.0001から0.02まで線形に増やします。ガウスノイズの和もガウス分布になるため、学習画像は閉形式で任意のステップtに直接移せます。xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noiseでは、alpha_bar_tは最初のtステップにおける1 - betaの積です。学習では画像、ランダムなステップ、新しいノイズを繰り返し選び、xtを構築して、平均二乗誤差を使ってネットワークにそのノイズを予測させます。これはDDPMの簡略化された目的関数です。このコードはDDPMのスケジュールを使って、写真に順方向プロセスを適用します。

import urllib.request

import numpy as np
import torch
from PIL import Image

# DDPM's linear schedule: 1,000 steps whose noise variance rises from 0.0001 to 0.02
betas = torch.linspace(1e-4, 0.02, 1000)
alpha_bar = torch.cumprod(1 - betas, dim=0)

# Load a photo and scale its pixels to [-1, 1]
image = Image.open(urllib.request.urlopen("https://ultralytics.com/images/bus.jpg")).convert("RGB").resize((256, 256))
x0 = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1) / 127.5 - 1

# Jump straight to step t in closed form: x_t = sqrt(alpha_bar) * x0 + sqrt(1 - alpha_bar) * noise
frames = []
for t in (0, 99, 299, 599, 999):
    xt = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * torch.randn_like(x0)
    frames.append(((xt.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).numpy())
    print(f"step {t + 1}: image weight {alpha_bar[t].sqrt():.3f}, noise weight {(1 - alpha_bar[t]).sqrt():.3f}")
Image.fromarray(np.hstack(frames)).save("forward_diffusion.jpg")

画像の重みは、ステップ1の1.000から、ステップ100で0.947、ステップ300で0.630、ステップ600で0.161、ステップ1,000で0.006へと低下し、ノイズの重みは1.000まで上昇します。保存されたforward_diffusion.jpgでは、バスが鮮明な状態から、粒状になり、薄くなって、最後に消えていく様子を確認できます。

生成では、プロセスを逆向きに実行します。完全なガウスノイズから始め、各ステップでネットワークが現在のサンプルに含まれるノイズを予測し、ソルバーがその一部を取り除き、わずかにノイズが減った結果を次のステップに渡します。テキストプロンプトやクラスラベルは条件としてネットワークに入力され、多くの場合、クロスアテンションを通じて反映されます。Jonathan HoとTim Salimansが提案した分類器なしガイダンスでは、1つのネットワークを条件ありと条件なしの両方で学習させ、サンプリング時に無条件予測から条件付き予測の方向へ外挿します。Stable Diffusion 1.xのパイプラインで既定値の7.5など、ガイダンススケールを高くすると、多様性を犠牲にしてプロンプトにより忠実になります。

サンプリングの高速化#

DDPMのサンプリングでは画像1枚につきネットワークを1,000回評価するため、分野の進歩の多くは評価回数の削減に向けられてきました。

• より優れたサンプラー：DDIMは、学習済みDDPMを決定論的なサンプラーで再利用し、品質の高いサンプルを実時間で10～50倍速く生成します。その後のODEソルバーでは、一般的な生成ステップ数が約20～50回に減りました。
• 潜在拡散：潜在拡散モデルはピクセルではなく、オートエンコーダーで圧縮した表現のノイズを除去します。Stable Diffusionのオートエンコーダーは512×512×3の画像を64×64×4の潜在表現に変換し、値の数を48分の1にします。
• 蒸留：一貫性モデルは、ノイズを含む任意のサンプルを直接きれいなサンプルに変換する方法を学習し、1ステップ生成を可能にします。潜在一貫性モデルは同じ考え方をStable Diffusionに適用し、2～4ステップで生成します。
• より直線的な経路：フローマッチングと整流フローは、ソルバーがより少ないステップでたどれる、ノイズからデータへの直線経路に沿って学習します。

ノイズ除去ネットワーク自体も、DDPMやStable Diffusion 1.xの畳み込み型U-Netから、現在の画像・動画生成器の大半で使われる拡散Transformer（DiT）へと移行しました。

拡散モデルとほかの生成モデルの比較#

• 敵対的生成ネットワーク（GANs）：GANは1回の順方向処理でサンプルを生成しますが、生成器と識別器による敵対的学習は不安定になりやすく、生成器がデータの一部しかカバーしないモード崩壊も起こりやすいです。拡散モデルの学習は安定した回帰問題で、データ分布をより網羅的にカバーできますが、サンプリングに多くのステップを要します。
• 変分オートエンコーダー（VAEs）：VAEは1回の処理で潜在ベクトルをサンプルにデコードしますが、生成サンプルはぼやける傾向があります。潜在拡散モデルは、VAE型のオートエンコーダーを画像の圧縮のみに使用し、生成は拡散モデルに任せます。
• 自己回帰モデル：言語モデルや視覚自己回帰モデルは、それまでの内容を条件として、一度に1つずつトークンやスケールを生成します。拡散モデルは各ステップでサンプル全体を並列に調整し、拡散言語モデルは同じノイズ除去をテキストに適用します。

実際のアプリケーション#

• 画像生成と編集：テキストから画像を生成するモデルはプロンプトから画像を作成し、画像インペインティングはマスクされた領域を周囲に合うように再生成します。ControlNetはエッジ、深度、ポーズのマップを使ってレイアウトを制約し、拡散アップスケーラーは超解像を実行します。
• 動画生成：動画生成モデルは多くのフレームをまとめて表す圧縮潜在表現のノイズを除去するため、各フレームは隣接フレームと一緒に生成されます。
• 合成学習データ：生成画像は、まれな欠陥、珍しい天候、まれな視点など、ビジョンデータセットの不足を補います。Ultralytics Platformのデータセットでは、類似画像を生成アクションにより、既存の画像を参考に1バッチあたり1～16枚の新しい画像を作成できます。採用した画像はラベルなしでデータセットに追加され、アノテーションできる状態になります。
• ロボティクス：拡散ポリシーは、カメラ画像を条件として、ノイズ除去によってロボットの動作シーケンスを生成します。
• 科学：AlphaFold 3は、タンパク質、核酸、小分子の複合体の構造を予測します。生の原子座標のノイズを除去する拡散モジュールを使用し、AlphaFold 2の構造モジュールを置き換えています。