拡散モデルを理解する：生成AIへの深い洞察

拡散モデルは、画像、音声、動画生成において最先端の結果を生み出す、生成AIにおける強力なアプローチとして台頭してきました。この詳細な技術記事では、拡散モデルの仕組み、その主要な革新点、そしてなぜこれほど成功したのかを探求します。このエキサイティングな新技術の数学的基礎、学習プロセス、サンプリングアルゴリズム、そして最先端の応用例について解説します。

拡散モデル入門

拡散モデルは、拡散過程を逆転させることでデータを徐々にノイズ除去することを学習する生成モデルの一種です。核となるアイデアは、純粋なノイズから始めて、ターゲット分布からの高品質なサンプルへと反復的に洗練することです。

このアプローチは、非平衡熱力学、特に構造を回復するための拡散の逆転過程に着想を得ています。機械学習の文脈では、データへのノイズの段階的な追加を逆転させることを学習すると考えることができます。

拡散モデルの主な利点には以下があります：

• 多くの場合でGANを凌駕する、最先端の画質
• 敵対的ダイナミクスなしの安定した学習
• 高度に並列化可能
• 柔軟なアーキテクチャ – 同じ次元の入出力をマッピングする任意のモデルを使用可能
• 強固な理論的基盤

拡散モデルの仕組みについて、さらに深く掘り下げてみましょう。

Source: Song et al.

拡散モデルにおける順方向および逆方向のプロセスは、確率微分方程式によって支配されます。順方向SDEはデータにノイズを加え、徐々にノイズ分布へと変換します。学習されたスコア関数によって導かれる逆方向SDEは、ノイズを段階的に除去し、ランダムノイズから現実的な画像を生成します。このアプローチは、連続状態空間において高品質な生成性能を達成するための鍵となります。

順方向拡散過程

順方向拡散過程は、実データ分布からサンプリングされたデータ点x₀から始まり、Tタイムステップにわたってガウシアンノイズを徐々に加え、次第にノイズの多いバージョンx₁, x₂, …, xTを生成します。

各タイムステップtにおいて、次の式に従って少量のノイズを追加します：

x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε

ここで：

• β_tは、各ステップで追加されるノイズの量を制御する分散スケジュールです
• εはランダムなガウシアンノイズです

このプロセスは、xTがほぼ純粋なガウシアンノイズになるまで続きます。

数学的には、これをマルコフ連鎖として記述できます：

q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)

ここで、Nはガウス分布を表します。

β_tスケジュールは通常、初期タイムステップでは小さく、時間とともに増加するように選択されます。一般的な選択肢には、線形、コサイン、またはシグモイドスケジュールがあります。

逆方向拡散過程

拡散モデルの目標は、この過程の逆を学習すること、つまり純粋なノイズxTから始めて、それを徐々にノイズ除去し、クリーンなサンプルx₀を回復することです。

この逆過程を次のようにモデル化します：

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))

ここで、μ_θとσ_θ^2は、θによってパラメータ化された学習済み関数（通常はニューラルネットワーク）です。

重要な革新点は、完全な逆分布を明示的にモデル化する必要がないことです。代わりに、既知である順方向過程の観点からパラメータ化することができます。

具体的には、最適な逆過程の平均μ*は次のように表せることを示せます：

μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))

ここで：

• α_t = 1 – β_t
• ε_θは学習されたノイズ予測ネットワークです

これにより、単純な目的関数が得られます。つまり、各ステップで追加されたノイズを予測するようにニューラルネットワークε_θを学習させます。

学習目的関数

拡散モデルの学習目的関数は、変分推論から導出できます。いくつかの簡略化の後、単純なL2損失に到達します：

L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]

ここで：

• tは1からTまで一様にサンプリングされます
• x₀は学習データからサンプリングされます
• εはサンプリングされたガウシアンノイズです
• x_tは、順方向過程に従ってx₀にノイズを加えて構築されます

言い換えれば、各タイムステップで追加されたノイズを予測するようにモデルを学習させています。

モデルアーキテクチャ

Source: Ronneberger et al.

U-Netアーキテクチャは、拡散モデルにおけるノイズ除去ステップの中心です。これは、再構成プロセス中に細部を保持するのに役立つスキップ接続を備えたエンコーダ-デコーダ構造を特徴としています。エンコーダは高レベルな特徴を捉えながら入力画像を段階的にダウンサンプリングし、デコーダはエンコードされた特徴をアップサンプリングして画像を再構成します。このアーキテクチャは、画像セグメンテーションのような正確な位置特定を必要とするタスクに特に効果的です。

ノイズ予測ネットワークε_θは、同じ次元の入出力をマッピングする任意のアーキテクチャを使用できます。U-Netスタイルのアーキテクチャは、特に画像生成タスクで人気のある選択肢です。

典型的なアーキテクチャは次のようになります：

 class DiffusionUNet(nn.Module): def __init__(self): super().__init__() # Downsampling self.down1 = UNetBlock(3, 64) self.down2 = UNetBlock(64, 128) self.down3 = UNetBlock(128, 256) # Bottleneck self.bottleneck = UNetBlock(256, 512) # Upsampling self.up3 = UNetBlock(512, 256) self.up2 = UNetBlock(256, 128) self.up1 = UNetBlock(128, 64) # Output self.out = nn.Conv2d(64, 3, 1) def forward(self, x, t): # Embed timestep t_emb = self.time_embedding(t) # Downsample d1 = self.down1(x, t_emb) d2 = self.down2(d1, t_emb) d3 = self.down3(d2, t_emb) # Bottleneck bottleneck = self.bottleneck(d3, t_emb) # Upsample u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb) u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb) u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb) # Output return self.out(u1) 

 def sample(model, n_samples, device): # Start with pure noise x = torch.randn(n_samples, 3, 32, 32).to(device) for t in reversed(range(1000)): # Add noise to create x_t t_batch = torch.full((n_samples,), t, device=device) noise = torch.randn_like(x) x_t = add_noise(x, noise, t) # Predict and remove noise pred_noise = model(x_t, t_batch) x = remove_noise(x_t, pred_noise, t) # Add noise for next step (except at t=0) if t > 0: noise = torch.randn_like(x) x = add_noise(x, noise, t-1) return x