Contrastive learning unlocks geometric insights for dataset pruning

Abstract

Dataset pruning aims at selecting a subset of the data so that the model trained on the subset performs comparably to the one trained on the full dataset. In the era of big data, unsupervised pruning of the dataset can alleviate the issue of the expensive labeling process from the beginning. Existing methods sort and select instances by well-designed importance metrics, while the unsupervised ones commonly regard representation learning as a black box employed to get embeddings, with its properties remaining insufficiently explored for dataset pruning. In this study, we revisit self-supervised Contrastive Learning by observing the learned embedding manifold, introducing Curvature Estimation to characterize the geometrical properties of the manifold. The statistical results reveal that the embedding distribution of instances on manifold surfaces is not uniform. Based on this observation, we propose an unsupervised dataset pruning strategy by performing downsampling in geometric areas with high instance density, namely KITTY sampling. Extensive experiments demonstrate that our proposed methods have achieved leading performances on CV dataset pruning compared to the baselines. Code is available at https://github.com/Frostland12138/KITTY.