Achieving deep clustering through the use of variational autoencoders and similarity-based loss

1.   Introduction

Clustering is a crucial and challenging research topic in machine learning, computer vision and data mining. The goal of clustering is to separate similar and unlabeled data into several independent subsets, each representing a class. Many clustering methods ^[1,2,3,4] have been proposed in the literature and are widely applied to various practical problems. However, their shortcomings are becoming more and more serious for high-dimensional data. The similarity measures based on the conventional clustering method are limited by the local relationship in the data space. They thus cannot capture the hidden and hierarchical dependences of the data in the underlying structure. To solve the curse of dimensionality caused by high-dimensional data, researchers first proposed dimensionality reduction and feature transformation methods such as principal component analysis (PCA) ^[5], independent component analysis ^[6], kernel methods ^[7] and non-negative matrix factorization ^[8]. Mapping the original high-dimensional data space to a new feature space significantly increases the density of the samples in this space and helps improve the efficiency of similarity measures. However, this feature space is not necessarily the most suitable for clustering tasks.

The rapid development of deep learning technology provides a new direction for clustering. Clustering methods with deep learning, i.e., deep clustering ^[9,10], can significantly benefit clustering tasks by transforming data into cluster-oriented representations. Practical experiments show that most deep clustering methods outperform traditional clustering methods. In most deep clustering methods, high-dimensional data are mapped to a low-dimensional feature representation via a powerful deep learning algorithm; then, the clustering algorithm is directly applied. Feature representation and clustering are usually handled independently. For example ^[11], stacked autoencoders (AEs) are used to extract nonlinear feature representations of graphs and then directly execute the k-means algorithm. The problem with such an approach is that they do not fully exploit the potential of feature representation in deep learning. In recent years, variational deep embedding (VaDE) ^[12] has been proposed, which exploits the advantages of variational AEs (VAEs) to yield an unsupervised generative clustering method that combines feature representation and clustering. VAE-based frameworks can accurately capture the data distribution in the latent space to represent the observed data and learn meaningful representations. The association of feature representations and clustering algorithms can enable the learning of cluster-oriented representations.

However, existing deep clustering methods still have many challenging issues worth investigating. For example, feature representation in VaDE focuses on the global structure of the data. Still, it ignores the local structure; additionally, some recent works ^[13,14] proved that similarity-based approaches in deep learning consider the local structure and thus lead to improved clustering performance.

This paper introduces a new deep clustering framework, namely, deep clustering through the use of VAEs and similarity-based loss (DVAE), that combines a VAE and a gaussian mixture model (GMM) to solve the issues above. The proposed method can optimize feature representation learning and clustering simultaneously in an end-to-end manner. Our method improves the latent space's intra-cluster compactness and inter-cluster separability during its training procedure. The main idea is to encourage the latent manifold's structure globally toward a more favorable state for clustering. The block diagram of our algorithm is shown in Figure 1. Specifically, we have added a new similarity-based graph embedding loss to the VAE's objective function. This dynamically creates a fully connected undirected graph in the latent space, treating the data as nodes on the graph and minimizing the weighted distances of their edges. We use cluster prediction to create weights for the edges and 2 - Wasserstein distances of the posterior distribution over the latent space as low-dimensional vectors representing relationships. Then, we conducted many experiments to evaluate the proposed method. In addition, we further discuss the ablation experiments, parameter sensitivity analysis and visualization of our method.

The main contributions of this paper include:

● The introduction of an end-to-end VAE-based clustering framework that can simultaneously optimize embedding representations and clustering. A combination of these two aspects can promote meaningful cluster-oriented representations.

● The introduction of a novel similarity-based loss as a graph embedding extension. We use the transition probabilities of cluster predictions to represent edge weights and the 2-Wasserstein distance of the posterior distribution to compute a low-dimensional representation of the relationship. The proposed loss can improve cluster-oriented representations by imposing local structural constraints, resulting in more compact clusters. To the best of our knowledge, this structure of graph embedding loss has not been applied in unsupervised clustering.

