[Paper Review] Deep $k$-Means: Jointly clustering with $k$-Means and learning representations

Maziar Moradi Fard, Thibaut Thonet+1 more|arXiv (Cornell University)|Jun 26, 2018

Gaussian Processes and Bayesian Inference39 references259 citations

TL;DR

DKM jointly learns data representations and k-Means clustering by a differentiable soft-assignment surrogate, achieving strong clustering without sacrificing reconstruction quality. It outperforms several deep clustering baselines across image and text datasets.

ABSTRACT

We study in this paper the problem of jointly clustering and learning representations. As several previous studies have shown, learning representations that are both faithful to the data to be clustered and adapted to the clustering algorithm can lead to better clustering performance, all the more so that the two tasks are performed jointly. We propose here such an approach for $k$-Means clustering based on a continuous reparametrization of the objective function that leads to a truly joint solution. The behavior of our approach is illustrated on various datasets showing its efficacy in learning representations for objects while clustering them.

Motivation & Objective

• Motivate and address the challenge of clustering in high-dimensional data by jointly learning representations and cluster assignments.
• Develop a differentiable end-to-end optimization that unifies representation learning with k-Means clustering.
• Show that a continuous relaxation of the k-Means objective enables scalable SGD-based training.
• Compare against state-of-the-art deep clustering baselines under controlled, fair conditions.
• Demonstrate the method’s effectiveness across diverse datasets (image and text).

Proposed method

• Formulate deep k-Means as minimizing a reconstruction loss plus a clustering loss in an embedding space learned by an auto-encoder.
• Introduce a continuous generalization of the k-Means objective using a differentiable surrogate Gk,f based on a parameterized softmax (Gk,f(hθ(x),α;R)).
• Use α as an inverse temperature and apply a deterministic annealing schedule to gradually harden assignments.
• Optimize θ (auto-encoder) and R (cluster centers) jointly via SGD with the surrogate loss, allowing end-to-end training.
• Discuss choices for Gk,f (including softmax-based surrogates) and conditions under which the surrogate converges to the true hard assignment.
• Consider pretraining vs. annealing as strategies for initialization and initialization of cluster centers.

Experimental results

Research questions

• RQ1 Can k-Means clustering and representation learning be truly optimized jointly in a differentiable, end-to-end framework?
• RQ2 Does a soft, differentiable surrogate for the argmin in k-Means (with annealing) achieve competitive or superior clustering without explicit discrete updates?
• RQ3 How does joint learning with a continuous relaxation compare to pretraining-based deep clustering methods across varied data modalities?
• RQ4 What is the impact of annealing α on clustering quality and representation learning?
• RQ5 Are the learned representations more discriminative for clustering than those produced by decoupled or traditional auto-encoder + k-Means pipelines?

Key findings

• DKM with annealing (DKMa) outperforms non-pretrained DCN across all datasets and approaches pretraining-based methods in clustering quality.
• With pretraining, DKMp achieves results competitive with or better than strong baselines, often surpassing IDEC and DCN on multiple datasets.
• Jointly learning representations and clustering via the differentiable surrogate yields state-of-the-art or near state-of-the-art performance among k-Means–related deep clustering methods on MNIST, USPS, 20NEWS, and RCV1.
• Annealing the surrogate parameter α serves as an effective alternative to pretraining for initializing cluster centers and achieving good local optima.
• The learned embedding space for DKM, especially with pretraining, yields well-separated clusters as visualized by t-SNE on USPS.
• Across datasets, the DKM variants consistently outperform standard KM and AE-KM baselines, illustrating the benefit of joint optimization.