技術的サマリー：ArcVQ-VAE

問題定義

ベクトル量子化変分オートエンコーダ（VQ-VAE）は、画像モデリングにおける離散表現の学習のための基盤的な枠組みとして確立され、圧縮、生成、理解タスクの中核を担っています。しかし、標準的な VQ-VAE モデルは、そのコードブックの有限な容量に起因する重大な制限に直面しています。主要な問題はコードブックの崩壊であり、量子化中にコードブックベクトルの小さな部分集合のみが積極的に利用され、大部分が非活性のまま残る現象です。

著者らは、潜在空間における幾何学的な不均衡を根本原因として特定しました。従来のトレーニングでは、頻繁に選択される（高有用性の）コードブックベクトルがエンコーダ出力に向かって更新されるにつれて大きな $\ell_2$ ノルムを蓄積する一方、未使用のベクトルは初期化付近に留まり、ノルムはほぼゼロのままとなります。これにより以下が生じます：

1. 冗長性：高有用性のベクトルが互いにクラスタリングし、幾何学的な多様性が低下する。
2. 不十分なカバレッジ：潜在空間が均一に覆われず、モデルが豊かで多様な表現を捉える能力が制限される。
3. 最適化されていない利用：コードブックの限られた容量が十分に活用されず、再構成精度とサンプル品質が制限される。

既存の解決策は、コードブック選択の頻度やメカニズム（例：オンライン K-means、確率的量子化）に焦点を当てる傾向がありますが、潜在空間の根本的な幾何学的不均衡には十分に対処できていません。

手法：ArcVQ-VAE

本論文は、球面角マージン事前分布（SAMP）を導入する新しいベクトル量子化フレームワークであるArcVQ-VAEを提案します。このフレームワークは、追加のネットワークコンポーネントを導入することなく統一された確率的設定内で動作し、コードブックと潜在ベクトルに適用される 2 つの主要なメカニズムに依存します。

1. 球体境界ノルム正則化（BBNR）

特定のベクトルが過度なノルム増大を通じて潜在空間を支配するのを防ぐため、著者らはすべてのコードブックベクトルを時間依存のユークリッド球内に収まるように制約します。

• 初期化：コードブックベクトルは $\ell_2$ 正規化され、単位超球面上に位置します。
• 動的制約：ノルムの上限 $M(t)$ は、トレーニングステップ $t$ とともに指数関数的に増加します（$M(t) = \exp(\alpha \cdot t)$）。
• メカニズム：各トレーニングバッチの後、ベクトルのノルムが $M(t)$ を超える場合、そのベクトルは境界まで再スケーリングされます。この設計は、早期のトレーニング中に厳格な制約を課して premature なクラスタリングを防ぎつつ、後期段階ではより表現力のあるベクトル大きさを可能にするために制約を徐々に緩和します。

2. アークコサイン加算マージン損失（ArcLoss）

顔認識における ArcFace に着想を得たこの損失関数は、単位超球面上の潜在ベクトル間の角分離性を高めることを促します。

• 正規化：エンコーダ出力（$z_e$）とコードブックベクトル（$e_j$）の両方が $\ell_2$ 正規化され、最近傍選択基準がユークリッド距離から角（測地線）距離に変換されます。
• 損失の定式化：この損失は、コードブックエントリに最も近い上位 $k$ 個の潜在トークンを正のペアとして扱い、それ以外を負のペアとして扱います。正のペアに対して加算的な角マージン $m$ を適用します：
  $$L_A = -\frac{1}{K} \sum_{j=1}^{K} \log \frac{\sum_{i \in N_j^{(k)}} e^{s \cos(\theta_{i,j} + m)}}{\sum_{i \in N_j^{(k)}} e^{s \cos(\theta_{i,j} + m)} + \sum_{i \notin N_j^{(k)}} e^{s \cos \theta_{i,j}}}$$
  ここで、$\theta_{i,j}$ は潜在ベクトルとコードブックベクトルの間の角度です。
• 効果：これにより、正のペアはより密に整列（角度が小さくなる）し、負のペアはより遠くへ押しやられ、より識別性が高く均一に分散した潜在空間が作成されます。この損失計算中、コードブックベクトルにはストップグラディエント演算子が適用され、コードブックが狭いバッチ駆動の分布に崩壊するのを防ぎます。

