EdVAE: Mitigating Codebook Collapse with Evidential Discrete Variational Autoencoders

要約

コードブックの崩壊は、ベクトル量子化変分オートエンコーダー (VQ-VAE) のような離散表現空間を使用して深い生成モデルをトレーニングする場合によく発生する問題です。
エンコーダがデータを表現するためにコードブックの埋め込み全体にわたる分布を直接学習する代替設計の離散変分オートエンコーダ (dVAE) でも同じ問題が発生することがわかりました。
ソフトマックス関数を使用して確率分布を取得すると、最も一致するコードブック要素に過信した確率が割り当てられるため、コードブックの崩壊が発生するという仮説を立てます。
この論文では、dVAE のコードブック崩壊問題に対処するために、ソフトマックスの代わりに証拠深層学習 (EDL) を組み込む新しい方法を提案します。
我々は、ソフトマックスの使用とは対照的に、コードブック埋め込み全体で確率分布を達成することの重要性を明らかに監視しています。
さまざまなデータセットを使用した実験では、EdVAE と呼ばれるモデルが、再構成パフォーマンスを向上させながらコードブックの崩壊を軽減し、dVAE および VQ-VAE ベースのモデルと比較してコードブックの使用率を向上させることがわかりました。

要約(オリジナル)

Codebook collapse is a common problem in training deep generative models with discrete representation spaces like Vector Quantized Variational Autoencoders (VQ-VAEs). We observe that the same problem arises for the alternatively designed discrete variational autoencoders (dVAEs) whose encoder directly learns a distribution over the codebook embeddings to represent the data. We hypothesize that using the softmax function to obtain a probability distribution causes the codebook collapse by assigning overconfident probabilities to the best matching codebook elements. In this paper, we propose a novel way to incorporate evidential deep learning (EDL) instead of softmax to combat the codebook collapse problem of dVAE. We evidentially monitor the significance of attaining the probability distribution over the codebook embeddings, in contrast to softmax usage. Our experiments using various datasets show that our model, called EdVAE, mitigates codebook collapse while improving the reconstruction performance, and enhances the codebook usage compared to dVAE and VQ-VAE based models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google