Year: 2021

Abstract:
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 efficient 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.

heading: 1 INTRODUCTION
text: 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); Buciluǎ 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 efficient 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 significantly 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.

heading: 2 METHOD

heading: 2.1 PROTOTYPICAL CONTRASTIVE PREDICTIVE CODING
text: Given data x with a random variable x ∼ X , let the teacher network fT and the student network fS , where they map x into RDT and RDS respectively. Also let T and S be random variables for representation fT (X) and fS(X) respectively. One can transfer the representational knowledge of T to S by maximizing the mutual information I(T ;S), where it is defined 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 zs ∼ S with a positive zt ≡ zt0 and N − 1 negatives {ztj}N−1j=1 sampled from T , i.e. (zt, zs) ∼ p(T, S) and {(ztj , zs)}N−1j=1 ∼ p(T )p(S), the following inequality holds for any critic h : RDT × RDS → R+:
I(T ;S) ≥ E [ log h(zt, zs)
1 N ∑N−1 j=0 h(ztj , zs)
] (1)
Previous works set the critic by the exponential of cosine similarity between two unit feature vectors, i.e. h(zt, zs) = exp(−zt · zs) where zt and zs are `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 zs and ztjs.
On the other hand, we project feature vectors into a probability space. To do that, we append a linear prototypes WT ∈ RDT×K and WS ∈ RDS×K at the top of fT and fS so that they have same output dimension of K. For brevity, let z̄s = WSzs and z̄t = WT zt and T̄ , S̄ be random variables for z̄t and z̄s respectively. Then we set probability of student ps by K-categorical distribution defined by the softmax operator on z̄s = WSzs with temperature τs > 0:
p(k)s = exp
( z̄ (k) s /τs )∑K k′=1 exp ( z̄ (k′) s /τs
) . (2) Similarly we define probability of teacher ptj with temperature τt > 0. Then we define the critic between z̄t and z̄s by the negative exponential of cross-entropy between ps and pt, i.e.,
h(z̄t, z̄s) = e −H(pt,ps) = e
∑K k=1 p (k) t log p (k) s . (3)
Then h is a positive bounded function and is maximized when pt matches with ps. Then by plugging Eq. 3 into Eq. 1, it follows that
I(T ;S) ≥ I(T̄ ; S̄) ≥ E [ log e−H(pt,ps)
1 N ∑N−1 j=0 e −H(ptj ,ps)
] (4)
= E [ log exp(pt · z̄s/τs)
1 N ∑N−1 j=0 exp ( ptj · z̄s/τs )] (5) ≥ E [ log
exp(pt · z̄s/τs) 1 N ∑N−1 j=0 ∑K k=1 p (k) tj exp ( z̄ (k) s /τs )] (6) = E [ log
exp(pt · z̄s/τs)∑K k=1 q (k) exp ( z̄ (k) s /τs )] , IProtoCPC, (7) where q(k) = 1N ∑N−1 j=0 ptj is a mean of teachers’ probability which we call a prior. The first 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 define Prototypical Contrastive Predictive Coding (ProtoCPC) objective IProtoCPC in Eq. 7. In addition, we define ProtoCPC loss LProtoCPC 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 LProtoCPC-Unif by the re-substitution entropy estimator of zs via a von-Mises Fisher kernel density estimation (vMF-KDE) (Ahmad & Lin, 1976):
LProtoCPC-Unif = Ezs∼S [ log K∑ k=1 q(k) exp(z̄(k)s /τs) ] = Ezs∼S [ log K∑ k=1 q(k) exp(wk · zs/τs) ] (8)
= Ezs∼S [log p̂vMF-KDE(zs)] + logZvMF = −Ĥ(zs) + logZvMF, (9)
where each wk is a k-th column of WS 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̂(zs). The ZvMF 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 ztj 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 sufficient 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:
q(k) ← mpq(k) + (1−mp) 1
N N∑ j=1 p (k) tj , (10)
where mp > 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:
argmax Pt
〈Pt, Zt〉+ τtH(Pt), s.t. Pt ∈ RN×K+ , Pt1K = 1N , P>t 1N = N
K 1K , (11)
where Zt is a matrix whose rows are ztj , H(Pt) = ∑N−1 j=0 ∑K k=1−p (k) t log p (k) 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 Pt into following form:
p (k) tj = βke z (k) tj /τt∑K
k′=1 βk′e z (k′) tj /τt
, (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.

heading: 2.2 APPLICATIONS OF PROTOCPC TO DISTILLATION TASKS
text: 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 fT , the goal of supervised model compression is to train a smaller student network fS 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 fS . The distillation loss is a cross-entropy between teacher’s probability output pt and student’s probability output ps. Formally, it is given by following:
min fS
E(x,y)∼(X,Y ) [ LSup(fS(x), y) + βτ2H(pt(x), ps(x)) ] , (13)
where y is a label, pt(x) and ps(x) are computed by the softmax operator on each fT (x) and fS(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 pt. Then our loss becomes following:
min fS
E(x,y)∼(X,Y ) [ LSup(fS(x), y) + βτ2LProtoCPC(pt(x), ps(x)) ] . (14)
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 gT , the goal of self-supervised model compression is to transfer the representational knowledge of gT to a smaller self-supervised student network. Fang et al. (2021) first 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 filled 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 gT be a composition of base encoder fT and the projection head hT . Then we train student network gS , where it is composed of smaller encoder fS and projection head hS of the same architecture as hT . Then we append prototypes WT and WS for each gT and gS 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 pt(x) and probability of student ps(x). For ps we use softmax operator with temperature τs > 0 and for pt we use SK operator with temperature τt > 0 over the batch of samples. Then the objective is given by following:
min gS ,WS
Ex∼X [ LProtoCPC(pt(x), ps(x)) ] . (15)
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 transformationM, then the self-supervised model compression with multi-crops is given by following:
min gS ,WS
Ex∼X,x̃j∼A(x) [ LProtoCPC(pt(x), ps(x)) + ∑ j LProtoCPC(pt(x), ps(x̃j)) ] . (16)
For prototypes of teacher network WT , we copy the parameters of WS to WT 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 WT . We present ablation studies on setting prototypes for WT .
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 first 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 (x1, x2) 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 pt(x1) and probability of student network ps(x2), and symmetrically between pt(x2) and ps(x1). The objectives are given by:
min θs
Ex∼X,x1,x2∼A(x) [ 1
2 LProtoCPC
( pt(x1), ps(x2) ) + 1
2 LProtoCPC
( pt(x2), ps(x1) )] , (17)
where A(·) is a data augmentation operator. Note that multi-crops data augmentation can be further applied.

heading: 3 RELATED WORK
text: 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 `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; Buciluǎ 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.

