MLエンジニアへの道 #54 - ノイズ除去拡散

このブログ記事では、コンピュータビジョンにおけるノイズ除去拡散確率モデルについて紹介します。

このシリーズの以前の記事では、画像生成のためのGANとVAEについて議論し、さらに複雑と思われるセグメンテーションとオブジェクト検出について議論しました。しかし、これまで議論してきた生成モデルはすべて、特に写真のようなリアルで高解像度の画像を生成する点で改善の余地がありました。そこで、本記事では、最近のモデルの印象的なパフォーマンスに貢献した画像生成モデルの改良点について説明します。

インセプションスコア & FID

しかし、生成モデルに入る前に、人間の観察に加えて画像生成を評価するための指標について議論する必要があります（これはVAEとGANの記事でその前提条件を扱っておらず、DVAEでは再構成よりも離散的な潜在表現の獲得に焦点を当てていたため、議論されてきませんでした）。主に使用される2つの指標としてインセプションスコア（IS）とフレシェ・インセプション距離（FID）があります。

インセプションスコアは、Inception-v3（ImageNetで事前訓練されたCNN）によって生成された画像に対して予測されたラベルの条件付き確率分布とラベルの周辺確率分布の間のKLダイバージェンスによって測定されます（ $\text{IS} = \text{exp}(\text{E}_{x \sim p_g}[D_{KL}(p(y|x) || p(y))])$ ）。生成モデルがInception-v3によって明確に分類できる画像を生成すること、つまり偏った条件付き確率分布 $p(y|x)$ ( $y$ がラベル、 $x$ が生成された画像)を得ることが望ましいです。同時に、生成モデルが多様な分類ラベルの画像を生成すること、つまり均一なラベルの周辺確率分布が得られることも重要です。したがって、条件付き確率がより偏り、周辺確率がより均一になるにつれて増加するKLダイバージェンスを使用することは理にかなっています。

ISはこの分野で広く使用されていますが、注意すべきいくつかの欠点があります。最も明白で重要な欠点は、先端の画像分類モデルではなく、 ImageNet（1000クラス）でのみ事前訓練され、入力サイズが限られている（ $300 \times 300$ ）Inception-v3に依存していることです。品質評価のために実際の画像と比較する代わりに条件付き確率分布の歪みを使用していることもあり、より広いクラスセットと複数のオブジェクトを持つ高解像度の写真のようなリアルな画像の品質を必ずしも捉えられない可能性があります。生成モデルは同じラベルに対してImageNetの画像に非常に似ている同じ画像を生成して高い多様性とISスコアを達成することもできます。

ISの品質評価の問題を軽減するために、FIDも使用します。これは、Inception-v3によって実際の画像と生成された画像から抽出された特徴（および線形層に渡されて埋め込みを作成する）の分布（正規分布と仮定）を比較します（ $\text{FID} = || \mu_X - \mu_Y ||^2 + \text{Tr}(\Sigma_X + \Sigma_Y - 2\sqrt{\Sigma_X \Sigma_Y})$ 、ここで $X$ と $Y$ は実際の画像と生成された画像であり、 $\text{Tr}$ はトレース関数で、すべての対角成分の合計を取ります）。 FIDでは実際の画像を使用して生成物がどれだけリアルかを品質評価に組み込むことができますが、依然としてImageNetで事前訓練されたInception-v3に依存しています。また、正規分布の仮定の妥当性を維持するために大きなサンプルが必要であり、計算コストが高くなるのにも関わらず、平均と共分散のみを使用するため、微妙な違いを捉えられない可能性があります。したがって、テキストや音声生成などの他の生成タスクと同様に、他の潜在的な指標を探求し続け、人間による評価を共に使用することが重要です。

階層的VAE

VAEは、ベイズ的視点（ここでは議論を避けますが、将来的に取り上げるかもしれません）と深層学習の視点の両方から興味深いものです。入力をディストリビューション潜在空間にマッピングすることで、データ圧縮と画像生成の両方を提供します。しかし、VAEには独自の課題があります。十分に大きなデコーダーが全ての画像を記憶し、潜在変数を無視してしまうことで、エンコーダーが有用な潜在表現を学習できなくなる可能性があります。また、潜在変数のサイズが様々な特徴を捉えるのに十分でない場合、重大な情報損失に悩まされることがあります。

この問題に対する直感的な解決策は、中間潜在変数を導入し、最終的な潜在変数（通常は標準ガウス分布）に到達する前に徐々に特徴を捉えることです。これにより帰納的バイアスを導入し、少なくとも理論上は学習を容易にすることが見込まれます。しかし、この階層的VAEアプローチは posterior collapse （事後崩壊）に悩まされます。事後崩壊とは、モデルが迅速に無意味な標準ガウス分布を生成することを学習し、そのサンプルが小さなデコーダーにさえも無用なガウスノイズとして無視されてしまう現象です。この事後崩壊のため、階層的VAEは実際には、残差接続などのトリックを使用して最後の潜在変数を再構成に利用するようモデルを促さない限り、通常のVAEよりも性能が低くなります。

