Muse: Text-To-Image Generation via Masked Generative Transformersを読んだメモ

昔に書きかけた記事の供養になります。一部書きかけですがご容赦ください。

Museとは

Muse¹は高速かつ高精細な画像を生成することができることを特徴とする画像生成手法です。
どの程度高速化と言えば、Stable Diffusion²（表中LDM）と比較して3～15倍高速でありながら、Stable Diffusionを上回る美麗な画像を生成することができます。

Museは、Non-diffusion modelかつNon-autoregressive modelな最近では珍しいタイプの画像生成モデルです。

また、従来のLDM（Latent Diffusion Model）やDiffusion modelにおいて、Stable DiffusionやImagenのような著名なモデルでは潜在変数の推論にU-Net³を用いることが多くありましたが、MuseではTransformerを用いて潜在変数を推論します。
これらの違いをまとめたものは以下になります。

	Muse	Stable Diffusion	Imagen
学習するモデルの出力	潜在変数	潜在変数	画像
潜在変数→画像へのデコーダ	VQ-VAE	VAE	なし
生成モデル	Transformer	U-Net	U-Net
誤差関数	Cross Entropy	Squared Error	Squared Error

本記事は、このMuseについて述べられているMuse: Text-To-Image Generation via Masked Generative Transformersを読んだ備忘録になります。
興味がある部分だけを読み込んでいるので、もしかしたら勘違い等あるかもしれませんが、その際はコメント等いただけますと幸いです。

モデル構造

論文中ではBase ModelとSuper-resolution Modelの2種類のモデルが紹介されていますが、今回説明するモデルはより低性能なBase Modelと呼ばれるモデルです。

Base modelの構成は以下のとおりです。
論文の図に少しだけ手を加えています。

Stable Diffusionとくらべてかなりシンプルなモデル構成ですね。
各モデルは以下の学習済みモデルを使用しています。
次節以降で各モデルについて簡単に説明していきます。

Model	学習済みモデル
Text Encoder	T5 Encoder
Base Transformer	学習済みのT5-XLL
SuperRes Transformer	学習済みのT5-XLL
VQTokenizer 1	256x256の画像で学習したVQGANのEncoder
VQTokenizer 2	512x512の画像で学習したVQGANのEncoder
VQ Decoder	VQGANのDecoder

VQGAN

VQGAN⁴とはVQ-VAE⁵を用いた画像生成モデルです。
今回主に使用するVQ-VAEにフォーカスして説明します。

VQ-VAEはVAEの1種でVector Quantization(ベクトル量子化)を伴う手法です。
Encoderは通常のVAEと同じくデータを低次の潜在変数へ圧縮する構成ですが、Encoder以降の取り扱いは大きく異なります。
VQ-vAEでデコードされる潜在変数は学習済みのEmbedding機構に保存されたデータのみをデコードします。

どういうことかと言えば、Embedding内ベクトルと潜在変数を比較して、最も差が小さいEmbedding内のベクトルをDecoderでデコードする構成になっています。

T5

T5⁶は Text-To-Text Transfer Transformer でテキストからテキストへの変換を学習する転移学習方法で、機械翻訳や文書要約に柔軟に対応することが可能であるとが示されています。
MuseではこのT5を使って学習されたTransformer⁷を用いてVQ-VAEの潜在変数を推論します。
厳密には、Transformerの推論値はVQ-VAEのEmbeddingに対するインデックスであり、このインデックスを使ってEmbeddingから潜在変数に変換。
VQ-VAEのDecoderを用いて潜在変数を画像へと変換します。

Transformer

今となっては当たり前に使われているモデルであるTransformerについても軽くご紹介します。
詳しくはTensorflowのチュートリアル[^8]がコードを交えて説明していてわかりやすいと思います。
基本的にTesnsorflowはNumPyと同じインタフェースなので、使ったことがなくても問題なく読めると思います。

それはさておき、Transformerの構造は以下の通りです。

