VQ-Diffusion: 離散空間における条件付き潜在拡散

VQ-Diffusion は、中国科学技術大学と Microsoft が開発した条件付き潜在拡散モデルです。ほとんどの一般的な拡散モデルが連続空間で動作するのに対し、VQ-Diffusion は離散的なベクトル集合で構成された量子化潜在空間上でノイズ付加と除去のプロセスを行います。

アーキテクチャと技術実装

VQ-Diffusion は、次元削減のために VQ-VAE を、拡散プロセスのためにエンコーダ‑デコーダ transformer を組み合わせて利用します。

VQ-VAE 潜在空間

画像はまず VQ-VAE エンコーダを用いて離散トークンまたは埋め込みベクトルにエンコードされます。このプロセスは画像をパッチに分割し、各パッチを固定サイズの語彙コードブックから最も近いエントリに置き換えることを含みます。VQ-Diffusion は特に Taming Transformers の VQGAN バリアントを使用し、拡散訓練中は凍結された事前学習済み VQ-VAE を利用します。

前方拡散プロセス

前方プロセスでは、潜在トークンがクリーンな状態からノイズ付加された状態へと遷移します。各トークンは次のいずれかを取ります。

そのまま残る。
異なる潜在ベクトルに再サンプリングされる（等確率）。
マスクされる。

トークンがマスクされると、以後はマスクされたままです。このプロセスはハイパーパラメータ $\alpha_{t}$、$\beta_{t}$、$\gamma_{t}$ によって制御され、$\gamma_{t}$ はトークンがマスクになる確率、$\alpha_{t} + \beta_{t}$ はトークンが変化しない確率です。個々の非マスク潜在ベクトルへの遷移確率は $\beta_{t}$ です。

逆プロセスの近似

ノイズを除去して画像を生成するために、エンコーダ‑デコーダ transformer がプロンプト $y$ に条件付けされた未ノイズ潜在 ($x_{0}$) のクラスを近似します。

Encoder: 重みが凍結された CLIP テキストエンコーダ。
Decoder: 全潜在ピクセルに対してマスクされていないグローバル注意を提供し、ベクトル埋め込み上のカテゴリ分布の対数確率を出力する transformer。

デコーダは単一のフォワードパスで未ノイズ潜在全体の分布を予測するため、現在の状態 $x_{t}$ に対するグローバル自己注意を実現します。

VQ-Diffusion と他の生成モデルの比較

VQ-Diffusion は、連続拡散モデルと自己回帰 (AR) モデルの両方に対する技術的代替手段を提供します。

連続拡散との比較

ほとんどの最新拡散モデルは連続的で、ガウスノイズを逐次的に加え、U‑Net を用いてそのノイズを予測することで逆プロセスを近似します。対照的に VQ-Diffusion は離散値上で動作し、$x_{0}$ の分布を直接予測することがノイズ予測よりも明確な目的となります。

自己回帰 (AR) モデルとの比較

VQ-Diffusion は、AR トランスフォーマーベースの画像モデルに共通する 3 つの主要な課題—解像度が上がるにつれて推論速度が線形に低下すること、誤差蓄積、方向バイアス—に対処します。

推論速度: AR モデルは画像確率をピクセルごとにラスタースキャン順で条件付けします。拡散ステップ数が潜在ピクセル数未満である限り、VQ-Diffusion の推論計算量は優れています。ITHQ データセット（32×32 潜在解像度）では、最大 100 ステップで訓練され、約 10 倍のビッグ O 改善が得られます。実際には、AR 手法より最大 15 倍高速で、かつ画像品質も向上します。
誤差訂正: AR モデルが教師強制を必要とするのに対し、VQ-Diffusion はマスクとランダムトークン置換の両方を訓練に組み込むことで、誤って予測されたトークンを修正する方法を学習します。
グローバルコンテキスト: VQ-Diffusion は $x_{t}$ 上でグローバルコンテキストを提供しながら $x_{t-1}$ を予測しますが、AR モデルは将来のピクセルを見ることを防ぐためにマスク注意に制限されます。

実装とさらなる最適化

VQ-Diffusion は diffusers ライブラリの VQDiffusionPipeline を通じて Hugging Face に統合されています。モデルは ITHQ、CUB-200、Oxford-102、MSCOCO、Conceptual Captions、LAION-400M、ImageNet など多様なデータセットで訓練されています。

さらなる最適化として、逆拡散ステップをスキップする時間ストライド $\Delta t$ を用いた高速推論戦略があります。また、"Improved Vector Quantized Diffusion Models" では、離散的な classifier‑free guidance によるサンプル品質向上と、結合分布問題に対処する代替推論戦略が提案されています。