Learning Robust Representation for Clustering through Locality Preserving Variational Discriminative Network

To model the global structure, we take advantage of Variational Deep Embedding (VaDE) (Jiang et al. 2017), which extends VAE with GMM to model the distribution of latent representations. Specifically, both VAE and VaDE maximize the log probability of the given data point:

$$\max_{\phi,\theta}\sum_{i=1}^{N}\log{p(x_{i};\phi,\theta)}$$

(2)

As the optimization is non-trivial, an evidence lower bound (ELBO) (Kingma and Welling 2014; Jiang et al. 2017) is proposed and maximized instead. In order to understand the ELBO, we further decompose it as:

	$\displaystyle L_{\text{ELBO}}(x_{i})=$	$\displaystyle E_{q(z,c|x_{i})}[\log{p(x_{i}|z)}]$		(3)
		$\displaystyle-D_{\text{KL}}(q(z,c|x_{i})||p(z,c))$

where $L_{\text{ELBO}}$ is the ELBO loss to be optimized. Term $x_{i}$ denotes the input data, $z$ represents the latent variables, $c$ is the GMM cluster in the latent layer. Eqn.3 shows that the ELBO loss consists of two terms. The first term corresponds to the reconstruction loss, which encourages the network to better approximate the original data. The second term is the regularization term, which requires the distribution of latent variables $z$ to be as close as Mixture-of-Gaussian distribution $p(\bm{z},c)$. Since ELBO preserves the global distribution of latent variables, our global structure term $L_{\text{G}}(x_{i};\phi,\theta)$ is set to the negative ELBO loss (since we minimize the objective), precisely, $L_{\text{G}}(x_{i};\phi,\theta)=-L_{\text{ELBO}}(x_{i})$.

In practice, following the derivation of Jiang et al. (2017), $L_{\text{G}}(x_{i};\phi,\theta)$ is computed as

		$\displaystyle{}L_{\text{G}}(x_{i};\phi,\theta)$		(4)
	$\displaystyle=$	$\displaystyle-\sum_{d=1}^{D}{x_{i}^{d}\log{\bm{\mu}_{x_{i}}|_{d}}}+(1-x_{i}^{d})log(1-\bm{\mu}_{x_{i}}|_{d})$
		$\displaystyle+\sum_{c=1}^{K}{\gamma_{ic}\sum_{j=1}^{J}{\big{(}\log{\bm{\sigma}_{c}^{2}|_{j}}}+\dfrac{\bm{\widetilde{\sigma}}_{i}^{2}|_{j}}{\bm{\sigma}_{c}^{2}|_{j}}+\dfrac{{(\bm{\widetilde{\mu}}_{i}|_{j}-\mu_{c}|_{j})^{2}}}{\bm{\sigma}_{c}^{2}|_{j}}\big{)}}$
		$\displaystyle-\sum_{c=1}^{K}{\gamma_{ic}\log{\dfrac{\pi_{ic}}{\gamma_{ic}}}}-\dfrac{1}{2}\sum_{j=1}^{J}(1+\log{\bm{\widetilde{\sigma}}_{i}^{2}|_{j}})$

where $D$ is the dimension of input $x_{i}$ and reconstructed outputs $\bm{\mu}_{x_{i}}$. $x_{i}^{d}$ is the $d$-th element of $x_{i}$, $*|_{j}$ denotes the $j$-th element of $*$. $\gamma_{ic}$ represents $q(c|z_{i})$, which is the GMM distribution for latent variable $z_{i}$. Moreover,

	$\displaystyle[\bm{\widetilde{\mu}}_{i},\log\bm{\widetilde{\sigma}}_{i}^{2}]$	$\displaystyle=g_{g}(x_{i};\phi)$		(5)
	$\displaystyle z_{i}$	$\displaystyle=\widetilde{\bm{\mu}}_{i}+\widetilde{\bm{\sigma}}_{i}\circ{\epsilon}_{i}$		(6)
	$\displaystyle\bm{\mu}_{x_{i}}$	$\displaystyle=f_{g}(z_{i};\theta)$		(7)

where $g_{g}(*;\phi)$ and $f_{g}(*;\theta)$ are two neural networks parameterized by $\phi$ and $\theta$, $z_{i}$ is generated from the latent distribution $\mathcal{N}(\bm{\widetilde{\mu}}_{i},\bm{\widetilde{\sigma}}_{i}^{2})$, $\epsilon_{i}$ is a vector with elements drawn from independent normal distribution and $\circ$ denotes the element-wise multiplication. $\bm{\mu}_{x_{i}}$ is the reconstructed data of $x_{i}$.

