Towards Green VAE: A Light Pixel-weighting Technique to Enhance Variational AutoEncoder

Abstract

Variational autoencoders (VAEs) has been a popular generative model for its effectiveness, mathematical foundation, and its impact to other approaches in deep generative learning. For its relatively light-weights and easiness for training, compared with Generative Adversarial Networks (GANs) or other large-scale model e.g. Diffusion, VAEs become a viable tool in cost-efficient generative applications especially considering the so-called green AI. However, in comparison to GANs, the performance of VAEs in generating realistic images is still inferior to the state-of-the-art generative adversarial network (GAN). In this paper, we argue that this problem is at least partly due to the irrational reconstruction in VAE that all pixels are equally weighted, which is harmful to the generating ability (or density estimate). Motivated by this, we propose to compute the weights of pixels. First, we formulate the problem of finding the appropriate weights into an optimal problem, and then give an analytical solution. Moreover, we propose a method to apply the computed weights of pixels into the training pipeline with almost no computation overhead which fits with the spirit of green AI for more energy-saving AI especially for deep learning models. Experiments on MNIST, Fashion-MNIST and CIFAR-10 show that our method can significantly improves VAE in terms of FID.