Uncertainty-Aware Contrastive Learning for deep clustering

Abstract

Deep clustering aims to group unlabeled data into meaningful clusters by learning discriminative feature representations. However, ambiguous features often lead to noisy representations and inconsistent semantics, limiting improvements in clustering performance. To address this issue, we propose an Uncertainty-Aware Contrastive Learning (UACL) method for deep clustering, which achieves robustness by adaptively restricting the learning of ambiguous features. Specifically, we model pairwise similarity evidence via subjective logic theory, formulating co-cluster probabilities as a Dirichlet distribution to quantify epistemic uncertainty from feature ambiguity. Guided by this uncertainty, we design a dynamic weight-updating strategy that progressively extracts information from potential positives, enhancing the model’s ability to learn discriminative representations and semantically consistent clusters. Furthermore, to enforce attribute consistency, we develop an Attribute Distribution Alignment module that aligns similarity and uncertainty. Extensive experiments on five benchmark datasets demonstrate UACL outperforms current state-of-the-art methods, with an improved ACC of 3.5% for CIFAR-100 and 7.0% for ImageNet-Dogs. The source code is available at: https://github.com/YL616/UACL.