In Eqn.3, it is essential to reconstruct the original data, which is achieved by minimizing the cross entropy loss or mean square error between raw data and the latent representations. However, such objectives are too strict, which would make the model vulnerable to the input noise. In order to improve the robustness of our model, we reformulate the reconstruction loss over all input data samples (for the simplicity, we omit the variable $c$),

		$\displaystyle E_{x\sim P(X)}\big{[}E_{q(z|x)}[\log(p(x|z))]\big{]}$		(8)
	$\displaystyle=$	$\displaystyle\int\int p(x)q(z|x)\log(p(x|z))dzdx$
	$\displaystyle=$	$\displaystyle\int\int p(x)q(z|x)\log\big{(}\frac{q(z|x)p(x)}{q(z)}\big{)}dzdx$
	$\displaystyle=$	$\displaystyle\int\int p(x)q(z|x)\log\big{(}\frac{q(z|x)p(x)}{q(z)p(x)}\big{)}dzdx$
		$\displaystyle+\int\int p(x)q(z|x)\log(p(x))dzdx$
	$\displaystyle=$	$\displaystyle D_{\text{KL}}(Q(Z|X)P(X)||Q(Z)P(X))-H(X)$
	$\displaystyle=$	$\displaystyle I(X,Z)-H(X)$

where $X=\{x_{1},...,x_{n}\}$ represent the set of all the input samples, $Z$ represents the corresponding latent variables. Since $H(X)$ solely depends on the input data, the maximization of the first term in Eqn.3 is equivalent to maximize the mutual information $I(X,Z)$.

To estimate the mutual information $I(X,Z)$, Hjelm et al. (2019) introduce a JS-divergence based approximation,

$$I^{\text{(JSD)}}(X,Z)=D_{\text{JS}}(Q(Z|X)P(X),Q(Z)P(X))$$

(9)

One way to compute the above JS-divergence in Eqn. 9 is using a discriminator $D$ (Nowozin, Cseke, and Tomioka 2016; Hjelm et al. 2019; Yang et al. 2019):

	$\displaystyle D_{\text{JS}}(P(X),Q(X))$	$\displaystyle=\frac{1}{2}\max_{D}\{E_{x\sim P(X)}[\log(\sigma(D(x)))]$		(10)
		$\displaystyle+E_{x\sim Q(X)}[\log(1-\sigma(D(x)))]\}$
		$\displaystyle+\log 2$