元々、NLPのために提案されたモデルで、InputsとOutputsにはカテゴリ値（NLPであれば単語のIndex）が来る場合を想定しています。
したがって、Embeddingでは、カテゴリ値をベクトルへと変換する処理になります。
Positoinal Encodingでは、単語の絶対位置を保存するために専用のベクトルを掛けることで埋め込みます。

上記の図でFeed Forwardは全結合層のみで構成されるネットワークで、Multi-Head Attentionは以下の構造のモデルです。

Multi-Head AttentionのQ, K, Vの各ヘッドには主に以下のパターンで使用されています。

• Q, K, V全てに同じ入力。
• Q, KはEncoderの出力，VにDecoderの処理中の値。

Museでも上記の2パターンのうち下側のパターンが使われていて、Encoderにテキスト、Decoderに画像を入力し、最終的にVQ-VAEの潜在変数のインデックスを多クラス分類によって得ます。

Museの学習

Museの学習は先述の通り最終的にVQ-VAEの潜在変数のインデックスを多クラス分類によって得るTransformerを学習することでなされます。
ただし、その学習の仕方にいくつか工夫があります。
本節ではこれらについて簡単に説明していきます。

モデルの階層化

書きかけの項目です。ご容赦ください。

データのマスキング

書きかけの項目です。ご容赦ください。

Museでは画像とテキスト両方から抽出した潜在変数をマスクして学習することで推論速度、安定性の向上など様々な効果を得ています。
具体的なマスキングの方策は以下のとおりです。

• テキスト
  • 方法：ランダムに10%のデータでテキストを使用せずに画像を生成
  • 効果：
• 画像
  • 方法：Mask gitで提案された以下のマスク処理をランダムに実施。マスク頻度はarccos分布に基づいて実施。
  • 効果：推論の高速化

おわりに

初見ではMuseもLDM（Latent Diffusion Model）かと思いましたが、そうではありませんね。
Diffusion Modelは何らか仮定したノイズ分布を除去する確率過程を踏むのに対して、Museはそれらの過程を踏まないためDiffusion Modelではありません。
また、様々な工夫を重ねることで、高速かつ高精細な画像を生成することができます。

実際にPyTorchで実装されたコードもあるので、コードを読むことが得意な方は実際に読んでみてもいいかもしれません。
https://github.com/lucidrains/muse-maskgit-pytorch

1. H. Chang et al., “Muse: Text-To-Image Generation via Masked Generative Transformers,” Jan. 2023, [Online]. Available: http://arxiv.org/abs/2301.00704. ↩

2. C. Saharia et al., “Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding,” 2022, [Online]. Available: http://arxiv.org/abs/2205.11487. ↩

3. W. Weng and X. Zhu, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” IEEE Access, vol. 9, pp. 16591–16603, 2021, doi: 10.1109/ACCESS.2021.3053408.
   [^8] Tensorflow チュートリアル https://www.tensorflow.org/text/tutorials/transformer ↩

4. P. Esser, R. Rombach, and B. Ommer, “Taming transformers for high-resolution image synthesis,” Proc. IEEE Comput. Soc. Conf. Comput. Vis. Pattern Recognit., pp. 12868–12878, 2021, doi: 10.1109/CVPR46437.2021.01268. ↩

5. A. Van Den Oord, O. Vinyals, and K. Kavukcuoglu, “Neural discrete representation learning,” Adv. Neural Inf. Process. Syst., vol. 2017-December, no. Nips, pp. 6307–6316, 2017, [Online]. Available: https://arxiv.org/pdf/1711.00937. ↩

6. C. Raffel et al., “Exploring the limits of transfer learning with a unified text-to-text transformer,” J. Mach. Learn. Res., vol. 21, pp. 1–67, 2020, [Online]. Available: https://arxiv.org/pdf/1910.10683.pdf. ↩

7. A. Vaswani et al., “Attention is all you need,” Adv. Neural Inf. Process. Syst., vol. 2017-Decem, no. Nips, pp. 5999–6009, 2017, [Online]. Available: https://arxiv.org/pdf/1706.03762.pdf. ↩