Matching aggregate posteriors in the variational autoencoder

The variational autoencoder (VAE) [22] is a well-studied, deep, latent-variable model (DLVM) that optimizes the variational lower bound of the log marginal data likelihood. However, the VAE's known failure to match the aggregate posterior often results unacceptable latent distribution, e.g. with pockets, holes, or clusters, that fail to adequately resemble the prior. The training of the VAE under different scenarios can also result in posterior collapse, which is associated with a loss of information in the latent space. This paper addresses these shortcomings in VAEs by reformulating the objective function to match the aggregate/marginal posterior distribution to the prior. We use kernel density estimate (KDE) to model the aggregate posterior. We propose an automated method to estimate the kernel and account for the associated kernel bias in our estimation, which enables the use of KDE in high-dimensional latent spaces. The proposed method is named the aggregate variational autoencoder (AVAE) and is built on the theoretical framework of the VAE. Empirical evaluation of the proposed method on multiple benchmark datasets demonstrates the advantages of the AVAE relative to state-of-the-art (SOTA) DLVM methods.