技術的概要：ワッサーシュタイン勾配流を用いた離散自己回帰事前分布の学習

1. 問題定義

離散画像トークナイザーは、通常、2 段階のレシピを用いて訓練される。まず、トークナイザーは画像の再構成（条件付き尤度に焦点を当てた）のために訓練され、次に、得られた固定化されたトークン系列に対して自己回帰（AR）事前分布が訓練される。この分離は根本的なミスマッチを生み出す。つまり、トークナイザーは画像情報を保持することを学習するが、AR 事前分布の帰納的バイアスについては認識しないままとなる。その結果、学習されたトークンは再構成には優れているものの、左から右へと順次予測する AR モデルにとっては予測が困難である可能性がある。

既存の手法は、デコーダーを修正すること（例えば、テイルドロップアウト、拡散デコーダー）によって、このミスマッチを緩和し、順次構造を促進しようとする。しかし、これらのアプローチは依然としてデコーダー中心の監督に依存しており、AR 事前分布のモデル化の難易度を直接利用してトークナイザーの学習を導くものではない。さらに、連続潜在空間向けに設計された標準的な分布レベルのマッチング手法（例えば、WAE-MMD、InfoVAE）は、量子化を通じた微分可能な経路が存在しないため、ハードな離散トークナイザーには直接適用できない。

2. 手法

理論的枠組み：三項変分整合性（TVC）

著者は、潜在変数学習を分析するために**三項変分整合性（Tripartite Variational Consistency: TVC）**を導入する。TVC は、生成目的をベイズの定理から導出された 3 つの整合性条件に分解する。

1. 条件付き尤度整合性: モデルのデコーダーと、データにアンカーされた事後分布との整合。
2. 事前分布整合性: モデルの事前分布と、エンコーダーの集積事後分布との整合。
3. 事後分布整合性: モデルの事後分布と、データにアンカーされた事後分布との整合。

著者は、標準的な 2 段階訓練が第 1 段階で尤度整合性を満たし、第 2 段階で事前分布整合性を満たす一方で、集積事後分布が本質的に「AR 友好型」であることを保証するための信号をトークナイザーに提供しないことを指摘する。一方、ELBO 型の訓練はインスタンスごとの事前分布整合性を強制するため、事後分布の崩壊（posterior collapse）を招くことが多い。

提案手法：WGF-DPD および wAR-Tok

このギャップを埋めるため、本論文は**分布事前分布発散（Distributional Prior Divergence: DPD）*を提案する。これは、個々のインスタンスを制約することなく、トークンの集積*分布を目標とする AR 事前分布に一致させる項である。離散トークンに対してこれを最適化するため、**ワッサーシュタイン勾配流（Wasserstein Gradient Flow: WGF）**を採用する。

中核となるメカニズムWGF-DPDは、以下のように動作する。

1. プロキシモデル: トークナイザー（$q_\phi$）が生成する現在の集積トークン分布を追跡するように、軽量な補助 AR モデル（$q_\psi$）を訓練する。
2. スコアベース更新: トークナイザーの更新方向は、集積事後分布と目標事前分布との間の KL 発散のワッサーシュタイン勾配から導出される。これにより、粒子レベルの勾配更新が得られる。
   $$\nabla_z \mathcal{L}_{DPD} \approx -\nabla_z \log p_\theta(x|z) + (\log q_\psi(z) - \log p_\theta(z))$$
   • 第一項は標準的な再構成勾配である。
   • 第二項は事前分布マッチングスコアであり、プロキシモデルと目標 AR 事前分布の対数確率の差を表す。
3. 効率的な最適化: 決定的に、この更新にはプロキシおよび目標 AR モデルを通じた順伝播（フォワードパス）のみが必要である。AR モデルを通じた逆伝播は不要である。勾配は、微分不可能な量子化ステップを跨いで、ストレートスルー推定器（STE）を介してエンコーダーパラメータに伝播される。

得られたトークナイザーwAR-Tokは、再構成目的と WGF-DPD シグナルと共に共同訓練される。

3. 主要な貢献

• TVC 枠組み: 潜在変数学習の理論的分解であり、「事前分布整合性」を明確に独立した目的として特定し、2 段階訓練がトークナイザーに事前分布のモデル化の難易度について認識させない理由を説明する。
• WGF-DPD アルゴリズム: ワッサーシュタイン勾配流を用いて集積トークン分布を AR 事前分布に整合させる、離散トークナイザー向けの新しい最適化手法。離散量子化を通じた逆伝播の非計算可能性や、敵対的マッチングの不安定性を回避する。
• wAR-Tok: WGF-DPD シグナルをトークナイザー訓練に統合する実用的な実装。エンコーダー側の整合シグナルを生成するために、補助 AR モデルを通じた順伝播のみを必要とする。

4. 実験結果

本手法は、1D トークナイゼーション、テイルドロップアウト、拡散デコーダー、およびタイドタイムステップを含むベースラインと比較し、CIFAR-10およびImageNet 256×256で評価された。

• CIFAR-10:
  • 再構成品質を一致させた場合（rFID $\approx$ 5.0）、wAR-Tok は生成 FID（gFID）で14.50、AR 生成損失（gLoss）で0.42を達成し、最強のベースライン（ImageFolder/One-D-Piece: gFID 17.60, gLoss 0.78）を大幅に上回った。
  • 訓練反復回数を一致させた場合、wAR-Tok は gFID 11.30、gLoss 0.34に達し、より速い収束と高品質を示した。
• ImageNet 256×256:
  • ベースラインと同じエンコーダー・デコーダーおよび AR 事前分布アーキテクチャを使用し、wAR-Tok は gFID を6.32 から 5.42へ、評価用 AR 損失を7.89 から 7.07へ削減した。再構成品質にはわずかなトレードオフ（rFID 2.05 $\to$ 2.35）があった。
• 訓練ダイナミクス: AR 訓練曲線は、ベースラインと比較して訓練全体を通じて wAR-Tok がより低いクロスエントロピー損失と高いトークン予測精度を維持することを示しており、学習されたトークンが本質的に AR 事前分布にとってモデル化しやすいことを示唆している。

5. 意義と主張

本論文は、wAR-Tokが、トークナイザー訓練中に分布レベルの事前分布マッチングシグナルを導入することで、離散トークナイザーと自己回帰事前分布の間のミスマッチを効果的に解決すると主張する。

• 効率性: この手法は計算効率が高く、訓練ステップごとにプロキシとターゲットの 2 回の順伝播を追加するだけで済み、AR モデルを通じた高価な逆伝播を回避する。
• 汎用性: AR 事前分布に焦点を当てているが、TVC 枠組みおよび WGF-DPD アプローチは事前分布に依存しないものとして提示されており、対数確率が計算可能な任意の離散事前分布に適用可能である。
• 謙虚さ: 著者は明示的に、WGF-DPD の効果を分離するために、最近のトークナイザーで一般的な補助技術（DINO 整合、蒸留、拡散デコーダー事後訓練など）を意図的に除外したため、絶対的な最先端（SOTA）性能を主張するものではないと述べている。結果は、制御された環境における提案された整合目的の有効性を示している。

この研究は、単なる再構成可能性だけでなく、下流の生成モデルによるトークンの「予測可能性」を明示的に最適化することが、より効率的で高品質な離散生成システムにつながることを示唆している。