自己回帰モデル

VAEが画像生成性能で苦戦する一方、自己回帰モデルはより良い結果を達成しました。ピクセル間の相関関係を活用することで、ピクセルの値の最初のサンプリングから祖先サンプリングを通じて一貫性のある画像を生成できます。しかし、代償として、潜在表現を学習することも、1回の反復で多くの隣接した相関のあるピクセルを生成することもできませんでした。これは、可能なピクセル値を絞り込み、すべての可能性の平均を生成することを避けるために条件付き確率(相関関係)に依存しているためです。これにより、画像を生成するために多数の反復（約 $10^6$ ）が必要となりました。

自己回帰モデルを階層的VAEのとして解釈すると、より良い性能につながった違いを見つけることができます。自己回帰モデルは固定エンコーダーを利用し、相関のないピクセルを暗黙的に徐々に除去して中間潜在変数（入力と同じサイズ）と最終潜在変数（すべてのピクセルが除去される）を生成します。デコーダーはこれらの潜在変数を反復的に処理して、元の画像を徐々に復元し、潜在変数間の条件付き確率や相関関係を捉えることによって元の画像を再構築することのみに焦点を当てています。これは、高度にパラメータ化され、圧縮された潜在変数を学習するエンコーダーが、画像生成に至っては不必要な複雑さをもたらしている可能性があるという以前の階層的VAEの潜在的な問題を示唆しています。また、モデルが1回の反復でできるだけ多くの相関のないピクセル値を処理できるようにすることで、性能と速度を大幅に向上させる可能性があると推測できます。

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

そこで、固定されたエンコーダーがガウスノイズサンプルを追加して(入力と同じサイズの)潜在変数が純粋な標準ガウスノイズになるまで進み、パラメータ化されたデコーダーが徐々に潜在変数のノイズを除去して入力を復元するというモデルを考えることができます。このモデルでは、すべてのピクセル値に加えられるノイズは独立にサンプリングされ相関がないため、モデルは潜在変数間の相関関係をできるだけ多く捉えることに集中でき、より少ない反復（〜 $10^2$ ）で自己回帰型画像生成が可能になります。このアプローチはノイズ除去拡散確率モデル（DDPMs）と呼ばれ、拡散確率モデル（DPMs）の中で最も一般的なアプローチと言えます。（スコアベースの生成モデルやフローベースの拡散モデルなど他のタイプもありますが、これらは将来的に取り上げるかもしれません。）

L = \text{E}_q[D_{KL}(q(x_T | x_0) || p(x_T)) + \sum_{t > 1} D_{KL}(q(x_{t-1} | x_t, x_0) || p(x_{t-1} | x_t)) - logp_{\theta}(x_0 | x_1)]

ここで、 $q$ はエンコーダーであり、 $q(x_t | x_{t-1}) := N(x_t; \sqrt{1-\beta_t}, \beta_t\text{I})$ のように分散スケジュール $\beta_t$ でノイズを追加し、 $p_{\theta}$ はパラメータ化されたデコーダーです。最初の項 $D_{KL}(q(x_T | x_0) || p(x_T))$ は、エンコーダーのノイズスケジュールが訓練可能な場合に、最後の潜在変数が標準ガウス分布 $p(x_T) \sim N(0,\text{I})$ に従うことを促します。2番目の項はデコーダーを訓練して中間潜在変数から次の潜在変数へのノイズ除去を行うことを促すもので、最後の項 $-logp_{\theta}(x_0 | x_1)$ は再構成損失であり、 $x$ とデコーダー出力間のMSEとして簡略化できます。表記 $\alpha_t = 1 - \beta_t$ と $\bar{\alpha}_t = \Pi_{s=1}^t \alpha_s$ を使用して $q(x_t | x_0)$ を $N(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)\text{I})$ と表現でき、これは $x_t(x_0, \epsilon) = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon$ （ここで $\epsilon \sim N(0, \text{I})$ ）として再パラメータ化できます。

デコーダーにこれらの中間潜在変数（ノイズの乗った画像）を生成させるのではなく、デコーダー $\epsilon_{\theta}(x_t)$ にノイズ因子 $\epsilon$ を予測させることで、それを中間潜在変数（前の反復で計算された）から差し引いて元の画像を復元できるようにすることが可能です。そうすることで、デコーダーのタスクはより単純になります（デコーダーは訓練中の各ステップで元の画像を復元することに集中でき、徐々に元の画像に近づくことを学習する必要がないため）。また、損失 $L$ は予測されたノイズ因子 $\epsilon_{\theta}(x_t)$ と実際のノイズ因子 $\epsilon$ 間のMSEによって簡単に計算できる様にもなります。（これはスケジュール $\beta$ が一定に保たれている場合です。 $\beta$ を学習可能にする場合は、記事の最後に引用されている元の論文で確認できる他の項を考慮する必要があります。）

