VQ-Diffusion

'VQ-Diffusion'

ベクトル量子化拡散（VQ-Diffusion）は、中国科学技術大学とMicrosoftによって開発された条件付き潜在拡散モデルです。一般的に研究されている拡散モデルとは異なり、VQ-Diffusionのノイジングとデノイジングのプロセスは量子化された潜在空間で動作します。つまり、潜在空間は離散的なベクトルの集合で構成されています。離散的な拡散モデルは、連続的な対応物と比較する興味深い比較対象を提供します。

• Hugging Faceモデルカード
• Hugging Face Spaces
• オリジナルの実装
• 論文

デモ

🧨 Diffusersを使用すると、わずか数行のコードでVQ-Diffusionを実行できます。

依存関係をインストールする

pip install 'diffusers[torch]' transformers ftfy

パイプラインをロードする

• 🤗変換器を使用した確率的な時系列予測
• Apple SiliconでのCore MLを使用した安定した拡散を利用する
• GPT2からStable Diffusionへ：Hugging FaceがElixirコミュニティに参入します

from diffusers import VQDiffusionPipeline

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")

FP16の重みを使用する場合

from diffusers import VQDiffusionPipeline
import torch

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16")

GPUに移動する

pipe.to("cuda")

パイプラインを実行する！

prompt = "プールで遊ぶテディベア。"

image = pipe(prompt).images[0]

アーキテクチャ

VQ-VAE

画像はVQ-VAEエンコーダを使用して、離散的な「トークン」または埋め込みベクトルのセットにエンコードされます。そのために、画像はパッチに分割され、各パッチは固定サイズの語彙を持つコードブックの中から最も近いエントリで置き換えられます。これにより、入力ピクセル空間の次元を削減します。VQ-DiffusionはTaming TransformersのVQGANバリアントを使用しています。VQ-VAEについては、このブログ投稿が理解を深めるための良いリソースです。

VQ-Diffusionは、拡散トレーニングプロセス中に凍結された事前学習済みのVQ-VAEを使用しています。

フォワードプロセス

フォワード拡散プロセスでは、各潜在トークンは同じままにするか、異なる潜在ベクトルに再サンプリングされる（等しい確率で）か、マスクされることができます。潜在トークンがマスクされると、マスクされた状態が続きます。α t \alpha_t α t ​ 、β t \beta_t β t ​ 、γ t \gamma_t γ t ​ は、ステップ t − 1 t-1 t − 1 からステップ t t t へのフォワード拡散プロセスを制御するハイパーパラメータです。γ t \gamma_t γ t ​ は、マスクされていないトークンがマスクされる確率です。α t + β t \alpha_t + \beta_t α t ​ + β t ​ は、マスクされていないトークンが同じままでいる確率です。トークンは、β t \beta_t β t ​ の確率で任意の個々の非マスクされた潜在ベクトルに遷移することができます。言い換えると、α t + K β t + γ t = 1 \alpha_t + K \beta_t + \gamma_t = 1 α t ​ + K β t ​ + γ t ​ = 1 です。ここで、 K K K は非マスクされた潜在ベクトルの数です。詳細については、論文のセクション4.1を参照してください。

逆プロセスの近似

エンコーダ-デコーダトランスフォーマは、プロンプト y y y に依存したノイズのない潜在変数 x 0 x_0 x 0 ​ のクラスを近似します。エンコーダは、凍結された重みを持つCLIPテキストエンコーダです。デコーダトランスフォーマは、すべての潜在ピクセルに対してマスクされていないグローバルなアテンションを提供し、ベクトル埋め込みのカテゴリ分布の対数確率を出力します。デコーダトランスフォーマは、1回のフォワードパスでノイズのない潜在変数の完全な分布を予測し、x t x_t x t ​ 上のグローバル自己アテンションを提供します。ノイズのない値に対する条件付きシーケンス・トゥ・シーケンスの問題としてフレーム化することは、エンコーダ-デコーダトランスフォーマが適している理由の一部を提供します。

ARモデルのセクションでは、VQ-Diffusionのアーキテクチャに関する追加のコンテキストがARトランスフォーマベースのモデルと比較して提供されます。

「Taming Transformers」は、画像データに対してトランスフォーマーが計算的に実行可能になるよう、生のピクセルを圧縮された潜在空間内の離散トークンに変換する方法についての良い議論を提供しています。

VQ-Diffusionの文脈で

