[論文レビュー] Vector Quantized Wasserstein Auto-Encoder

Tung-Long Vuong, Trung Le+5 more|arXiv (Cornell University)|Feb 12, 2023

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

本稿では、離散コード語系列の分布とデータ分布の間の Wasserstein 距離を最小化することによって、深層離散潜在表現を学習するための新規手法、Vector Quantized Wasserstein Auto-Encoder (VQ-WAE) を提案する。Wasserstein 距離の理論的性質を活用し、最適化可能な形に目的関数を再定式化することで、VQ-VAE や SQ-VAE の変種と比較して、より優れたコードブック利用効率、画像再構成および生成性能、およびより分離され、クラスタリングされた潜在表現を達成した。

ABSTRACT

Learning deep discrete latent presentations offers a promise of better symbolic and summarized abstractions that are more useful to subsequent downstream tasks. Inspired by the seminal Vector Quantized Variational Auto-Encoder (VQ-VAE), most of work in learning deep discrete representations has mainly focused on improving the original VQ-VAE form and none of them has studied learning deep discrete representations from the generative viewpoint. In this work, we study learning deep discrete representations from the generative viewpoint. Specifically, we endow discrete distributions over sequences of codewords and learn a deterministic decoder that transports the distribution over the sequences of codewords to the data distribution via minimizing a WS distance between them. We develop further theories to connect it with the clustering viewpoint of WS distance, allowing us to have a better and more controllable clustering solution. Finally, we empirically evaluate our method on several well-known benchmarks, where it achieves better qualitative and quantitative performances than the other VQ-VAE variants in terms of the codebook utilization and image reconstruction/generation.

研究の動機と目的

• 既存の VQ-VAE 変種におけるコードブック利用効率と表現品質の限界を是正すること。
• 最適輸送理論を用いた生成的アプローチから深層離散表現を学習することの探求。
• Wasserstein 距離とベクトル量子化を結びつける理論的根拠に基づくフレームワークの構築。
• コードブック内のコード語の均一かつ意味のある使用を保証することで、学習された表現の品質を向上させること。
• Wasserstein 距離の最小化が、より優れた再構成、生成、および下流の表現性能をもたらすことを実証すること。

提案手法

• 生成モデルを定式化し、デコーダがコード語系列の離散分布を、Wasserstein 距離を最小化することでデータ分布へマッピングする。
• 再構成項と潜在表現とコード語分布の間の Wasserstein 距離を含む等価な目的関数を導出する。
• Wasserstein 距離とクラスタリングを結びつける理論的枠組みを導入し、潜在空間の構造的制御を可能にする。
• 確定的エンコーダとデコーダを用い、潜在空間における最近傍割当により定義されるベクトル量子化を実装する。
• 訓練の安定化を図り、バックプロパゲーション中にコードブックを固定しつつエンコーダへの勾配伝播を保証するため、ストップ勾配操作を適用する。
• 画像生成を可能にするために、簡単な離散潜在系列の事前分布（例：自己回帰モデル）を採用する。

Figure 1: Reconstruction results for the FFHQ dataset.

実験結果

リサーチクエスチョン

• RQ1 VAE に類似したフレームワークにおいて、Wasserstein 距離を効果的に活用して離散表現学習を改善できるか？
• RQ2 潜在コード分布とデータ分布の間の Wasserstein 距離を最小化することで、より優れたコードブック利用効率と表現品質が得られるか？
• RQ3 Wasserstein 距離の理論的性質を活用して、より構造的でクラスタリングされた潜在表現を実現できるか？
• RQ4 VQ-WAE は、画像再構成、生成、および下流の分類において、VQ-VAE や SQ-VAE と比較してどのように異なるか？
• RQ5 提案されたフレームワークは、最適化の改善により、より多様で意味のあるコード語を自然に導くことができるか？

主な発見

• VQ-WAE は、すべてのデータセットにおいて、1層分類精度が最高を記録した。CIFAR10 では 50.19%、MNIST では 95.62%、SVHN では 38.38% を達成し、VQ-VAE や SQ-VAE を上回った。
• 画像生成において、VQ-WAE は FID スコアが最低であった。CIFAR10 では非条件付きで 87.73、条件付きで 88.51、MNIST では非条件付きで 8.21、条件付きで 3.88 を記録し、優れた画像品質を示した。
• VQ-WAE ではコードブック利用効率が顕著に向上し、より多くのコード語が活用され、VQ-VAE や SQ-VAE と比較して、死活コード語や未使用コード語が著しく減少した。
• 潜在空間の可視化結果から、VQ-WAE の表現は各クラスごとに明確に分離され、凝集したクラスタを形成している一方、VQ-VAE や SQ-VAE の表現はクラス間で散らばり、重複していた。
• アブレーションスタディにより、VQ-WAE はコードブックを最適に利用しており、より多様で意味のあるコード語と、より構造化された潜在多様体を実現していることが確認された。
• 理論的分析により、Wasserstein 距離の定式化における最適カップリング（定理 2.3 の OP）は、コード語系列の連合分布を学習する道筋を示唆しており、これにより生成モデルのさらなる向上が可能であると考えられた。

Figure 2: Latent distribution over the codebook on test-set.