● Experiments on five publicly available data sets demonstrate that our method is highly competitive and superior to several state-of-the-art clustering methods.

The rest of this paper is structured as follows: Section 2 reviews related work. Section 3 details the framework of the VAE, including generative models, inference models and the lower bounds of evidence. At the same time, we introduce a new loss function based on the idea of a VAE. In Section 4, we conduct comprehensive experiments to evaluate the proposed method, including comparing several clustering methods and performing ablation studies, parameter sensitivity analysis and an experimental visualization. Finally, the work of this paper is summarized in Section 5.

2.   Related works

2.1.   Deep clustering

The goal of clustering tasks is to partition a given data set without prior knowledge and group data with similar structures or patterns into the same cluster, with different clusters having high heterogeneity. Clustering is an important and challenging research direction in machine learning; it has been studied by many scholars, resulting in many clustering algorithms. GMMs and k-means are the most widely used partitioning algorithms based on optimizing a specific cost function that allows separation between clusters. Most traditional clustering algorithms have achieved great success on low-dimensional data. Nonetheless, they cannot achieve good clustering results in the high-dimensional space ^[15] when the data volume and dimensionality grow exponentially. The reason is that there is much noisy information in high-dimensional data sets, and traditional distance measures are not suitable for high-dimensional data. To solve this problem, researchers have proposed dimensionality reduction-based clustering. Clustering based on dimensionality reduction refers to mapping high-dimensional data to low-dimensional space through the use of some dimensionality reduction technique, including PCA ^[5], linear discriminant analysis ^[16] and kernel-based PCA ^[17]. Clustering methods based on dimensionality-reduced data have good performance on a limited number of tasks ^[18,19,20]. Furthermore, the model proposed in ^[21] entails the use of a spectral clustering-like approach that enables the learning of new Laplacian matrices and attainment of more robust data clustering through the use of the manifold structure and local discriminative information.

However, with the continuous improvement of practical application requirements, the traditional clustering methods based on dimensionality reduction are becoming more and more limited by the bottleneck of not being able to mine the underlying structural information of the data. Researchers have introduced deep learning to extract representations to overcome the limitations of shallow linear or nonlinear feature extraction algorithms. Deep learning allows for the mining of more complex nonlinear features and relationships. It integrates the feature extraction work into the model building process, avoids manual operations, and significantly improves the clustering effect ^[22]. A typical example is presented in ^[11], wherein the model first learns feature representations in the reconstruction task through a sparse AE-based network before k-means is applied for clustering in the feature space. AEs and traditional clustering techniques can extract more cluster-oriented features with excellent performance. The model proposed in ^[23] adopts a similar approach, completing representation learning through deep ternary networks. Subsequently, adversarial AE (AAE) ^[24] combines the discriminator from generative adversarial networks with AE, resulting in a more powerful model.

Deep embedding for clustering (DEC) ^[25] is a well-studied deep clustering method that uses an AE as the network framework and cluster assignment reinforcement loss as regularization. DEC ignores the problem that clustering loss cannot guarantee local structure; however, improved deep embedded clustering (IDEC) ^[26] improves this. IDEC uses an incomplete AE. By fusing clustering loss and the AE loss, IDEC jointly assigns cluster labels and learns features suitable for clustering and preserving the data structure. K-Autoencoders (K-DAE) ^[27] is an AE-based deep clustering method that extends the k-means algorithm. K-DAE uses both the reconstruction loss of an AE and the loss function of k-means to train the network.

2.2.   Variational autoencoder

The VAE is a deep generative model ^[28]. VAEs requires the following assumptions. The low-dimensional latent variable ${\bf z}$ is generated from a prior distribution $p({\bf z})$, for example, a multivariate standard Gaussian distribution. The high-dimensional observation data points ${\bf x}$ are automatically generated from the conditional distribution $p({\bf x}|{\bf z})$ given the sampling of ${\bf z}$. These two structures of VAEs are called the inference model and generative model, respectively. VAEs aim to find the true posterior $p({\bf x}|{\bf z})$ to infer the latent distribution from the observed samples. Using the Bayes theorem, we can obtain this posterior distribution:

