Improve variational autoEncoder with auxiliary softmax multiclassifier

TL;DR

This paper introduces VAE-AS, an improved variational autoencoder that uses an auxiliary softmax classifier to maintain mutual information, effectively addressing posterior collapse and enhancing image reconstruction quality.

Contribution

The paper proposes a novel auxiliary softmax multi-classifier to improve VAE training by preserving mutual information and reducing posterior collapse.

Findings

VAE-AS effectively maintains mutual information during training.

VAE-AS reduces the posterior collapse problem.

VAE-AS improves image reconstruction quality.

Abstract

As a general-purpose generative model architecture, VAE has been widely used in the field of image and natural language processing. VAE maps high dimensional sample data into continuous latent variables with unsupervised learning. Sampling in the latent variable space of the feature, VAE can construct new image or text data. As a general-purpose generation model, the vanilla VAE can not fit well with various data sets and neural networks with different structures. Because of the need to balance the accuracy of reconstruction and the convenience of latent variable sampling in the training process, VAE often has problems known as "posterior collapse". images reconstructed by VAE are also often blurred. In this paper, we analyze the main cause of these problem, which is the lack of mutual information between the sample variable and the latent feature variable during the training process. To maintain mutual information in model training, we propose to use the auxiliary softmax multi-classification network structure to improve the training effect of VAE, named VAE-AS. We use MNIST and Omniglot data sets to test the VAE-AS model. Based on the test results, It can be show that VAE-AS has obvious effects on the mutual information adjusting and solving the posterior collapse problem.