[論文レビュー] SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization

Yuhta Takida, Takashi Shibuya+8 more|arXiv (Cornell University)|May 16, 2022

Domain Adaptation and Few-Shot Learning被引用数 13

ひとこと要約

本稿では、確率的量子化と自己自己冷却トレーニングを備えた変分オートエンコーダーであるSQ-VAEを提案する。この手法は、ストップ・グラデントやコードブックリセットといったヒューリスティクスに依存せずに、コードブックの利用を向上させる。トレーニング可能な確率的デ量子化と量子化を導入することで、SQ-VAEはトレーニング中に自然に決定論的量子化へと自己自己冷却され、視覚および音声タスクにおいてVQ-VAEや標準VAEよりも優れた再構成性能と生成性能を達成する。

ABSTRACT

One noted issue of vector-quantized variational autoencoder (VQ-VAE) is that the learned discrete representation uses only a fraction of the full capacity of the codebook, also known as codebook collapse. We hypothesize that the training scheme of VQ-VAE, which involves some carefully designed heuristics, underlies this issue. In this paper, we propose a new training scheme that extends the standard VAE via novel stochastic dequantization and quantization, called stochastically quantized variational autoencoder (SQ-VAE). In SQ-VAE, we observe a trend that the quantization is stochastic at the initial stage of the training but gradually converges toward a deterministic quantization, which we call self-annealing. Our experiments show that SQ-VAE improves codebook utilization without using common heuristics. Furthermore, we empirically show that SQ-VAE is superior to VAE and VQ-VAE in vision- and speech-related tasks.

研究の動機と目的

• 決定論的量子化とヒューリスティクスに基づくトレーニングスキームにより、コードブックベクトルの一部しか使用されないVQ-VAEにおけるコードブックコラプスの問題に対処すること。
• ストップ・グラデント、EMA、またはコードブックリセットに依存しない、標準的な変分ベイズの枠組み内で動作するトレーニングフレームワークの開発。
• 確率的量子化の確率性を時間経過とともに減少させる自己自己冷却プロセスにより、コードブックのより高い利用を実現すること。
• カテゴリカルデータに特化した変種を用いて、データタイプに一般化可能な手法の設計。特に、von Mises–Fisher分布を用いたカテゴリカルデータ向けの最適化。
• ハイパーパramータチューニングや補助技術を用いずに、視覚および音声ベンチマークにおいて再構成と生成の両面で優れた性能を示すこと。

提案手法

• 連続的な潜在変数を、微分可能でトレーニング可能な確率分布を用いて、離散的コードブックベクトルへマッピングする確率的デ量子化プロセスを導入する。
• コードブックベクトル上でのトレーニング可能な後方カテゴリカル分布を用いた確率的量子化を適用し、離散的選択の勾配伝播を可能にする。
• 量子化プロセスの確率性がトレーニング中に減少する自己自己冷却メカニズムを用いることで、自然に決定論的量子化へ収束する。
• ストレートスラッシュ推定を必要としない、標準的なバックプロパゲーションで最適化可能な、証拠下限界（ELBO）としての目的関数を定式化する。
• 2つの変種を提案する：一般データ向けのガウス型SQ-VAEと、カテゴリカルデータ向けのvMF SQ-VAE。後者は、離散分布における性能向上を図るためにvon Mises–Fisher分布を用いる。
• Gumbel-Softmaxまたはコンcreteリラクゼーションのアプローチに従い、唯一の主要ハイパーパramータ（確率性の温度）のみを用いてモデルをトレーニングする。

実験結果

リサーチクエスチョン

• RQ1 VAEフレームワークにおける確率的量子化は、決定論的VQ-VAEと比較してコードブックの利用を向上させることができるか？
• RQ2 SQ-VAEにおける自己自己冷却プロセスは、ヒューリスティクスによる干渉なしに、より優れた表現学習と再構成性能をもたらすか？
• RQ3 SQ-VAEは、標準VAEおよびVQ-VAEと比較して、視覚および音声タスクでどのように性能を発揮するか？
• RQ4 vMF SQ-VAEは、画像ピxlsなどのカテゴリカルデータ分布において、ガウス型SQ-VAEを上回る性能を発揮するか？
• RQ5 コードブックサイズの拡大は、VQ-VAEではしばしばコラプスを引き起こすが、SQ-VAEでは性能向上をもたらすか？

主な発見

• SQ-VAEは、ストップ・グラデント、EMA、コードブックリセットを一切使用せず、コードブックの利用度をVQ-VAEを上回り、コードブックの有効な探索が可能であることを示している。
• SQ-VAEにおける自己自己冷却プロセスは、量子化における確率性を段階的に低減させ、安定した収束とより優れた表現学習を実現している。
• MNISTおよびGray-CelebAでは、VAEおよびVQ-VAEと比較して、SQ-VAEがより優れた再構成性能を示し、再構成誤差が低くなっている。
• CelebA-HQ 256×256では、SQ-VAEが高精細な再構成および生成サンプルを生成しており、大規模データセットへのスケーラビリティを示している。
• vMF SQ-VAEは、カテゴリカルデータにおいてガウス型SQ-VAEを顕著に上回り、離散分布へのvon Mises–Fisher分布の有効性を裏付けている。
• SQ-VAEではコードブックサイズの拡大が性能向上をもたらす一方、VQ-VAEでは大きなコードブックがしばしばコラプスを引き起こし、効果が得られない。