Combining clustering and representation learning is one of the most promising approaches for unsupervised learning of deep neural networks. However, doing so naively leads to ill posed learning problems with degenerate solutions. In this paper, we propose a novel and principled learning formulation that addresses these issues. The method is obtained by maximizing the information between labels and input data indices. We show that this criterion extends standard crossentropy minimization to an optimal transport problem, which we solve efficiently for millions of input images and thousands of labels using a fast variant of the Sinkhorn-Knopp algorithm. The resulting method is able to self-label visual data so as to train highly competitive image representations without manual labels. Our method achieves state of the art representation learning performance for AlexNet and ResNet-50 on SVHN, CIFAR-10, CIFAR-100 and ImageNet.
Summary:
The paper proposes a self-supervised learning procedure to train deep neural networks within an unsupervised learning setting. The authors build their work on a pretext task that consists in maximizing the information between the input data samples and labels that are basically obtained by a self-labeling procedure that clusters them in K distinct classes. This is similar to what was done in DeepCluster, a previous self-supervised algorithm that self-labels samples through clustering. However, differently from that approach, the current method does not introduce any additional clustering cost functions. Instead it implicitly achieves self-labeling by simply adding the constraint that label assignments equally partition the dataset. This constraint acts as a "regularizer" that allows the authors to minimize the cross-entropy loss between inputs and pseudo-labels while avoiding the degenerate trivial solution where all samples are assigned to the same pseudo-label. As a result, the authors are able to derive a self-labeling method that optimizes the same cross-entropy loss as the classification task. That is done by remarking that minimizing the loss function over the pseudo-label assignments (under the equal partition constraint) can be formulated as an optimal transport problem that can be solved efficiently with a fast version of the Sinkhorn-Knopp algorithm. In practice, what the authors do at training time is to alternate between 1) minimizing the average cross-entropy loss by training the neural network (feature extractor + linear classification head) given a fixed pseudo-labels assignment, and 2) optimizing the pseudo-label assignments implemented as a matrix Q of posterior distribution of labels given the sample index, which, as said, can be done efficiently with a KL-regularized version of Sinkhorn. Moreover, this last step can be carried out simultaneously for multiple distinct classification heads (with possibly different number of labels), each sharing the same feature extractor but inducing a different matrix Q. At this point, the number of classification heads can be treated as a hyperparameter of the algorithm.
The authors then go on to show that this new algorithm is competitive with current state of the art method with several architectures in terms of providing a good feature extractor for downstream image classification, detection and segmentation tasks. They for instance consistently beat DeepCluster, the main direct competitor, on classification and detection tasks.
They also conduct ablation studies that provide interesting insights on the functioning of their algorithms and the effects of the multiple classification heads, the number of clusters, and the quality of the learned assignment. Intriguingly, on ImageNet they find for example that AlexNet obtains better performance at validation when it's trained from scratch on labels obtained with their self-labeling procedure, as opposed to the original labels.
Decision:
In my opinion this paper should be a clear accept. The paper is well written, presents an elegant idea in a clear and straight-forward manner, and is solidly built on top of the current literature on self-supervised learning for image processing, which is also very well summarized.
The feature extractors obtained with the proposed algorithm are convincingly tested and validated on several downstream tasks (like classification on ImageNet, PascalVOC classification, detection and segmentation), and that is done for several base architectures, obtaining performances that are competitive with state-of-the-art. In addition, a series of careful ablation studies help in gleaning some scientific understanding on the method.
Minor comments:
- The authors refer to traditional clustering like k-means as being "generative", which a little confusing. Clustering algorithms can be derived within a probabilistic framework by positing a generative model of the data, however, strictly speaking, k-means is not by itself is not a generative approach. It's a minor point, but it could be helpful to be more precise about this, in order to avoid possible confusion. The main point that the authors want to make in this regard is that their framework eschews having to posit an additional clustering cost function.