ViT-VQGAN：画像量子化技術を再考して画像生成と画像理解の性能を向上(2/2)

１．ViT-VQGAN：画像量子化技術を再考して画像生成と画像理解の性能を向上(2/2)まとめ

・VIMは犬や猫など与えられた特定のクラスの画像を合成するクラス条件付き生成が可能
・画像理解タスクでも分類精度を60.3%(iGPT-L)から73.2%に向上させることに成功
・画像量子化の改良により画像生成と理解の両方において優れた結果を示すことができた

２．VIMとは？

以下、ai.googleblog.comより「Vector-Quantized Image Modeling with Improved VQGAN」の意訳です。元記事は2022年5月18日、Jiahui YuさんとJing Yu Kohさんによる投稿です。

アイキャッチ画像のクレジットはPhoto by Bonnie Kittle on Unsplash

学習済ViT-VQGANでは、画像は整数で表される個々のトークンに符号化され、それぞれのトークンは入力画像の8×8範囲を含みます。これらのトークンを用いて、デコーダのみのTransformerを学習し、一連の画像トークンを自己回帰的に予測します。この2段階モデルであるVIMは、Transformerの出力ソフトマックス分布からトークンごとにサンプリングするだけで、条件無しに画像生成を行うことが可能です。

また、VIMは与えられたクラス(例えば、犬や猫)の特定の画像を合成するような、クラス条件付き生成を行うことができます。私達は、学習時とサンプリング時の両方で、画像トークンの前にクラスIDトークンを付加することにより、条件無し生成をクラス条件付き生成に拡張します。

ImageNetで学習したクラス条件付き画像生成をした未加工の犬画像
条件付けしたクラスは、アイリッシュテリア、ノーフォークテリア、ノーウィッチテリア、ヨークシャーテリア、ワイヤーヘアードフォックステリア、レイクランドテリア。

また、VIMの画像理解能力を検証するために、画像理解能力測定の標準ベンチマークであるImageNet分類用の線形射影レイヤーの微調整を行いました。ImageGPTと同様に、特定のブロックでのレイヤー出力を取り、トークン特徴をシーケンス上で平均化し(凍結)、平均化された特徴表現をクラス予測に投影するソフトマックス層(学習可能)を挿入します。これにより、特徴表現学習に有用な情報を提供する中間特徴を捉えることができます。

実験結果

全ViT-VQGANモデルを128のCloudTPUv4コアに分散した256の学習バッチサイズで学習させます。すべてのモデルは、256×256解像度の入力画像で学習させます。事前に学習したViT-VQGAN画像量子化器の上に、条件無しおよびクラス条件付き画像合成のためのTransformerモデルを学習し、従来の研究結果と比較します。

私達は、広く用いられているImageNetベンチマークを用いて、クラス条件付き画像合成と教師なし特徴表現学習に対する提案手法の性能を測定しました。

以下の表では、Fréchet Inception Distance(FID)によって測定されたクラス条件付き画像合成の性能を示しています。先行研究と比較して、VIMはFIDを3.07(低いほど良い)に改善し、VQGANモデル(FID 7.35)と比較して58.6%の相対的な改善を達成しました。また、VIMは画像理解能力も向上させ、Inception Score(IS)は188.6から227.4とVQGANに対して20.6%の向上が見られました。

Model	Acceptance Rate	FID	IS
Validation data	1	1.62	235
DCTransformer	1	36.5	N/A
BigGAN	1	7.53	168.6
BigGAN-deep	1	6.84	203.6
IDDPM	1	12.3	N/A
ADM-G, 1.0 guid.	1	4.59	186.7
VQVAE-2	1	~31	~45
VQGAN	1	17.04	70.6
VQGAN	0.5	10.26	125.5
VQGAN	0.25	7.35	188.6
ViT-VQGAN (Ours)	1	4.17	175.1
ViT-VQGAN (Ours)	0.5	3.04	227.4

クラス条件付き画像合成のためのFréchet Inception Distance(FID)と画像理解のためのInception Score(IS)の異なるモデル間の比較、いずれも解像度256×256のImageNet上での比較。受容率は、VQGANでの処理と同様に、ResNet-101分類モデルでフィルタリングした結果を示しています。

生成モデルを学習した後、線形層を微調整することで学習した画像特徴表現をテストし、画像理解能力を測定するための標準的なベンチマークであるImageNet分類を実行しました。本モデルは、画像理解タスクにおいて従来の生成モデルを凌駕し、線形プロービング(linear probing、すなわち、モデルの残りの部分を凍結したまま、単一の線形分類層を学習する)により分類精度を60.3%(iGPT-L)から73.2%に向上させることに成功しました。これらの結果は、VIMの強力な生成能力と画像表現学習能力を示しています。

まとめ

私達は、ベクトル量子化画像モデリング(VIM:Vector-quantized Image Modeling)を提案しました。これは、画像トークンを自己回帰的に予測するためにTransformerを事前学習させるもので、離散画像トークンは改良されたViT-VQGAN画像量子化器から生成されます。

画像量子化の改良により、画像の生成と理解の両方において優れた結果を示すことができました。本研究結果が、画像生成と理解のためのより統一的なアプローチに向けた今後の研究の刺激となることを期待しています。

謝辞

Xin Li, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, Yonghui WuにVIMの論文作成において感謝します。Wei Han、Yuan Cao、Jiquan Ngiam、Vijay Vasudevan、Zhifeng Chen、Claire Cuiの有益な議論とフィードバック、および、このプロジェクトを通してサポートしてくれたGoogle Research と Brain Teamの他の皆さんに感謝します。