2025-11-10T02:44:53.419690

Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models

Zheng, Li
Vector Quantized Variational Autoencoders (VQ-VAEs) leverage self-supervised learning through reconstruction tasks to represent continuous vectors using the closest vectors in a codebook. However, issues such as codebook collapse persist in the VQ model. To address these issues, existing approaches employ implicit static codebooks or jointly optimize the entire codebook, but these methods constrain the codebook's learning capability, leading to reduced reconstruction quality. In this paper, we propose Group-VQ, which performs group-wise optimization on the codebook. Each group is optimized independently, with joint optimization performed within groups. This approach improves the trade-off between codebook utilization and reconstruction performance. Additionally, we introduce a training-free codebook resampling method, allowing post-training adjustment of the codebook size. In image reconstruction experiments under various settings, Group-VQ demonstrates improved performance on reconstruction metrics. And the post-training codebook sampling method achieves the desired flexibility in adjusting the codebook size.
academic

ベクトル量子化モデルにおける自己拡張可能コードブックのグループ単位最適化

基本情報

  • 論文ID: 2510.13331
  • タイトル: Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models
  • 著者: Hong-Kai Zheng, Piji Li(南京航空航天大学)
  • 分類: cs.CV
  • 発表時期/会議: ICLR 2026
  • 論文リンク: https://arxiv.org/abs/2510.13331

要約

Vector Quantized Variational Autoencoders (VQ-VAEs)は、再構成タスクを通じた自己教師あり学習を行い、コードブック内の最も近いベクトルを使用して連続ベクトルを表現します。しかし、VQモデルではコードブック崩壊などの問題が依然として存在します。これらの問題に対処するため、既存の方法は暗黙的な静的コードブックまたはコードブック全体の共同最適化を採用していますが、これらの方法はコードブックの学習能力を制限し、再構成品質の低下をもたらします。本論文ではGroup-VQを提案し、コードブックをグループ単位で最適化します。各グループは独立に最適化され、グループ内では共同最適化が行われます。このアプローチは、コードブック利用率と再構成性能の間のトレードオフを改善します。さらに、訓練不要なコードブック再サンプリング方法を導入し、訓練後のコードブックサイズの調整を可能にします。様々な設定での画像再構成実験において、Group-VQは再構成指標で改善されたパフォーマンスを示しています。

研究背景と動機

問題記述

Vector Quantization (VQ)は、連続特徴を離散トークンにマッピングする技術であり、VQ-VAEで広く応用されています。しかし、従来のVQ訓練はコードブック利用率の低さという問題に直面しており、コードベクトルの一部のみが使用・更新され、「コードブック崩壊」をもたらし、モデルの符号化能力を制限します。

既存方法の限界

  1. Vanilla VQ: 各コードベクトルが独立に更新され、コードブック崩壊を引き起こしやすい
  2. Joint VQ方法(SimVQ、VQGAN-LCなど): 共有パラメータを通じてコードブック全体を共同最適化し、100%の利用率を達成できますが、コードブックの学習能力を制限します

研究動機

著者は実験を通じて、Joint VQは100%のコードブック利用率に迅速に達することができますが、同じ利用率の下では、その再構成品質はVanilla VQより劣っていることを発見しました。これはコードブック利用率と再構成性能の間にトレードオフが存在することを示しており、より良いバランス戦略が必要です。

核心的貢献

  1. Group-VQ方法の提案: VQモデルにおける利用率と再構成性能のバランスを取るグループベースのコードブック最適化方法
  2. Joint VQ方法の一般化: 共有パラメータの観点からJoint VQを再解釈し、訓練後のコードブックサンプリング方法を導入
  3. 訓練不要なコードブック調整: 訓練後にモデルを再訓練することなく、柔軟にコードブックサイズを調整可能
  4. 包括的な実験検証: 画像再構成タスクでGroup-VQとコードブック再サンプリングの有効性を検証

方法の詳細

タスク定義

画像 IRH×W×3I \in \mathbb{R}^{H \times W \times 3} が与えられたとき、VQ-VAEはまずエンコーダを使用して特徴マップ ZRh×w×dZ \in \mathbb{R}^{h \times w \times d} を取得し、その後量子化器を通じて各特徴ベクトル zRdz \in \mathbb{R}^d をコードブック C={qiqiRd,i=0,1,...,n1}C = \{q_i | q_i \in \mathbb{R}^d, i = 0,1,...,n-1\} 内の最も近いコードベクトルに置き換えます:

