Self-labelling via simultaneous clustering and representation learning

1 INTRODUCTION . Learning from unlabelled data can dramatically reduce the cost of deploying machine learning algorithms to new applications , thus amplifying their impact in the real world . Self-supervision is an increasingly popular framework for learning without labels . The idea is to define pretext learning tasks that can be constructed from raw data alone , but that still result in neural networks that transfer well to useful applications . Much of the research in self-supervision has focused on designing new pretext tasks . However , given supervised data such as ImageNet ( Deng et al. , 2009 ) , the standard classification objective of minimizing the cross-entropy loss still results in better or at least as good pre-training than any of such methods ( for a given amount of data and for a given model complexity ) . This suggests that the task of classification is sufficient for pre-training networks , provided that suitable data labels are available . In this paper , we thus focus on the problem of obtaining the labels automatically by designing a self-labelling algorithm . Learning a deep neural network together while discovering the data labels can be viewed as simultaneous clustering and representation learning . The latter can be approached by combining cross-entropy minimization with an off-the-shelf clustering algorithm such as K-means . This is precisely the approach adopted by the recent DeepCluster method ( Caron et al. , 2018 ) , which achieves excellent results in unsupervised representation learning . However , combining representation learning , which is a discriminative task , with clustering is not at all trivial . In particular , we show that the combination of cross-entropy minimization and K-means as adopted by DeepCluster can not be described as the optimization of an overall learning objective ; instead , there exist degenerate solutions that the algorithm avoids via particular implementation choices . In order to address this technical shortcoming , in this paper , we contribute a new principled formulation for simultaneous clustering and representation learning . The starting point is to minimize a single loss , the cross-entropy loss , for learning the deep network and for estimating the data labels . This 1https : //github.com/yukimasano/self-label is often done in semi-supervised learning and multiple instance learning . However , when applied naively to the unsupervised case , it immediately leads to a degenerate solution where all data points are mapped to the same cluster . We solve this issue by adding the constraint that the labels must induce an equipartition of the data , which we show maximizes the information between data indices and labels . We also show that the resulting label assignment problem is the same as optimal transport , and can therefore be solved in polynomial time by linear programming . However , since we want to scale the algorithm to millions of data points and thousands of labels , standard transport solvers are inadequate . Thus , we also propose to use a fast version of the Sinkhorn-Knopp algorithm for finding an approximate solution to the transport problem efficiently at scale , using fast matrix-vector algebra . Compared to methods such as DeepCluster , the new formulation is more principled and allows to more easily demonstrate properties of the method such as convergence . Most importantly , via extensive experimentation , we show that our new approach leads to significantly superior results than DeepCluster , achieving the new state of the art for representation learning approaches . In fact , the method ’ s performance surpasses others that use a single type of supervisory signal for self-supervision , and is on par or better than very recent contributions as well ( Tian et al. , 2019 ; He et al. , 2019 ; Misra & van der Maaten , 2019 ; Oord et al. , 2018 ) . 2 RELATED WORK . Our paper relates to two broad areas of research : ( a ) self-supervised representation learning , and ( b ) more specifically , training a deep neural network using pseudo-labels , i.e . the assignment of a label to each image . We discuss closely related works for each . Self-supervised learning : A wide variety of methods that do not require manual annotations have been proposed for the self-training of deep convolutional neural networks . These methods use various cues and proxy tasks namely , in-painting ( Pathak et al. , 2016 ) , patch context and jigsaw puzzles ( Doersch et al. , 2015 ; Noroozi & Favaro , 2016 ; Noroozi et al. , 2018 ; Mundhenk et al. , 2017 ) , clustering ( Caron et al. , 2018 ; Huang et al. , 2019 ; Zhuang et al. , 2019 ; Bautista et al. , 2016 ) , noise-as-targets ( Bojanowski & Joulin , 2017 ) , colorization ( Zhang et al. , 2016 ; Larsson et al. , 2017 ) , generation ( Jenni & Favaro , 2018 ; Ren & Lee , 2018 ; Donahue et al. , 2017 ; Donahue & Simonyan , 2019 ) , geometry ( Dosovitskiy et al. , 2016 ) , predicting transformations ( Gidaris et al. , 2018 ; Zhang et al. , 2019 ) and counting ( Noroozi et al. , 2017 ) . Most recently , contrastive methods have shown great performance gains , ( Oord et al. , 2018 ; Hénaff et al. , 2019 ; Tian et al. , 2019 ; He et al. , 2019 ) by leveraging augmentation and adequate losses . In ( Feng et al. , 2019 ) , predicting rotation ( Gidaris et al. , 2018 ) is combined with instance retrieval ( Wu et al. , 2018 ) and multiple tasks are combined in ( Doersch & Zisserman , 2017 ) . Pseudo-labels for images : In the self-supervised domain , we find a spectrum of methods that either give each data point a unique label ( Wu et al. , 2018 ; Dosovitskiy et al. , 2016 ) or train on a flexible number of labels with K-means ( Caron et al. , 2018 ) , with mutual information ( Ji et al. , 2018 ) or with noise ( Bojanowski & Joulin , 2017 ) . In ( Noroozi et al. , 2018 ) a large network is trained with a pretext task and a smaller network is trained via knowledge transfer of the clustered data . Finally , ( Bach & Harchaoui , 2008 ; Vo et al. , 2019 ) use convex relaxations to regularized affine-transformation invariant linear clustering , but can not scale to larger datasets . Our contribution is a simple method that combines a novel pseudo-label extraction procedure from raw data alone and the training of a deep neural network using a standard cross-entropy loss . 3 METHOD . We will first derive our self-labelling method , then interpret the method as optimizing labels and targets of a cross-entropy loss and finally analyze similarities and differences with other clusteringbased methods . 3.1 SELF-LABELLING . Neural network pre-training is often achieved via a supervised data classification task . Formally , consider a deep neural network x = Φ ( I ) mapping data I ( e.g . images ) to feature vectors x ∈ RD . The model is trained using a dataset ( e.g . ImageNet ) of N data points I1 , . . . , IN with corresponding labels y1 , . . . , yN ∈ { 1 , . . . , K } , drawn from a space of K possible labels . The representation is followed by a classification head h : RD → RK , usually consisting of a single linear layer , converting the feature vector into a vector of class scores . The class scores are mapped to class probabilities via the softmax operator : p ( y = ·|xi ) = softmax ( h ◦ Φ ( xi ) ) . The model and head parameters are learned by minimizing the average cross-entropy loss E ( p|y1 , . . . , yN ) = − 1 N N∑ i=1 log p ( yi|xi ) . ( 1 ) Training with objective ( 1 ) requires a labelled dataset . When labels are unavailable , we require a self-labelling mechanism to assign the labels automatically . In semi-supervised learning , self-labelling is often achieved by jointly optimizing ( 1 ) with respect to the model h ◦ Φ and the labels y1 , . . . , yN . This can work if at least part of the labels are known , thus constraining the optimization . However , in the fully unsupervised case , it leads to a degenerate solution : eq . ( 1 ) is trivially minimized by assigning all data points to a single ( arbitrary ) label . To address this issue , we first rewrite eq . ( 1 ) by encoding the labels as posterior distributions q ( y|xi ) : E ( p , q ) = − 1 N N∑ i=1 K∑ y=1 q ( y|xi ) log p ( y|xi ) . ( 2 ) If we set the posterior distributions q ( y|xi ) = δ ( y−yi ) to be deterministic , the formulations in eqs . ( 1 ) and ( 2 ) are equivalent , in the sense that E ( p , q ) = E ( p|y1 , . . . , yN ) . In this case , optimizing q is the same as reassigning the labels , which leads to the degeneracy . To avoid this , we add the constraint that the label assignments must partition the data in equally-sized subsets . Formally , the learning objective objective2 is thus : min p , q E ( p , q ) subject to ∀y : q ( y|xi ) ∈ { 0 , 1 } and N∑ i=1 q ( y|xi ) = N K . ( 3 ) The constraints mean that each data point xi is assigned to exactly one label and that , overall , the N data points are split uniformly among the K classes . The objective in eq . ( 3 ) is combinatorial in q and thus may appear very difficult to optimize . However , this is an instance of the optimal transport problem , which can be solved relatively efficiently . In order to see this more clearly , let Pyi = p ( y|xi ) 1N be the K ×N matrix of joint probabilities estimated by the model . Likewise , let Qyi = q ( y|xi ) 1N be K ×N matrix of assigned joint probabilities . Using the notation of ( Cuturi , 2013 ) , we relax matrix Q to be an element of the transportation polytope U ( r , c ) : = { Q ∈ RK×N+ | Q1 = r , Q > 1 = c } . ( 4 ) Here 1 are vectors of all ones of the appropriate dimensions , so that r and c are the marginal projections of matrix Q onto its rows and columns , respectively . In our case , we require Q to be a matrix of conditional probability distributions that split the data uniformly , which is captured by : r = 1 K · 1 , c = 1 N · 1 . With this notation , we can rewrite the objective function in eq . ( 3 ) , up to a constant shift , as E ( p , q ) + logN = 〈Q , − logP 〉 , ( 5 ) where 〈·〉 is the Frobenius dot-product between two matrices and log is applied element-wise . Hence optimizing eq . ( 3 ) with respect to the assignments Q is equivalent to solving the problem : min Q∈U ( r , c ) 〈Q , − logP 〉 . ( 6 ) 2We assume for simplicity that K divides N exactly , but the formulation is easily extended to any N ≥ K by setting the constraints to either bN/Kc or bN/Kc+ 1 , in order to assure that there is a feasible solution . This is a linear program , and can thus be solved in polynomial time . Furthermore , solving this problem always leads to an integral solution despite having relaxed Q to the continuous polytope U ( r , c ) , guaranteeing the exact equivalence to the original problem . In practice , however , the resulting linear program is large , involving millions of data points and thousands of classes . Traditional algorithms to solve the transport problem scale badly to instances of this size . We address this issue by adopting a fast version ( Cuturi , 2013 ) of the Sinkhorn-Knopp algorithm . This amounts to introducing a regularization term min Q∈U ( r , c ) 〈Q , − logP 〉+ 1 λ KL ( Q‖rc > ) , ( 7 ) where KL is the Kullback-Leibler divergence and rc > can be interpreted as a K ×N probability matrix . The advantage of this regularization term is that the minimizer of eq . ( 7 ) can be written as : Q = diag ( α ) Pλ diag ( β ) ( 8 ) where exponentiation is meant element-wise and α and β are two vectors of scaling coefficients chosen so that the resulting matrix Q is also a probability matrix ( see ( Cuturi , 2013 ) for a derivation ) . The vectors α and β can be obtained , as shown below , via a simple matrix scaling iteration . For very large λ , optimizing eq . ( 7 ) is of course equivalent to optimizing eq . ( 6 ) , but even for moderate values of λ the two objectives tend to have approximately the same optimizer ( Cuturi , 2013 ) . Choosing λ trades off convergence speed with closeness to the original transport problem . In our case , using a fixed λ is appropriate as we are ultimately interested in the final clustering and representation learning results , rather than in solving the transport problem exactly . Our final algorithm ’ s core can be described as follows . We learn a model h◦Φ and a label assignment matrix Q by solving the optimization problem eq . ( 6 ) with respect to both Q , which is a probability matrix , and the model h ◦ Φ , which determines the predictions Pyi = softmaxy ( h ◦ Φ ( xi ) ) . We do so by alternating the following two steps : Step 1 : representation learning . Given the current label assignments Q , the model is updated by minimizing eq . ( 6 ) with respect to ( the parameters of ) h ◦ Φ . This is the same as training the model using the common cross-entropy loss for classification . Step 2 : self-labelling . Given the current model h ◦ Φ , we compute the log probabilities P . Then , we find Q using eq . ( 8 ) by iterating the updates ( Cuturi , 2013 ) ∀y : αy ← [ Pλβ ] −1y ∀i : βi ← [ α > Pλ ] −1i . Each update involves a single matrix-vector multiplication with complexityO ( NK ) , so it is relatively quick even for millions of data points and thousands of labels and so the cost of this method scales linearly with the number of images N . In practice , convergence is reached within 2 minutes on ImageNet when computed on a GPU . Also , note that the parameters α and β can be retained between steps , thus allowing a warm start of Step 2 .

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.