拡散モデル

現代の拡散モデルは主に連続的です。前方プロセスでは、連続的な拡散モデルは反復的にガウスノイズを追加します。逆プロセスは、p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)) p θ ​ ( x t − 1 ​ ∣ x t ​ ) = N ( x t − 1 ​ ; μ θ ​ ( x t ​ , t ) , Σ θ ​ ( x t ​ , t ) ) によって近似されます。より単純な場合のDDPMでは、共分散行列が固定され、U-Netが x t x_t x t ​ のノイズを予測するためにトレーニングされ、x t − 1 x_{t-1} x t − 1 ​ はノイズから導かれます。

近似的な逆プロセスは、離散逆プロセスと構造的に似ています。ただし、離散の場合、x t x_t x t ​ のノイズを予測するための明確な類似点はありませんし、x 0 x_0 x 0 ​ の分布を直接予測することがより明確な目標です。

連続的な拡散モデルに比べて、離散拡散モデルに関する文献は少ないです。不平衡熱力学を使用した深層非教師付き学習は、二項分布上の拡散モデルを紹介しています。Argmax FlowsとMultinomial Diffusionは、多項分布に対する離散拡散を拡張し、言語モデリングタスクのためにノイズのない分布を予測するためにトランスフォーマーをトレーニングします。Structured Denoising Diffusion Models in Discrete State-Spacesは、代替のノイジングプロセス（一様、吸収、離散化ガウス、トークン埋め込み距離）によって多項拡散を一般化します。連続的な拡散モデルでも代替のノイジングプロセスは可能ですが、論文で指摘されているように、加法ガウスノイズのみが注目を集めています。

自己回帰モデル

VQ-DiffusionをARモデルと比較する方が興味深いかもしれません。なぜなら、ARモデルでは、トランスフォーマーを使用して離散分布を予測することがより頻繁に行われます。トランスフォーマーはARモデリングで成功を収めていますが、画像の解像度が増加するにつれて推論速度が線形に減少し、エラーが蓄積し、方向バイアスが発生するという問題があります。VQ-Diffusionはこれらの3つの課題を改善しています。

AR画像生成モデルは、各ピクセルがラスタスキャン順序（左から右へ、上から下へ）で前のピクセルに依存するように画像の確率を分解することを特徴としています。つまり、p ( x ) = ∏ i p ( x i ∣ x i − 1 , x i − 2 , . . . x 2 , x 1 ) p(x) = \prod_i p(x_i | x_{i-1}, x_{i-2}, … x_{2}, x_{1}) p ( x ) = ∏ i ​ p ( x i ​ ∣ x i − 1 ​ , x i − 2 ​ , . . . x 2 ​ , x 1 ​ ) です。その結果、モデルは対数尤度を直接最大化することでトレーニングすることができます。さらに、実際のピクセル（潜在ではない）値で動作するARモデルは、離散多項分布からチャンネル値を予測します。つまり、まず赤チャンネルの値が256通りのソフトマックスからサンプリングされ、次に緑チャンネルの予測は赤チャンネルの値に依存します。

AR画像生成モデルは、トランスフォーマーを計算的に実行可能にするための多くの作業によってアーキテクチャが進化してきました。トランスフォーマーベースのモデル以前には、PixelRNN、PixelCNN、PixelCNN++が最先端の技術でした。

「Image Transformer」では、トランスフォーマーベースのモデルへの移行と、トランスフォーマーベースのモデル以外のモデルについての良い議論が提供されています（省略された引用については論文を参照）。

小さな画像の各ピクセルを順次予測するために再帰ニューラルネットワークをトレーニングすることは、計算上非常に困難です。そのため、畳み込みニューラルネットワークを使用する並列化可能なモデル（PixelCNNなど）が最近注目を集め、PixelRNNを品質で超えるようになりました。

畳み込みニューラルネットワークに比べてRNNの一般的には比較的限定された受容野があります。これは、対称性や遮蔽などの画像に一般的な長距離の現象をモデル化する能力に悪影響を及ぼす可能性があります。特に、層の数が少ない場合にはさらに影響が出ます。受容野を拡大することは、品質を著しく向上させることが示されていますが（Salimans et al.）、その代わりにパラメータ数とそれに伴う計算性能の大幅なコストがかかり、このようなモデルのトレーニングをより困難にする可能性があります。

