Cooperation in the Latent Space: The Benefits of Adding Mixture Components in Variational Autoencoders
This work addresses the problem of making VAEs more flexible and effective for researchers and practitioners in machine learning, though it is incremental as it builds on existing VAE and importance sampling methods.
The paper tackles improving variational autoencoders (VAEs) by adding mixture components, showing that this monotonically increases the ELBO and enhances latent representations across image and single-cell datasets, with some Mixture VAEs achieving state-of-the-art log-likelihood results on MNIST and FashionMNIST.
In this paper, we show how the mixture components cooperate when they jointly adapt to maximize the ELBO. We build upon recent advances in the multiple and adaptive importance sampling literature. We then model the mixture components using separate encoder networks and show empirically that the ELBO is monotonically non-decreasing as a function of the number of mixture components. These results hold for a range of different VAE architectures on the MNIST, FashionMNIST, and CIFAR-10 datasets. In this work, we also demonstrate that increasing the number of mixture components improves the latent-representation capabilities of the VAE on both image and single-cell datasets. This cooperative behavior motivates that using Mixture VAEs should be considered a standard approach for obtaining more flexible variational approximations. Finally, Mixture VAEs are here, for the first time, compared and combined with normalizing flows, hierarchical models and/or the VampPrior in an extensive ablation study. Multiple of our Mixture VAEs achieve state-of-the-art log-likelihood results for VAE architectures on the MNIST and FashionMNIST datasets. The experiments are reproducible using our code, provided here: https://github.com/lagergren-lab/mixturevaes.