TOWARDS EFFICIENT VARIATIONAL AUTO-ENCODER USING WASSERSTEIN DISTANCE

VAE, or variational auto-encoder, compresses data into latent attributes and generates new data of different varieties. VAE with KL divergence loss has been considered an effective technique for data augmentation. In this paper, we propose using Wasserstein distance as a measure of distributional similarity for the latent attributes and show its superior theoretical lower bound (ELBOW) compared with that of KL divergence (ELBOKL) under mild conditions. Using multiple experiments, we demonstrate that the new loss function converges faster and generates better quality data to aid image classification tasks. We also propose implementing a dynamically changing hyper-parameter tuning schedule to avoid the potential overfitting of ELBOW.