…セルフアテンションは、必然的に順次処理する必要があるPixelRNNのほぼ無制限の受容野と、より並列化可能なPixelCNNとそのさまざまな拡張の限られた受容野とのトレードオフをより良く実現できます。

Image Transformerは、ピクセルの近傍領域に対してセルフアテンションを制限することで、トランスフォーマーを使用しています。

Taming TransformersとDALL-E 1は、畳み込みとトランスフォーマーを組み合わせています。両方のモデルは、離散的な潜在空間を学習するためにVQ-VAEをトレーニングし、その後トランスフォーマーを圧縮された潜在空間でトレーニングします。トランスフォーマーの文脈はグローバルですが、マスクされています。これは、以前に予測された潜在ピクセル全体に対してアテンションが提供されるためですが、モデルはまだAR（オートリグレッション）なので、まだ予測されていないピクセルにはアテンションを提供することはできません。

ImageBARTは、畳み込み、トランスフォーマー、拡散プロセスを組み合わせています。モデルは、短い多項式拡散プロセスでさらに圧縮された離散的な潜在空間を学習します。その後、別々のエンコーダ・デコーダトランスフォーマーが拡散プロセスの各ステップを逆に予測するようにトレーニングされます。エンコーダトランスフォーマーはx t ​のグローバルなコンテキストを提供し、デコーダトランスフォーマーはx t − 1 ​の潜在ピクセルをオートリグレッシブに予測します。その結果、各ピクセルはよりノイズの多い画像においてグローバルなクロスアテンションを受けます。より複雑なデータセットでは、2〜5つの拡散ステップが使用されます。

ARモデルは大きな進歩を遂げていますが、高解像度の画像ほど推論速度が線形に低下し、エラーの蓄積や方向バイアスの問題があります。同じサイズのARトランスフォーマーモデルに対して、VQ-Diffusionの推論のビッグOは、拡散ステップの数が潜在ピクセルの数よりも少ない限り、より良い結果となります。ITHQデータセットでは、潜在解像度は32×32であり、モデルは約100の拡散ステップまでトレーニングされています。これにより、約10倍のビッグOの改善が実現されます。実際には、「VQ-Diffusionは、より良い画質を実現しながら、AR手法よりも15倍高速化することができます」（詳細は論文を参照）とされています。また、VQ-Diffusionは教師強制を必要とせず、誤って予測されたトークンを修正することを学習します。トレーニング中、ノイズのある画像はマスクされ、潜在ピクセルはランダムなトークンで置き換えられます。VQ-Diffusionは、x t ​を予測する際にもx t − 1 ​をグローバルな文脈として提供することができます。

VQ-Diffusionと 🧨 Diffusersのさらなるステップ

これまで、ITHQデータセットでトレーニングされたVQ-Diffusionモデルの移植のみを行いました。CUB-200、Oxford-102、MSCOCO、Conceptual Captions、LAION-400M、ImageNetでトレーニングされたVQ-Diffusionモデルもリリースされています。

VQ-Diffusionは、より高速な推論戦略もサポートしています。ネットワークの再パラメータ化は、ノイズのない画像に条件付けられた拡散プロセスの事後分布に依存しています。時間ストライドΔtを使用する場合、類似の式が適用されます。これにより、逆拡散ステップの数をスキップします。p θ ( x t − Δ t ∣ x t , y ) = ∑ x ~ 0 = 1 K q ( x t − Δ t ∣ x t , x ~ 0 ) p θ ( x ~ 0 ∣ x t , y ) p_\theta (x_{t – \Delta t } | x_t, y) = \sum_{\tilde{x}_0=1}^{K}{q(x_{t – \Delta t} | x_t, \tilde{x}_0)} p_\theta(\tilde{x}_0 | x_t, y) p θ ​ ( x t − Δ t ​ ∣ x t ​ , y ) = ∑ x ~ 0 ​ = 1 K ​ q ( x t − Δ t ​ ∣ x t ​ , x ~ 0 ​ ) p θ ​ ( x ~ 0 ​ ∣ x t ​ , y ) 。

Improved Vector Quantized Diffusion Modelsは、離散のクラシファイアフリーガイダンスと「共分布の問題」に対処するための代替推論戦略を使用して、VQ-Diffusionのサンプル品質を向上させます。離散のクラシファイアフリーガイダンスはdiffuserに統合されますが、代替推論戦略はまだ追加されていません。

貢献は歓迎です！

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!