In the denominator of the above equation, the marginalization over the latent variables of the high-dimensional data distribution $p({\bf x})$ is complicated. Therefore, the VAE introduces a new easy-to-solve approximate distribution $q({\bf z}|{\bf x})$ to complete the function of the inference model and use this variational distribution to approximate the true posterior and minimize the Kullback-Leibler (KL) divergence between them, which can be written as

Since $p({\bf x})$ and ${\bf z}$ are irrelevant, we can derive the core formula of VAE:

The term to the right of the equal sign is called the evidence lower bound (ELBO) for ${\bf x}$. According to Jensen's inequality ^[29], KL divergence is non-negative. So, $log(p({\bf x}))$ is always greater than the ELBO. We can see that minimizing the distance between the variational and true distribution is equal to maximizing the ELBO. At this point, the VAE transforms the probabilistic inference problem of summing variables into the problem of optimizing the ELBO. We can write the loss of the VAE as

VAEs has flexibility, theory, scalability, and robust feature extraction capabilities. Researchers have combined VAEs with classical clustering methods and succeeded in recent years. The model proposed in ^[30] first tries to use the mixed Gaussian distribution as the approximate posterior of the VAE. Realigning each point in the VAE latent space to its nearest class neighbor may improve the clustering effect.

The main work related to our proposed method is VaDE. VaDE and our model can learn latent feature representations based on the VAE framework and GMM priors. However, the differences are significant; specifically, our method uses a similarity-based latent space metric to preserve local data structure. On the other hand, our proposed method emphasizes that if the samples are close in the latent space, promoting shortening their distance leads to better clustering performance. Our proposed method enables the learning of more compact feature representations. The experimental results show that the proposed method outperforms VaDE.

3.   Method

This section first details the clustering process for the improved VAE and its framework, including the generative model, the inference model, and the ELBO. After that, we introduce the details of the proposed similarity-based loss function.

3.1.   Variational autoencoder

Let us begin with a set of $D$-dimensional original data points $X = \{{\bf x}_i \in \mathbb{R}^D \}_{i = 1}^N$ derived from $C$ nonlinear manifolds, where $N$ is the number of data points and $C$ is the number of corresponding classes.

3.1.1.   Generative model

In our model, the observation variable ${\bf x}$ is generated by the latent variables ${\bf z}$. Instead of the single Gaussian prior, we treat the GMM as the prior latent space to capture the distribution of data. We model the distribution of $p_{\theta}({\bf x}, {\bf z}, {\bf c})$ which can be written as

where $p({\bf c})$ is the categorical distribution parameterized by $\pi \in \mathbb R_{+}^C$, $\sum_{k = 1}^C \pi_k = 1, \, {\bf c} \in \{ 0, 1\}$, $p({\bf z}|{\bf c})$ is the $k{\rm{th}}$ Gaussian component parameterized by the mean $\boldsymbol \mu_k$ and the variance $\boldsymbol \sigma_k^2$, $k = \{1, 2, ..., C\}$, $p({\bf x}|{\bf z})$ is the Bernoulli distribution or Gaussian distribution parameterized by $\boldsymbol \mu _x$, $\boldsymbol \mu _x$ is the mean characterized by the deep neural network $f$ with the parameters $\theta$, which can be written as $[\boldsymbol \mu_x] = f({\bf z}; \theta)$, $\lambda$ is a predefined parameter and ${\bf I}$ is an identity matrix.

3.1.2.   Inference model

To apply variational inference to $p_{\theta}({\bf x}, {\bf z}, {\bf c})$, we use the variational distribution $q_{\phi}({\bf z}| {\bf x})$ to approximate the true posterior $p({\bf z}| {\bf x})$. We use the same strategy as VaDE, i.e., assuming that $q({\bf z}, {\bf c}|{\bf x})$ is a mean-field distribution, but using a different factorization method, mathematically, follows:

