Balancing Reconstruction Quality and Regularisation in ELBO for VAEs

Abstract

A trade-off exists between reconstruction quality and the prior regularisation in the Evidence Lower Bound (ELBO) loss that Variational Autoencoder (VAE) models use for learning. There are few satisfactory approaches to deal with a balance between the prior and reconstruction objective, with most methods dealing with this problem through heuristics. In this paper, we show that the noise variance (often set as a fixed value) in the Gaussian likelihood p(x|z) for real-valued data can naturally act to provide such a balance. By learning this noise variance so as to maximise the ELBO loss, we automatically obtain an optimal trade-off between the reconstruction error and the prior constraint on the posteriors. This variance can be interpreted intuitively as the necessary noise level for the current model to be the best explanation of the observed dataset. Further, by allowing the variance inference to be more flexible it can conveniently be used as an uncertainty estimator for reconstructed or generated samples. We demonstrate that optimising the noise variance is a crucial component of VAE learning, and showcase the performance on MNIST, Fashion MNIST and CelebA datasets. We find our approach can significantly improve the quality of generated samples whilst maintaining a smooth latent-space manifold to represent the data. The method also offers an indication of uncertainty in the final generative model.

arXiv comment8 pages for main contents and 15 pages for supplemental materials that include data pre-processing, model architectures and more results

References in corpus (10)

• Auto-Encoding Variational Bayes
• Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms
• GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium
• Variational Inference with Normalizing Flows
• Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders
• Understanding disentangling in $\beta$-VAE
• Importance Weighted Autoencoders
• Improving Variational Inference with Inverse Autoregressive Flow
• Improving Variational Auto-Encoders using Householder Flow
• Inference Suboptimality in Variational Autoencoders

Cited by in corpus (3)

• Generating Multivariate Load States Using a Conditional Variational Autoencoder
• Re-parameterizing VAEs for stability
• LaDDer: Latent Data Distribution Modelling with a Generative Prior