heading: 4 EXPERIMENT

heading: 4.1 SUPERVISED MODEL COMPRESSION
text: 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.
Results on ImageNet In Table 3, 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. Remark that our method is simpler in its implementation and shows better performance than CRD.

heading: 4.2 SELF-SUPERVISED MODEL COMPRESSION
text: 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 classification (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 classification. 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).
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.

heading: 4.3 SELF-DISTILLATION
text: 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.
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 classification.
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 mp and observe that non-zero mp results in better representation learning.

heading: 4.4 ABLATION STUDY
text: 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.

heading: 5 CONCLUSION
text: 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.

heading: A DETAILS IN EXPERIMENTS

heading: A.1 PSEUDOCODE FOR PROTOCPC LOSS
text: The PyTorch style pseudo-code for our ProtoCPC is demonstrated in Algorithm 1.
Algorithm 1 ProtoCPC loss PyTorch-style pseudocode.
# tps, tpt: student and teacher temperatures # m: prior momentum rate prior = torch.ones(1, K) # initialize prior with uniform def ProtoCPC(zt, zs):
zt = zt.detach() pt = SK(zt / tpt) zs = zs / tps
prior = m*prior + (1-m)* K * torch.mean(pt, dim=0) # sum of prior is always K
loss_align = -torch.sum(pt * zs, dim=1) loss_unif = torch.logsumexp(zs + torch.log(prior), 1) loss = loss_align + loss_unif return loss.mean()

heading: A.2 SUPERVISED MODEL COMPRESSION
text: Many of followings are borrowed from CRD (Tian et al., 2019).

heading: A.2.1 OTHER METHODS
text: We compare to the following other state-of-the-art methods from the literature:
1. 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)

