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