[Paper Review] Simple and Effective VAE Training with Calibrated Decoders

Oleh Rybkin, Kostas Daniilidis+1 more|arXiv (Cornell University)|Jun 23, 2020

Digital Media Forensic DetectionComputer Science65 references35 citations

TL;DR

The paper analyzes calibrated decoders for VAEs, introduces a sigma-VAE with analytic variance estimation, and shows it eliminates beta hyperparameter tuning while improving generation quality across image and video datasets.

ABSTRACT

Variational autoencoders (VAEs) provide an effective and simple method for modeling complex distributions. However, training VAEs often requires considerable hyperparameter tuning to determine the optimal amount of information retained by the latent variable. We study the impact of calibrated decoders, which learn the uncertainty of the decoding distribution and can determine this amount of information automatically, on the VAE performance. While many methods for learning calibrated decoders have been proposed, many of the recent papers that employ VAEs rely on heuristic hyperparameters and ad-hoc modifications instead. We perform the first comprehensive comparative analysis of calibrated decoder and provide recommendations for simple and effective VAE training. Our analysis covers a range of image and video datasets and several single-image and sequential VAE models. We further propose a simple but novel modification to the commonly used Gaussian decoder, which computes the prediction variance analytically. We observe empirically that using heuristic modifications is not necessary with our method. Project website is at https://orybkin.github.io/sigma-vae/

Motivation & Objective

• Assess how calibrated decoders affect VAE performance without manual beta tuning.
• Identify decoder parameterizations that yield well-calibrated uncertainty and stable training.
• Develop a simple, analytic method to set decoder variance and compare to gradient-based learning.
• Evaluate calibrated decoders on a range of image and video datasets and model types.

Proposed method

• Review and compare various calibrated decoder architectures for Gaussian and discrete decoders.
• Propose a Gaussian decoder with a single shared variance, and an analytic optimal-variance formulation (the sigma-VAE).
• Relate calibrated decoders to beta-VAE and show how calibration corresponds to accounting for decoder uncertainty.
• Derive the objective L = D ln sigma + (D/(2 sigma^2)) MSE(x, x̂) + KL(q(z|x)||p(z)).
• Explore per-pixel, per-image, and shared-variance decoders and analyze stability and impact on MI and prior match.
• Empirically evaluate on SVHN, CelebA, CIFAR, and BAIR SVG with single-image and sequential VAE models.

Experimental results

Research questions

• RQ1 Does calibrating the decoder remove the need to tune the KL weight beta in VAEs across datasets and architectures?
• RQ2 Which decoder parameterizations yield well-calibrated uncertainty and stable training for VAEs?
• RQ3 Can an analytic solution for decoder variance improve learning speed and sample quality?
• RQ4 How do calibrated decoders impact latent variable information content and prior alignment?
• RQ5 What are the practical trade-offs of per-pixel vs shared variance in terms of ELBO, FID, and sample quality?

Key findings

• Calibrated decoders can match or exceed beta-VAE performance without manual beta tuning and improve ELBO and sample quality.
• A Gaussian decoder with a shared variance (sigma-VAE) often outperforms unit-variance decoders and tuned beta-VAE setups.
• An analytic optimal-variance solution (optimal sigma-VAE) yields faster convergence and higher log-likelihoods than gradient-learned variances.
• Per-pixel variance decoders can hurt sample quality and prior alignment compared to shared or per-image variations.
• Optimal sigma-based methods achieve strong ELBO and sample quality across multiple datasets and model types.