where the mathematical forms of the two distributions to the right of the equal sign are

where

To choose a suitable $q_{\rho}({\bf c}| {\bf z})$, we assume that $q({\bf z}|{\bf x})$ is a Gaussian distribution and $q({\bf c}|{\bf z})$ is a multinomial distribution, where $\{ \tilde{\boldsymbol \mu}, \tilde{\boldsymbol \sigma} \}$ and $\tilde{ \pi}$ are some arbitrarily determined value of Network $g$ with the parameter $\phi$ and Network $h$ with the parameter $\rho$.

Since the model adopts a random sampling strategy, the VAE introduces a reparameterization trick to allow backpropagation and obtain ${\bf z}$, as described by the following formula:

where $\circ$ denotes the element-wise product, and $\epsilon$ is sampled from the distribution $\mathcal{N}(0, {\bf I})$.

This paper uses the classifier $h$ to distinguish different clusters. Once we have trained the variational encoder, we use the pre-trained model to extract features from the data and use k-means to cluster the features to obtain pseudo-labels. Since no labels are available in the clustering task, we use these pseudo labels to train the classifier.

3.1.3.   Evidence lower bound

By minimizing the $KL$ distance between $p({\bf z}, {\bf c}|{\bf x})$ and $q({\bf z}, {\bf c}|{\bf x})$, the loss function of the VAE with the GMM prior can be written as

According to Eqs 3.1 and 3.2, we take the example that ${\bf x}$ is binary and apply the reparameterization trick; then, Eq 3.5 can be rewritten as

where $*|_i$ denotes the $i$th element of $*$. Here

where $f, g$ and $h$ are the generative model, inference model and classifier, respectively.

3.2.   Differences from traditional variational autoencoders

Our proposed method is a generative clustering model based on VAEs. Our optimization follows steps of a VAE, using stochastic gradient variational Bayes and reparameterization techniques to optimize the ELBO. The difference between our method and the traditional VAE is reflected in two ways: 1) Our proposed method introduces a GMM to construct the model generation process. That is, $p(\boldsymbol{z})$ is a GMM instead of the normal distribution in traditional VAE; 2) Given the problem of ignoring the capture of local structures in traditional VAEs, we have added a constraint-based on the Wasserstein distance to the loss function of the VAE to help the model deal with complex data clustering problems.

3.3.   Similarity-based loss

In this section, we propose a novel similarity-based loss function to preserve the local structure of the data in the latent space. Our proposed loss can be seen as an extension of graph embedding. We dynamically create a fully connected undirected graph in the latent space, treating the data as nodes and the correspondence between their distributions as edges on the graph.

We shall now discuss the construction of a weighted graph in the latent space from the samples. It is a fully connected graph with an adjacency matrix ${\bf W} \in \mathbb{R}^{N \times N}$. Each non-negative element of the real symmetric matrix $W$ represents the similarity between a pair of vertices on the graph. This matrix can be formed using various similarity measures, such as a Gaussian kernel function. Inspired by prior class information ^[31] in supervised learning and mixed class information ^[32] in semi-supervised learning, we propose a new cluster label similarity measure from latent space distance and generalize it to unsupervised study. The authors of ^[32] proved that, ideally, data of the same class would cluster together to form a compact cluster. Therefore, there is a positive transition probability between data points in the same cluster. We use a formula variant for transition probability to describe the extent of the data similarity; it is defined as

where ${\bf w}_{ij}$ represents an element of ${\bf W}$, representing the distance between the data points ${\bf x}_i$ and ${\bf x}_j$ in the category space.

To better characterize the similarity relationship for the data points, we propose a low-dimensional embedding relationship based on the distance of the data distribution in the low-dimensional latent space as a graph embedding. At this point, we can write out the form of the loss function in the VAE framework with graph embedding:

