How to train your VAE

TL;DR

This paper investigates the role of KL Divergence in VAEs, proposing a new ELBO formulation with a Gaussian mixture posterior, regularization, and adversarial training to improve face generation quality.

Contribution

It introduces a novel ELBO redefinition using Gaussian mixtures, adds regularization to prevent variance collapse, and employs adversarial training for enhanced realism in VAE outputs.

Findings

Generated realistic face images

Improved texture realism with PatchGAN discriminator

Enhanced VAE training stability

Abstract

Variational Autoencoders (VAEs) have become a cornerstone in generative modeling and representation learning within machine learning. This paper explores a nuanced aspect of VAEs, focusing on interpreting the Kullback-Leibler (KL) Divergence, a critical component within the Evidence Lower Bound (ELBO) that governs the trade-off between reconstruction accuracy and regularization. Meanwhile, the KL Divergence enforces alignment between latent variable distributions and a prior imposing a structure on the overall latent space but leaves individual variable distributions unconstrained. The proposed method redefines the ELBO with a mixture of Gaussians for the posterior probability, introduces a regularization term to prevent variance collapse, and employs a PatchGAN discriminator to enhance texture realism. Implementation details involve ResNetV2 architectures for both the Encoder and Decoder. The experiments demonstrate the ability to generate realistic faces, offering a promising solution for enhancing VAE-based generative models.