Towards Better Data Augmentation using Wasserstein Distance in Variational Auto-encoder

TL;DR

This paper introduces Wasserstein distance into variational auto-encoders to improve data augmentation, resulting in better convergence and more effective synthetic data for image classification.

Contribution

The paper proposes replacing KL divergence with Wasserstein distance in VAE, showing improved theoretical bounds and practical performance for data augmentation.

Findings

Wasserstein-based VAE has a superior ELBO lower bound compared to KL-based VAE.

The new loss function improves convergence properties.

Generated images better support image classification tasks.

Abstract

VAE, or variational auto-encoder, compresses data into latent attributes, and generates new data of different varieties. VAE based on KL divergence has been considered as an effective technique for data augmentation. In this paper, we propose the use of Wasserstein distance as a measure of distributional similarity for the latent attributes, and show its superior theoretical lower bound (ELBO) compared with that of KL divergence under mild conditions. Using multiple experiments, we demonstrate that the new loss function exhibits better convergence property and generates artificial images that could better aid the image classification tasks.