Augmented ELBO regularization for enhanced clustering in variational autoencoders

Abstract

With significant advances in deep neural networks, various new algorithms have emerged that effectively model latent structures within data, surpassing traditional clustering methods. Each data point is expected to belong to a single cluster in a typical clustering algorithm. However, clustering based on variational autoencoders (VAEs) represents the expectation of the overall clusters, denoted as $c=1,\dots ,K$ in the KL divergence term. Consequently, the latent embedding $z$ can be learned to exist across multiple clusters with relatively balanced probabilities, rather than being strongly associated with a specific cluster. This study introduces an additional regularizer to encourage the latent embedding $z$ to have a strong affiliation with specific clusters. We introduce optimization methods to maximize the ELBO that includes the newly added regularization term and explore methods to eliminate computationally challenging terms. The positive impact of this regularization on clustering accuracy was verified by examining the variance of the final cluster probabilities. Furthermore, an enhancement in the clustering performance was observed when regularization was introduced.