Following Eqn. 9 and Eqn.10, we utilize a neural network to serve as the robust embedding discriminator,

		$\displaystyle L_{\text{MI}}(x_{i};\phi,\rho)=-\log[\sigma(D(x_{i},z_{i};\rho)]$		(11)
		$\displaystyle-E_{(x_{j},z_{i})\sim p(z_{i})p(x_{j})}[\log(1-\sigma(D(x_{j},z_{i};\rho)))]$

where $D(*;\rho)$ is a neural network parameterized by $\rho$. Term $z_{i}$ is the latent variable generated from $x_{i}$ by Eqn.5 and Eqn.6. Discriminator $D(*;\rho)$ outputs a binary value indicating whether $z$ is generated from $x$. The second term in Eqn.LABEL:eq:LMI is negative sampling estimation to train the discriminator. In practice, we randomly draw samples from the datasets to form the negative pairs with $z_{i}$.³³3In the appendix, we further provide a theoretical proof for the effectiveness of the proposed robust discriminator.

One thing should be noted is that $L_{\text{MI}}(x_{i};\phi,\rho)$ cannot replace the original reconstruction loss in Eqn.3. This is because Eqn.9 is valid only when $P(X)$ and $Q(X)$ are close. As we maximize the mutual information $I(X,Z)$ during training, $P(X)$ and $Q(X)$ will be gradually varied and Eqn.9 will no longer hold. In order to avoid this problem, we combine the reconstruction term and the robust embedding discriminator in our method to achieve better clustering performance (Eqn.1).

While the global structure models the latent space based on clustering distribution (e.g. GMM), such methods may generate unsatisfactory results for the data around the cluster boundary. To solve the problem, we introduce a local structure model that captures the pairwise relationships by minimizing the KL-divergence on the distance distribution.

Inspired by Maaten and Hinton (2008), we define the following pairwise distance distribution $p_{ij}$ in the latent space:

$$p_{j|i}=\frac{\exp(-\|\bm{\widetilde{\mu}}_{i}-\bm{\widetilde{\mu}}_{j}\|^{2}/2\eta_{i}^{2})}{\sum_{k\not=i}\exp(-\|\bm{\widetilde{\mu}}_{k}-\bm{\widetilde{\mu}}_{i}\|^{2}/2\eta_{i}^{2})}$$

(12)

$$p_{ij}=\frac{p_{j|i}+p_{i|j}}{2n}$$

(13)

where $\widetilde{\bm{\mu}}_{i}$ and $\widetilde{\bm{\mu}}_{j}$ are the mean of the latent variables obtained from Eqn. 5. Here, we select $\widetilde{\bm{\mu}}$ instead of $z$ because it contains much less noise. Term $\eta_{i}$ is a hyper-parameter adjusted according to the data density. One way to compute $\eta_{i}$ is from a pre-defined perplexity value:

$$\text{Perplexity}(P_{i})=2^{-\sum_{j}p_{j|i}\log_{2}p_{j|i}}$$

(14)

In order to encode the pairwise information, we apply a neural network to minimize the KL-divergence of the distance distribution. To be more precise, we first apply a locality mapping network that converts $\widetilde{\bm{\mu}}_{i}$ to $\bm{\widetilde{\bm{o}}}^{\prime}_{i}$:

$$\bm{\widetilde{\bm{o}}}^{\prime}_{i}=f_{lp}(\bm{\widetilde{\mu}}_{i};\psi)$$

(15)

where $f_{lp}(*;\psi)$ is a neural network parameterized by $\psi$. Once $\bm{\widetilde{\bm{o}}}^{\prime}_{i}$ is obtained, the pairwise distribution $q_{ij}$ between $\bm{\widetilde{\bm{o}}}^{\prime}_{i}$ and $\bm{\widetilde{\bm{o}}}^{\prime}_{j}$ is defined using a student t-distribution with one degree of freedom,

$$q_{ij}=\frac{(1+||\bm{\widetilde{\bm{o}}}^{\prime}_{i}-\bm{\widetilde{\bm{o}}}^{\prime}_{j}||^{2})^{-1}}{\sum_{k}{\sum_{k\not=l}(1+||\bm{\widetilde{\bm{o}}}^{\prime}_{k}-\bm{\widetilde{\bm{o}}}^{\prime}_{l}||^{2})^{-1}}}$$

(16)

t-distribution is applied due to its computational efficiency and its property to mitigate crowding problems (Maaten and Hinton 2008).

Finally, the local structure modeling $L_{\text{LP}}(x_{i},x_{j};\phi,\psi)$ is given by measuring the KL-divergence between the two joint probabilities $p_{ij}$ and $q_{ij}$:

$$L_{\text{LP}}(x_{i},x_{j};\phi,\psi)=\sum_{i\not=j}p_{ij}\log\frac{p_{ij}}{q_{ij}}$$

(17)

An intuitive interpretation for $L_{\text{LP}}(x_{i};\phi,\psi)$ is that it encourages $\bm{\widetilde{\bm{o}}}^{\prime}_{i}$ and $\bm{\widetilde{\bm{o}}}^{\prime}_{j}$ to be closer when the pairwise distance of $\widetilde{\bm{\mu}}_{i}$ and $\widetilde{\bm{\mu}}_{j}$ are small, and vice versa. As a result, the final output of $\bm{\widetilde{\bm{o}}}^{\prime}$ not only inherits the global structure information from $\bm{\widetilde{\mu}}$ but also encodes local structure information from the pairwise relationship, thus can be used as an effective representation for clustering.

To test the ability of our model to solve general clustering tasks, we do experiments on both vision and textual datasets.

For both vision and textual datasets, we combine the training and test sets together because of the unsupervised setting. Furthermore, for vision data, we concatenate the pixels to form the input vector without preprocessing. Following the settings of Xie, Girshick, and Farhadi (2016); Jiang et al. (2017), we set the encoder of global structure as $D-500-500-2000-10$ and decoder as $10-2000-500-500-D$, where $D$ is the input dimension. The robust discriminator uses the structure of $(10+D)-256-1$, where $10$ is the dimension of the latent representation. For the local structure model, we adopt the network with $10-256-256-256-10$. All layers are fully connected using ReLU as the activation function. We set $\alpha_{0}=1$ and $\alpha_{1}=10^{-4}$ in MNIST and Fashion-MNIST and $\alpha_{0}=1$ and $\alpha_{1}=10^{-3}$ in REUTERS-10k and REUTERS. The batch size is 1,000. We use Adam optimizer (Kingma and Ba 2015). We apply K-means on the latent representations (i.e. $\bm{\widetilde{\bm{o}}}^{\prime}$) to do clustering.

We compare the results of our model with several state-of-the-art baseline models. All the K-means algorithms mentioned below use 10 different K-means++ initialization (Arthur and Vassilvitskii 2007) and finally select the settings with the lowest average distances.

• •

  K-means: apply the K-means algorithm to the original data.

• •

  AE + K-means: A pipeline method that trains an Autoencoder first and then applies the K-means algorithm to the latent representations produced by the Autoencoder.

• •

  VAE + K-means: A pipeline method that trains a Variational Autoencoder first and then applies K-means to the latent representations.

• •

  DEC (Xie, Girshick, and Farhadi 2016): Following the experiment in the paper, DEC is first pre-trained with the autoencoder setting, and then iteratively optimized by minimizing the KL-divergence between the distribution of the current soft cluster assignment and its proposed auxiliary target distribution.

• •

  ClusterGAN (Mukherjee et al. 2019): introduces a variation of GAN with an inverse-mapping network and a clustering-specific loss, resulting in superior results among GAN-based algorithms in the clustering task.

• •

  VaDE (Jiang et al. 2017): a generative clustering framework combining VAE with GMM prior. It is reported that VaDE achieves the best performance compared with other general AE/VAE based clustering methods.

The settings of the encoder and decoder in AE and VAE are all the same as our proposed model.

While evaluating the real performance of clustering is arguably a non-trivial problem, we follow the mainstream practices and use the unsupervised clustering accuracy (ACC), Normalized Mutual Information (NMI) and Adjusted Rand Index (ARI) as the metrics to evaluate the models. In our experiments, we set the number of clusters the same as the ground-truth labels.

In this section, we explain the experiment results in details on vision and textual datasets.

In Table 3, we show the results of the ablation study on both MNIST and REUTERS datasets. As our model contains three parts ($L_{\text{G}}$,$L_{\text{MI}}$,$L_{\text{LP}}$), we design the experiment to study the contribution of each component. We use the method that only models the global structure $L_{G}$ as the baseline to compare, which is almost equivalent to VaDE. From the experiment results, it can be seen that the local structure component ($L_{\text{LP}}$) is crucial and improves the clustering performance significantly, while the robust embedding discriminator $L_{\text{MI}}$ only improves the performance by a slight margin. This result is in-line with our expectation because the embedding discriminator is better at improving the robustness of the system rather than the accuracy.

To testify the robustness of our model, we do experiments on the MNIST dataset following Gilmer et al. (2019). Particularly, we gradually corrupt the pixels in the input images with independent Gaussian Noise ($N(0,\sigma^{2})$). In Fig.3, we show that the proposed robust embedding discriminator ($L_{MI}$) performs significantly better than the methods only using the global structure ($L_{G}$ only). In addition, the experiment results also show that the local structure ($L_{LP}$) contributes to the robustness, especially when the image is seriously corrupted. We can also observe that VaDE and ClusterGAN are fragile to the input noise, especially ClusterGAN, which collapses even under small pollution ratio.

In Fig. 4, we show a concrete example by performing an image retrieval task using the latent representations. In particular, we use the representation of an image as the query to retrieve top $K$ images based on euclidean distances. Specifically, in the task we extract digit “4” form the dataset (this is a difficult example as it looks like “9”). The top 10 nearest neighbors found by our model are all correct, while the ones obtained using VaDE are a mixture of “4” and “9”. This is because the global structure modeling usually focuses on the cluster centroids and modeling the data distribution around those centroids, thus it is difficult to differentiate the data around the cluster boundary. On the other hand, since our model encodes the local pairwise relationships, it can incorporate the information of the boundary points, therefore performs better under such situations.

In Fig. 5 we examine the accuracy of our model on MNIST by adjusting the embedding dimension. It can be observed that the results are quite stable with various dimensions. One interesting observation is that the cluster accuracy is almost kept when we reduce the dimension to 2, which makes our model promising to be served as a data visualization tool. In Fig. 6 we show the visualization for the latent representations on the test datasets of MNIST and REUTERS-10k, which all contain 10,000 samples. Different colors represent the gold labels. It shows that our proposed LPVDN creates a cohesive structure for the same clusters, while clearly separates the latent representations of different clusters.