eXact-Prior Variational Autoencoder (X-VAE): Learning Data-Adaptive Gaussian Mixture Priors for Latent Distributions

問題提起

標準的な変分オートエンコーダ（VAE）は、潜在空間に対して固定された等方性ガウス事前分布 $p(z) = \mathcal{N}(0, I)$ に依存しています。これは、証拠下界（ELBO）におけるカルバック・ライブラー（KL）ダイバージェンスの計算には数学的に便利ですが、この仮定は実世界のデータセットにおける潜在表現の真の複雑な構造を捉えることができない場合が多く、以下の問題を引き起こします。

• 事前分布と事後分布のミスマッチ（Prior-Posterior Mismatch）: 学習された潜在分布が標準正規分布から大きく逸脱し、再構成精度やサンプルの品質を制限します。
• 事後崩壊（Posterior Collapse）: モデルが潜在変数を無視してしまい、生成能力が低下します。
• 表現力の限界: 硬直的な事前分布が潜在空間を制約し、モデルが多峰性（マルチモーダル）や構造的に制約されたデータを効果的に表現することを妨げます。

VampPrior、フローベースの事前分布、あるいは階層型モデルといった既存の解決策は、柔軟性を向上させますが、多くの場合、追加の学習可能パラメータ、計算量の増加、および最適化手順の複雑化を招きます。

手法：X-VAEフレームワーク

著者らは、事前分布の推定と変分学習を分離する eXact-Prior Variational Autoencoder (X-VAE) を提案しています。X-VAEは、複雑な事前分布ネットワークを学習する代わりに、**学習済み決定論的オートエンコーダ（AE）**の潜在表現から、データ適応型の事前分布を導出します。

手法は以下の2つのステージで進行します。

1. 事前分布の推定（事前学習ステージ）
再構成誤差を最小化するように決定論的AEを学習させます。このAEによって生成された潜在コード $\{z_i^{AE}\}_{i=1}^N$ を用いて、事前分布の統計量を推定します。著者らは、事前分布 $p_{AE}(z)$ のための2つの構成を提案しています。

• 単一ガウス事前分布: AEの潜在コードの経験的な平均（$\mu_{AE}$）と共分散（$\Sigma_{AE}$）によって、単一のガウス事前分布を定義します：$p_{AE}(z) = \mathcal{N}(\mu_{AE}, \Sigma_{AE})$。
• ガウス混合モデル（GMM）事前分布: AEの潜在コードに対して、$K$ 個のコンポーネントを持つ対角ガウス混合モデルを適合させます（例：期待値最大化法による）。この混合モデルは固定され、事前分布として機能します：
  $$p_{AE}(z) = \sum_{k=1}^K \pi_k^p \mathcal{N}(z; \mu_k^p, \text{diag}((\sigma_k^p)^2))$$
  ここで、混合重み $\pi^p$、平均 $\mu^p$、および標準偏差 $\sigma^p$ は、AEから導出された固定パラメータであり、VAEの学習中には学習されません。

2. ルーテッド・トランスポートを用いた変分学習
VAEはこの固定された事前分布を用いて学習されます。近似事後分布 $q_\phi(z|x)$ を混合事前分布に適合させるために、著者らは**ルーテッド・トランスポート（routed-transport）**メカニズムを導入しています。

• カテゴリカル・ルーティング: 各潜在座標 $k$ について、事前分布の重み $\pi^p$（または事後分布の重み $\pi^q$）から混合コンポーネント $c_k$ が選択されます。
• トランスポート演算子: 事前分布の選択されたコンポーネントの分布からノイズサンプル $\epsilon_k$ が抽出されます。このノイズは、線形写像を用いてエンコーダの事後分布へと「輸送」されます：
  $$T_k(\epsilon_k) = \mu_k^q + \sigma_k^q \frac{1}{\sigma_k^p} (\epsilon_k - \mu_k^p)$$
  これにより、もし $\epsilon_k \sim \mathcal{N}(\mu_k^p, (\sigma_k^p)^2)$ であれば、$T_k(\epsilon_k) \sim \mathcal{N}(\mu_k^q, (\sigma_k^q)^2)$ となり、バックプロパゲーションのための再パラメータ化トリックが維持されます。
