S EED: S ELF - SUPERVISED D ISTILLATION F OR V ISUAL R EPRESENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper is concerned with self-supervised learning for small models.
The problem is motivated by our empirical studies that while the widely used contrastive self-supervised learning method has shown great progress on large model training, it does not work well for small models.
To address this problem, we propose a new learning paradigm, named SE lf-Sup E rvised D istillation ( S EED), where we leverage a larger network (as Teacher) to transfer its representational knowledge into a smaller architecture (as Student) in a self-supervised fashion.
Instead of directly learning from unlabeled data, we train a student encoder to mimic the similarity score distribution inferred by a teacher over a set of instances.
We show that S EED dramatically boosts the performance of small networks on downstream tasks.
Compared with self-supervised baselines, S EED improves the top-1 accuracy from 42.2% to 67.6% on EfﬁcientNet-B0 and from 36.3% to 68.2% on MobileNetV3-Large on the ImageNet-1k dataset.
INTRODUCTION
[Figure]
The burgeoning studies and success on self-supervised learning ( SSL ) for visual representation are mainly marked by its extraordinary potency of learning from unlabeled data at scale.
Accompanying with the SSL is its phenomenal beneﬁt of obtaining task-agnostic representations while allowing the training to dispense with prohibitively expensive data labeling.
Major ramiﬁcations of visual SSL include pretext tasks (Noroozi & Favaro, 2016; Zhang et al., 2016; Gidaris et al., 2018; Zhang et al., 2019;
Feng et al., 2019), contrastive representation learning (Wu et al., 2018;
He et al., 2020; Chen et al., 2020a), online/ofﬂine clustering (Yang et al., 2016; Caron et al., 2018; Li et al., 2020; Caron et al., 2020;
Grill et al., 2020), etc.
Among them, several recent works (He et al., 2020; Chen et al., 2020a; Caron et al., 2020) have achieved comparable or even better accuracy than the supervised pre-training when transferring to downstream tasks, e.g. semi-supervised classiﬁcation, object detection.
The aforementioned top-performing SSL algorithms all involve large networks ( e.g. , ResNet-50 (He et al., 2016) or larger), with, however, little attention on small networks.
Empirically, we ﬁnd that existing techniques like contrastive learning do not work well on small networks.
For instance, the linear probe top-1 accuracy on ImageNet using MoCoV2 (Chen et al., 2020c) is only 36.3% with MobileNet-V3-Large (see Figure 1), which is much lower compared with its supervised training accuracy 75.2% (Howard et al., 2019).
For EfﬁcientNet-B0, the accuracy is 42.2% compared with its supervised training accuracy 77.1% (Tan & Le, 2019).
We conjecture that this is because smaller models with fewer parameters cannot effectively learn instance level discriminative representation with large amount of data.
To address this challenge, we inject knowledge distillation (KD)
(Buciluˇa et al., 2006; Hinton et al., 2015) into self-supervised learning and propose self-supervised distillation (dubbed as S EED) as a new learning paradigm.
That is, train the larger, and distill to the smaller both in self-supervised manner.
Instead of directly conducting self-supervised training on a smaller model, S EED ﬁrst trains a large model (as the teacher) in a self-supervised way, and then distills the knowledge to the smaller model (as the student).
Note that the conventional distillation is for supervised learning, while the distillation here is in the self-supervised setting without any labeled data.
Supervised distillation can be formulated as training a student to mimic the probability mass function over classes predicted by a teacher model.
In unsupervised knowledge distillation setting, however, the distribution over classes is not directly attainable.
Therefore, we propose a simple yet effective self-supervised distillation method.
Similar to (He et al., 2020; Wu et al., 2018), we maintain a queue of data samples.
Given an instance, we ﬁrst use the teacher network to obtain its similarity scores with all the data samples in the queue as well as the instance itself.
Then the student encoder is trained to mimic the similarity score distribution inferred by the teacher over these data samples.
The simplicity and ﬂexibility that S EED brings are self-evident.
It does not require any clustering/prototypical computing procedure to retrieve the pseudo-labels or latent classes.
The teacher model can be pre-trained with any advanced SSL approach, e.g. , MoCo-V2 (Chen et al., 2020c),
SimCLR (Chen et al., 2020a), S WAV (Caron et al., 2020).
The knowledge can be distilled to any target small networks (either shallower, thinner, or totally different architectures).
To demonstrate the effectiveness, we comprehensively evaluate the learned representations on series of downstream tasks, e.g. , fully/semi-supervised classiﬁcation, object detection, and also assess the transferability to other domains.
For example, on ImageNet-1k dataset, S EED improves the linear probe accuracy of EfﬁcientNet-B0 from 42.2% to 67.6% (a gain over 25%), and MobileNet-V from 36.3% to 68.2% (a gain over 31%) compared to MoCo-V2 baselines, as shown in Figure 1 and Section 4.
Our contributions can be summarized as follows:
• We are the ﬁrst to address the problem of self-supervised visual representation learning for small models.
• We propose a self-supervised distillation ( S EED) technique to transfer knowledge from a large model to a small model without any labeled data.
• With the proposed distillation technique ( S EED), we signiﬁcantly improve the state-of-theart SSL performance on small models.
• We exhaustively compare a variety of distillation strategies to show the validity of S EED under multiple settings.
R ELATED W ORK
Among the recent literature in self-supervised learning, contrastive based approaches show prominent results on downstream tasks.
Majority of the techniques along this direction are stemming from noise-contrastive estimation (Gutmann & Hyvärinen, 2010) where the latent distribution is estimated by contrasting with randomly or artiﬁcially generated noises.
Oord et al.
ﬁrst proposed Info-NCE to learn image representations by predicting the future using an auto-regressive model for unsupervised learning.
Follow-up works include improving the efﬁciency (Hénaff et al., 2019), and using multi-view as positive samples (Tian et al., 2019b).
As these approaches can only have the access to limited negative instances, Wu et al.
(2018) designed a memory-bank to store the previously seen random representations as negative samples, and treat each of them as independent categories (instance discrimination).
However, this approach also comes with a deﬁciency that the previously stored vectors are inconsistent with the recently computed representations during the earlier stage of pre-training.
Chen et al.
(2020a) mitigate this issue by sampling negative samples from a large batch.
Concurrently, He et al.
(2020) improve the memory-bank based method and propose to use the momentum updated encoder for the remission of representation inconsistency.
Other techniques include Misra & Maaten (2020) that combines the pretext-invariant objective loss with contrastive learning, and Wang & Isola (2020) that decomposes contrastive loss into alignment and uniformity objectiveness.
Knowledge distillation (Hinton et al., 2015) aims to transfer knowledge from a cumbersome model to a smaller one without losing too much generalization power, which is also well investigated in model compression (Buciluˇa et al., 2006).
Instead of mimicking the teacher’s output logit, attention transfer (Zagoruyko & Komodakis, 2016) formulates knowledge distillation on attention maps.
Similarly, works in (Ahn et al., 2019; Yim et al., 2017; Koratana et al., 2019; Huang & Wang, 2017) have utilized different learning objectives including consistency on feature maps, consistency on probability mass function, and maximizing the mutual information.
CRD (Tian et al., 2019a), which is derived from CMC (Tian et al., 2019b), optimizes the student network by a similar objective to Oord et al.
(2018) using a derived lower bound on mutual information.
However, the aforementioned efforts all focus on task-speciﬁc distillation ( e.g. , image classiﬁcation) during the ﬁne-tuning phase rather than a task-agnostic distillation in the pre-training phase for the representation learning.
Several works on natural language pre-training proposed to leverage knowledge distillation for a smaller yet stronger small models. For instances, DistillBert (Sanh et al., 2019), TinyBert (Jiao et al., 2019), and MobileBert (Sun et al., 2020), have used knowledge distillation for model compression and shown their validity on multiple downstream tasks.
Similar works also emphasize the value of smaller and faster models for language representation learning by leveraging knowledge distillation (Turc et al., 2019; Sun et al., 2019).
These works all demonstrate the effectiveness of knowledge distillation for language representation learning in small models, while are not extended to the pre-training for visual representations.
S EED closely relates to the above techniques but aims to facilitate visual representation learning during pre-training phase for small models , which as far as we know has not yet been investigated.
M ETHOD 3. P RELIMINARY ON K NOWLEDGE D ISTILLATION
Knowledge distillation (Hinton et al., 2015; Bucilu ˇ a et al., 2006) is an effective technique to transfer knowledge from a strong teacher network to a target student network. The training task can be generalized as the following formulation: [Equation]
where x i is an image, y i is the corresponding annotation, θ S is the parameter set for the student network, and θ T is the set for the teacher network.
The loss L sup is the alignment error between the network prediction and the annotation. For example in image classiﬁcation task (Mishra &
Marr, 2017; Shen & Savvides, 2020; Polino et al., 2018; Cho & Hariharan, 2019), it is normally a cross entropy loss.
For object detection (Liu et al., 2019; Chen et al., 2017), it includes bounding box regression as well. The loss of L distill is the mimic error of the student network towards a pre-trained teacher network.
For example in (Hinton et al., 2015), the teacher signal comes from the softmax prediction of multiple large-scale networks and the loss is measured by the Kullback–Leibler divergence.
In (Romero et al., 2014), the task is to align the intermediate feature map values and to minimize the squared l 2 distance.
The effectiveness has been well demonstrated in the supervised setting with labeled data, but remains unknown for the unsupervised setting, which is our focus.
3. S ELF -S UPERVISED D ISTILLATION FOR V ISUAL R EPRESENTATION
Different from supervised distillation, S EED aims to transfer knowledge from a large model to a small model without requiring labeled data, so that the learned representations in small model can be used for downstream tasks.
Inspired by contrastive SSL , we formulate a simple approach for the distillation on the basis of instance similarity distribution over a contrastive instance queue.
Similar to He et al. , we maintain an instance queue for storing data samples’ encoding output from the teacher.
Given a new sample, we compute its similarity scores with all the samples in the queue using both the teacher and the student models.
We require that the similarity score distribution computed by the student matches with that computed by the teacher, which is formulated as minimizing the cross entropy between the student and the teacher’s similarity score distributions (as illustrated in Figure 2). [Figure]
𝑧
[Figure] 𝑧 𝑇 SoftMax Cross
[Figure]
Entropy
[Figure]
Teacher Probability
SoftMax
Pre
[Figure]
trained and Frozen
L
[Figure]
Norm
L2-Norm …
[Figure] Student [Figure] Speciﬁcally, for a randomly augmented view x i of an image, it is ﬁrst mapped and normalized into feature vector representations
z Ti = f Tθ ( x i ) / ||
f Tθ ( x i ) || 2 , and z Si =
f
Sθ ( x i ) /
||
f Sθ ( x i )
|| 2 , where z Ti , z Si ∈ R D , and f Tθ and f Sθ denote the teacher and student encoders, respectively.
Let D =
[ d 1 ... d K ] denote the instance queue where K is the queue length and d j is the feature vector obtained from the teacher encoder.
Similar to the contrastive learning framework, D is progressively updated under the “ ﬁrst-in ﬁrst-out ” strategy as distillation proceeds.
That is, we en-queue the visual features of the current batch inferred by the teacher and de-queue the earliest seen samples at the end of iteration.
Note that the maintained samples in queue D are mostly random and irrelevant to the target instance x i .
Minimizing the cross entropy between the similarity score distribution computed by the student and teacher based on D softly contrasts x i with randomly selected samples, without directly aligning with the teacher encoder.
To address this problem, we add the teacher’s embedding ( z Ti ) into the queue and form D + =
[ d 1 ... d K , d K +1 ] with d K +1 = z Ti . Let p T ( x i ; θ T ; D + ) denote the similarity score between x i and d j ’s ( j = 1 , ..., K + 1 ) computed by the teacher model.
p T ( x i ; θ T ; D + ) is deﬁned as
[Equation]
and τ T is a temperature parameter for the teacher.
Note, we use () T to represent the feature from the teacher network and use ( · ) to represent the inner product between two features.
Similarly let p S ( x i ; θ S , D + ) denote the similarity score computed by the student model, which is deﬁned as
[Equation]
and τ S is a temperature parameter for the student.
Our self-supervised distillation can be formulated as minimizing the cross entropy between the similarity scores of the teacher, p T ( x i ; θ T , D + ) , and the student, p S ( x i ; θ S , D + ) , over all the instances x i , that is,
[Equation]
Since the teacher network is pre-trained and frozen, the queued features are consistent during training w.r.t. the student network.
The higher the value of p T j is, the larger weight will be laid on p S j .
Due to the l 2 normalization, similarity score between z Ti and d K +1 remains constant 1 before softmax normalization, which is the largest among p Tj . Thus, the weight for p SK +1 is the largest and can be adjusted solely by tuning the value of τ T .
By minimizing the loss, the feature of z Si can be aligned with z Ti and meanwhile contrasts with other unrelated image features in D . We further discuss the relation of these two goals with our learning objective in Appendix 1.5.
Relations with Info-NCE loss. When τ T → 0 , the softmax function for p T smoothly approaches to a one-hot vector, where p TK +1 equals 1 and all others 0. In this extreme case, the loss becomes
[Equation]
which is similar to the widely-used Info-NCE loss (Oord et al., 2018) in contrastive-based SSL (see discussion in Appendix 1.6).
E XPERIMENT 4. P RE - TRAINING Self-Supervised Pre-training of Teacher Network.
By default, we use MoCo-V2 (Chen et al., 2020c) to pre-train the teacher network.
Following (Chen et al., 2020a), we use ResNet as the network backbone with different depths/widths and append a multi-layer-perceptron (MLP) layer (two linear layers and one ReLU (Nair & Hinton, 2010) activation layer in between) at the end of the encoder after average pooling.
The dimension of the last feature dimension is 128.
All teacher networks are pre-trained for 200 epochs due to the computational limitation unless explicitly speciﬁed. As our distillation is independent with the teacher pre-training algorithm, we also show results with other self-supervised pre-trained models for teacher network, e.g. , S WAV (Caron et al., 2020), SimCLR (Chen et al., 2020a).
Self-Supervised Distillation on Student Network.
We choose multiple smaller networks with fewer learnable parameters as the student network: MobileNet-v3-Large (Howard et al., 2017), EfﬁcientNet-B0 (Tan & Le, 2019), and smaller ResNet with fewer layers (ResNet-18, 34).
Similar to the pre-training for teacher network, we add one additional MLP layer on the basis of the student network.
Our distillation is trained with a standard SGD optimizer with momentum 0.9 and a weight decay parameter of 1e-4 for 200 epochs.
The initial learning rate is set as 0.03 and updated by a cosine decay scheduler (Nair & Hinton, 2010) with 5 warm-up epochs and batch size 256.
In Eq. 4, the teacher temperature is set as τ T = 0 .
01 and the student temperature is τ S = 0 .
The queue size of K is 65,536.
In the following subsections and appendix, we also show results with different hyper-parameter values, e.g. , for τ T and K .
4. F INE - TUNING AND EVALUATION
In order to validate the effectiveness of self-supervised distillation, we choose to assess the performance of representations of the student encoder on several downstream tasks.
We ﬁrst report its performances of linear evaluation and semi-supervised linear evaluation on the ImageNet {+ILSVRC
+} (Deng et al., 2009) dataset.
To measure the feature transferability brought by distillation, we also conduct evaluations on other tasks, which include object detection and segmentation on the VOC07 (Everingham et al.) and MS-COCO (Lin et al., 2014) datasets.
At the end, we compare the transferability of the features learned by distillation with ordinary self-supervised contrastive learning on the tasks of linear classiﬁcation on datasets from different domains.
Linear and K NN Evaluation on ImageNet.
We conduct the supervised linear classiﬁcation on ImageNet-1K, which contains ∼ 1.3M images for training, and 50,000 images for validation, spanning 1,000 categories.
Following previous works in (He et al., 2020; Chen et al., 2020a), we train a single linear layer classiﬁer on top of the frozen network encoder after self-supervised pretraining/distillation.
SGD optimizer is used to train the linear classiﬁer for 100 epochs with weight decay to be 0.
The initial learning rate is set as 30 and is then reduced by a factor of 10 at 60 andepochs (similar as in Tian et al.
(2019a)).
Notably, when training the linear classiﬁer for MobileNet and EfﬁcientNet, we reduce the initial learning rate to 3.
The results are reported in terms of Topand Top-5 accuracy.
We also perform classiﬁcation using K -Nearest Neighbors ( K NN) based on the learned 128d vector from the last MLP layer.
The sample is classiﬁed by taking the most frequent label of its K ( K = 10 ) nearest neighbors.
Table 1 shows the results with various teacher networks and student networks.
We list the baseline of contrastive self-supervised pre-training using MoCo-V2 (Chen et al., 2020c) in the ﬁrst row for each student architecture. We can see clearly that smaller networks perform rather worse.
For example,
[Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
+
[Figure] [Figure] [Figure]
+
[Figure]
[Figure] [Figure]
+8.
[Figure]
+5.
[Figure] [Figure]
MobileNet-V3 can only reach 36.3%.
This aligns well with previous conclusions from (Chen et al., 2020a;b) that bigger models are desired to perform better in contrastive-based self-supervised pretraining.
We conjecture that this is mainly caused by the inability of smaller network to discriminate instances in a large-scale dataset.
The results also clearly demonstrate that the distillation from a larger network helps boosting the performances of small networks, and show obvious improvement.
For instance, with MoCo-V2 pre-trained ResNet-152 (for 400 epochs) as the teacher network, the Topaccuracy of MobileNet-V3-Large can be signiﬁcantly improved from 36.3% to 61.4%. Furthermore, we use ResNet-50 × 2 (provided in Caron et al.
as the teacher network and adopt the multi-crop trick (see Appendix 1.3 for details).
The accuracy can be further improved to 68.2% (last row of
Table 1) for MobileNet-V3-Large with 800 epochs of distillation.
We note that the gain beneﬁted from distillation becomes more distinct on smaller architectures and we further study the effect of various teacher models in ablations.
Semi-Supervised Evaluation on
ImageNet
Following (Oord et al., 2018; Kornblith et al., 2019;
Kolesnikov et al., 2019), we evaluate the representation on the semi-supervised task, where a ﬁxed 1% or 10% subsets of ImageNet training data (Chen et al., 2020a) are provided with the annotations.
After the self-supervised learning with and without distillation, we also train a classiﬁer on top of the representation.
The results are shown in Figure 3, where the baseline without distillation is depicted when teacher parameters are 0.
As we can see, the accuracy is also improved remarkably with S EED distillation, and a stronger teacher network with more parameters leads to a better performed student network.
Transferring to Classiﬁcation.
To further study whether the improvement of the learned representations by distillation is conﬁned to ImageNet, we evaluate on additional classiﬁcation datasets to study the generalization and transferability of the feature representation.
We strictly follow the linear evaluation and ﬁne-tuning settings from (Kornblith et al., 2019; Chen et al., 2020a; Grill et al., 2020), that a linear layer is trained on the basis of frozen features.
We report Top-1 Accuracy of models before and after distillation from various architectures on CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009), SUN-397 (Xiao et al., 2010) datasets (see Figure 4).
More details regarding pre-processing and training can be found in Appendix 1.2.
Notably, we observe that our distillation surpasses contrastive self-supervised pre-training consistently on all benchmarks, verifying the effectiveness of S EED.
T o pA cc u r a c y ( % )
Number of Parameters of the Teacher (Millions)
CIFAR55 .57 .
CIFAR
SUNResNet-
EﬃcientNet
[Figure]
Obj
[Figure]
COCO Obj.
Det.
COCO Inst.
Segm.
AP bb
AP bb
AP bb
AP bb
AP bb
AP bb
AP mk
AP mk
AP mk
R- (cid:55) 46.
[Table]
This also proves the generalization ability of the learned representations from distillation to a wide range of data domain and different classes.
Transferring to Detection and Segmentation.
We conduct two downstream tasks here.
The ﬁrst is
Faster R-CNN (Ren et al., 2015) model for object detection trained on VOC-07+12 train+val set and evaluated on VOC-07 test split.
The second is Mask R-CNN (
He et al., 2017) model for the object detection and instance segmentation on COCO 2017 dataset (Lin et al., 2014).
The pre-trained model serves as the initial weight and following He et al.
, we ﬁne-tune all the layers of the model.
More experiment settings can be found in Appendix 1.2.
The results are illustrated in Table 2.
As we can see, on VOC, the distilled pre-trained model achieves a large improvement.
With ResNet-152 as the teacher network, the Resnet18-based Faster R-CNN model shows +0.7 point improvement on AP, +1.4 improvement on AP 50 and +1.6 on AP 75 .
On COCO, the improvement is relatively minor and the reason could be that COCO training set has ∼ 118k training images while VOC has only ∼ 16.5k training images.
A larger training set with more ﬁne-tuning iterations reduces the importance of the initial weights.
4. A BLATION S TUDY
We now explore the effects of distillation using different Teacher architectures, Teacher Pre-training algorithms, various distillation strategies and hyper-parameters.
Different Teacher Networks.
Figure 5 summarizes the accuracy of ResNet-18 and EfﬁcientNet-B distilled from wider and deeper ResNet architectures.
We see clear performance improvement as depth and width of teacher network increase: compared to ResNet-50, deeper (ResNet-101) and wider (ResNet-50 × 2) substantially improve the accuracy.
However, further architectural enlargement has relatively limited effects, and we suspect the accuracy might be limited by the student network capacity in this case. Different Teacher Pre-training Algorithms. In Table 3, we show the Top-1 accuracy of ResNetdistilled from ResNet-50 with different pre-training algorithms, i.e. , MoCo-V1 (He et al., 2020), MoCo-V2 (Chen et al., 2020c), SimCLR (Chen et al., 2020a), and S WAV (Caron et al., 2020)).
Notably, the aforementioned methods all unanimously adopt contrastive-based pre-training except S WAV, which is based upon online clustering.
We ﬁnd that our S EED is agnostic to pre-training approaches, making it easy to use any self-supervised models (including clustering-based approach like S WAV) in self-supervised distillation. In addition, we observe that more training epochs for both teacher SSL and distillation epochs can bring beneﬁcial gain. [Figure]
[Figure]
[Table] [Table] Other Distillation Strategies.
We explore several alternative distillation strategies. l 2 -Distance : where the l2 -distance of teacher student’s embeddings are minimized, motivated by (Romero et al., 2014).
K -Means : we exploit K -Means clustering to assign a pseudo-label based on the teacher network’s representation.
Online Clustering : we continuously update the clustering centers during distillation for pseudo-label generation. Binary Contrastive Loss : we adopt an Info-NCE alike loss for contrastive distillation (Tian et al., 2019a). We provide details for other strategies in Appendix 1.4.
Table 4 shows the results for each method on ResNet-18 (student) distilled from ResNet50.
From the results, the simple l2 -distance minimizing approach can achieve a decent accuracy, which demonstrates the effectiveness of applying the distillation idea to the self-supervised learning. Meanwhile, our proposed SEED outperforms these alternatives with highest accuracy, which shows the superiority of aligning the student towards the teacher and contrasting with the irrelevant samples.
Other Hyper-Parameters. Table 5 summarizes the distillation performances on multiple datasets using different temperature τ T .
We observe a better performance when decreasing τ T to 0.01 for ImageNet-1k and CIFAR-10 dataset, and to 1e-3 for CIFAR-100 datasets.
When τ is large, the softmax-normalized similarity score of p Tj between z Ti and instance d j in the queue D + also becomes large, which means the student’s feature should be less discrminative with the features of other images to some extent.
When τ T is 0, the teacher model will generate a one-hot vector, which only treats z Ti as a positive instance and all others in the queue as negative.
Thus, the best τ is a trade-off depending on the data distribution.
We further compare effect of different hyper-parameters in Appendix 1.8.
C ONCLUSIONS
Self-Supervised Learning is acknowledged for its remarkable ability in learning from unlabeled, and large scale data. However, a critical impedance for the SSL pre-training on smaller architecture comes from its low capacity of discriminating enormous number of instances.
Instead of directly learning from unlabeled data, we proposed S EED as a novel self-supervised learning paradigm, which learns representation by self-supervised distillation from a bigger SSL pre-trained model.
We show in extensive experiments that S EED effectively addresses the weakness of self-supervised learning for small models and achieves state-of-the-art results on various benchmarks of small architectures.