Training β-VAE by Aggregating a Learned Gaussian Posterior with a Decoupled Decoder

要約

変分オートエンコーダー (VAE) における再構成損失とカルバック ライブラー ダイバージェンス (KLD) 損失はしばしば拮抗的な役割を果たし、$\beta$-VAE における KLD 損失の重みを調整して 2 つの損失のバランスをとることは、
トリッキーでデータセット固有のタスク。
その結果、重み $\beta$ が慎重に調整されていない場合、VAE トレーニングの現在の慣行では、再構成の忠実度と潜在空間の連続性 $/$disentanglement との間にトレードオフが生じることがよくあります。
この論文では、直感と2つの損失の拮抗メカニズムの慎重な分析を提示し、洞察に基づいて、VAEをトレーニングするためのシンプルで効果的な2段階の方法を提案します。
具体的には、この方法は、新しい条件付き分布 $p_{\phi} (x
入力データ $x$ の |z)$。
実験的に、集約された VAE が潜在空間に関するガウスの仮定を最大限に満たしながら、潜在空間が $\mathcal{N}(\mathbf{0},I) によって大まかに正則化されている場合に匹敵する再構成エラーを達成することを示します。
$.
提案されたアプローチでは、一般的な VAE トレーニング プラクティスで必要とされる特定のデータセットが与えられた場合、ハイパーパラメーター (つまり、KLD 重み $\beta$) の調整は必要ありません。
3D 頭蓋骨の再構成と形状の完成を目的とした医療データセットを使用してこの方法を評価し、その結果は、提案された方法を使用してトレーニングされた VAE の有望な生成能力を示しています。
さらに、潜在変数のガイド付き操作を通じて、既存のオートエンコーダー (AE) ベースのアプローチと、形状完成問題に対する VAE などの生成的アプローチとの間の接続を確立します。
コードと事前トレーニング済みの重みは、https://github.com/Jianningli/skullVAE で入手できます。

要約(オリジナル)

The reconstruction loss and the Kullback-Leibler divergence (KLD) loss in a variational autoencoder (VAE) often play antagonistic roles, and tuning the weight of the KLD loss in $\beta$-VAE to achieve a balance between the two losses is a tricky and dataset-specific task. As a result, current practices in VAE training often result in a trade-off between the reconstruction fidelity and the continuity$/$disentanglement of the latent space, if the weight $\beta$ is not carefully tuned. In this paper, we present intuitions and a careful analysis of the antagonistic mechanism of the two losses, and propose, based on the insights, a simple yet effective two-stage method for training a VAE. Specifically, the method aggregates a learned Gaussian posterior $z \sim q_{\theta} (z|x)$ with a decoder decoupled from the KLD loss, which is trained to learn a new conditional distribution $p_{\phi} (x|z)$ of the input data $x$. Experimentally, we show that the aggregated VAE maximally satisfies the Gaussian assumption about the latent space, while still achieves a reconstruction error comparable to when the latent space is only loosely regularized by $\mathcal{N}(\mathbf{0},I)$. The proposed approach does not require hyperparameter (i.e., the KLD weight $\beta$) tuning given a specific dataset as required in common VAE training practices. We evaluate the method using a medical dataset intended for 3D skull reconstruction and shape completion, and the results indicate promising generative capabilities of the VAE trained using the proposed method. Besides, through guided manipulation of the latent variables, we establish a connection between existing autoencoder (AE)-based approaches and generative approaches, such as VAE, for the shape completion problem. Codes and pre-trained weights are available at https://github.com/Jianningli/skullVAE

arxiv情報

提供元, 利用サービス

arxiv.jp, Google