where ${\bf Q}_{ij}$ is a simplified notation of ${\bf Q}(q({\bf z}|\boldsymbol x_i)||q({\bf z}|\boldsymbol x_j))$, indicating the distance between the two distributions. $d_{\rm{Wasserstein}}$ is a variant of the 2 - Wasserstein distance. Since $q({\bf z}|\boldsymbol x)$ is in Gaussian form, there is a closed-form solution here, given by

Given Eq 3.11, we can further simplify the third term on the right side of the equal sign as

Substituting Eq 3.12 into Eq 3.11, we get

Unlike existing similarity-based methods, which consider the similarity of the original data in the Euclidean space, our proposed loss function focuses on the relationship between the posterior prediction of the data and the similarity in the latent space. Similarity increases when data in the latent space is close, and vice versa. This loss function encourages data in the same cluster to cluster together in the latent space, and data in different clusters are mutually exclusive to promote separation.

3.4.   Variational autoencoder with latent clusters

We can write the final loss as

where $\beta \geq 0$ is a weight coefficient that controls the proposed loss function. For our proposed loss to always have a positive effect on the model during the learning phase, we have the following settings: the weight $\beta$ is 0 for the first $T$ iterations and takes appropriate values thereafter. It works for all tasks, and the model performs best when $T$ is 50.

4.   Experiments

The experiment proves the superiority of our proposed method. To make a fair comparison, we use the same network structure and parameters as VaDE. Like VaDE, we stack the AE into a deep VAE and then pre-train the generative and inference models to eliminate the initial VAE's problems. We use k-means clustering to complete the initialization of the classifier and Gaussian model by generating pseudo-labels in the latent space clustering to train the classifier network and determining the initialization of the mean and variance of the GMM. The pseudo-code of our algorithm is as follows:

4.1.   Data sets

The data sets used in this experiment include MNIST, HHAR, REUTERS-10K and USPS. Specifically,

● MNIST is a database containing handwritten digital images from 250 different people. It includes a training set of 60,000 examples and a test set of 10,000 samples. MNIST consists of 10 types of handwritten digital pictures and digital labels composed of 28 $\times$ 28 pixels.

● HHAR is a heterogeneous human activity recognition data set acquired from smartphones and smartwatches, and it contains 10,200 sensor records. The number of dimensions recorded by the sensor is 561 dimensions, which are divided into six categories of human activities, including Biking, Sitting, Standing, Walking, Stair Up and Stair down.

● Reuters is an English text data set from 11,228 press articles in Reuters, labeled with 46 topics. Reuters-10k ^[33] is a subset of 10,000 samples randomly sampled from Reuters, and it includes four categories: corporate/industrial, government/social, markets and economics.

● STL-10 is a database of 13,000 labeled and 100,000 unlabeled color images. Compared to MNIST, it has a higher resolution (96 × 96 pixels) and is highly complex (10 real-world classes). We applied ResNet-50 to reduce the image dimensions to 2048.

● USPS is a data set containing 9,298 digital images from envelopes automatically scanned by the United States Postal Service. It consists of 16 $\times$ 16 pixel grayscale samples, and each image in the data set has been centered and normalized.

The detailed information is shown in Table 1.

4.2.   Experimental settings

For the proposed method, we use the same network structure as VaDE, with $D-500-500-2000-10$ as the encoder network structure, $10-2000-500-500-D$ as the decoder and $10-C (C = 10)$ or $10-C-C (C < 10)$ for the classifier, where $D$ represents the number of dimensions of the data and $C$ represents the number of classes. All layers are fully connected in our model, and ReLU is used as an activation function. We use a batch size of 300 for all data sets. MNIST, HHAR, REUTERS-10K, STL-10 and USPS have an initial learning rate of 0.002, reduced every 10 epochs with a decay rate of 0.9. During the first 50 epochs, the value of $\beta$ is 0.

4.3.   Evaluation metrics

We measured the performances of our proposed method and the benchmarks using clustering accuracy (ACC) and normalized mutual information (NMI).

ACC is defined as

