Prototypical Contrastive Predictive Coding

1 INTRODUCTION . 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 can not 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 . 2 METHOD . 2.1 PROTOTYPICAL CONTRASTIVE PREDICTIVE CODING . 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 .

Summary: 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.