Regularizing Variational Autoencoder with Diversity and Uncertainty Awareness

TL;DR

This paper introduces DU-VAE, a novel regularization method for Variational Autoencoders that enhances latent space diversity and reduces uncertainty, leading to improved generative and classification performance without altering training procedures.

Contribution

The paper proposes DU-VAE, a new approach that implicitly regularizes the latent space using Dropout and Batch-Normalization, improving diversity and reducing uncertainty in VAEs.

Findings

Outperforms state-of-the-art baselines on likelihood estimation

Achieves better latent space diversity and lower uncertainty

Improves classification accuracy on benchmark datasets

Abstract

As one of the most popular generative models, Variational Autoencoder (VAE) approximates the posterior of latent variables based on amortized variational inference. However, when the decoder network is sufficiently expressive, VAE may lead to posterior collapse; that is, uninformative latent representations may be learned. To this end, in this paper, we propose an alternative model, DU-VAE, for learning a more Diverse and less Uncertain latent space, and thus the representation can be learned in a meaningful and compact manner. Specifically, we first theoretically demonstrate that it will result in better latent space with high diversity and low uncertainty awareness by controlling the distribution of posterior's parameters across the whole data accordingly. Then, without the introduction of new loss terms or modifying training strategies, we propose to exploit Dropout on the variances and Batch-Normalization on the means simultaneously to regularize their distributions implicitly. Furthermore, to evaluate the generalization effect, we also exploit DU-VAE for inverse autoregressive flow based-VAE (VAE-IAF) empirically. Finally, extensive experiments on three benchmark datasets clearly show that our approach can outperform state-of-the-art baselines on both likelihood estimation and underlying classification tasks.