R EPRESENTATION D ISTILLATION BY P ROTOTYPICAL C ONTRASTIVE P REDICTIVE C ODING
Anonymous authors
Paper under double-blind review
A BSTRACT
Transferring representational knowledge of a model to another is a wide-ranging topic in machine learning.
Those applications include the distillation of a large supervised or self-supervised teacher model to a smaller student model or selfsupervised learning via self-distillation.
Knowledge distillation is an original method to solve these problems, which minimizes a cross-entropy loss between the prototypical probabilistic outputs of teacher and student networks.
On the other hand, contrastive learning has shown its competency in transferring representations as they allow students to capture the information of teacher representations.
In this paper, we amalgamate the advantages of knowledge distillation and contrastive learning by modeling the critic of a contrastive objective by the prototypical probabilistic discrepancy between two features.
We refer to it as prototypical contrastive predictive coding and present efﬁcient implementation using the proposed objective for three distillation tasks: supervised model compression, self-supervised model compression, and self-supervised learning via self-distillation.
Through extensive experiments, we validate the effectiveness of our method compared to various supervised and self-supervised knowledge distillation baselines.
I NTRODUCTION
In machine learning, knowledge distillation (KD) is a problem that aims to transfer the knowledge from one network (a teacher) to another one (a student).
The original method introduced by Hinton et al. (2015); Bucilu ˇ a et al.
(2006) minimizes the cross-entropy between the probabilistic outputs of teacher and student networks.
Even though the simplicity of its implementation, many other distillation methods cannot easily outperform KD.
Therefore the concept of knowledge distillation has been expanded to various machine learning tasks other than supervised model compressions, such as self-supervised learning (Caron et al., 2020; 2021) or self-supervised model compression (Fang et al., 2021).
Those methods rely on minimizing the discrepancy of probabilistic outputs of teacher and student networks, thus require a function that maps a data or feature into a probability space.
It is straightforward to use a linear mapping for it, which we refer to a prototypes .
On the other hand, contrastive learning has its merits in capturing correlations and high-order dependencies over teacher representations as they are trained by pulling the positives and pushing over the negatives.
Especially, contrastive objectives such as NCE (Gutmann & Hyvärinen, 2010) or
CPC (as known as InfoNCE)
(Oord et al., 2018; Bachman et al., 2019; Hjelm et al., 2018) are proven to be a lower bound to the mutual information allowing student network to capture the information from the teacher representations.
While many methods (Chen et al., 2020b;c) used those objectives for representation learning, Tian et al.
(2019) proposed contrastive representation distillation (CRD) and demonstrated the effectiveness of contrastive learning in transferring knowledge from one network to another.
However, the current contrastive learning requires large negative samples with careful sampling.
To ameliorate, we combine the prototypical method and contrastive objective to inherit their advantages.
To that end, we propose prototypical contrastive predictive coding (ProtoCPC), which utilizes the prototypes to generate probability distribution and model the critic by the discrepancy between two probabilistic outputs.
Furthermore, our ProtoCPC objective is a lower bound to the mutual information between teacher and student representations alike CPC.
But unlike CPC, it does not require commodious negative samples or careful sampling policies.
Given the ProtoCPC objective, we present efﬁcient applications on three distillation tasks: supervised model compression, self-supervised model compression, and self-supervised learning via self-distillation.
While ProtoCPC is simple in its implementation, we observe that it signiﬁcantly boosts the performance of representation distillation.
Our method outperforms KD, and partially outperform CRD on CIFAR-100 and ImageNet.
Moreover, our method achieves state-of-the-art performance on transferring various self-supervised teacher representations to a small self-supervised model.
Lastly, by construing self-supervised learning as a self-distillation, we validate the effectiveness of ProtoCPC in representation learning.
Our contributions are following:
• We propose a novel prototypical contrastive objective for transferring representational knowledge between models.
• We apply the proposed objective to three distillation tasks: supervised model compression, self-supervised model compression, and self-supervised learning via self-distillation.
• Experiments show the effectiveness of our method in various representation distillation benchmarks, especially achieving state-of-the-art performance in supervised/self-supervised model compression.
M ETHOD 2. P ROTOTYPICAL C ONTRASTIVE P REDICTIVE C ODING
Given data x with a random variable x ∼ X , let the teacher network f T and the student network f S , where they map x into R D T and R D S respectively.
Also let T and S be random variables for representation f T ( X ) and f S ( X ) respectively.
One can transfer the representational knowledge of
T to S by maximizing the mutual information
I ( T ; S ) , where it is deﬁned by the KL-divergence between the joint distribution p ( T, S ) and the product of marginal distribution p ( T ) p ( S ) .
However, as estimation and optimization of mutual information is challenging, many approaches count on maximizing variational lower bound to the mutual information (Poole et al., 2019).
The contrastive predictive coding (Oord et al., 2018), or as known as InfoNCE is a guaranteed lower bound to the mutual information (Oord et al., 2018; Tian et al., 2020; Bachman et al., 2019; Song & Ermon, 2020), and has shown its competency in both representation learning (He et al., 2020;
Chen et al., 2020b;c) and representation distillation (Tian et al., 2019).
Formally, given a z s ∼ S with a positive z t ≡ z t 0 and N − 1 negatives { z tj } N − 1 j =1 sampled from T , i.e. ( z t , z s ) ∼ p
( T, S ) and { ( z tj , z s ) }
N − 1 j =1 ∼ p ( T ) p ( S ) , the following inequality holds for any critic h : R D T ×
R D S
→ R + :
[Equation]
Previous works set the critic by the exponential of cosine similarity between two unit feature vectors, i.e. h ( z t , z s ) =
exp( − z t · z s ) where z t and z s are (cid:96) 2 -normalized.
Theoretically, the lower bound becomes tighter as N → ∞ .
In practice, the CPC objective requires using extremely large batch size or memory buffer that stores the negatives as it requires pairwise computation between z s and z tj s.
On the other hand, we project feature vectors into a probability space.
To do that, we append a linear prototypes W T ∈ R D T × K and W S ∈ R D S × K at the top of f T and f S so that they have same output dimension of K .
For brevity, let ¯
z s =
W S z s and ¯ z t =
W T z t and ¯ T, ¯ S be random variables for ¯ z t and ¯ z s respectively.
Then we set probability of student p s by K -categorical distribution deﬁned by the softmax operator on ¯
z s = W S z s with temperature τ s > 0 :
[Equation]
Similarly we deﬁne probability of teacher p tj with temperature τ t > 0 .
Then we deﬁne the critic between ¯ z t and ¯ z s by the negative exponential of cross-entropy between p s and p t , i
[Equation]
e.,
[Equation]
Then h is a positive bounded function and is maximized when p t matches with p s .
Then by plugging
Eq. 3 into Eq.
1, it follows that
[Equation]
where q ( k ) =
1 N (cid:80) N − 1 j =
p tj is a mean of teachers’ probability which we call a prior .
The ﬁrst inequality is from data processing inequality, the second equality is from crossing out the constant term, and third inequality is from Jensen’s inequality.
We deﬁne Prototypical Contrastive Predictive
Coding (ProtoCPC) objective I ProtoCPC in Eq.
In addition, we deﬁne ProtoCPC loss L ProtoCPC by the negative of ProtoCPC objective, thus minimizing ProtoCPC loss is equivalent to a variational maximization of mutual information between student and teacher representations.
Relationship with CPC While many lower bounds to the mutual information were proposed,
Tschannen et al.
(2019) observe that the tightness of bound does not necessarily imply a better representation learning performance.
From then, many works focused on analyzing the components of contrastive objective itself which are responsible for the empirical success.
Wang & Isola (2020) argued that the contrastive loss is composed of alignment and uniformity loss, where alignment loss accounts for the similarity of two positive features, and uniformity loss measures how the features are scattered in the unit hypersphere and show that both losses are important in contrastive learning.
We draw an analogy on ProtoCPC by dissecting into alignment and uniformity losses.
Since the alignment loss is straightforward, we focus on the uniformity loss.
We show that one can interpret the uniformity loss L ProtoCPC-Unif by the re-substitution entropy estimator of z s via a von-Mises Fisher kernel density estimation (vMF-KDE) (Ahmad & Lin, 1976):
[Equation]
where each w k is a k -th column of W S and acts as a mean direction of k -th vMF distribution and q ( k ) acts as a prior for each k -th vMF distribution.
The ˆ p vMF-KDE is thus the mixture of K vMF distribution with prior q ( k ) and then the uniformity loss is a re-substitution entropy ˆ H ( z s ) .
The Z vMF is a normalizing constant for vMF distribution.
Remark that the uniformity loss of CPC objective is also a re-substitution entropy with vMF-KDE, but the mean directions are given by negative samples z tj and the prior is uniform.
It shows that the ProtoCPC objective allows modeling of complex mixture of vMF distribution by exploiting prior term and using prototypes remove the dependency on negative samples.
Prior momentum Since ProtoCPC is contrastive, it requires sufﬁcient negatives to perform learning.
However, unlike CPC, ProtoCPC only requires prior q that accounts for the negatives.
Therefore, we use exponential moving average (EMA) on prior q for better estimation.
At each iteration, we update q by following update rule:
[Equation]
where m p > 0 is a momentum rate.
The prior momentum allows better estimation of prior term regardless of the size of negative samples.
Assignment of teacher probability While KD used softmax operator for both probabilities of teacher and student networks, many self-supervised methods (Asano et al., 2019; Caron et al., 2020) reported that the softmax operator can lead to collapse, i.e. every representation fall into the same one.
To compromise, many prototypical self-supervised methods resort on Sinkhorn-Knopp iterative algorithm by formulating the assignment of teacher probability as an optimal transport problem:
[Equation]
where Z t is a matrix whose rows are z tj , H ( P t ) =
(cid:80) N − 1 j
[Equation]
0 (cid:80)
Kk
[Equation]
1 − p
[Equation]
k
[Equation]
t log p
[Equation]
t is a entropy and τ t > 0 is a temperature that controls the smoothness of distribution.
Then the Eq.
11 can be solved by only few steps of Sinkhorn-Knopp iteration (Cuturi, 2013; Altschuler et al., 2017) which iteratively projects P t into following form:
[Equation]
(12) where β k is a normalizing constant.
We refer this to Sinkhorn-Knopp (SK) operator.
The SK operator allows each prototype to be non-zero, i.e. q ( k ) is non-zero for each k .
In practice, we only conduct a few steps of SK iterations.
2. A PPLICATIONS OF P ROTO CPC TO DISTILLATION TASKS
Fig.
1 depicts explanation on how we structure prototypical contrastive learning for three distillation tasks we consider: supervised model compression, self-supervised model compression and selfdistillation.
Supervised model compression Given a (supervised) pre-trained teacher network f T , the goal of supervised model compression is to train a smaller student network f S by utilizing the representational knowledge from the teacher.
To do that, Hinton et al.
(2015) proposed to minimize the knowledge distillation loss, where it is a weighted sum of the supervised loss and the distillation loss.
Here, the supervised loss is a conventional cross-entropy loss between the one-hot label and the output of f S .
The distillation loss is a cross-entropy between teacher’s probability output p t and student’s probability output p s .
Formally, it is given by following:
[Equation]
(13) where y is a label, p t ( x ) and p s ( x ) are computed by the softmax operator on each f T ( x ) and f S ( x ) with temperature τ > 0 , and β is a balancing weight.
Our approach is to replace the distillation loss by the ProtoCPC loss and use SK operator for p t .
Then our loss becomes following:
[Equation]
Remark that since the supervised networks already contain prototypes at the last layer, we use prototypes for both supervised learning and distillation at the same time.
Note that CRD (Tian et al., 2019) uses binary contrastive loss for distillation loss by maintaining a large memory buffer of negative samples.
Also, CRD requires an additional linear layer to implement inner-product critic.
However, our method conducts contrastive learning without additional training of embeddings or storing a memory buffer of negative samples.
Self-supervised model compression Given a large self-supervised pre-trained teacher network g T , the goal of self-supervised model compression is to transfer the representational knowledge of g T to a smaller self-supervised student network.
Fang et al.
ﬁrst proposed SEED, where the student network is trained by minimizing the cross-entropy loss between the similarity score of teacher and student.
The similarity score is computed by the pairwise computation over a memory buffer ﬁlled with features of the teacher network.
Our approach uses prototypes to generate the probabilistic output of teacher and student representations, and use ProtoCPC loss for effective distillation.
Let the pre-trained teacher network g T be a composition of base encoder f T and the projection head h T .
Then we train student network g S , where it is composed of smaller encoder f S and projection head h S of the same architecture as h T .
Then we append prototypes W T and W S for each g T and g S to ensure that they have the same output of dimension K .
Given a data x , we train the student network by minimizing the ProtoCPC loss between the probability of teacher p t ( x ) and probability of student p s ( x ) .
For p s we use softmax operator with temperature τ s > 0 and for p t we use SK operator with temperature τ t > 0 over the batch of samples.
Then the objective is given by following:
[Equation]
We can further use multi-crops data augmentation (Caron et al., 2020), where it augments multiple small crops of a data during the training to expedite the training.
Let ˜ x
j ∼ M ( x ) be small local crops with transformation M , then the self-supervised model compression with multi-crops is given by following:
[Equation]
For prototypes of teacher network W T , we copy the parameters of W S to W T at each iteration.
This allows our method to apply to any self-supervised teacher networks.
Note that if the teacher network is trained by prototypical methods such as DINO (Caron et al., 2021) or SwAV (Caron et al., 2020), we can re-use the pre-trained prototypes for W T .
We present ablation studies on setting prototypes for W T .
Self-distillation Additionally, we show that ProtoCPC loss can enhance the representation learning of a self-supervised model, where we interpret it as a self-distillation method.
We built on
DINO (Caron et al., 2021), which is the state-of-the-art method in self-supervised learning that ﬁrst regarded the self-supervised learning as a knowledge distillation problem.
Let the student network be the composition of a base encoder f θ S , we append projection head h θ S consists of MLP layers and a prototypical layer W θ S .
Then we set the teacher by the momentum encoder (He et al., 2020;
Chen et al., 2020b; Grill et al., 2020; Caron et al., 2021) on the student network, where it is a composition of encoder
f θ T , projection head h
θ T , and prototype W θ T .
We use
EMA for the weights of teacher network, where it is update by θ t → λθ t + (1 − λ ) θ s , where λ is a momentum rate.
Within the training, we generate a pair of views ( x 1 , x 2 ) from a data x ∼ X and they are passed to each student and teacher networks.
Then we compute ProtoCPC loss between the probability of teacher network p t ( x 1 ) and probability of student network p s ( x 2 ) , and symmetrically between p t ( x 2 ) and p s ( x 1 ) .
The objectives are given by:
[Equation]
(17) where A ( · ) is a data augmentation operator.
Note that multi-crops data augmentation can be further applied.
R ELATED W ORK
Contrastive learning There is a large body of works on contrastive objectives for representation learning.
The contrastive objectives such as CPC (Oord et al., 2018; Song & Ermon, 2020; Bachman et al., 2019; Hjelm et al., 2018) and multi-view data augmentation (Tian et al., 2020) are keys for the success of contrastive representation learning.
With sophisticated implementation such as momentum encoder (He et al., 2020;
Chen et al., 2020b; Grill et al., 2020), or caching negative samples (Wu et al., 2018;
Chen et al., 2020b;c) had dramatically reduced the gap between selfsupervised and supervised models.
While the contrastive objectives are proven to be a lower bound to the mutual information (Poole et al., 2019; Song & Ermon, 2020), some studies attribute that the success of representation learning to the contrastive objective itself rather than the tightness of mutual information estimation (Bachman et al., 2019).
Then many works focus on analyzing the contrastive objective itself (Wang & Isola, 2020) or count on the metric-learning perspective of contrastive learning that deviates into different objectives (Grill et al., 2020; Chen & He, 2021; Zbontar et al., 2021;
Bardes et al., 2021).
The key difference between many contrastive objectives and our method is that their discrepancy measure is given by the (cid:96) 2 -norm of two features, while our method projects feature into a probability space and characterize the discrepancy with probabilistic measure.
Our work is also related to methods that use clustering for representation learning (Tian et al., 2017;
Zhuang et al., 2019;
Li et al., 2020), which have been developed into prototypical methods (Asano et al., 2019; Li et al., 2020; Caron et al., 2020; 2021).
Knowledge distillation
The idea that transferring the large pre-trained model’s representation into a smaller one has been embodied by knowledge distillation (Hinton et al., 2015; Bucilu ˇ a et al., 2006).
Then various distillation criteria or utilization of intermediate feature maps were studied to enhance the performance (Romero et al., 2014; Zagoruyko & Komodakis, 2016a; Tung & Mori, 2019; Peng et al., 2019;
Ahn et al., 2019; Park et al., 2019; Passalis & Tefas, 2018; Heo et al., 2019; Kim et al., 2018;
Yim et al., 2017; Huang & Wang, 2017).
However, many of them were not able to outperform
KD by far.
Recently, Tian et al.
(2019) introduced contrastive representation distillation (CRD) which utilizes binary contrastive objective for distillation loss and shows the empirical superiority on various knowledge distillation benchmarks.
Moreover, recent studies focus on distilling the knowledge of large self-supervised models to the smaller ones using the concept of knowledge distillation (Fang et al., 2021;
Shen et al., 2021;
Noroozi et al., 2018; Chen et al., 2020a).
Fang et al.
(2021) proposed self-supervised representation distillation (SEED) which used features of teacher networks to compute probability distribution for distillation.
One can interpret SEED as the prototypical method where prototypes are given by the queue of teacher features.
E XPERIMENT 4. S UPERVISED M ODEL C OMPRESSION
Setup We experiment on CIFAR-100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) with various teacher-student combinations such as ResNet (He et al., 2016) and Wide
ResNet (WRN) (Zagoruyko & Komodakis, 2016b).
We compare with various distillation baselines such as KD (Hinton et al., 2015) and CRD (Tian et al., 2019).
Results on CIFAR-100 Table 1 and Table 2 compare top-1 accuracies of our method on supervised model compression on CIFAR-100.
Table 1 investigates teachers and students of the same architectural style, while Table 2 is focused on students and teachers from different architectures.
We observe that our ProtoCPC consistently outperforms KD and other distillation methods, yet ProtoCPC is on par with CRD where it also uses a contrastive learning method for distillation.
Remark that
CRD requires a large memory buffer of negative samples with careful sampling according to the class information and additional training of embeddings.
However, our ProtoCPC does not require massive negative samples and additional training of linear embeddings since we use the last linear layer as prototypes for our ProtoCPC.
[Table]
[Table]
Results on ImageNet In Table
[Table]
we compare our ProtoCPC to different distillation methods such as AT (Zagoruyko & Komodakis, 2016a) and Online-KD (Lan et al., 2018) on ImageNet, where pre-trained ResNet-34 is used for teacher and ResNet-18 is a student.
We observe that ProtoCPC outperforms various knowledge distillation baselines such as KD, and is slightly better than CRD
[Table]
Remark that our method is simpler in its implementation and shows better performance than CRD.
[Table]
4. S ELF - SUPERVISED MODEL COMPRESSION
Setup We experiment distillation of various self-supervised networks to ResNet-18 (He et al., 2016) on ImageNet (Deng et al., 2009) without class labels.
We consider following self-supervised teacher networks: MoCo-v2 (Chen et al., 2020b) pre-trained ResNet-50, SwAV (Caron et al., 2020) pre-trained ResNet-50 and DINO (Caron et al., 2021) pre-trained ResNet-50 and vision transformer (Dosovitskiy et al., 2020).
We train for 100 epochs and we additionally conduct experiments on using multi-crops data augmentation for SwAV and DINO pre-trained ResNet-50 networks.
For evaluation, we follow linear evaluation protocol which conducts supervised learning on the linear layer appended at the top of the frozen feature and k -nearest neighbor classiﬁcation ( k -NN).
Main results Table 4 show the main results of our self-supervised model compression compared to self-supervised learning (SSL) of itself.
We observe that ProtoCPC outperforms SSL with a large margin, especially showing superior performance in k -NN classiﬁcation.
Note that our method works well for various self-supervised teacher networks and even works well when the teacher and student networks are of different architectures (vision transformer teacher to ResNet student).
[Table]
[Table]
Comparison with SEED Table 5 compares our method with SEED (Fang et al., 2021), an original method for self-supervised model compression.
We observe that our method consistently outperforms SEED in both k -NN and linear evaluation with the same teacher network applied.
Although training for shorter epochs, our method achieves outperforms SEED.
When the teacher is pre-trained by SwAV, our method outperforms SEED without using multi-crops data augmentation, and the gap becomes larger when we use multi-crops as well.
When using
DINO self-supervised teacher and multi-crops data augmentation, we achieve the best results in representation learning of ResNet-18.
4. S ELF - DISTILLATION
Setup We evaluate our ProtoCPC objective in self-supervised learning, where we cast it as a selfdistillation.
We follow DINO (Caron et al., 2021) for experimental procedures, where the teacher network is set by the momentum encoder, and linear prototypes are attached to the features to distill the knowledge.
[Table] [Table]
We change the loss function from cross-entropy loss to ProtoCPC loss without any adjustment in other settings.
We use a prior momentum rate of 0.9 for ProtoCPC loss.
We evaluate on both ResNet-50 and DeiT-S/16, where DINO performs well on both architectures.
For evaluation, we report the results of linear evaluation protocol and k -nn classiﬁcation.
Results Table 6 compare the performance of representation learning of our method and
DINO.
We observe that ProtoCPC enhances the performance of representation learning without the adjustment of hyper-parameters.
Also, in
Table 7, we demonstrate the effect of prior momentum rate m p and observe that non-zero m p results in better representation learning.
4. A BLATION STUDY Cross-entropy v.s. ProtoCPC and Softmax
v.s. Sinkhorn-Knopp
The original KD is implemented by cross-entropy loss between the softmax output of the teacher and student network.
Our method uses ProtoCPC loss and SK operator for teacher network.
We provide ablation studies on the choice of the loss function and probability assignment operator on the subset of supervised model compression tasks.
In Table 8, we observe that using the SK operator boosts the performance of KD, and using ProtoCPC loss further enhances the performance.
Teacher’s prototypes for self-supervised model compression For prototypical self-supervised teachers, one can use the pre-trained prototypes for distillation.
We present an ablation between using pre-trained prototypes and copying student’s prototypes for teachers.
Table 9 show the results on distillation from SwAV ResNet-50 and DINO ResNet-50 teachers.
We observe that when distilling from SwAV teacher, using pre-trained prototypes performs better, while distilling from DINO teacher, copying from student’s prototypes performs better.
We suspect that as SwAV used single prototypes throughout pre-training, the pre-trained prototypes contain representational knowledge.
On the other hand, the pre-trained prototypes of DINO lack such representational knowledge
[Table] [Table]
pre-trained prototypes
C ONCLUSION
In this paper, we propose prototypical contrastive predictive coding, a simple yet effective method for the distillation of a network by combining the prototypical method and contrastive learning.
Our experiments show the effectiveness of our objective on various applications such as supervised/ self-supervised model compression and self-supervised learning by self-distillation.
A D ETAILS IN EXPERIMENTS A. P SEUDOCODE FOR P ROTO CPC LOSS
The PyTorch style pseudo-code for our ProtoCPC is demonstrated in Algorithm 1.
Algorithm 1 ProtoCPC loss PyTorch-style pseudocode.
[Table]
A. S UPERVISED MODEL COMPRESSION
Many of followings are borrowed from CRD (Tian et al., 2019)
. A.2. O THER M ETHODS
We compare to the following other state-of-the-art methods from the literature:
Knowledge Distillation (KD) (Hinton et al., 2015) 2.
Fitnets:
Hints for thin deep nets (Romero et al., 2014) 3.
Attention Transfer (AT) (Zagoruyko & Komodakis, 2016a) 4.
Similarity-Preserving Knowledge Distillation (SP) (Tung & Mori, 2019); 5.
Correlation Congruence (CC) (Peng et al., 2019) 6.
Variational information distillation for knowledge transfer (VID) (Ahn et al., 2019) 7.
Relational Knowledge Distillation (RKD) (Park et al., 2019) 8.
Learning deep representations with probabilistic knowledge transfer (PKT) (Passalis & Tefas, 2018) 9.
Knowledge transfer via distillation of activation boundaries formed by hidden neurons (AB) (Heo et al., 2019) 10.
Paraphrasing complex network: Network compression via factor transfer (FT) (Kim et al., 2018) 11.
A gift from knowledge distillation: Fast optimization, network minimization and transfer learning (FSP) (Yim et al., 2017) 12.
Like what you like: Knowledge distill via neuron selectivity transfer (NST) (Huang & Wang, 2017) 13.
Contrastive representation distillation (CRD) (Tian et al., 2019)
A.2. N ETWORK A RCHITECTURES
• Wide Residual Network (WRN) (Zagoruyko & Komodakis, 2016b).
WRN-d-w represnets wide resnet with depth d and width factor w .
• resnet (He et al., 2016).
We use resnet-d to represent cifar -style resnet with 3 groups of basic blocks, each with 16 , 32 , and 64 channels respectively.
In our experiments, resnet8 x and resnet32 x4 indicate a 4 times wider network (namely, with 64 , 128 , and 256 channels for each of the block)
• ResNet (He et al., 2016).
ResNet-d represents ImageNet -style ResNet with Bottleneck blocks and more channels.
• MobileNetV2 Sandler et al.
In our experiments, we use a width multiplier of 0 .
• vgg (Simonyan & Zisserman, 2014).
the vgg net used in our experiments are adapted from its original ImageNet counterpart.
• ShufﬂeNetV1 (Zhang et al., 2018), ShufﬂeNetV2 (Ma et al., 2018).
ShufﬂeNets are proposed for efﬁcient training and we adapt them to input of size 32x32.
A.2. I MPLEMENTATION D ETAILS
All methods evaluated in our experiments use SGD.
• For CIFAR-100, we initialize the learning rate as 0.05, and decay it by 0.1 every 30 epochs after the ﬁrst 150 epochs until the last 240 epoch.
For MobileNetV2, ShufﬂeNetV1 and ShufﬂeNetV2, we use a learning rate of 0.01 as this learning rate is optimal for these models in a grid search, while 0.05 is optimal for other models.
• For ImageNet, we follow the standard PyTorch practice but train for 10 more epochs.
Batch size is 64 for CIFAR-100 or 256 for ImageNet.
The student is trained by a combination of cross-entropy classiﬁcation objective and a knowledge distillation objective, shown as follows:
[Equation]
For the weight balance factor β , we directly use the optimal value from the original paper if it is speciﬁed, or do a grid search with teacher WRN-40-2 and student WRN-16-2.
This results in the following list of β used for different objectives:
KD (Hinton et al., 2015): α = 0 .
1 , β = 0 .
9 and T =2.
Fitnets (Romero et al., 2014): α = 1 , β =3.
AT (Zagoruyko & Komodakis, 2016a): α = 1 , β =4.
SP (Tung & Mori, 2019): α = 1 , β =5.
CC (Peng et al., 2019): α = 1 , β = 0 .6.
VID (Ahn et al., 2019): α = 1 , β =7.
RKD (Park et al., 2019): α = 1 , β 1 = 25 for distance and β 2 = 50 for angle.
For this loss, we combine both term following the original paper.
PKT (Passalis & Tefas, 2018): α = 1 , β =9.
AB (Heo et al., 2019): β = 0 , distillation happens in a separate pre-training stage where only distillation objective applies.
FT (Kim et al., 2018): α = 1 , β =11.
FSP (Yim et al., 2017): β = 0 , distillation happens in a separate pre-training stage where only distillation objective applies.
NST (Huang & Wang, 2017): α = 1 , β =13.
CRD (Tian et al., 2019): α = 1 , β = 0 .
Similar to KD, our ProtoCPC used α = 1 , β = 1 . 75 and temperature T = 4 .
For ablation study in
Table 8, we used α = 0 .
1 , β = 0 .
9 for CE losses and α = 1 , β = 1 .
75 for ProtoCPC losses.
A. S ELF - SUPERVISED MODEL COMPRESSION A.3. P RE - TRAINED TEACHERS
• MoCo-v2 (Chen et al., 2020b) uses InfoNCE loss and momentum encoder.
Since InfoNCE requires large negatives, they pertain a large queue which updates by ﬁrst-in-ﬁrst-out rule with features of momentum encoder.
We used the MoCo-v2 ResNet-50 (He et al., 2016) trained for 800 epochs.
• SwAV (Caron et al., 2020) is a prototypical method that generates probability by adopting prototypes.
They generate probability of a given feature by computing similarity with respect to prototypes and minimize the cross-entropy between probability outputs of two different views of an image.
They used Sinkhorn-Knopp iteration for probability output.
Also, they ﬁrst proposed multi-crops strategy, which additionally use small crops of an image to expedite the training.
We used the SwAV ResNet-50 trained for 800 epochs with multi-crops applied.
• DINO (Caron et al., 2021) is a prototypical method which uses momentum encoder.
The training progress is similar to SwAV except that they use momentum encoder on the prototypes and use online centering before the softmax operator.
They showed that their method is effective in training vision transformer as well as convnet such as ResNet-50.
We used the DINO ResNet-50 and DeiT-S/16 (Dosovitskiy et al., 2020) trained for 800 epochs with multi-crops applied.
We archive the checkpoints of teacher models from the author’s original implementation.
For each method, we also conduct self-supervised learning on ResNet-18 for fair comparison.
We used same hyper-parameters that were used to train ResNet-50 except that we trained for 100 epochs.
In Table 4, we report the results of self-supervised learning on ResNet-18.
A.3. N ETWORK ARCHITECTURES
We set the projection heads of student network to be same as the teacher network.
When teacher is
DINO DeiT, the teacher network do not contain batch normalization, but we add batch normalization to projection heads when training ResNet-18 student.
Remark that the projection heads of MoCo and SwAV have output dimension of 128, and the projection head of DINO has output of 256.
Then we set the number of prototypes to be K = 65536 throughout the experiments.
For SwAV, since we use pre-trained prototypes, the number of prototypes is 3000.
Every features are normalized before computation with prototypes, and prototypes are normalized during the training.
A.3. T RAINING HYPERPARAMETERS
For probability of teacher, we use SK operator with 3 steps of iteration and τ t = 0 .
For probability of student, we set τ s = 0 .
The prior momentum for ProtoCPC loss is 0.9.
We use SGD optimizer with batch size 512 and weight decay is 1e-4.
The learning rate is 0.6 and is decayed by cosine learning rate schedule to 1e-6.
A.3. E VALUATION
For evaluation, we use both linear evaluation protocol and k -nearest neighbor classiﬁcation.
For linear evaluation protocol, we freeze the trained weight and train a linear classiﬁer at the top of the frozen feature.
We train with SGD optimizer with batch size 256 and use learning rate of 0.3 with 100 epochs.
We use cosine learning rate decay schedule and we don’t use weight decay.
For k -NN classiﬁcation, we follow weighted k -NN with τ = 0 .
07 as done in (Wu et al., 2018).
A. S ELF - DISTILLATION
Most of our implementation is borrowed from DINO (Caron et al., 2021).
For self-distillation experiments, we only changed cross-entropy loss to ProtoCPC loss.
Detailed hyper-parameters are following:
• ResNet-50: we used 2 layers MLP projection head with output dimension 65536.
We trained for 100 epochs with 2x224 + 6x96 multi-crops, batch size 512, learning rate 0.6 with cosine annealing schedule, weight decay 1e-4, teacher temperature τ t = 0 .
04 , student temperature τ s = 0 .
1 , prior momentum m p =0.9.
• DeiT-S/16: we used 3 layers MLP projection head with output dimension 65536.
We trained for 100 epochs with 2x224 + 8x96 multi-crops, batch size 512, learning rate 0.001 with cosine annealing schedule, weight decay 1e-4, teacher temperature τ t = 0 .
04 , student temperature τ s = 0 .
1 , prior momentum m p =0.9.