MGVQ: Could VQ-VAE Beat VAE - A Generalizable Tokenizer with Multi-group Quantization

著者 Mingkai Jia, Wei Yin, Xiaotao Hu, Jiaxin Guo, Xiaoyang Guo, Qian Zhang, Xiao-Xiao Long, Ping Tan
所属 The Hong Kong University of Science and Technology, Horizon Robotics, The Chinese University of Hong Kong, Nanjing University
投稿日 2025年7月12日
カテゴリ cs.LG, cs.CV
arXiv 2507.07997v1 ↗

MGVQ: Could VQ-VAE Beat VAE - A Generalizable Tokenizer with Multi-group Quantization

基本情報

  • arXiv ID: 2507.07997v1 (https://arxiv.org/abs/2507.07997)
  • 著者: Mingkai Jia, Wei Yin, Xiaotao Hu, Jiaxin Guo, Xiaoyang Guo, Qian Zhang, Xiao-Xiao Long, Ping Tan
  • 所属: The Hong Kong University of Science and Technology, Horizon Robotics, The Chinese University of Hong Kong, Nanjing University
  • 投稿日: 2025年7月12日
  • カテゴリ: cs.LG, cs.CV

簡単に説明すると

MGVQは、画像を離散的なトークンに変換する新しいVector Quantized VAE（VQ-VAE）手法である。従来のVQ-VAEはVAEと比較して再構成品質が劣っていた。MGVQはマルチグループ量子化という新しいアプローチにより、ImageNetでのrFIDをSD-VAEの0.91から0.49に改善した。具体的には、潜在次元を保持したまま複数のサブコードブックを使用することで、コードブックの表現能力を10億倍以上に拡張し、情報損失を最小化している。

関連リンクは以下の通りである。

1. 研究概要

1.1 背景と動機

Vector Quantization（VQ）は、連続的な視覚データを離散的なトークンの集合に変換する基本技術である。高次元データを学習されたコードブック埋め込みの有限集合にマッピングすることで、VQはコンパクトな潜在空間を作成する。この離散表現は計算複雑性を削減するだけでなく、重要な意味情報を保持し、画像圧縮や特徴抽出などのタスクを可能にする。

しかし、現在のトークナイザーは、コードブックの表現能力が限られているため、依然として大きな情報損失に悩まされている。その結果、VQ-VAEはVAEと比較して劣ったエンコーディング・デコーディング性能を示す傾向がある。デコードされた画像には大きなエラーが生じることが多い。

既存の手法では、コードブックの崩壊を避けるために潜在次元を積極的に削減する戦略が取られてきた。例えば、LlamaGenとEmu3は潜在次元をそれぞれ8と4に削減し、コードブックサイズを16384と32768に拡張している。しかし、潜在次元の削減は、エンコードされた特徴からの情報損失をもたらし、最終的にパフォーマンスを低下させる。

1.2 主要な貢献

MGVQの主要な貢献は以下の通りである。

  • 離散コードブックの表現能力を10億倍以上に拡張する新しい手法の提案により、VAEとVQ-VAE間の再構成ギャップを縮小
  • 潜在次元を保持し、単一の統一されたコードブックの代わりに、サブコードブックのグループを使用し、ネストされたマスキング訓練戦略を適用する。これによりコードブックの最適化を容易にし、表現能力を向上させ、情報損失を削減
  • 512pと2k解像度のゼロショットベンチマークを構築し、再構成の汎化能力を包括的に評価。ImageNet 256pベンチマークと7つのゼロショットHDデータセットでSOTA性能を実現（rFID 0.49でSD-VAEの0.91から46%改善）

2. 提案手法

2.1 手法の概要

MGVQは、離散コードブックの表現能力を向上させるための新しい手法である。従来の手法が次元削減によってコードブックの崩壊を回避しようとしたのに対し、MGVQは潜在次元を保持して表現力のある特徴を維持しながら、情報損失を回避する。

具体的には、潜在埋め込みをサブトークンに分割し、それぞれを独立したサブコードブックにマッピングする。これにより、最適化を容易にし、より大きな表現能力を実現する。また、訓練中にネストされたマスキング戦略を採用し、エンコーダーが画像を順序立てて潜在空間に圧縮するよう強制する。

2.2 技術的詳細

MGVQの技術的な核心は、マルチグループ量子化メカニズムにある。

マルチグループ量子化：
各潜在埋め込み z(x) ∈ R^{C_l} を、G個のサブトークン {z(x_i)}_{i=1}^G に均等に分割する。各サブトークンは同じ次元 C_l/G を持つ。対応して、G個の非共有サブコードブックを設定し、それぞれが同じサイズKを持つ。各サブトークンは、関連するサブコードブックを使用して個別に量子化される。

最終的に、すべての量子化されたサブトークンを連結して再構成に使用する。例えば、潜在次元C_lを32、グループサイズGを4、総コードブックサイズを最大32768に設定した場合、コードブック次元は8と非常に小さく、サブコードブックサイズ（8192）は既存手法よりもはるかに小さくなる。これによりコードブックの崩壊が緩和され、8192^4という大きな表現能力が実現される。

ネストマスキング：
サブコードブックが同じものに収束するのを避けるため、訓練中にいくつかのサブトークンをネストされた方法でマスクすることを提案する。具体的には、量子化されたサブトークンが与えられた場合、保持するユニット数M_keep ∈ {1, ..., G}をランダムにサンプリングし、最後のG-M_keepトークンをマスクして削除する。この訓練方法により、エンコーダーは画像を順序立てて潜在空間に圧縮し、デコーダーは粗から細へと画像を再構成することを学習する。

2.3 新規性

MGVQの新規性は以下の点にある。

  • 次元削減ではなく次元保持による情報損失の回避という逆転の発想
  • 単一の大きなコードブックではなく、複数の小さなサブコードブックの組み合わせによる表現能力の指数的拡張
  • ネストマスキングによる順序付けられた表現学習の導入

3. 実験結果

3.1 実験設定

評価は以下のデータセットとメトリクスで実施された。

  • ImageNet 50k検証セット（256×256解像度）
  • ゼロショットベンチマーク（512×512および2560×1440解像度）
  • UHDBench：2k解像度の2293画像を含む新しいゼロショットベンチマーク
  • 評価指標：PSNR、rFID（reconstruction-FID）、SSIM、LPIPS

3.2 主要な結果

ImageNetベンチマークでの結果：
16倍ダウンサンプリングの場合、MGVQ-G8は以下の優れた性能を示した。

  • rFID: 0.49（SD-VAEの0.91から46%改善）
  • PSNR: 24.70
  • SSIM: 0.787

8倍ダウンサンプリングの場合はさらに良好な結果を達成した。

  • rFID: 0.27
  • PSNR: 29.96
  • SSIM: 0.918

ゼロショット評価：
512×512解像度の6つのベンチマークでMGVQは最高性能を達成した。具体的にはCelebA、DAVIS、TextOCR、VFHQ、Spring、ENeRFのすべてで最高スコアを記録した。特に注目すべきは、PSNRにおいてSD-VAEを一貫して上回った点である。

2560×1440解像度のUHDBenchとDAVISでも、PSNRで最高性能を記録した。

3.3 既存手法との比較

MGVQは既存のVQ-VAE手法（VQGAN、LlamaGen、OpenMagvit2、VAR）をすべて上回った。特に重要なのは、連続トークナイザーであるSD-VAEやSDXL-VAEと比較しても優れた性能を示した点である。これらの連続トークナイザーは数十億枚の画像で訓練されている。それにもかかわらず、MGVQはより良い再構成品質を達成した。

また、コードブック使用率においても100%を達成し、デッドポイント（未使用のコード）の問題を解決している。

4. 実用性評価

4.1 実装の容易性

MGVQの実装は比較的シンプルである。主な変更点は以下の通りである。

  • 既存のVQ-VAEアーキテクチャに対して、量子化層のみを変更
  • 複数のサブコードブックの管理は追加のインデックス処理のみで実現可能
  • ネストマスキングは訓練時のみ必要で、推論時には不要

GitHubでコードが公開予定であり、実装の詳細を確認できる。

4.2 計算効率

計算効率の観点では以下の特徴がある。

  • 推論時の計算コストは従来のVQ-VAEとほぼ同等
  • メモリ使用量は複数のサブコードブックにより増加する。しかし各サブコードブックが小さいため、全体的な増加は限定的
  • 訓練時はネストマスキングにより若干の計算オーバーヘッドが発生

4.3 応用可能性

MGVQは以下の応用分野で活用可能である。

  • 画像生成：LlamaGenの生成フレームワークにMGVQトークナイザーを適用した実験では、gFIDとISの両方で改善を確認
  • 画像圧縮：高い再構成品質により、画像品質を保ちながらデータサイズを削減するシステムの構築が可能
  • ビデオ処理：フレーム間の一貫性を保ちながら高品質な圧縮を実現
  • マルチモーダル学習：高品質な離散表現により、視覚情報とテキスト情報を統合的に扱うモデルの性能向上に貢献

5. まとめと所感

5.1 論文の意義

MGVQは、VQ-VAEの長年の課題であった「VAEとの性能差」を46%縮小した画期的な研究である。この研究の意義は以下の点にある。

第一に、「次元削減によるコードブック崩壊の回避」という従来の常識を覆し、「次元保持による情報損失の最小化」という新しいパラダイムを提示した。これは、VQ-VAEの設計思想に大きな影響を与える可能性がある。

第二に、単純ながら効果的なマルチグループ量子化により、表現能力を指数的に拡張することに成功した。これにより、離散表現でも連続表現に匹敵する品質を実現できることを実証した。

第三に、包括的なゼロショットベンチマークを構築し、手法の汎化能力を厳密に評価した。これは今後のVQ-VAE研究の標準的な評価プロトコルとなる可能性がある。

5.2 今後の展望

MGVQの今後の発展可能性として以下が考えられる。

  1. より大規模なグループ数への拡張：現在はG=4,8での実験だが、より大きなグループ数での性能向上の可能性がある。

  2. 適応的グループ分割：画像の複雑さに応じてグループ数を動的に調整する手法の開発。

  3. 階層表現の学習：ネストマスキングを拡張し、より洗練された階層的表現を学習する手法の探求。

  4. 大規模ビジョン-言語モデルへの統合：高品質な離散表現を活用した、計算コストを抑えたマルチモーダルモデルの構築。

総じて、MGVQはVQ-VAEの新しい可能性を切り開く重要な研究であり、画像処理と生成モデルの発展に大きく貢献する可能性を秘めている。