• 損失関数: ELBOは、閉形式のKLダイバージェンスを用いて最適化されます。GMMの場合、KL項はカテゴリカル項（コンポーネント割り当ての整合）と、コンポーネントごとのガウス項の和によって抑えられます：
  $$D_{KL}(q_\phi(z|x) \parallel p(z)) \leq D_{KL}(\pi^q \parallel \pi^p) + \sum_{k=1}^K \pi_k^q D_{KL}(\mathcal{N}_k^q \parallel \mathcal{N}_k^p)$$

3. 生成のための分散スケーリング
生成フェーズにおいて、著者らは事前分布の標準偏差に適用されるグローバルなスケーリング因子 $\alpha$ を導入しています。
$$z_{gen} \sim \mathcal{N}(\mu_{AE}, \text{diag}((\alpha \sigma_{AE})^2))$$
これにより、サンプルの忠実度（$\alpha$ が小さい場合、経験的な多様体に近いサンプル）と多様性（$\alpha$ が大きい場合、頻度の低い領域を探索するサンプル）のトレードオフを、再学習なしで明示的に制御することが可能になります。

主な貢献

1. データ適応型事前分布: 等方性ガウス分布を、学習済みAEの潜在統計量から推定された事前分布に置き換えることで、VAEを経験的なデータ分布に根付かせます。
2. 解析的な扱いやすさ: 提案された事前分布に対して閉形式のKLダイバージェンスを導出し、複雑な学習済み事前分布ネットワークやフローベースの変換を回避します。
3. ルーテッド・トランスポート・メカニズム: エンコーダの事後分布を混合コンポーネントにマッピングする新しいサンプリング戦略を提供し、再パラメータ化の勾配を維持しながら、潜在空間をAEのクラスター構造に固定します。
4. 制御可能な多様性: 推論時に多様性と忠実度のトレードオフを調整できる、単純な分散スケーリング機構（$\alpha$）を提供します。
5. パラメータ効率性: 事前分布は固定バッファとして格納されるため、生成モデルに学習可能なパラメータをゼロで追加します。

実験結果

著者らはMNIST、CelebA、および合成クラスターデータセットを用いてX-VAEを評価しました。

• 性能: CelebAにおいて、最良のX-VAE構成は、標準的なVAEやVampPrior、階層型VAEを上回る、最も低いテスト目的関数（164.65）と最高のFréchet Inception Distance（FID）80.40を達成しました。MNISTでは、X-VAE構成は標準的なVAE（47.21）に対してFIDを改善し（49.67）、VampPriorと同等またはわずかに上回る結果を示しました。
• 合成データ: 3つのクラスターを持つ合成データセットにおいて、X-VAEはデータの多峰構造を保持することに成功しましたが、標準的なVAEはモードを単一のガウス分布へと崩壊させました。
• 再構成: X-VAEは、ベースラインと同等の再構成品質を維持しながら、より現実的で多様なサンプルを生成しました。
• 効率性: 本手法は事前分布のために追加の学習可能パラメータを必要とせず、GM-VAEや階層型VAEのような学習済み事前分布を用いるアプローチと比較して計算効率が高いです。

意義と主張

本論文は、X-VAEが複雑な学習済み事前分布フレームワークに対するシンプルかつ効果的な代替案であることを主張しています。学習済みAEの経験的な潜在構造に事前分布を固定することで、X-VAEは以下を実現します。

• 標準的なVAEに固有の事前分布と事後分布のミスマッチを軽減します。
• 標準的なVAEの解析的な扱いやすさと計算効率を維持します。
• 生成された解が厳格な構造的制約を満たしつつ多様性を維持する必要がある、エンジニアリングや工業デザインなどの制約付き生成モデリングのための実用的なフレームワークを提供します。
• 高品質な生成モデリングには、必ずしも複雑でパラメータの多い事前分布ネットワークが必要ではなく、学習済み表現の統計的構造を活用することで、優れた整合性と制御性が得られることを示しています。

著者らは、学習済みAEの品質への依存や、AEとVAEの潜在空間間の残留するミスマッチの可能性といった限界も認めています。また、対角GMMの仮定がモード内の相関を捉えきれない可能性や、高次元において混合コンポーネントの数が自然なカテゴリ数を過剰に指定してしまう可能性についても述べています。今後の課題として、非ガウス事前分布、階層的拡張、および結合学習戦略の探索が示唆されています。