ここでは、直接ノイズを除去したものを予測するのではなく、ノイズを予測して差し引くことで間接的にノイズを除去します。これはもし最終潜在変数である標準ガウスノイズから直接ノイズを除去した画像を予測するようにデコーダーを訓練すると、推論時にすべての訓練画像の平均を予測することになり（利用できる有用な相関関係がない）、無効でぼやけた結果になってしまうためです。一方、ノイズサンプルの平均は依然として有効なノイズであり、自己回帰的に適用して有効なノイズ除去済みの画像に到達することができます。エンコーダーを比較的固定して画像生成に集中させ、ガウスノイズサンプルを追加してノイズを除去することで反復ごとに調整される無相関のピクセル値の数を最大化し、デコーダーにノイズ因子を予測させることで、DDPMsは階層的VAEや自己回帰モデルよりもパフォーマンスと速度の面で大きく改善されました。

U-Net & LDM & DiT

DDPMsでは、入力サイズと同じサイズの潜在変数のノイズを除去するようにデコーダーを訓練します。つまり、デコーダーブロックは同じ入力と出力サイズを持つ必要があります。このようなアーキテクチャは、セマンティックセグメンテーションに関する記事ですでに見たことがありますーU-Netです。安定した勾配のためのスキップ接続を持つU-Netアーキテクチャは、単純なノイズ除去タスクとDDPMsのデコーダーブロックとして強力なパフォーマンスを示すことがわかっています。

U-Netアーキテクチャは畳み込み層に限定されず、自己注意層やシフトウィンドウ注意層でも使用できます。実際、高い画像生成品質で認識されたStability AIによるオリジナルのStable Diffusionは、ノイズ除去ネットワークとして注意機構を備えた時間条件付きU-Net（位置埋め込みの上に正弦波時間埋め込みを付加する）を利用しました。 Stable Diffusionからのもう一つの貢献は、潜在拡散モデル（LDM）の導入です。入力と同じサイズの潜在変数を扱うため、モデルは巨大になる傾向があり、訓練が難しく推論が遅くなっていました。そこでLDMsは、ノイズ除去を適用するより小さな潜在表現を生成するための事前訓練されたVAEを導入し、パフォーマンスを維持しながらモデルを小さくしました。

U-Net以外にも、私たちは同じ入力サイズと出力サイズを持つ別のアーキテクチャであるトランスフォーマーを広範囲に研究してきました。 U-Netとは異なりダウンサンプリングやボトルネックを使用しないViTを使用し、ノイズのガウス分布の平均(とオプションで分散)を出力させることができます。これは拡散トランスフォーマー（DiT）と呼ばれています。大量のパラメータを気にしない場合（大規模な研究機関は計算リソースが豊富なため、しばしばこれを気にしません）、DDPMにDiTを使用することは理にかなっています。パラメータ数を気にする場合でも、 Stable Diffusion 3のように、LDMで潜在空間のサイズを縮小し、より少ないパラメータを実現しながら使用することができます。

結論

この記事では、画像生成の評価によく使用される指標であるISとFIDについて取り上げ、ディープラーニングの観点から（主に機械的な観点であり、ベイズ的観点や情報理論的観点からではなく）、階層的VAEや自己回帰モデルからノイズ除去拡散確率モデルへの自然な発展について説明しました。また、LDMについても簡単に議論し、ノイズ除去ネットワークに用いる明らかに適切なアーキテクチャとしてU-NetとDiTについて触れました。次の記事ではDDPMをより有用に活用できる様にする、条件付き生成ついて議論します。

リソース

Algorithm Simplicity. 2024. Why Does Diffusion Work Better than Auto-Regression?. YouTube.
Ho, J. et al. 2020. Denoising Diffusion Probabilistic Models. ArXiv.
mm_0824. 2021. Frechet Inception Distance(FID)を理解する. 楽しみながら学ぶ機械学習 / 自然言語処理入門.
mm_0824. 2021. Inception Scoreを理解する. 楽しみながら学ぶ機械学習 / 自然言語処理入門.
Rombach, R. et al. 2022. High-Resolution Image Synthesis with Latent Diffusion Models. ArXiv.
Tomczak, M. J. n.d. Diffusion-based Models. jmtomczak.
Tomczak, M. J. n.d. Hierarchical VAEs. jmtomczak.