Deep mean shift clustering

Abstract

We use Mean Shift clustering in the latent space of an auto-encoder to have a better representation of the data and a more structured latent space. Instead of just using the mode of the distribution calculated using kernel density estimates, we use trajectories of data points leading to the modes to better model the basin of attraction of each mode. This helps in better structuring of the latent space and results in a more inferential model. Since mean-shift can be modelled as an RNN-block [16] our method is end-to-end trainable. Tuning the bandwidth of mean-shift gives us the flexibility of clustering the latent space on different hierarchical levels. We modify the original trajectory based LSTM model by incorporating a discounting mechanism. We modified the mean shift implementation by using a fixed kernel for the mean shift iteratiosn. We also apply a new loss (Support Set Loss) to penalize the clusters made on the latent space. This uses the trajectories of the points segregated into groups which ended up in the same mode and those which didn’t. We have used this loss function in both semi-supervised and unsupervised fashion. In the end, we also propose a model which uses Contrastive Predictive Coding loss, given in [23] in the latent space as well as a regularizer for the encoding network model.