Shuffle to Learn: Self-supervised learning from permutations via differentiable ranking

1 INTRODUCTION . Supervised learning has achieved important successes on large annotated datasets ( Deng et al. , 2009 ; Amodei et al. , 2016 ) . However , most available data , whether images , audio , or videos are unlabelled . For this reason , pre-training representations in an unsupervised way , with subsequent fine-tuning on labelled data , has become the standard to extend the performance of deep architectures to applications where annotations are scarce , such as understanding medical images ( Rajpurkar et al. , 2017 ) , recognizing speech from under-resourced languages ( Rivière et al. , 2020 ; Conneau et al. , 2020 ) , or solving specific language inference tasks ( Devlin et al. , 2018 ) . Among unsupervised training schemes , self-supervised learning focuses on designing a proxy training objective , that requires no annotation , such that the representations incidentally learned will generalize well to the task of interest , limiting the amount of labeled data needed for fine-tuning . Such “ pretext ” tasks , a term coined by Doersch et al . ( 2015 ) , include learning to colorize an artificially gray-scaled image ( Larsson et al. , 2017 ) , inpainting removed patches ( Pathak et al. , 2016 ) or recognizing with which angle an original image was rotated ( Gidaris et al. , 2018 ) . Other approaches for self-supervision include classification to original images after data augmentation ( Chen et al. , 2020 ) and clustering ( Caron et al. , 2018 ) . In this work , we consider the pretext task of reordering patches of an input , first proposed for images by Noroozi & Favaro ( 2016 ) , the analogue of solving a jigsaw puzzle . In this setting , we first split an input into patches and shuffle them by applying a random permutation . We train a neural network to predict which permutation was applied , taking the shuffled patches as inputs . We then use the inner representations learned by the neural network as input features to a low-capacity supervised classifier ( see Figures 1 and 2 for illustration ) . We believe that permutations provide a promising avenue for self-supervised learning , as they are conceptually general enough to be applied across a large range of modalities , unlike colorization ( Larsson et al. , 2017 ) or rotations ( Gidaris et al. , 2018 ) that are specific to images . The idea of using permutations was also explored in Santa Cruz et al . ( 2018 ) where they use a bi level optimization scheme which leverages sinkhorn iterations to learn visual reconstructions . Their method resorts to approximating the permuation matrix with such continuous methods . Our method relies on no such approximations and can efficiently represent all possible permutations . Moreover , the encouraging results of Noroozi & Favaro ( 2016 ) when transferring learned image features for object detection and image retrieval inspire us to advance this method a step forward . However , including permutations into an end-to-end differentiable pipeline is challenging , as permutations are a discontinuous operation . Noroozi & Favaro ( 2016 ) circumvent this issue by using a fixed set of permutations and casting the permutation prediction problem as a classification one . Given that the number of possible permutations of n patches is n ! , this approach can not scale to exploiting the full set of permutations , even when n is moderately small . In this work , we leverage recent advances in differentiable ranking ( Berthet et al. , 2020 ; Blondel et al. , 2020b ) to integrate permutations into end-to-end neural training . This allows us to solve the permutation inversion task for the entire set of permutations , removing a bottleneck that was heretofore sidestepped in manners that could deteriorate downstream performance . Moreover , we successfully demonstrate for the first time the effectiveness of permutations as a pretext task on multiple modalities with minimal modality-specific adjustments . In particular , we improve music understanding by learning to reorder spectrogram frames , over the time and frequency axes . We also improve video understanding by reordering video frames along time . To summarize , we make the following two contributions . - We integrate differentiable ranking into end-to-end neural network training for representation learning . This provides an efficient manner to learn in reordering tasks for all permutations , for larger numbers of patches . We show that this drastic increase in the number of permutations improves the quality of learned representations for downstream tasks . - We successfully demonstrate for the first time the effectiveness of permutations as a general purpose self-supervision method , efficient on multiple modalities with extremely minimal modifications to the network . Additionally , the pre-trained representations perform well across diverse tasks of the same modality . We purposefully divert from domain-specific transformations , often inspired by data augmentation techniques , as predicting the permutation applied to input patches is not restricted either to a domain , nor a modality or a dimensionality . This also creates opportunities for applications beyond the scope of what we illustrate in our work . The rest of the paper is organized as follows . In Section 2 we present the problem formulation and the methods used . In Section 3 we demonstrate the effectiveness of our experiments on audio , video , and image tasks . Further details can be found in the Appendix . 2 METHODS . 2.1 GENERAL METHODOLOGY . To efficiently leverage the existence of large amounts of unlabeled data , we present a self-supervised pretext task that predicts the permutation applied to patches of an input . We do so in a manner that removes an existing bottleneck , and allows to use all possible permutations as targets during training . This pretext task is performed upstream and the internal representation learned by the pretext neural network can then be transferred and used on secondary downstream tasks – see Figures 1 and 2 . In the upstream task , for each data point , n patches , sub-parts x1 , . . . , xn of of identical dimensions d are extracted . Their exact structure depend naturally on the modality : e.g . horizontal bands of an image , frequency bands of a spectrogram , frames of a video . Accordingly , the dimensions in d can represent height , width , channels , etc . These patches are then permuted randomly , and organized in a tensor Xi of dimension n× d ( see Figure 2 ) , which is paired to the applied permutation as a label yi ( see Section 2.2 for details on permutation encoding ) . We then train the weights w of a neural network to minimize the loss between its output fw ( Xi ) , of size n , and the encoding yi ∈ Rn of the permutation applied to the n patches ( see Figure 2 ) . Note that the last operation of the network is a differentiable ranking operator y∗ε . This operator , the encoding of the permutations for the labels , and the loss used in this upstream task are detailed in Section 2.2 below . The network , and the details of the data-processing pipeline generating the patches , are detailed in Section 2.3 . After upstream training on the initial dataset , the upstream network weights can be used to generate representations . By truncating the network , removing some of the last layers , we can extract an embedding of any input vector . These representations can be used in a downstream task to train a new network , with its own weights , minimizing a loss ( typically classification or regression ) between its output and the downstream task labels ( see Figures 1 and 2 ) . We mostly evaluate our methods on downstream performance : the accuracy after training of the downstream network , on a task that is unknown during upstream training . However , the pretext reordering task can be of interest in and of itself , as in learning-to-rank problems ( Liu , 2011 ) , and we also report generalization performance in this task . As an aside , in the Jigsaw puzzle reassembly task , Noroozi & Favaro ( 2016 ) show that the choice of permutation set matters when it comes to performance . They make use of the Hamming distance to choose a permutation set that it maximally separated . This permutation set is diverse but is not close to covering the entire permutation space . By supporting the full set of permutations , our approach does not suffer from this issue . The downstream tasks are dataset-dependent . However , the network in the upstream reordering task requires minimal modification across tasks and modalities ( see Section 3 for details ) . In this work , we demonstrate the effectiveness of our method on audio , video , and images , across several classification and regression tasks . 2.2 DIFFERENTIABLE RANKING METHODOLOGY . Our methodology for representation learning relies importantly on the ability to incorporate ordering or ranking operations in an end-to-end differentiable pipeline . During training for the upstream task , the last two layers of the network consist of : a vector of score values θw ( X ) ∈ Rn , and network outputs fw ( X ) = y∗ε ( θw ( X ) ) ∈ Rn , using a differentiable ranking operator y∗ε that we detail here . The goal of the upstream task is to find good parameters w such that the network outputs fw ( X ) correctly recover the label y representing the permutation applied to the patches in X . In earlier works using permutations as pretext task ( Noroozi & Favaro , 2016 ; Lee et al. , 2017 ) , training with permutation labels is achieved by reducing the permutation inversion task to classification . More specifically , it is encoding a fixed number of permutations L n ! as classes . Each class is represented by a one-hot vector , and network outputs are logits θw ( X ) of size L , leading to a prediction by taking a softmax among these L classes . This approach is obviously limited : representing all the permutations requires in principle n ! classes , which is quickly not manageable , even for small values of n. Further , this does not take into account the similarity between permutations : with this encoding , permutations are orthogonal , no matter how similar they are . We address these issues by having network outputs θw ( X ) of size only n , and interpreting their relative orders as the predicted permutation ( e.g . y = ( 0 , 1 , 2 , 3 ) if they are in decreasing order , predicting the identity permutation ) . The pretext labels also encode the permutations in this manner . This gives a unique encoding to each permutation , operates in dimension n , and can be computed with sorting operations in time O ( n log n ) . Further , small distances in these encodings naturally represent similar permutations . However , applying directly a ranking operator y∗ on the last layer of the network would not allow for backpropagation of gradients : the function of the weights w 7→ L ( y∗ ( θw ( X ) ) ; y ) is piece-wise constant . Small changes in the weights w induce either large jumps or no change in value at all , its gradients are 0 almost everywhere , and undefined otherwise . In order to overcome this matter , we consider instead two differentiable ranking operations , one using stochastic perturbations , introduced in Berthet et al . ( 2020 ) and another one using regularization ( Blondel et al. , 2020b ) . These operations , denoted here by y∗ε , map any vector of k values to a point in the convex hull of permutation encodings in dimension k ( e.g . ( 0.1 , 0.9 , 2.2 , 2.8 ) over 4 elements ) . They can be thought of analogous to the softmax operator for ranking . They share some of its properties : good approximation of the original function , differentiability in the input values θ with non-zero derivatives everywhere , ease of computation , and tuning by a temperature parameter ε ( see Appendix A.1 for further details ) . Adjusting the network parameters w requires a notion of loss function between y∗ε ( θ ) and y . For the version of y∗ε ( θ ) based on stochastic perturbations ( Berthet et al. , 2020 ) , we use the associated Fenchel–Young loss ( Blondel et al. , 2020a ) , that act directly on θ = θw ( X ) ( outputs of the vector , and inputs of the sorting operations ) , written here as LFY ( θ ; y ) ( see Appendix A.1 ) . This loss is convex in θ , smooth , equal to 0 if and only if y∗ε ( θ ) = y . Its gradients are given by ∇θLFY ( θ ; y ) = y∗ε ( θ ) − y . We call this loss “ Perturbed F-Y ” in our empirical results . For the regularized version of y∗ε ( θ ) ( Blondel et al. , 2020b ) , we use 1 2 ‖y∗ε ( θ ) − y‖2 . We call this loss “ Fast Soft Ranking ( FSR ) ” in our empirical results . We opt for these two losses for their good theoretical properties and O ( n log n ) complexity . Other choices ( Mena et al. , 2018 ; Cuturi et al. , 2019 ; Vlastelica et al. , 2019 ; Rolı́nek et al. , 2020 ; Grover et al. , 2019 ) are also possible , potentially with higher computational cost , or regions with zero gradient .

Summary: This paper presents a self-supervised learning task of shuffling input patches and demanding the network to learn to unshuffle. A related prior work, Noorozi and Favaro (2016) uses a fixed set of permutations to do this task for a given number of patches, and the current paper argues to expand this idea for the full set of permutations. To this end, the paper encodes a permutation as a number tuple, with the goal of the network to learn to produce the correct tuple that has the numbers in order. As the numbers are discrete and thus non-differentiable, the paper suggests differentiable soft-variants using stochastic perturbations and regularizations (Fenchel-Young loss). Experiments are provided on audio and video tasks and show promise over the method of Noorozi and Favaro (2016).