where $l_n$ and $c_n$ represent the true label and cluster assignment generated by the input $x_n$, respectively. $m$ covers all possible one-to-one mappings between the predicted labels produced by the clustering and ground-truth labels.

NMI is defined as

where ${ I} (\cdot; \cdot)$ and $H(\cdot)$ represent the mutual information and the entropy, respectively.

4.4.   Comparison with several clustering approaches

We compared several methods to verify the effectiveness of our proposed method, including traditional clustering methods, i.e., a GMM and k-means algorithm, and deep clustering methods, i.e., AE+GMM and K-DAE ^[27], IDEC ^[26] and VaDE algorithms ^[12].

We present the ACC and NMI results obtained for each method in Tables 2 and 3 and highlight the best scores. From these two tables, we have the following observations. 1) Deep clustering performs better than traditional clustering. The excellent performance of the GMM and k-means algorithm on STL10 is attributable to the data being processed by ResNet-50, which is equivalent to the effect of applying traditional clustering after deep learning. 2) The VAE was is more stable and effective than the AE on these five data sets. On MNIST, HHAR, Reuters-10K and STL-10, VaDE demonstrated relatively competitive performance compared with other existing methods. 3) Our proposed method achieved the highest ACC for all data sets. On USPS, in particular, it far outperformed any clustering method.

4.5.   Ablation study

To verify the effectiveness of the similarity-based loss term introduced in our method, we further conducted ablation studies on four data sets, i.e., HHAR, Reuters-10k, STL-10K and USPS. Specifically, the network is learned by removing the similarity-based loss term and studying the network's performance. The experimental results are shown in Table 4, from which we can make the following observations. We can see that adding similarity-based terms benefits the clustering task. On HHAR and STL-10, the clustering performance improved by 2.48% and 0.55%, respectively, and the term of similarity was relatively mild for STL-10, although there was still a slight improvement. On Reuters-10k and USPS, adding similarity items significantly affected the clustering effect. The clustering performance improved by 5.33% on Reuters-10k and by an astonishing 45.71% on USPS, which strongly proves the effect of the similarity-based loss term.

4.6.   Parameter sensitivity

The parameter $\beta$ affects the contribution of our proposed similarity term to the total loss. To choose appropriate weights, we took {0.001, 0.005, 0.01, 0.05, 0.1, 1} and report the effect of $\beta$ on the NMI and ACC for the four data sets in Figure 2, respectively. The experiments show that taking the highest $\beta$ value in most cases will yield the best ACC and NMI values. For example, in the cases of HHAR and STL-10, $\beta = 1$ achieved the best results. In the case of MNIST, $\beta = 0.1$ achieved the best results because DVAE suffers from the same problem as VaDE when $\beta = 1$ gets stuck in undesired local minima or saddle points. The best way at the moment is to use pre-training to overcome that problem. At the same time, in the case of USPS, although $\beta$ shows a downward trend on the graph, the magnitude is minimal. To sum up, we recommend that $\beta$ takes the biggest value that the model can sustain.

4.7.   Experimental visualization

We made the following two visualizations. First, we specifically compared the performances of the DVAE and VaDE on the MNIST data sets; as shown in Figure 3, the abscissa is the number of iterations, and the ordinate is the value of ACC. As the number of iterations increased, the DVAE consistently outperformed VaDE. Second, we used the t-SNE to visualize the output of the encoder of the VAE to show the degree of separation of the clusters in the latent space in Figure 4. As shown in the figure, as the number of iterations increased, clusters of the same class eventually became more compact.

5.   Conclusions

We have proposed a new loss function for unsupervised clustering. This loss function enhances clustering performance by encouraging the global structure of the data manifold in the latent space toward a more compact direction. Specifically, this loss function is the cross-entropy of the predicted cluster transition probabilities and the Wasserstein distance of the predicted posterior distribution. We evaluated our method on four widely used data sets, demonstrating that our method is more robust and effective.

Conflict of interest

The authors declared that they have no conflicts of interest to this work.