全体の目的関数は、標準的な VQ-VAE 損失と ArcLoss を組み合わせ、初期トレーニングでは角の制約を強制し、後期では再構成精度を優先するために減衰係数 $\gamma(t)$ で重み付けされます。

主要な貢献

1. 新しいフレームワーク：コードブック利用の構造的限界を解決するために、潜在空間の幾何学を Spherical Angular-Margin Prior（SAMP）を通じて取り入れた ArcVQ-VAE の導入。
2. 効率性：この手法は実質的に追加の計算コストを伴わず、追加のネットワークコンポーネントを必要としません。ノルムベースのスケーリングと軽量な損失項の適用によって単純に実装されます。
3. 性能：包括的な実験により、ArcVQ-VAE は比較的小さなコードブックサイズであっても、再構成精度、表現の多様性、サンプル品質においてベースライン手法を一貫して上回ることが示されました。

実験結果

著者らは、ArcVQ-VAE を MNIST、CIFAR-10、ImageNet-1K で評価し、VQ-VAE、VQGAN、CVQ-VAE、RQ-VAE などの最先端のベースラインと比較しました。

• 再構成品質：ImageNet-1K において $16^2$ トークンマップと $K=1024$ のコードブックサイズを用いた場合、ArcVQ-VAE は rFID 1.63 を達成し、特定の構成において VQGAN（7.94）や CVQ-VAE（1.57）を上回りました。また、$32^2$ トークンマップでは 0.95 の rFID を達成しました。データセット全体を通じて、一貫して低い再構成誤差と高い SSIM/PSNR を示しました。
• コードブック利用率：ArcVQ-VAE は、ImageNet（$16^2$ トークン）で $K=1024$ の場合、**99%**のコードブック利用率を達成し、$32^2$ トークンの場合は 88% を達成しました。これは VQGAN（44%）よりも大幅に改善されたものであり、はるかに大きなコードブック（例：10 万エントリの VQGAN-LC）を使用する手法と同等の性能です。
• 生成品質：Latent Diffusion Model（LDM）のトークナイザーとして使用された場合、ArcVQ-VAE は ImageNet で FID 6.79、Inception Score 172.3 を達成し、MoVQ や CVQ-VAE などの他の VQ ベースのトークナイザーを上回りました。
• 定性的分析：可視化（t-SNE およびペアワイズ距離行列）により、ArcVQ-VAE がより均一に分散したコードブックエントリを生成し、コードワード間の距離が大きくなり、冗長性が減少していることが確認されました。定性的な再構成では、VQGAN に比べて微細な詳細（顔の特徴、テクスチャなど）の保存が優れていることが示されました。
• アブレーション研究：著者らは、BBNR と ArcLoss の組み合わせが改善の源泉であることを確認しました。どちらかのコンポーネントを削除すると性能が低下しました。また、この手法はハイパーパラメータの選択（上位 $k$、マージン $m$、スケーリング $s$）に対して頑健であり、LDM と PixelCNN の両方の異なる生成事前分布において効果的であることが示されました。

意義と主張

本論文は、ArcVQ-VAE が、潜在空間における均一性と識別性を促進する幾何学的制約を課すことで、コードブック崩壊問題に対するシンプルかつ効果的な解決策を提供すると主張しています。その意義は以下の点にあります：

• 効率性の向上：同様の利用率を達成するために巨大なコードブック（例：$K=100,000$）を必要とする手法と比較して、小さなコードブック予算（$K=1024$）で優れた性能を達成すること。
• 汎用性：このアプローチは特定の生成事前分布に依存せず、拡散モデルと自己回帰モデルの両方で恩恵を示すこと。
• 幾何学的洞察：潜在空間の幾何学的な不均衡（ノルムと角の分離）に対処することが、離散表現の有用性を最大化するために不可欠であることを強調すること。

著者らは謙虚な立場を維持し、この手法が最先端の実験結果を達成している一方で、特定の設計選択（上位 $k$ の選択や損失スケジューリングなど）はヒューリスティックであることを認めています。彼らは ArcVQ-VAE を、一般的な超球面モデリングの代替としてではなく、VQ ベースの離散トークナイザーを強化するために特別に設計された幾何学を考慮した正則化戦略として位置づけています。