q=argminqiCzqi,i=0,1,...,n1q = \arg\min_{q_i \in C} \|z - q_i\|, i = 0,1,...,n-1

モデルアーキテクチャ

Group-VQの設計

Group-VQはコードブック CCkk 個の互いに素なグループ(サブコードブック)に分割します:

C=j=0k1Gj,GjGj= if jjC = \bigcup_{j=0}^{k-1} G_j, \quad G_j \cap G_{j'} = \emptyset \text{ if } j \neq j'

各グループ GjG_j は独立に更新され、グループ内では共同最適化が採用されます。コードベクトル qjtGjq_{jt} \in G_j の勾配更新は以下の通りです:

qjtLcmt=qjtLj\nabla_{q_{jt}} L_{cmt} = \nabla_{q_{jt}} L_j

これにより、各グループはそのグループ内のコードベクトルから生成される勾配のみの影響を受けることが保証されます。

コードブックのパラメータ化

各グループ GjG_j は共有パラメータを通じてパラメータ化されます:

Gj=G^jWj+bjG_j = \hat{G}_j W_j + b_j

ここで:

  • G^jRnj×rj\hat{G}_j \in \mathbb{R}^{n_j \times r_j}: コードブックコア(固定分布からサンプリング)
  • WjRrj×dW_j \in \mathbb{R}^{r_j \times d}: プロジェクタ(学習可能)
  • bjRdb_j \in \mathbb{R}^d: バイアスベクトル

技術的革新点

1. 統一的視点での分析

  • Vanilla VQ: k=nk = n、各コードベクトルが1つのグループ
  • Joint VQ: k=1k = 1、コードブック全体が1つのグループ
  • Group-VQ: 1kn1 \leq k \leq n、両極端のケースのバランスを取ります

2. コードブック再サンプリング機構

生成的コードブックの特性を利用して、訓練後にコードブックコアを再サンプリングできます:

q~=v^Wj,v^N(0,I)q̃ = v̂ W_j, \quad v̂ \sim \mathcal{N}(0, I)

2つのモードをサポートします:

  • 再サンプリング: コードブックを完全に置き換え
  • 自己拡張: 元のコードブックに基づいて新しいコードベクトルを追加

実験設定

データセット

  • ImageNet-1k: メインデータセット
  • MS-COCO: 補足検証
  • 入力解像度: 128×128、ダウンサンプリング係数f=8

評価指標

  • rFID(再構成FID): 再構成画像と元画像の分布距離
  • LPIPS(VGG16): 知覚的類似性
  • PSNR: ピーク信号対雑音比
  • SSIM: 構造類似性指数

比較方法

  • VQGAN、ViT-VQGAN、VQGAN-FC
  • FSQ、LFQ(固定コードブック方法)
  • VQGAN-LC、SimVQ(Joint VQ方法)

実装の詳細

  • 学習率: 1×10⁻⁴
  • オプティマイザ: Adam (β₁=0.5, β₂=0.9)
  • バッチサイズ: 32/GPU
  • ハードウェア: NVIDIA A5000 GPU

実験結果

主要な結果

ImageNet-1kでのパフォーマンス比較(コードブックサイズ65,536):

方法グループ数利用率rFID↓LPIPS↓PSNR↑SSIM↑
VQGAN65,5361.4%3.740.1722.200.706
SimVQ1100.0%1.990.1224.340.788
Group-VQ6499.9%1.860.1124.370.787

Group-VQはすべての指標で最高のパフォーマンスを達成し、ベースライン方法を大幅に上回っています。

アブレーション実験

異なるグループ数の影響:

グループ数13264128512
利用率100%100%100%95.6%78.8%
rFID↓6.456.056.096.116.28

実験は32~64グループが最適な選択であることを示しており、コードブック利用率と再構成性能のバランスが取れています。

コードブック再サンプリング実験

コードブックサイズ調整結果:

方法コードブックサイズrFID↓PSNR↑
Group-VQ65,5361.8724.32
+ ダウンサンプリング32,7682.1624.02
+ アップサンプリング131,0721.7924.49
+ 自己拡張131,0721.7624.51

結果はコードブック再サンプリング方法の有効性を検証し、コードブックサイズを柔軟に調整でき、予想されるパフォーマンス変化を得られることを示しています。

可視化分析

ランダム投影を使用してコードベクトルを2D空間に可視化すると、以下が発見されました:

  1. 異なるグループは異なる特徴分布を学習
  2. グループ内のコードベクトルは相対的に類似し、グループ間には大きな差異
  3. 各グループの統計特性(平均値、分散、使用頻度)には明らかな差異

関連研究

VQ改善方法の分類

  1. Straight-Through Estimator改善: 勾配伝播の最適化
  2. マルチインデックス量子化: RQ-VAE、Product Quantizationなど
  3. コードブック改善: 本論文の重点関心領域

Joint VQ方法

  • VQGAN-LC: 事前訓練特徴の初期化+投影層
  • SimVQ: ランダム初期化+行列再パラメータ化
  • LFQ/FSQ: 固定コードブックで崩壊を回避

本論文はこれらの方法を「共有パラメータを通じて実現されるJoint VQ」として統一し、この基礎の上でグループ単位の最適化戦略を提案しています。

結論と考察

主要な結論

  1. コードブック利用率と再構成品質の間にはトレードオフが存在: 100%の利用率が必ずしも最高の再構成効果をもたらすわけではありません
  2. グループ単位の最適化は効果的なバランス戦略: Group-VQはグループ数の調整を通じて柔軟な制御を実現
  3. コードブック再サンプリングは実用的価値を提供: 訓練後にコードブックサイズを柔軟に調整可能

限界

  1. 生成タスクでの検証がない: 再構成タスクのみでテストされ、生成モデルの検証が不足
  2. グループ数の選択には調整が必要: 最適なグループ数は具体的なタスクとデータセットに依存
  3. 計算複雑性: 複数グループの最適化は訓練時間を増加させる可能性

今後の方向性

  1. 生成モデル(自己回帰モデルなど)でのGroup-VQの効果検証
  2. 適応的グループ数選択戦略の探索
  3. Group-VQと他のVQ改善方法の組み合わせの研究

深い評価

利点

  1. 理論的貢献が明確: グループ最適化の観点から既存のVQ方法を統一的に理解し、新しい分析視点を提供
  2. 方法がシンプルで効果的: Group-VQの設計は直感的で、実装と理解が容易
  3. 実験が充分: 複数データセット、複数アーキテクチャでの包括的検証、詳細なアブレーション実験
  4. 実用的価値が高い: コードブック再サンプリング方法は実際の応用における柔軟性の必要性を解決

不足点

  1. 理論分析が十分でない: グループ単位の最適化がなぜより効果的なのかについての理論的説明が不足
  2. 適用範囲が限定的: 主に画像再構成に焦点を当てており、他のモダリティとタスクでの効果は未知
  3. 計算オーバーヘッド分析が不足: 複数グループ最適化の計算コストについて詳細な分析がない

影響力

  1. 学術的価値: VQ研究に新しい最適化思想を提供し、後続研究を刺激する可能性
  2. 実用的価値: コードブック再サンプリング方法は実際の展開で非常に価値がある
  3. 再現性: 著者がコード公開を約束しており、方法の普及に有利

適用シーン

  1. 画像/ビデオ符号化: 高品質再構成が必要な圧縮タスク
  2. マルチモーダル学習: 汎用的なベクトル量子化コンポーネントとして
  3. 生成モデル: 生成モデルに離散表現を提供するトークナイザーとして

参考文献

本論文は主に以下の重要な研究に基づいています:

  1. Van Den Oord et al. (2017) - VQ-VAE原始論文
  2. Zhu et al. (2024b) - SimVQ方法
  3. Yu et al. (2023) - LFQ方法
  4. Mentzer et al. (2023) - FSQ方法

まとめ: これはVQ領域に重要な貢献をもたらした論文です。Group-VQ方法はシンプルで効果的であり、コードブック最適化に新しい思想をもたらします。コードブック再サンプリング方法は非常に強い実用的価値を持っています。理論分析と適用範囲でさらに改善の余地がありますが、全体的には高品質の研究成果です。