Clustering via Self-Supervised Diffusion

TL;DR

This paper introduces CLUDI, a novel self-supervised clustering framework that leverages diffusion models and pre-trained Vision Transformers to improve clustering accuracy and robustness in high-dimensional data.

Contribution

The paper presents a new diffusion-based clustering method combining generative diffusion models with Vision Transformer features, using a teacher-student paradigm for stable and accurate clustering.

Findings

Achieves state-of-the-art clustering performance on challenging datasets.

Demonstrates robustness and adaptability to complex data distributions.

Introduces a stochastic diffusion-based data augmentation strategy.

Abstract

Diffusion models, widely recognized for their success in generative tasks, have not yet been applied to clustering. We introduce Clustering via Diffusion (CLUDI), a self-supervised framework that combines the generative power of diffusion models with pre-trained Vision Transformer features to achieve robust and accurate clustering. CLUDI is trained via a teacher-student paradigm: the teacher uses stochastic diffusion-based sampling to produce diverse cluster assignments, which the student refines into stable predictions. This stochasticity acts as a novel data augmentation strategy, enabling CLUDI to uncover intricate structures in high-dimensional data. Extensive evaluations on challenging datasets demonstrate that CLUDI achieves state-of-the-art performance in unsupervised classification, setting new benchmarks in clustering robustness and adaptability to complex data distributions. Our code is available at https://github.com/BGU-CS-VIL/CLUDI.