[論文レビュー] Jigsaw-VAE: Towards Balancing Features in Variational Autoencoders

Saeid Asgari Taghanaki, Mohammad Havaei+4 more|arXiv (Cornell University)|May 12, 2020

Generative Adversarial Networks and Image Synthesis参考文献 41被引用数 6

ひとこと要約

本稿では、確率的空間的並べ替えと混合事前分布を適用することで、潜在表現における特徴のアンバランスを軽減する変分オートエンコーダー、Jigsaw-VAEを提案する。これにより、主な特徴と希少な特徴の両方のバランスの取れた学習が可能となり、分布シフト下の下流クラスタリングタスクで性能が向上し、ベースラインVAEと比較して生成の多様性と現実性が向上する。

ABSTRACT

The latent variables learned by VAEs have seen considerable interest as an unsupervised way of extracting features, which can then be used for downstream tasks. There is a growing interest in the question of whether features learned on one environment will generalize across different environments. We demonstrate here that VAE latent variables often focus on some factors of variation at the expense of others - in this case we refer to the features as ``imbalanced''. Feature imbalance leads to poor generalization when the latent variables are used in an environment where the presence of features changes. Similarly, latent variables trained with imbalanced features induce the VAE to generate less diverse (i.e. biased towards dominant features) samples. To address this, we propose a regularization scheme for VAEs, which we show substantially addresses the feature imbalance problem. We also introduce a simple metric to measure the balance of features in generated images.

研究の動機と目的

• VAEにおける特徴のアンバランス問題に対処すること。具体的には、色のような支配的特徴に過剰適合し、形状のような希少な特徴を損なう問題を解消すること。
• テストデータの分布が学習時と異なる際の下流タスク（例：クラスタリング）におけるVAEベースの表現の一般化性能を向上させること。
• 潜在空間における複数の変動要因（例：色と構造）のバランスの取れた学習を促進する手法を開発すること。
• 生成されたサンプルがトレーニングデータの特徴をどれだけ保持しているかを定量的に評価するための新しい指標、特徴存在度指標（FPM）を導入すること。
• 並べ替えに基づく正則化が、再構成性能やサンプリング品質を損なわせることなく、VAEにおける表現バイアスを効果的に低減できることを示すこと。

提案手法

• 入力画像に対してランダムな空間的タイリングとシャッフルを適用する確率的ジャイガーパーティションレイヤーを導入し、VAEが局所的な構造的関係を学習するように強制する。
• 潜在空間に混合事前分布を適用することで、尤度項とKL正則化項の衝突を緩和し、事後分布の整合性を向上させる。
• 各サンプルごとに並べ替え操作を独立して適用し、各順方向パスで画像パッチの配置が異なるようにする。
• 並べ替えに不変な再構成損失と混合事前分布を組み合わせた修正されたVAE目的関数を用い、特徴のバランスの取れた学習を促進する。
• 新しい正則化成分を組み込んだ標準的な下界の下界（ELBO）を用いて、モデルをエンドツーエンドで訓練する。
• VAEが学習した潜在コードにクラスタリングヘッドを訓練することで、下流のクラスタリングタスクに適応し、正規化相互情報量（NMI）を用いて特徴のバランスを評価する。

実験結果

リサーチクエスチョン

• RQ1 VAEにおける確率的入力並べ替えは、画像データにおける色のような支配的特徴へのバイアスを低減できるか？
• RQ2 提案されたJigsaw-VAEは、標準VAEや他の分離化手法と比較して、潜在空間における学習特徴のバランスを改善できるか？
• RQ3 テスト分布が1つの特徴（例：すべての数字が黄色に着色）に偏った場合、モデルは下流のクラスタリングタスクでどの程度の性能を示すか？
• RQ4 提案された特徴存在度指標（FPM）は、生成されたサンプルにトレーニングデータの特徴がどれだけ保持されているかを効果的に定量化できるか？
• RQ5 Jigsaw-VAEは、β-VAE や d-VAE などの既存のVAE変種と比較して、より多様で現実的なサンプルを生成できるか？

主な発見

• Jigsaw-VAEは、単一色のMNISTテストセットでNMIスコア0.9473を達成し、β-VAE（0.8766）とVAE（0.8872）を大きく上回り、色と形状の特徴のバランスが良好であることを示している。
• 単一色テストセットでは、Jigsaw-VAEの中央NMIは0.379であったのに対し、Mixup-VAEは0.114、β-VAEは7.8e-6にとどまり、特徴分布シフトに対して高い耐性を示している。
• 再構成タスクにおいて、Jigsaw-VAEは入力の数字が一様に着色された場合でも、色と形状の両方を生成サンプルに保持することができた。これに対して、非Jigsawモデルは支配的クラスタに基づいて誤った色を割り当てた。
• 提案された特徴存在度指標（FPM）は、標準設定下ではβ-VAEと同等のMSEおよびFPMスコアを示したが、Jigsaw-VAEがトレーニングデータの特徴を生成画像により多く保持していることを示した。
• Jigsaw-VAEは、サンプル間の補間が滑らかで、特徴の急激な変化を避けており、潜在空間における分離性と連続性が優れていることを示している。
• CelebAデータセットにおいて、Jigsaw-VAEはβ-VAEよりもより現実的で多様なサンプルを生成した。β-VAEはFPMとMSEスコアが高くても、視覚的に現実性に欠ける画像を生成していた。