heading: A.2.2 NETWORK ARCHITECTURES
text: • 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 x4 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. (2018). In our experiments, we use a width multiplier of 0.5. • vgg (Simonyan & Zisserman, 2014). the vgg net used in our experiments are adapted from
its original ImageNet counterpart. • ShuffleNetV1 (Zhang et al., 2018), ShuffleNetV2 (Ma et al., 2018). ShuffleNets are proposed
for efficient training and we adapt them to input of size 32x32.
A.2.3 IMPLEMENTATION DETAILS
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 first 150 epochs until the last 240 epoch. For MobileNetV2, ShuffleNetV1 and ShuffleNetV2, 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 classification objective and a knowledge distillation objective, shown as follows:
L = αLsup + βLdistill (18)
For the weight balance factor β, we directly use the optimal value from the original paper if it is specified, 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:
1. KD (Hinton et al., 2015): α = 0.1, β = 0.9 and T = 4 2. Fitnets (Romero et al., 2014): α = 1, β = 100 3. AT (Zagoruyko & Komodakis, 2016a): α = 1, β = 1000 4. SP (Tung & Mori, 2019): α = 1, β = 3000 5. CC (Peng et al., 2019): α = 1, β = 0.02 6. VID (Ahn et al., 2019): α = 1, β = 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. 8. PKT (Passalis & Tefas, 2018): α = 1, β = 30000 9. AB (Heo et al., 2019): β = 0, distillation happens in a separate pre-training stage where
only distillation objective applies. 10. FT (Kim et al., 2018): α = 1, β = 500 11. FSP (Yim et al., 2017): β = 0, distillation happens in a separate pre-training stage where
only distillation objective applies. 12. NST (Huang & Wang, 2017): α = 1, β = 50 13. CRD (Tian et al., 2019): α = 1, β = 0.8
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.

heading: A.3 SELF-SUPERVISED MODEL COMPRESSION

heading: A.3.1 PRE-TRAINED TEACHERS
text: • 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 first-in-first-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 first 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.

heading: A.3.2 NETWORK ARCHITECTURES
text: 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.

heading: A.3.3 TRAINING HYPERPARAMETERS
text: For probability of teacher, we use SK operator with 3 steps of iteration and τt = 0.04. For probability of student, we set τs = 0.1. 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.

heading: A.3.4 EVALUATION
text: For evaluation, we use both linear evaluation protocol and k-nearest neighbor classification. For linear evaluation protocol, we freeze the trained weight and train a linear classifier 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 classification, we follow weighted k-NN with τ = 0.07 as done in (Wu et al., 2018).

heading: A.4 SELF-DISTILLATION
text: 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 mp=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 mp=0.9.

Review summaries:
This paper studies a new contrastive loss function for knowledge distillation and self-supervised representation learning. The paper demostrates results on transferring knowledge from a teacher to student using both supervised and self-supervised objectives, and on representation learning based on self-distillation. As prior work has already applied contrastive learning to each of these problems the main novelty of the paper is the exact form of the contrastive loss. The loss is computed over projections of the embeddings into a k-dimensional simplex (probability space) where cross-entropy can be used as the similarity function. This allows for encoding "negatives" into a distribution over k prototypes that are updated with EMA over iterations of training rather than requiring a large batch of negatives during each iteration of training. The results appear to slightly outperform SOTA methods.

This paper provides a method for combining contrastive learning and clustering (prototypical probabilities) for three knowledge distillation tasks - supervised model compression, self-supervised model compression, and self-supervised learning with self-distillation.  Traditional contrastive learning methods rely on the similarity of feature vectors from the teacher and student networks using cosine similarity, log bilinear model (van den Oord et al), etc. This paper proposes to project the feature vectors into a probability space via linear prototypes (layer) of the same dimension (i.e, cluster the feature vectors). The proposed method overcomes the disadvantage of having large negative examples for contrastive learning by maintaining an EMA of prior that accounts for the negative examples.  The main contribution of this paper is the ProtoCPC loss that combines the advantages of contrastive predictive coding loss using prototypes over the teacher and student representation.   The authors empirically show that the proposed method is useful for various distillation tasks on image classification tasks on CIFAR-100 and Imagenet datasets.

This paper proposes to combine knowledge distillation (KD) and contrastive learning for distillation tasks. Concretely, it models the critic of a contrastive objective by the prototypical probabilistic discrepancy between two features. The authors then carry out extensive experiments on supervised model compression, self-supervised model compression, and self-supervised learning through self-distillation. The empirical results show that the proposed method outperforms other strong baselines.